fix: cancel active /workOrder futures on query cancellation

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). -->
@@ -842,6 +843,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>

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);
+  }
+}

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();
+        }
+      }
+    }
+  }
+}