diff --git a/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/TaskedSpaceBasedStep.java b/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/TaskedSpaceBasedStep.java index eedd8835f9..738c2ed91a 100644 --- a/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/TaskedSpaceBasedStep.java +++ b/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/TaskedSpaceBasedStep.java @@ -38,6 +38,7 @@ import com.here.xyz.jobs.JobClientInfo; import com.here.xyz.jobs.steps.execution.LambdaBasedStep.LambdaStepRequest.ProcessUpdate; import com.here.xyz.jobs.steps.execution.StepException; +import com.here.xyz.jobs.steps.execution.db.Database; import com.here.xyz.jobs.steps.impl.SpaceBasedStep; import com.here.xyz.jobs.steps.impl.transport.tasks.TaskPayload; import com.here.xyz.jobs.steps.impl.transport.tasks.TaskProgress; @@ -58,10 +59,16 @@ import java.io.IOException; import java.math.BigDecimal; import java.math.RoundingMode; +import java.sql.Array; +import java.sql.ResultSet; import java.sql.SQLException; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -83,7 +90,12 @@ public abstract class TaskedSpaceBasedStep extends SpaceBasedStep { private static final Logger logger = LogManager.getLogger(); + /** Marker output-set key/file prefix used to indicate successful step finalization. */ public static final String FINALIZATION_MARKER = "finalized_marker"; + /** Number of consecutive unknown running-query checks before a task is considered retryable. */ + public static final Integer MAX_UNKNOWN_TASK_QUERY_CHECKS = 3; + /** Maximum number of retry attempts, for a server side killed single, before failing retry handling. */ + public static final Integer MAX_TASK_RETRY_ATTEMPTS = 3; private TaskedSpaceBasedQueryBuilder taskedSpaceBasedQueryBuilder; { @@ -491,8 +503,9 @@ private void startTask(TaskProgress taskProgressAndItem) throws TooManyResour prepareTaskQuery(taskProgressAndItem.getTaskId()); - runReadQueryAsync(buildTaskQuery(taskProgressAndItem.getTaskId(), (I) taskProgressAndItem.getTaskInput(), failureCallback), - queryRunsOnWriter() ? dbWriter() : dbReader(), 0d/*perItemAcus.doubleValue()*/, false); + runReadQueryAsync(buildTaskQuery(taskProgressAndItem.getTaskId(), (I) taskProgressAndItem.getTaskInput(), failureCallback) + .withLabel(getId() + "#taskId", String.valueOf(taskProgressAndItem.getTaskId())), + queryRunsOnWriter() ? dbWriter() : dbReader(), 0d/*perItemAcus.doubleValue()*/, false); } } @@ -696,7 +709,7 @@ public AsyncExecutionState getExecutionState() throws UnknownStateException { return AsyncExecutionState.SUCCEEDED; try { - TaskProgress taskProgress = getTaskProgress(); + TaskProgress taskProgress = getTaskProgress(); return evaluateExecutionState(taskProgress); } catch (SQLException e) { @@ -709,13 +722,57 @@ public AsyncExecutionState getExecutionState() throws UnknownStateException { } } - private AsyncExecutionState evaluateExecutionState(TaskProgress taskProgress) throws UnknownStateException { + private AsyncExecutionState evaluateExecutionState(TaskProgress taskProgress) + throws UnknownStateException, SQLException, WebClientException, TooManyResourcesClaimed { if (taskProgress.isComplete()) { //Only log and return RUNNING to avoid double success handling; infoLog(STEP_ON_STATE_CHECK, "All tasks finalized!"); - }if (taskProgress.hasRunningTasks()) - // Check if the expected queries are still running. - return super.getExecutionState(); + } + if (taskProgress.hasRunningTasks()) { + Set expectedTaskIds = taskProgress.getStartedNotFinalizedTaskIds(); + Set expectedTaskIdStrings = expectedTaskIds.stream() + .map(String::valueOf) + .collect(Collectors.toSet()); + + if (expectedTaskIdStrings.isEmpty()) { + infoLog(STEP_ON_STATE_CHECK, "Started tasks are reported but started_not_finalized_task_ids is empty."); + throw new UnknownStateException("No started-not-finalized task ids available for step: " + getGlobalStepId() + " !"); + } + + boolean runsOnWriter = queryRunsOnWriter(); + Database database = runsOnWriter ? dbWriter() : dbReader(); + //Check if all expected TaskQueries areRunning. Collect all taskIds where we are not able + //to find a running query. + Set runningTaskIds = SQLQuery.areRunning( + requestResource(database, 0d), database.getRole() != WRITER, + getId() + "#taskId", expectedTaskIdStrings + ); + + Set unknownStateTaskIds = expectedTaskIds.stream() + .filter(taskId -> !runningTaskIds.contains(String.valueOf(taskId))) + .collect(Collectors.toSet()); + + if (unknownStateTaskIds.isEmpty()) + return AsyncExecutionState.RUNNING; + + Set exceededTaskIds = incrementUnknownQueryStateForTasks(unknownStateTaskIds); + for(int exceededTaskId : exceededTaskIds){ + infoLog(STEP_ON_STATE_CHECK, "Unknown queryState of taskId " + exceededTaskId + " has exceeded unknown query state threshold. Retry the task now!"); + try { + TaskProgress retryTaskProgress = resetTaskForRetry(exceededTaskId); + if (retryTaskProgress == null) { + throw new StepException("Retry threshold exceeded for taskId " + exceededTaskId + "."); + } + startTask(retryTaskProgress); + } catch (StepException e1){ + throw e1; + } catch (Exception e2){ + throw new StepException("Not able to retry taskId " + exceededTaskId, e2); + } + } + + return AsyncExecutionState.RUNNING; + } if (taskProgress.hasNoRunningTasks()) { infoLog(STEP_ON_STATE_CHECK, "No running tasks detected. StartedTasks: " + taskProgress.getStartedTasks() + "," + " FinalizedTasks: " + taskProgress.getFinalizedTasks() + " !"); @@ -783,6 +840,40 @@ private void updateQueryTaskItemOutput(SpaceBasedTaskUpdate update) throws WebCl , db(WRITER), 0); } + private Set incrementUnknownQueryStateForTasks(Set unknownStateTaskIds) + throws WebClientException, SQLException, TooManyResourcesClaimed { + if (unknownStateTaskIds == null || unknownStateTaskIds.isEmpty()) + return Set.of(); + + SQLQuery query = getQueryBuilder().buildIncrementUnknownQueryStateStatement(unknownStateTaskIds, MAX_UNKNOWN_TASK_QUERY_CHECKS); + + return runReadQuerySync(query, db(WRITER), 0, rs -> { + Set exceededTaskIds = new java.util.LinkedHashSet<>(); + while (rs.next()) + exceededTaskIds.add(rs.getInt("task_id")); + return exceededTaskIds; + }); + } + + private TaskProgress resetTaskForRetry(int taskId) throws WebClientException, SQLException, TooManyResourcesClaimed { + return runReadQuerySync(getQueryBuilder().buildResetTaskForRetryStatement(taskId), db(WRITER), 0, rs -> { + if (!rs.next()) + throw new SQLException("Task for retry not found: taskId=" + taskId); + + int retryAttempt = rs.getInt("retry_attempts"); + if (retryAttempt > MAX_TASK_RETRY_ATTEMPTS) + return null; + + try { + return new TaskProgress<>(rs.getInt("task_id"), + XyzSerializable.deserialize(rs.getString("task_input"), new TypeReference() {})); + } + catch (JsonProcessingException e) { + throw new StepException("Can not deserialize task_input for retry of taskId=" + taskId + "!", e); + } + }); + } + private TaskProgress executeTaskItemQuery(SQLQuery query) throws WebClientException, SQLException, TooManyResourcesClaimed { TaskProgress taskProgress; try { @@ -816,10 +907,20 @@ private TaskProgress getTaskProgress() throws WebClientException, SQLException, rs -> { if (!rs.next()) return null; - return new TaskProgress(rs.getInt("total"), rs.getInt("started"), rs.getInt("finalized")); + return new TaskProgress(rs.getInt("total"), rs.getInt("started"), rs.getInt("finalized"),getIntegerList(rs, "started_not_finalized_task_ids")); }); } + private Set getIntegerList(ResultSet rs, String columnName) throws SQLException { + Array sqlArray = rs.getArray(columnName); + + return sqlArray == null + ? Set.of() + : Arrays.stream((Object[]) sqlArray.getArray()) + .map(value -> ((Number) value).intValue()) + .collect(Collectors.toSet()); + } + private SQLQuery resetTaskItemWhichAreNotFinalized() { infoLog(STEP_EXECUTE, "Reset task items for restart."); return getQueryBuilder().buildResetTaskItemWhichAreNotFinalizedStatement(); diff --git a/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/tasks/TaskProgress.java b/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/tasks/TaskProgress.java index cbb287363c..d985810317 100644 --- a/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/tasks/TaskProgress.java +++ b/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/tasks/TaskProgress.java @@ -18,12 +18,15 @@ */ package com.here.xyz.jobs.steps.impl.transport.tasks; +import java.util.Set; + public class TaskProgress { private int totalTasks; private int startedTasks; private int finalizedTasks; private Integer taskId; private I taskInput; + private Set startedNotFinalizedTaskIds = Set.of(); public TaskProgress() {} @@ -31,6 +34,11 @@ public TaskProgress(Integer taskId) { this.taskId = taskId; } + public TaskProgress(Integer taskId, I taskInput) { + this.taskId = taskId; + this.taskInput = taskInput; + } + public TaskProgress(int totalTasks, int startedTasks, int finalizedTasks) { this.totalTasks = totalTasks; this.startedTasks = startedTasks; @@ -45,6 +53,13 @@ public TaskProgress(int totalTasks, int startedTasks, int finalizedTasks, Intege this.taskInput = taskInput; } + public TaskProgress(int totalTasks, int startedTasks, int finalizedTasks, Set startedNotFinalizedTaskIds) { + this.totalTasks = totalTasks; + this.startedTasks = startedTasks; + this.finalizedTasks = finalizedTasks; + this.startedNotFinalizedTaskIds = startedNotFinalizedTaskIds; + } + public int getTotalTasks() { return totalTasks; } @@ -85,6 +100,14 @@ public void setTaskInput(I taskInput) { this.taskInput = taskInput; } + public Set getStartedNotFinalizedTaskIds() { + return startedNotFinalizedTaskIds; + } + + public void setStartedNotFinalizedTaskIds(Set startedNotFinalizedTaskIds) { + this.startedNotFinalizedTaskIds = startedNotFinalizedTaskIds; + } + public boolean isComplete() { return totalTasks == finalizedTasks; } @@ -109,6 +132,7 @@ public String toString() { "totalTasks=" + totalTasks + ", startedTasks=" + startedTasks + ", finalizedTasks=" + finalizedTasks + + ", startedNotFinalizedTaskIds=" + startedNotFinalizedTaskIds + ", taskId=" + taskId + ", taskInput=" + taskInput + '}'; diff --git a/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/tools/TaskedSpaceBasedQueryBuilder.java b/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/tools/TaskedSpaceBasedQueryBuilder.java index ef6ee61172..5d0525d5bb 100644 --- a/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/tools/TaskedSpaceBasedQueryBuilder.java +++ b/xyz-jobs/xyz-job-steps/src/main/java/com/here/xyz/jobs/steps/impl/transport/tools/TaskedSpaceBasedQueryBuilder.java @@ -23,6 +23,8 @@ import com.here.xyz.models.hub.Space; import com.here.xyz.util.db.SQLQuery; +import java.util.Set; + import static com.here.xyz.events.ContextAwareEvent.SpaceContext; import static com.here.xyz.jobs.steps.impl.transport.TaskedSpaceBasedStep.SpaceBasedTaskUpdate; @@ -42,10 +44,13 @@ public SQLQuery buildTaskTableStatement() { task_output JSONB, started BOOLEAN DEFAULT false, finalized BOOLEAN DEFAULT false, + unknown_query_state_occurrences INTEGER DEFAULT 0, + retry_attempts INTEGER DEFAULT 0, + started_at TIMESTAMP DEFAULT NULL, + updated_at TIMESTAMP DEFAULT NULL, CONSTRAINT ${primaryKey} PRIMARY KEY (task_id) ); """) - //TODO: CHECK CONSTRAINT!! .withVariable("table", getTemporaryJobTableName()) .withVariable("schema", schema) .withVariable("primaryKey", getTemporaryJobTableName() + "_primKey"); @@ -54,15 +59,17 @@ public SQLQuery buildTaskTableStatement() { public SQLQuery buildUpdateTaskItemOutputStatement(SpaceBasedTaskUpdate update) { return new SQLQuery(""" UPDATE ${schema}.${table} - SET task_output = ( - COALESCE(task_output, '{}'::JSONB) || #{taskUpdate}::JSONB - ) || jsonb_build_object( - 'taskOutput', - COALESCE(task_output->'taskOutput', '{}'::JSONB) - || COALESCE((#{taskUpdate}::JSONB)->'taskOutput', '{}'::JSONB) - ) - WHERE task_id = #{taskId}; - """) + SET + updated_at = now(), + task_output = ( + COALESCE(task_output, '{}'::JSONB) || #{taskUpdate}::JSONB + ) || jsonb_build_object( + 'taskOutput', + COALESCE(task_output->'taskOutput', '{}'::JSONB) + || COALESCE((#{taskUpdate}::JSONB)->'taskOutput', '{}'::JSONB) + ) + WHERE task_id = #{taskId}; + """) .withVariable("schema", schema) .withVariable("table", getTemporaryJobTableName()) .withNamedParameter("taskId", update.taskId) @@ -72,7 +79,8 @@ public SQLQuery buildUpdateTaskItemOutputStatement(SpaceBasedTaskUpdate update) public SQLQuery buildResetTaskItemWhichAreNotFinalizedStatement() { return new SQLQuery(""" UPDATE ${schema}.${table} t - SET started = false + SET started = false, + updated_at = now() WHERE started = true AND finalized = false; """) .withVariable("schema", schema) @@ -89,7 +97,8 @@ public SQLQuery retrieveTaskStatisticsQuery() { return new SQLQuery(""" SELECT COUNT(1) as total, SUM((started = true)::int) as started, - SUM((finalized = true)::int) as finalized + SUM((finalized = true)::int) as finalized, + ARRAY_AGG(task_id) FILTER (WHERE started = true AND finalized = false) as started_not_finalized_task_ids FROM ${schema}.${table}; """) .withVariable("schema", schema) @@ -144,6 +153,42 @@ public SQLQuery buildTemporaryJobTableDropStatement() { .withVariable("schema", schema); } + public SQLQuery buildIncrementUnknownQueryStateStatement(Set unknownStateTaskIds, int maxUnknownTaskQueryChecks) { + return new SQLQuery(""" + WITH updated AS ( + UPDATE ${schema}.${table} + SET unknown_query_state_occurrences = COALESCE(unknown_query_state_occurrences, 0) + 1, + updated_at = now() + WHERE task_id IN ( + SELECT value::INT + FROM jsonb_array_elements_text(#{missingTaskIds}::JSONB) + ) + RETURNING task_id, unknown_query_state_occurrences + ) + SELECT task_id + FROM updated + WHERE unknown_query_state_occurrences >= #{maxUnknownTaskQueryChecks}; + """) + .withVariable("schema", schema) + .withVariable("table", getTemporaryJobTableName()) + .withNamedParameter("missingTaskIds", XyzSerializable.serialize(unknownStateTaskIds)) + .withNamedParameter("maxUnknownTaskQueryChecks", maxUnknownTaskQueryChecks); + } + + public SQLQuery buildResetTaskForRetryStatement(int taskId) { + return new SQLQuery(""" + UPDATE ${schema}.${table} + SET unknown_query_state_occurrences = 0, + updated_at = now(), + retry_attempts = retry_attempts + 1 + WHERE task_id = #{taskId} + RETURNING task_id, task_input, retry_attempts; + """) + .withVariable("schema", schema) + .withVariable("table", getTemporaryJobTableName()) + .withNamedParameter("taskId", taskId); + } + private String getTemporaryJobTableName(String stepId) { return JOB_DATA_PREFIX + stepId; } diff --git a/xyz-jobs/xyz-job-steps/src/main/resources/jobs/transport.sql b/xyz-jobs/xyz-job-steps/src/main/resources/jobs/transport.sql index 1ff7242185..11936fce4e 100644 --- a/xyz-jobs/xyz-job-steps/src/main/resources/jobs/transport.sql +++ b/xyz-jobs/xyz-job-steps/src/main/resources/jobs/transport.sql @@ -569,7 +569,7 @@ BEGIN IF task_item.task_id IS NOT NULL THEN EXECUTE format( 'UPDATE %1$s C - SET started = true + SET started = true, started_at = now() WHERE C.task_id = %2$L;', get_table_reference(ctx->>'schema', ctx->>'stepId' ,'JOB_TABLE'), task_item.task_id @@ -663,7 +663,7 @@ BEGIN COALESCE(t.task_output->''taskOutput'', ''{}''::JSONB) || COALESCE((%2$L::JSONB)->''taskOutput'', ''{}''::JSONB) ), - finalized = %3$L + updated_at = now(), finalized = %3$L WHERE task_id = %4$L;', get_table_reference(ctx->>'schema', ctx->>'stepId', 'JOB_TABLE'), p_task_output::TEXT, diff --git a/xyz-util/src/main/java/com/here/xyz/util/db/SQLQuery.java b/xyz-util/src/main/java/com/here/xyz/util/db/SQLQuery.java index 3b6f492ddd..050ed9ec9c 100644 --- a/xyz-util/src/main/java/com/here/xyz/util/db/SQLQuery.java +++ b/xyz-util/src/main/java/com/here/xyz/util/db/SQLQuery.java @@ -40,6 +40,7 @@ import java.util.Collection; import java.util.Collections; import java.util.HashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Map.Entry; @@ -910,13 +911,24 @@ private static void killByLabel(String labelIdentifier, String labelValue, long } private static SQLQuery buildLabelMatchQuery(String labelIdentifier, String labelValue) { - return new SQLQuery("strpos(query, '/*labels(') > 0 AND substring(query, " - + "strpos(query, '/*labels(') + 9, " - + "strpos(query, ')*/') - 9 - strpos(query, '/*labels('))::json->>#{labelIdentifier} = #{labelValue}") - .withNamedParameter("labelIdentifier", labelIdentifier) + return new SQLQuery("strpos(query, '/*labels(') > 0 AND ${{extractedLabelValue}} = #{labelValue}") + .withQueryFragment("extractedLabelValue", buildLabelValueExtraction(labelIdentifier)) .withNamedParameter("labelValue", labelValue); } + /** + * Builds the SQL fragment that extracts the value of the label with the specified identifier from the {@code /*labels(...)*}{@code /} + * comment prefix of a query in {@code pg_stat_activity}. + * @param labelIdentifier The label identifier whose value should be extracted (e.g. taskId). + * @return The SQL fragment resolving to the label value as text. + */ + private static SQLQuery buildLabelValueExtraction(String labelIdentifier) { + return new SQLQuery("substring(query, " + + "strpos(query, '/*labels(') + 9, " + + "strpos(query, ')*/') - 9 - strpos(query, '/*labels('))::json->>#{labelIdentifier}") + .withNamedParameter("labelIdentifier", labelIdentifier); + } + public static boolean isRunning(DataSourceProvider dataSourceProvider, boolean useReplica, String queryId) throws SQLException { return isRunning(dataSourceProvider, useReplica, QUERY_ID, queryId); } @@ -938,6 +950,59 @@ public static boolean isRunning(DataSourceProvider dataSourceProvider, boolean u .run(dataSourceProvider, rs -> rs.next(), useReplica); } + /** + * Checks which of the provided values are currently running for one label identifier. + * + * @param labelIdentifier The label identifier to check (e.g. taskId). + * @param labelValues The set of values to check for the given label identifier. + * @return The subset of provided label values that are currently running. + */ + public static Set areRunning(DataSourceProvider dataSourceProvider, boolean useReplica, + String labelIdentifier, Set labelValues) throws SQLException { + if (labelValues == null || labelValues.isEmpty()) + return Set.of(); + + return new SQLQuery(""" + WITH expected_values AS (${{expectedValues}}), + active_values AS ( + SELECT ${{labelValue}} AS label_value + FROM pg_stat_activity + WHERE state = 'active' + AND pid != pg_backend_pid() + AND strpos(query, '/*labels(') > 0 + ) + SELECT DISTINCT e.label_value + FROM expected_values e + JOIN active_values a + ON a.label_value = e.label_value + """) + .withQueryFragment("expectedValues", buildLabelValuesQuery(labelValues)) + .withQueryFragment("labelValue", buildLabelValueExtraction(labelIdentifier)) + .withLoggingEnabled(false) + .run(dataSourceProvider, rs -> { + Set runningValues = new LinkedHashSet<>(); + while (rs.next()) + runningValues.add(rs.getString("label_value")); + return runningValues; + }, useReplica); + } + + private static SQLQuery buildLabelValuesQuery(Set labelValues) { + if (labelValues == null || labelValues.isEmpty()) + return new SQLQuery("SELECT NULL::text AS label_value WHERE FALSE"); + + List values = new ArrayList<>(); + int i = 0; + for (String labelValue : labelValues) { + String valueParam = "labelValue" + i; + values.add(new SQLQuery("SELECT #{" + valueParam + "} AS label_value") + .withNamedParameter(valueParam, labelValue)); + i++; + } + + return SQLQuery.join(values, " UNION ALL "); + } + private static String getClashing(Map map1, Map map2) { if (map1 == null || map2 == null) return null; diff --git a/xyz-util/src/test/java/com/here/xyz/test/sql/base/SQLQueryIT.java b/xyz-util/src/test/java/com/here/xyz/test/sql/base/SQLQueryIT.java index b887d5ffb1..e194cf4f23 100644 --- a/xyz-util/src/test/java/com/here/xyz/test/sql/base/SQLQueryIT.java +++ b/xyz-util/src/test/java/com/here/xyz/test/sql/base/SQLQueryIT.java @@ -27,6 +27,8 @@ import com.here.xyz.util.db.datasource.DataSourceProvider; import java.sql.SQLException; import java.util.Map; +import java.util.Set; + import org.junit.jupiter.api.Test; public class SQLQueryIT extends SQLITBase { @@ -64,6 +66,70 @@ public void startAndKillQuery() throws Exception { } } + @Test + public void areRunningCheckMultipleValuesForSameLabel() throws Exception { + try (DataSourceProvider dsp = getDataSourceProvider()) { + SQLQuery longRunningQuery1 = new SQLQuery("SELECT pg_sleep(500)") + .withLabel("taskId", "1"); + SQLQuery longRunningQuery2 = new SQLQuery("SELECT pg_sleep(500)") + .withLabel("taskId", "2"); + + new Thread(() -> { + try { + longRunningQuery1.run(dsp); + } + catch (SQLException e) { + //Ignore + } + }).start(); + + new Thread(() -> { + try { + longRunningQuery2.run(dsp); + } + catch (SQLException e) { + //Ignore + } + }).start(); + + //Wait until both queries are visible in pg_stat_activity. + long deadline = System.currentTimeMillis() + 5_000; + while (System.currentTimeMillis() < deadline && SQLQuery.areRunning(dsp, false, "taskId", Set.of("1", "2")).size() < 2) + Thread.sleep(100); + + //Check that both taskIds are currently running + assertEquals(Set.of( + "1", + "2" + ), SQLQuery.areRunning(dsp, false, "taskId", Set.of( + "1", + "2" + ))); + + assertEquals(Set.of( + "1" + ), SQLQuery.areRunning(dsp, false, "taskId", Set.of( + "1" + ))); + + assertEquals(Set.of( + "2" + ), SQLQuery.areRunning(dsp, false, "taskId", Set.of( + "2" + ))); + + //Kill the long-running queries + longRunningQuery1.kill(); + longRunningQuery2.kill(); + + //Check that the original queries are not running anymore + assertFalse(SQLQuery.isRunning(dsp, false, longRunningQuery1.getQueryId())); + assertFalse(SQLQuery.isRunning(dsp, false, longRunningQuery2.getQueryId())); + assertEquals(Set.of(), SQLQuery.areRunning(dsp, false, "taskId", Set.of( + "1","2"))); + } + } + @Test public void runAsyncQuery() throws Exception { SQLQuery longRunningAsyncQuery = new SQLQuery("SELECT pg_sleep(10)").withAsync(true);