fix: stabilize task queue and ingestion CI tests

embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/IngestionSmokeTest.java

@@ -24,6 +24,7 @@
 import org.apache.druid.common.utils.IdUtils;
 import org.apache.druid.data.input.impl.CsvInputFormat;
 import org.apache.druid.data.input.impl.TimestampSpec;
+import org.apache.druid.indexer.TaskStatusPlus;
 import org.apache.druid.indexing.common.task.CompactionTask;
 import org.apache.druid.indexing.common.task.IndexTask;
 import org.apache.druid.indexing.common.task.NoopTask;
@@ -33,13 +34,17 @@
 import org.apache.druid.indexing.kafka.simulate.KafkaResource;
 import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorSpec;
 import org.apache.druid.indexing.overlord.Segments;
+import org.apache.druid.indexing.overlord.TaskMaster;
+import org.apache.druid.indexing.overlord.TaskRunner;
+import org.apache.druid.indexing.overlord.TaskRunnerWorkItem;
 import org.apache.druid.indexing.overlord.supervisor.SupervisorStatus;
 import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.metadata.storage.postgresql.PostgreSQLMetadataStorageModule;
 import org.apache.druid.query.DruidMetrics;
 import org.apache.druid.query.http.SqlTaskStatus;
+import org.apache.druid.rpc.indexing.OverlordClient;
 import org.apache.druid.segment.metadata.Metric;
 import org.apache.druid.tasklogs.TaskLogStreamer;
 import org.apache.druid.testing.embedded.EmbeddedBroker;
@@ -67,6 +72,7 @@
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -145,7 +151,51 @@ protected EmbeddedDruidCluster addServers(EmbeddedDruidCluster cluster)
   @AfterEach
   public void cleanUp()
   {
-    markSegmentsAsUnused(dataSource);
+    try {
+      final List<SupervisorStatus> supervisors = new ArrayList<>();
+      cluster.callApi().onLeaderOverlord(OverlordClient::supervisorStatuses).forEachRemaining(supervisors::add);
+      for (final SupervisorStatus supervisor : supervisors) {
+        if (dataSource.equals(supervisor.getId())) {
+          cluster.callApi().onLeaderOverlord(o -> o.terminateSupervisor(supervisor.getId()));
+        }
+      }
+
+      cluster.callApi()
+             .waitForResult(this::cancelTasksForCurrentTest, Set::isEmpty)
+             .withTimeoutMillis(60_000)
+             .go();
+    }
+    finally {
+      markSegmentsAsUnused(dataSource);
+    }
+  }
+
+  private Set<String> cancelTasksForCurrentTest()
+  {
+    final Set<String> taskIds = new HashSet<>();
+    for (final String state : List.of("running", "pending", "waiting")) {
+      for (final TaskStatusPlus task : cluster.callApi().getTasks(dataSource, state)) {
+        taskIds.add(task.getId());
+      }
+    }
+
+    // A task can be complete in storage while still occupying a worker slot.
+    // Docker subclasses may use an external leader, so its runner is not available here.
+    final Optional<TaskRunner> taskRunner = overlord.bindings().getInstance(TaskMaster.class).getTaskRunner();
+    if (taskRunner.isPresent()) {
+      final List<TaskRunnerWorkItem> runnerTasks = new ArrayList<>(taskRunner.get().getPendingTasks());
+      runnerTasks.addAll(taskRunner.get().getRunningTasks());
+      for (final TaskRunnerWorkItem task : runnerTasks) {
+        if (dataSource.equals(task.getDataSource())) {
+          taskIds.add(task.getTaskId());
+        }
+      }
+    }
+
+    for (final String taskId : taskIds) {
+      cluster.callApi().onLeaderOverlord(o -> o.cancelTask(taskId));
+    }
+    return taskIds;
   }
 
   protected int markSegmentsAsUnused(String dataSource)
@@ -363,7 +413,7 @@ public void test_streamLogs_ofCancelledTask() throws Exception
     final String taskId = IdUtils.getRandomId();
     final long runDurationMillis = 100_000L;
     cluster.callApi().onLeaderOverlord(
-        o -> o.runTask(taskId, new NoopTask(taskId, null, null, runDurationMillis, 0L, null))
+        o -> o.runTask(taskId, new NoopTask(taskId, null, dataSource, runDurationMillis, 0L, null))
     );
 
     eventCollector.latchableEmitter().waitForEvent(

embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/KafkaBoundedSupervisorTest.java

@@ -29,11 +29,18 @@
 import org.apache.druid.query.DruidMetrics;
 import org.apache.druid.testing.embedded.EmbeddedDruidCluster;
 import org.apache.druid.testing.embedded.StreamIngestResource;
+import org.apache.druid.testing.embedded.tools.EventSerializer;
+import org.apache.druid.testing.embedded.tools.JsonEventSerializer;
+import org.apache.druid.testing.embedded.tools.StreamGenerator;
+import org.apache.druid.testing.embedded.tools.WikipediaStreamEventStreamGenerator;
+import org.apache.kafka.clients.producer.ProducerRecord;
 import org.joda.time.Period;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -203,12 +210,26 @@ private KafkaSupervisorSpec createBoundedKafkaSupervisor(
         .build(dataSource, topic);
   }
 
+  private void publishRecordsToBothPartitions(String topic)
+  {
+    final EventSerializer serializer = new JsonEventSerializer(overlord.bindings().jsonMapper());
+    final StreamGenerator generator = new WikipediaStreamEventStreamGenerator(serializer, 100, 100);
+    final List<byte[]> records = generator.generateEvents(10);
+    final List<ProducerRecord<byte[], byte[]>> producerRecords = new ArrayList<>();
+    // Fixed per-partition end offsets require data in both partitions; Kafka's default partitioner need not balance it.
+    for (int i = 0; i < records.size(); i++) {
+      producerRecords.add(new ProducerRecord<>(topic, i % 2, null, records.get(i)));
+    }
+    kafkaServer.produceRecordsWithoutTransaction(producerRecords);
+    Assertions.assertEquals(Map.of("0", 500L, "1", 500L), kafkaServer.getPartitionOffsets(topic));
+  }
+
   @Test
   public void test_boundedSupervisor_withMismatchedMetadata_is_unhealthy()
   {
     final String topic = IdUtils.getRandomId();
     kafkaServer.createTopicWithPartitions(topic, 2);
-    publish1kRecords(topic, false);
+    publishRecordsToBothPartitions(topic);
 
     // Get the current end offsets for all partitions
     Map<String, Long> currentOffsets = kafkaServer.getPartitionOffsets(topic);
@@ -282,7 +303,7 @@ public void test_boundedSupervisor_doesNotSilentlyCompleteWhenStaleOffsetExceeds
   {
     final String topic = IdUtils.getRandomId();
     kafkaServer.createTopicWithPartitions(topic, 2);
-    publish1kRecords(topic, false);
+    publishRecordsToBothPartitions(topic);
 
     // Run 1: ingest up to offset 100 on each partition and complete.
     Map<String, Long> startOffsets1 = new HashMap<>();

indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueScaleTest.java

@@ -179,12 +179,8 @@ public void doMassLaunchAndExit() throws Exception
       taskQueue.add(testTask);
     }
 
-    // in theory we can get a race here, since we fetch the counts at separate times
+    // The running, pending, and waiting counters are separate snapshots and cannot be summed while tasks transition.
     Assertions.assertEquals(numTasks, taskQueue.getTasks().size(), "all tasks should be known");
-    long runningTasks = taskQueue.getRunningTaskCount().values().stream().mapToLong(Long::longValue).sum();
-    long pendingTasks = taskQueue.getPendingTaskCount().values().stream().mapToLong(Long::longValue).sum();
-    long waitingTasks = taskQueue.getWaitingTaskCount().values().stream().mapToLong(Long::longValue).sum();
-    Assertions.assertEquals(numTasks, (runningTasks + pendingTasks + waitingTasks), "all tasks should be known");
 
     // Wait for all tasks to finish.
     final TaskLookup.CompleteTaskLookup completeTaskLookup =
@@ -194,12 +190,18 @@ public void doMassLaunchAndExit() throws Exception
       Thread.sleep(100);
     }
 
-    Thread.sleep(100);
+    // Completion is persisted before cleanup finishes. The test timeout bounds this wait.
+    while (!taskStorage.getActiveTasks().isEmpty()
+           || taskQueue.getRunningTaskCount().values().stream().anyMatch(count -> count != 0)
+           || taskQueue.getPendingTaskCount().values().stream().anyMatch(count -> count != 0)
+           || taskQueue.getWaitingTaskCount().values().stream().anyMatch(count -> count != 0)) {
+      Thread.sleep(100);
+    }
 
     Assertions.assertEquals(0, taskStorage.getActiveTasks().size(), "no tasks should be active");
-    runningTasks = taskQueue.getRunningTaskCount().values().stream().mapToLong(Long::longValue).sum();
-    pendingTasks = taskQueue.getPendingTaskCount().values().stream().mapToLong(Long::longValue).sum();
-    waitingTasks = taskQueue.getWaitingTaskCount().values().stream().mapToLong(Long::longValue).sum();
+    final long runningTasks = taskQueue.getRunningTaskCount().values().stream().mapToLong(Long::longValue).sum();
+    final long pendingTasks = taskQueue.getPendingTaskCount().values().stream().mapToLong(Long::longValue).sum();
+    final long waitingTasks = taskQueue.getWaitingTaskCount().values().stream().mapToLong(Long::longValue).sum();
     Assertions.assertEquals(0, runningTasks, "no tasks should be running");
     Assertions.assertEquals(0, pendingTasks, "no tasks should be pending");
     Assertions.assertEquals(0, waitingTasks, "no tasks should be waiting");

server/src/main/java/org/apache/druid/segment/metadata/AbstractSegmentMetadataCache.java

@@ -591,7 +591,7 @@ public void removeSegment(final DataSegment segment)
       segmentsNeedingRefresh.remove(segment.getId());
       unmarkSegmentAsMutable(segment.getId());
 
-      segmentMetadataInfo.compute(
+      final ConcurrentSkipListMap<SegmentId, AvailableSegmentMetadata> remainingSegments = segmentMetadataInfo.compute(
           segment.getDataSource(),
           (dataSource, segmentsMap) -> {
             if (segmentsMap == null) {
@@ -605,7 +605,11 @@ public void removeSegment(final DataSegment segment)
               }
               removeSegmentAction(segment.getId());
               if (segmentsMap.isEmpty()) {
-                tables.remove(segment.getDataSource());
+                // Emit the removal action only if this call actually removed the table, so that a concurrent
+                // refresh which also finds the datasource gone cannot report the same removal twice.
+                if (tables.remove(segment.getDataSource()) != null) {
+                  removeDataSourceAction(segment.getDataSource());
+                }
                 log.info("dataSource [%s] no longer exists, all metadata removed.", segment.getDataSource());
                 return null;
               } else {
@@ -615,6 +619,9 @@ public void removeSegment(final DataSegment segment)
             }
           }
       );
+      if (remainingSegments == null) {
+        dataSourcesNeedingRebuild.remove(segment.getDataSource());
+      }
 
       lock.notifyAll();
     }
@@ -625,6 +632,16 @@ public void removeSegment(final DataSegment segment)
    */
   protected abstract void removeSegmentAction(SegmentId segmentId);
 
+  /**
+   * Called under the cache lock after the last segment of a datasource has been removed and its table was actually
+   * removed from {@link #tables} by that removal. It is not called when no table existed for the datasource, so a
+   * single datasource removal triggers this action at most once even if a refresh observes the removal concurrently.
+   */
+  protected void removeDataSourceAction(String dataSource)
+  {
+    // No additional action by default.
+  }
+
   @VisibleForTesting
   public void removeServerSegment(final DruidServerMetadata server, final DataSegment segment)
   {

sql/src/main/java/org/apache/druid/sql/calcite/schema/BrokerSegmentMetadataCache.java

@@ -251,11 +251,11 @@ public void refresh(final Set<SegmentId> segmentsToRefresh, final Set<String> da
       final RowSignature rowSignature = buildDataSourceRowSignature(dataSource);
       if (rowSignature == null) {
         log.info("datasource [%s] no longer exists, all metadata removed.", dataSource);
-        tables.remove(dataSource);
-        emitMetric(
-            Metric.DATASOURCE_REMOVED,
-            1,
-            ServiceMetricEvent.builder().setDimension(DruidMetrics.DATASOURCE, dataSource));
+        // The last-segment callback may already have removed the table and emitted the metric while this refresh
+        // was in flight. Only emit if this refresh is the one that actually removed the table.
+        if (tables.remove(dataSource) != null) {
+          emitDataSourceRemoved(dataSource);
+        }
         continue;
       }
 
@@ -264,12 +264,11 @@ public void refresh(final Set<SegmentId> segmentsToRefresh, final Set<String> da
         // and a new datasource is added
         log.info("datasource [%s] schema has not been initialized yet, "
                  + "check coordinator logs if this message is persistent.", dataSource);
-        // this is a harmless call
-        tables.remove(dataSource);
-        emitMetric(
-            Metric.DATASOURCE_REMOVED,
-            1,
-            ServiceMetricEvent.builder().setDimension(DruidMetrics.DATASOURCE, dataSource));
+        // Usually there is no table to remove here. If there was one, a concurrent last-segment callback may have
+        // removed it and emitted the metric already, so only emit if this refresh actually removed the table.
+        if (tables.remove(dataSource) != null) {
+          emitDataSourceRemoved(dataSource);
+        }
         continue;
       }
 
@@ -307,6 +306,22 @@ protected void removeSegmentAction(SegmentId segmentId)
     // noop, no additional action needed when segment is removed.
   }
 
+  @Override
+  protected void removeDataSourceAction(String dataSource)
+  {
+    // The last-segment callback can remove the table without another schema refresh.
+    emitDataSourceRemoved(dataSource);
+  }
+
+  private void emitDataSourceRemoved(String dataSource)
+  {
+    emitMetric(
+        Metric.DATASOURCE_REMOVED,
+        1,
+        ServiceMetricEvent.builder().setDimension(DruidMetrics.DATASOURCE, dataSource)
+    );
+  }
+
   private Set<String> queryDataSources()
   {
     Set<String> dataSources = new HashSet<>();