This is an automated email from the ASF dual-hosted git repository.
kfaraz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 8d634fdf8d9 fix: cancel active `/workOrder` futures on query
cancellation (#20313)
8d634fdf8d9 is described below
commit 8d634fdf8d9ac14b763db30636157407261675fc
Author: Sasha Syrotenko <[email protected]>
AuthorDate: Thu Oct 15 17:13:56 2026 +0300
fix: cancel active `/workOrder` futures on query cancellation (#20313)
Changes:
- Track in-flight requests in `DartWorkerClientImpl`
- Cancel in-flight requests when the client is closed
---
embedded-tests/pom.xml | 6 +
.../testing/embedded/msq/BaseDartQueryTest.java | 131 ++++++++++++++
.../msq/EmbeddedDartQueryCancellationTest.java | 195 +++++++++++++++++++++
.../embedded/msq/EmbeddedDartReportApiTest.java | 103 +----------
.../msq/dart/controller/DartWorkerManager.java | 7 +
.../msq/dart/worker/DartWorkerClientImpl.java | 78 ++++++++-
.../org/apache/druid/msq/exec/ControllerImpl.java | 2 +-
.../msq/dart/controller/DartWorkerManagerTest.java | 43 +++++
.../msq/dart/worker/DartWorkerClientImplTest.java | 99 +++++++++++
9 files changed, 559 insertions(+), 105 deletions(-)
diff --git a/embedded-tests/pom.xml b/embedded-tests/pom.xml
index 3928027576a..31a89e0b921 100644
--- a/embedded-tests/pom.xml
+++ b/embedded-tests/pom.xml
@@ -401,6 +401,7 @@
<artifactId>simple-client-sslcontext</artifactId>
<version>${project.parent.version}</version>
<scope>test</scope>
+
</dependency>
<!-- Pin protobuf-java-util to the version from druid-protobuf-extensions
to satisfy
the enforcer's RequireUpperBoundDeps rule (google-cloud-storage
brings a lower version). -->
@@ -830,6 +831,11 @@
<artifactId>docker-java-transport</artifactId>
<scope>test</scope>
</dependency>
+ <dependency>
+ <groupId>org.eclipse.jetty</groupId>
+ <artifactId>jetty-server</artifactId>
+ <scope>test</scope>
+ </dependency>
</dependencies>
<build>
diff --git
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/BaseDartQueryTest.java
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/BaseDartQueryTest.java
new file mode 100644
index 00000000000..e3c6560fe79
--- /dev/null
+++
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/BaseDartQueryTest.java
@@ -0,0 +1,131 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.testing.embedded.msq;
+
+import org.apache.druid.common.utils.IdUtils;
+import org.apache.druid.guice.SleepModule;
+import org.apache.druid.indexing.common.task.IndexTask;
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.sql.http.GetQueryReportResponse;
+import org.apache.druid.testing.embedded.EmbeddedBroker;
+import org.apache.druid.testing.embedded.EmbeddedClusterApis;
+import org.apache.druid.testing.embedded.EmbeddedCoordinator;
+import org.apache.druid.testing.embedded.EmbeddedDruidCluster;
+import org.apache.druid.testing.embedded.EmbeddedHistorical;
+import org.apache.druid.testing.embedded.EmbeddedIndexer;
+import org.apache.druid.testing.embedded.EmbeddedOverlord;
+import org.apache.druid.testing.embedded.auth.EmbeddedBasicAuthResource;
+import org.apache.druid.testing.embedded.indexing.MoreResources;
+import org.apache.druid.testing.embedded.junit5.EmbeddedClusterTestBase;
+import org.junit.jupiter.api.BeforeAll;
+
+/**
+ * Shared embedded cluster setup for Dart query tests.
+ */
+public abstract class BaseDartQueryTest extends EmbeddedClusterTestBase
+{
+ private static final int MAX_RETAINED_REPORT_COUNT = 10;
+
+ protected final EmbeddedBroker broker1 = new EmbeddedBroker();
+ protected final EmbeddedBroker broker2 = new EmbeddedBroker();
+ protected final EmbeddedIndexer indexer = new EmbeddedIndexer();
+ protected final EmbeddedOverlord overlord = new EmbeddedOverlord();
+ protected final EmbeddedHistorical historical = new EmbeddedHistorical();
+ protected final EmbeddedCoordinator coordinator = new EmbeddedCoordinator();
+
+ protected EmbeddedMSQApis msqApis;
+ protected String ingestedDataSource;
+
+ private void configureBroker(final EmbeddedBroker broker, final int port)
+ {
+ broker.addProperty("druid.msq.dart.controller.heapFraction", "0.5")
+ .addProperty("druid.msq.dart.controller.maxRetainedReportCount",
String.valueOf(MAX_RETAINED_REPORT_COUNT))
+ .addProperty("druid.query.default.context.maxConcurrentStages", "1")
+ .addProperty("druid.sql.planner.enableSysQueriesTable", "true")
+ .addProperty("druid.plaintextPort", String.valueOf(port));
+ }
+
+ @Override
+ protected EmbeddedDruidCluster createCluster()
+ {
+ coordinator.addProperty("druid.manager.segments.useIncrementalCache",
"always");
+ overlord.addProperty("druid.manager.segments.pollDuration", "PT0.1s");
+
+ configureBroker(broker1, 7082);
+ configureBroker(broker2, 7083);
+
+ historical.addProperty("druid.msq.dart.worker.heapFraction", "0.5")
+ .addProperty("druid.msq.dart.worker.concurrentQueries", "1");
+
+ indexer.setServerMemory(400_000_000)
+ .addProperty("druid.segment.handoff.pollDuration", "PT0.1s")
+ .addProperty("druid.processing.numThreads", "2")
+ .addProperty("druid.worker.capacity", "4");
+
+ return EmbeddedDruidCluster.withEmbeddedDerbyAndZookeeper()
+ .addCommonProperty("druid.msq.dart.enabled",
"true")
+ .addResource(new EmbeddedBasicAuthResource())
+ .useLatchableEmitter()
+ .addServer(coordinator)
+ .addServer(overlord)
+ .addServer(broker1)
+ .addServer(broker2)
+ .addServer(indexer)
+ .addServer(historical)
+ .addExtension(SleepModule.class);
+ }
+
+ @BeforeAll
+ protected void setupData()
+ {
+ msqApis = new EmbeddedMSQApis(cluster, overlord);
+
+ ingestedDataSource = EmbeddedClusterApis.createTestDatasourceName();
+ final String taskId = IdUtils.getRandomId();
+ final IndexTask task =
MoreResources.Task.BASIC_INDEX.get().dataSource(ingestedDataSource).withId(taskId);
+ cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, task));
+ cluster.callApi().waitForTaskToSucceed(taskId, overlord);
+
+ cluster.callApi().waitForAllSegmentsToBeAvailable(ingestedDataSource,
coordinator, broker1);
+ cluster.callApi().waitForAllSegmentsToBeAvailable(ingestedDataSource,
coordinator, broker2);
+ }
+
+ /**
+ * Polls the report API on {@link #broker1} until a report is available.
+ */
+ protected GetQueryReportResponse waitForReport(final String sqlQueryId)
+ {
+ final long timeout = 30_000;
+ final long deadline = System.currentTimeMillis() + timeout;
+ while (System.currentTimeMillis() < deadline) {
+ final GetQueryReportResponse report =
msqApis.getDartQueryReport(sqlQueryId, broker1);
+ if (report != null) {
+ return report;
+ }
+ try {
+ Thread.sleep(100);
+ }
+ catch (InterruptedException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ throw new ISE("Timed out after[%,d] ms waiting for query to be in RUNNING
state", timeout);
+ }
+}
diff --git
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedDartQueryCancellationTest.java
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedDartQueryCancellationTest.java
new file mode 100644
index 00000000000..d769587f4cb
--- /dev/null
+++
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedDartQueryCancellationTest.java
@@ -0,0 +1,195 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.testing.embedded.msq;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.common.concurrent.Execs;
+import org.apache.druid.query.QueryContexts;
+import org.eclipse.jetty.server.Connector;
+import org.eclipse.jetty.server.NetworkConnector;
+import org.eclipse.jetty.server.Server;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.io.OutputStream;
+import java.net.ServerSocket;
+import java.net.Socket;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Embedded tests for Dart query cancellation.
+ */
+public class EmbeddedDartQueryCancellationTest extends BaseDartQueryTest
+{
+ @Test
+ @Timeout(value = 180, unit = TimeUnit.SECONDS)
+ public void test_cancelDartQuery_cancelsWorkOrderRetries() throws Exception
+ {
+ final String sqlQueryId = UUID.randomUUID().toString();
+ final String sql = StringUtils.format(
+ "SELECT SLEEP(TIMESTAMP_TO_MILLIS(__time) * 0 + 60) FROM \"%s\"",
+ ingestedDataSource
+ );
+ final CountDownLatch workOrderAttempts = new CountDownLatch(2);
+ final CountDownLatch workOrderAfterCancellation = new CountDownLatch(1);
+ final AtomicBoolean cancellationCompleted = new AtomicBoolean();
+
+ // Drop worker responses so the controller keeps retrying /workOrder.
+ try (HistoricalBlackhole blackhole = new HistoricalBlackhole(
+ false,
+ true,
+ workOrderAttempts,
+ workOrderAfterCancellation,
+ cancellationCompleted
+ )) {
+ final ListenableFuture<String> queryFuture = msqApis.submitDartSqlAsync(
+ sql,
+ Map.of(QueryContexts.CTX_SQL_QUERY_ID, sqlQueryId),
+ broker1
+ );
+ try {
+ waitForReport(sqlQueryId);
+
+ Assertions.assertTrue(
+ workOrderAttempts.await(15, TimeUnit.SECONDS),
+ "Dart controller did not retry /workOrder"
+ );
+
+ // Cancel while the /workOrder request is in flight.
+ final ListenableFuture<Boolean> cancellation =
cluster.callApi().onTargetBrokerAsync(
+ broker1,
+ broker -> broker.cancelSqlQuery(sqlQueryId)
+ );
+
+ Assertions.assertTrue(cancellation.get(30, TimeUnit.SECONDS));
+ cancellationCompleted.set(true);
+
+ Assertions.assertFalse(
+ workOrderAfterCancellation.await(5, TimeUnit.SECONDS),
+ "Dart controller sent /workOrder after cancellation"
+ );
+ }
+ finally {
+ queryFuture.cancel(true);
+ }
+ }
+ }
+
+ /**
+ * Replaces the historical HTTP connector with a socket that selectively
acknowledges worker requests.
+ */
+ private class HistoricalBlackhole implements AutoCloseable
+ {
+ private final boolean acknowledgeStop;
+ private final boolean acknowledgeWorkOrder;
+ private final CountDownLatch workOrderAttempts;
+ private final CountDownLatch workOrderAfterCancellation;
+ private final AtomicBoolean cancellationCompleted;
+ private final Connector connector;
+ private final ServerSocket socket;
+ private final ExecutorService executor;
+ private boolean restored;
+
+ private HistoricalBlackhole(
+ final boolean acknowledgeWorkOrder,
+ final boolean acknowledgeStop,
+ final CountDownLatch workOrderAttempts,
+ final CountDownLatch workOrderAfterCancellation,
+ final AtomicBoolean cancellationCompleted
+ ) throws Exception
+ {
+ this.acknowledgeWorkOrder = acknowledgeWorkOrder;
+ this.acknowledgeStop = acknowledgeStop;
+ this.workOrderAttempts = workOrderAttempts;
+ this.workOrderAfterCancellation = workOrderAfterCancellation;
+ this.cancellationCompleted = cancellationCompleted;
+ this.connector =
historical.bindings().getInstance(Server.class).getConnectors()[0];
+
+ final int port = ((NetworkConnector) connector).getLocalPort();
+ // Keep the Historical announced and replace only its HTTP endpoint.
+ connector.stop();
+
+ this.socket = new ServerSocket(port);
+ this.executor =
Execs.singleThreaded("EmbeddedDartQueryCancellationTest-blackhole-%s");
+ this.executor.submit(this::acceptRequests);
+ }
+
+ private void acceptRequests()
+ {
+ while (!socket.isClosed()) {
+ try (Socket request = socket.accept()) {
+ request.setSoTimeout(1_000);
+ final String requestLine = new BufferedReader(
+ new InputStreamReader(request.getInputStream(),
StandardCharsets.UTF_8)
+ ).readLine();
+
+ final boolean workOrder = requestLine != null &&
requestLine.contains("/workOrder");
+ final boolean stop = requestLine != null &&
requestLine.contains("/stop");
+
+ if (workOrder) {
+ if (cancellationCompleted.get()) {
+ workOrderAfterCancellation.countDown();
+ } else {
+ workOrderAttempts.countDown();
+ }
+ }
+
+ if ((workOrder && acknowledgeWorkOrder) || (stop &&
acknowledgeStop)) {
+ final OutputStream output = request.getOutputStream();
+ output.write(
+ "HTTP/1.1 202 Accepted\r\nContent-Length: 0\r\nConnection:
close\r\n\r\n"
+ .getBytes(StandardCharsets.UTF_8)
+ );
+ output.flush();
+ }
+ }
+ catch (final IOException e) {
+ if (!socket.isClosed()) {
+ // The client can close a request while cancellation is racing
with a retry.
+ }
+ }
+ }
+ }
+
+ @Override
+ public void close() throws Exception
+ {
+ if (!restored) {
+ restored = true;
+ socket.close();
+ executor.shutdownNow();
+ if (!connector.isRunning()) {
+ connector.start();
+ }
+ }
+ }
+ }
+}
diff --git
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedDartReportApiTest.java
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedDartReportApiTest.java
index 4e933d7d77e..16ff27cb4ca 100644
---
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedDartReportApiTest.java
+++
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedDartReportApiTest.java
@@ -25,15 +25,11 @@ import com.google.common.collect.ImmutableMap;
import com.google.common.util.concurrent.ListenableFuture;
import io.netty.handler.codec.http.HttpMethod;
import io.netty.handler.codec.http.HttpResponseStatus;
-import org.apache.druid.common.utils.IdUtils;
import org.apache.druid.data.input.impl.CsvInputFormat;
import org.apache.druid.error.DruidException;
import org.apache.druid.error.ThrowableMatcher;
-import org.apache.druid.guice.SleepModule;
import org.apache.druid.indexer.TaskState;
import org.apache.druid.indexer.report.TaskReport;
-import org.apache.druid.indexing.common.task.IndexTask;
-import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.http.client.CredentialedHttpClient;
import org.apache.druid.java.util.http.client.HttpClient;
@@ -55,16 +51,8 @@ import org.apache.druid.sql.http.GetQueryReportResponse;
import org.apache.druid.sql.http.QueryInfo;
import org.apache.druid.sql.http.ResultFormat;
import org.apache.druid.testing.embedded.EmbeddedBroker;
-import org.apache.druid.testing.embedded.EmbeddedClusterApis;
-import org.apache.druid.testing.embedded.EmbeddedCoordinator;
-import org.apache.druid.testing.embedded.EmbeddedDruidCluster;
-import org.apache.druid.testing.embedded.EmbeddedHistorical;
-import org.apache.druid.testing.embedded.EmbeddedIndexer;
-import org.apache.druid.testing.embedded.EmbeddedOverlord;
import org.apache.druid.testing.embedded.auth.EmbeddedBasicAuthResource;
import org.apache.druid.testing.embedded.auth.HttpUtil;
-import org.apache.druid.testing.embedded.indexing.MoreResources;
-import org.apache.druid.testing.embedded.junit5.EmbeddedClusterTestBase;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
@@ -81,85 +69,20 @@ import java.util.stream.Collectors;
* Embedded test for the Dart report API at {@code
/druid/v2/sql/queries/{id}/reports}.
* Uses batch ingestion to avoid dependency on Kafka/Docker.
*/
-public class EmbeddedDartReportApiTest extends EmbeddedClusterTestBase
+public class EmbeddedDartReportApiTest extends BaseDartQueryTest
{
- private static final int MAX_RETAINED_REPORT_COUNT = 10;
-
// Authentication constants - use shared constants from
EmbeddedBasicAuthResource where available
private static final String REGULAR_USER = "regularUser";
private static final String REGULAR_PASSWORD = "helloworld";
- private final EmbeddedBroker broker1 = new EmbeddedBroker();
- private final EmbeddedBroker broker2 = new EmbeddedBroker();
- private final EmbeddedIndexer indexer = new EmbeddedIndexer();
- private final EmbeddedOverlord overlord = new EmbeddedOverlord();
- private final EmbeddedHistorical historical = new EmbeddedHistorical();
- private final EmbeddedCoordinator coordinator = new EmbeddedCoordinator();
-
- private EmbeddedMSQApis msqApis;
- private String ingestedDataSource;
private HttpClient adminClient;
private HttpClient regularUserClient;
- private void configureBroker(EmbeddedBroker broker, int port)
- {
- broker.addProperty("druid.msq.dart.controller.heapFraction", "0.5")
- .addProperty("druid.msq.dart.controller.maxRetainedReportCount",
String.valueOf(MAX_RETAINED_REPORT_COUNT))
- .addProperty("druid.query.default.context.maxConcurrentStages", "1")
- .addProperty("druid.sql.planner.enableSysQueriesTable", "true")
- .addProperty("druid.plaintextPort", String.valueOf(port));
- }
-
- @Override
- protected EmbeddedDruidCluster createCluster()
- {
- coordinator.addProperty("druid.manager.segments.useIncrementalCache",
"always");
- overlord.addProperty("druid.manager.segments.pollDuration", "PT0.1s");
-
- // Enable Dart with report retention on both brokers, with different ports
- configureBroker(broker1, 7082);
- configureBroker(broker2, 7083);
-
- historical.addProperty("druid.msq.dart.worker.heapFraction", "0.5")
- .addProperty("druid.msq.dart.worker.concurrentQueries", "1");
-
- indexer.setServerMemory(400_000_000)
- .addProperty("druid.segment.handoff.pollDuration", "PT0.1s")
- .addProperty("druid.processing.numThreads", "2")
- .addProperty("druid.worker.capacity", "4");
-
- return EmbeddedDruidCluster.withEmbeddedDerbyAndZookeeper()
- .addCommonProperty("druid.msq.dart.enabled",
"true")
- .addResource(new EmbeddedBasicAuthResource())
- .useLatchableEmitter()
- .addServer(coordinator)
- .addServer(overlord)
- .addServer(broker1)
- .addServer(broker2)
- .addServer(indexer)
- .addServer(historical)
- .addExtension(SleepModule.class);
- }
-
@BeforeAll
- protected void setupData()
+ protected void setupAuth()
{
- msqApis = new EmbeddedMSQApis(cluster, overlord);
-
- // Set up HTTP clients for admin and regular user
setupAdminClient();
setupRegularUserAndClient();
-
- // Ingest test data once, using batch ingestion.
- ingestedDataSource = EmbeddedClusterApis.createTestDatasourceName();
- final String taskId = IdUtils.getRandomId();
- final IndexTask task =
MoreResources.Task.BASIC_INDEX.get().dataSource(ingestedDataSource).withId(taskId);
- cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, task));
- cluster.callApi().waitForTaskToSucceed(taskId, overlord);
-
- // Wait for segments to be available on both brokers
- cluster.callApi().waitForAllSegmentsToBeAvailable(ingestedDataSource,
coordinator, broker1);
- cluster.callApi().waitForAllSegmentsToBeAvailable(ingestedDataSource,
coordinator, broker2);
}
private void setupAdminClient()
@@ -579,28 +502,6 @@ public class EmbeddedDartReportApiTest extends
EmbeddedClusterTestBase
return sqlQueryIds;
}
- /**
- * Polls the report API on {@link #broker1} until a report is available.
- */
- private GetQueryReportResponse waitForReport(String sqlQueryId)
- {
- final long timeout = 30_000;
- final long deadline = System.currentTimeMillis() + timeout;
- while (System.currentTimeMillis() < deadline) {
- final GetQueryReportResponse report =
msqApis.getDartQueryReport(sqlQueryId, broker1);
- if (report != null) {
- return report;
- }
- try {
- Thread.sleep(100);
- }
- catch (InterruptedException e) {
- throw new RuntimeException(e);
- }
- }
- throw new ISE("Timed out after[%,d] ms waiting for query to be in RUNNING
state", timeout);
- }
-
/**
* Polls the report API on the specified brokers until completed reports are
available from all of them.
*/
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartWorkerManager.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartWorkerManager.java
index 5df0f347f59..76962b1974d 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartWorkerManager.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartWorkerManager.java
@@ -215,6 +215,13 @@ public class DartWorkerManager implements WorkerManager
}
stopFuture.set(null);
+ } else if (state.get() == State.STOPPED) {
+ try {
+ FutureUtils.getUnchecked(stopFuture, false);
+ }
+ catch (Throwable ignored) {
+ // Suppress errors.
+ }
}
}
}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java
index 804daa42f7f..9357a874483 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java
@@ -21,7 +21,9 @@ package org.apache.druid.msq.dart.worker;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes;
+import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
import com.google.errorprone.annotations.concurrent.GuardedBy;
import io.netty.handler.codec.http.HttpMethod;
import it.unimi.dsi.fastutil.Pair;
@@ -30,6 +32,7 @@ import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.java.util.http.client.response.HttpResponseHandler;
import org.apache.druid.msq.dart.worker.http.DartWorkerResource;
import org.apache.druid.msq.exec.WorkerClient;
+import org.apache.druid.msq.kernel.WorkOrder;
import org.apache.druid.msq.rpc.BaseWorkerClientImpl;
import org.apache.druid.query.QueryContexts;
import org.apache.druid.rpc.FixedServiceLocator;
@@ -44,8 +47,12 @@ import org.apache.druid.utils.CloseableUtils;
import javax.annotation.Nullable;
import java.io.Closeable;
import java.net.URI;
+import java.util.ArrayList;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
/**
* Dart implementation of {@link WorkerClient}. Uses the same {@link
BaseWorkerClientImpl} as the task-based engine.
@@ -68,6 +75,10 @@ public class DartWorkerClientImpl extends
BaseWorkerClientImpl implements DartWo
@GuardedBy("clientMap")
private boolean closed;
+ private final Set<ListenableFuture<?>> activeRequests =
ConcurrentHashMap.newKeySet();
+ private final Set<ListenableFuture<Void>> activeWorkOrders =
ConcurrentHashMap.newKeySet();
+ private volatile boolean isStopping = false;
+
/**
* Create a worker client.
*
@@ -127,6 +138,13 @@ public class DartWorkerClientImpl extends
BaseWorkerClientImpl implements DartWo
public void close()
{
synchronized (clientMap) {
+ closed = true;
+
+ // Cancel requests before closing locators so in-flight requests do not
continue retrying.
+ final List<ListenableFuture<?>> requests = new
ArrayList<>(activeRequests);
+ activeRequests.clear();
+ requests.forEach(request -> request.cancel(true));
+
for (Map.Entry<String, Pair<ServiceClient, Closeable>> entry :
clientMap.entrySet()) {
CloseableUtils.closeAndSuppressExceptions(
entry.getValue().right(),
@@ -135,13 +153,34 @@ public class DartWorkerClientImpl extends
BaseWorkerClientImpl implements DartWo
}
clientMap.clear();
- closed = true;
}
}
@Override
- public ListenableFuture<?> stopWorker(String workerId)
+ public ListenableFuture<Void> postWorkOrder(final String workerId, final
WorkOrder workOrder)
+ {
+ synchronized (clientMap) {
+ if (isStopping) {
+ return Futures.immediateCancelledFuture();
+ }
+
+ final ListenableFuture<Void> future = super.postWorkOrder(workerId,
workOrder);
+
+ activeWorkOrders.add(future);
+ future.addListener(() -> activeWorkOrders.remove(future),
MoreExecutors.directExecutor());
+
+ return future;
+ }
+ }
+
+ @Override
+ public ListenableFuture<?> stopWorker(final String workerId)
{
+ synchronized (clientMap) {
+ isStopping = true;
+ activeWorkOrders.forEach(future -> future.cancel(true));
+ }
+
return getClient(workerId).asyncRequest(
new RequestBuilder(HttpMethod.POST, "/stop"),
IgnoreHttpResponseHandler.INSTANCE
@@ -165,7 +204,7 @@ public class DartWorkerClientImpl extends
BaseWorkerClientImpl implements DartWo
client = baseClient;
}
- return Pair.of(client, locator);
+ return Pair.of(new RequestTrackingClient(client, activeRequests), locator);
}
private Pair<ServiceClient, Closeable> getClientAndLocator(final String
workerIdString)
@@ -184,6 +223,39 @@ public class DartWorkerClientImpl extends
BaseWorkerClientImpl implements DartWo
}
}
+ private static class RequestTrackingClient implements ServiceClient
+ {
+ private final ServiceClient delegate;
+ private final Set<ListenableFuture<?>> activeRequests;
+
+ private RequestTrackingClient(
+ final ServiceClient delegate,
+ final Set<ListenableFuture<?>> activeRequests
+ )
+ {
+ this.delegate = delegate;
+ this.activeRequests = activeRequests;
+ }
+
+ @Override
+ public <IntermediateType, FinalType> ListenableFuture<FinalType>
asyncRequest(
+ final RequestBuilder requestBuilder,
+ final HttpResponseHandler<IntermediateType, FinalType> handler
+ )
+ {
+ final ListenableFuture<FinalType> future =
delegate.asyncRequest(requestBuilder, handler);
+ activeRequests.add(future);
+ future.addListener(() -> activeRequests.remove(future),
MoreExecutors.directExecutor());
+ return future;
+ }
+
+ @Override
+ public ServiceClient withRetryPolicy(final ServiceRetryPolicy retryPolicy)
+ {
+ return new RequestTrackingClient(delegate.withRetryPolicy(retryPolicy),
activeRequests);
+ }
+ }
+
/**
* Service client that adds the {@link
DartWorkerResource#HEADER_CONTROLLER_HOST} header.
*/
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
index bd87333e532..3e8d1f8ba18 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
@@ -437,6 +437,7 @@ public class ControllerImpl implements Controller
MultiStageQueryContext.getSketchEncoding(querySpec.getContext())
);
closer.register(workerSketchFetcher::close);
+ closer.register(() -> workerManager.stop(true));
// Execution-related: run the multi-stage QueryDefinition.
final InputSpecSlicerFactory inputSpecSlicerFactory =
makeInputSpecSlicerFactory(
@@ -2542,7 +2543,6 @@ public class ControllerImpl implements Controller
log.debug("Query [%s] starting task launcher.", queryDef.getQueryId());
workerTaskLauncherFuture = workerManager.start();
- closer.register(() -> workerManager.stop(true));
workerTaskLauncherFuture.addListener(
() ->
diff --git
a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartWorkerManagerTest.java
b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartWorkerManagerTest.java
index a685a1457b1..8f8639b79cf 100644
---
a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartWorkerManagerTest.java
+++
b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartWorkerManagerTest.java
@@ -23,11 +23,13 @@ import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
import it.unimi.dsi.fastutil.ints.IntSet;
import org.apache.druid.common.guava.FutureUtils;
import org.apache.druid.error.DruidException;
import org.apache.druid.indexer.TaskState;
import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.concurrent.Execs;
import org.apache.druid.msq.dart.worker.DartWorkerClientImpl;
import org.apache.druid.msq.dart.worker.WorkerId;
import org.apache.druid.msq.exec.WorkerManager;
@@ -43,6 +45,11 @@ import org.mockito.MockitoAnnotations;
import java.util.Collections;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
public class DartWorkerManagerTest
{
@@ -199,4 +206,40 @@ public class DartWorkerManagerTest
// Ensure the future from start() resolves.
Assertions.assertNull(FutureUtils.getUnchecked(future, true));
}
+
+ @Test
+ public void test_concurrentStop_waitsForFirstStop() throws Exception
+ {
+ final SettableFuture<Void> stopFuture = SettableFuture.create();
+ final CountDownLatch stopRequested = new CountDownLatch(1);
+
Mockito.when(workerClient.stopWorker(WORKERS.get(0))).thenAnswer(invocation -> {
+ stopRequested.countDown();
+ return stopFuture;
+ });
+
Mockito.when(workerClient.stopWorker(WORKERS.get(1))).thenReturn(Futures.immediateFuture(null));
+
+ workerManager.start();
+ final ExecutorService stopExecutor = Execs.multiThreaded(2,
"dart-worker-manager-stop-test-%s");
+
+ try {
+ final Future<?> firstStop = stopExecutor.submit(() ->
workerManager.stop(true));
+ Assertions.assertTrue(stopRequested.await(10, TimeUnit.SECONDS));
+
+ final Future<?> concurrentStop = stopExecutor.submit(() ->
workerManager.stop(true));
+ Assertions.assertThrows(
+ TimeoutException.class,
+ () -> concurrentStop.get(1, TimeUnit.SECONDS),
+ "Concurrent stop returned before the first stop completed"
+ );
+
+ stopFuture.set(null);
+ firstStop.get(10, TimeUnit.SECONDS);
+ concurrentStop.get(10, TimeUnit.SECONDS);
+ }
+ finally {
+ stopFuture.set(null);
+ stopExecutor.shutdownNow();
+ Assertions.assertTrue(stopExecutor.awaitTermination(10,
TimeUnit.SECONDS));
+ }
+ }
}
diff --git
a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartWorkerClientImplTest.java
b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartWorkerClientImplTest.java
index df38d020818..dc829267944 100644
---
a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartWorkerClientImplTest.java
+++
b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartWorkerClientImplTest.java
@@ -19,10 +19,15 @@
package org.apache.druid.msq.dart.worker;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
import org.apache.druid.error.DruidException;
import org.apache.druid.java.util.common.concurrent.Execs;
import org.apache.druid.java.util.http.client.HttpClient;
+import org.apache.druid.msq.kernel.WorkOrder;
import org.apache.druid.rpc.ServiceClient;
+import org.apache.druid.rpc.ServiceClientFactory;
import org.apache.druid.rpc.ServiceClientFactoryImpl;
import org.apache.druid.rpc.ServiceClosedException;
import org.apache.druid.segment.TestHelper;
@@ -106,6 +111,100 @@ public class DartWorkerClientImplTest
Assertions.assertThrows(DruidException.class, () ->
workerClient.getClient(WORKER_ID.toString()));
}
+ @Test
+ public void test_stopWorker_cancelsActiveWorkOrder()
+ {
+ final SettableFuture<Void> workOrderFuture = SettableFuture.create();
+ final ServiceClient serviceClient = Mockito.mock(ServiceClient.class);
+
Mockito.<ListenableFuture<Void>>when(serviceClient.asyncRequest(Mockito.any(),
Mockito.any()))
+ .thenReturn(workOrderFuture, Futures.immediateFuture(null));
+
+ final ServiceClientFactory clientFactory =
Mockito.mock(ServiceClientFactory.class);
+ Mockito.when(clientFactory.makeClient(Mockito.any(), Mockito.any(),
Mockito.any())).thenReturn(serviceClient);
+
+ final DartWorkerClientImpl client = new DartWorkerClientImpl(
+ QUERY_ID,
+ clientFactory,
+ TestHelper.makeSmileMapper(),
+ "localhost:8080"
+ );
+
+ try {
+ client.postWorkOrder(WORKER_ID.toString(),
Mockito.mock(WorkOrder.class));
+ client.stopWorker(WORKER_ID.toString());
+
+ Assertions.assertTrue(workOrderFuture.isCancelled());
+ }
+ finally {
+ client.close();
+ }
+ }
+
+ @Test
+ public void test_postWorkOrder_afterStop_isCancelled()
+ {
+ final SettableFuture<Void> workOrderFuture = SettableFuture.create();
+ final ServiceClient serviceClient = Mockito.mock(ServiceClient.class);
+
Mockito.<ListenableFuture<Void>>when(serviceClient.asyncRequest(Mockito.any(),
Mockito.any()))
+ .thenReturn(Futures.immediateFuture(null), workOrderFuture);
+
+ final ServiceClientFactory clientFactory =
Mockito.mock(ServiceClientFactory.class);
+ Mockito.when(clientFactory.makeClient(Mockito.any(), Mockito.any(),
Mockito.any())).thenReturn(serviceClient);
+
+ final DartWorkerClientImpl client = new DartWorkerClientImpl(
+ QUERY_ID,
+ clientFactory,
+ TestHelper.makeSmileMapper(),
+ "localhost:8080"
+ );
+
+ try {
+ client.stopWorker(WORKER_ID.toString());
+ final ListenableFuture<Void> returned = client.postWorkOrder(
+ WORKER_ID.toString(),
+ Mockito.mock(WorkOrder.class)
+ );
+
+ Assertions.assertTrue(returned.isCancelled());
+ Assertions.assertFalse(workOrderFuture.isCancelled());
+ }
+ finally {
+ client.close();
+ }
+ }
+
+ @Test
+ public void test_closeCancelsActiveRequests()
+ {
+ final SettableFuture<Void> finishFuture = SettableFuture.create();
+ final SettableFuture<Void> stopFuture = SettableFuture.create();
+ final ServiceClient serviceClient = Mockito.mock(ServiceClient.class);
+
Mockito.<ListenableFuture<Void>>when(serviceClient.asyncRequest(Mockito.any(),
Mockito.any()))
+ .thenReturn(finishFuture, stopFuture);
+
+ final ServiceClientFactory clientFactory =
Mockito.mock(ServiceClientFactory.class);
+ Mockito.when(clientFactory.makeClient(Mockito.any(), Mockito.any(),
Mockito.any())).thenReturn(serviceClient);
+
+ final DartWorkerClientImpl client = new DartWorkerClientImpl(
+ QUERY_ID,
+ clientFactory,
+ TestHelper.makeSmileMapper(),
+ "localhost:8080"
+ );
+
+ try {
+ client.postFinish(WORKER_ID.toString());
+ client.stopWorker(WORKER_ID.toString());
+ client.close();
+
+ Assertions.assertTrue(finishFuture.isCancelled());
+ Assertions.assertTrue(stopFuture.isCancelled());
+ }
+ finally {
+ client.close();
+ }
+ }
+
/**
* Verify that a request to {@link #WORKER_ID} fails immediately, rather
than retrying, due to its client
* being closed.
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
