From ab8b89af3111dd1616c26c740dc5f8ec48e4a9ad Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Sun, 23 Mar 2025 10:58:58 +0530 Subject: [PATCH 01/10] Improve concurrency in TaskQueue --- .../kafka/supervisor/KafkaSupervisorTest.java | 141 +++--- .../supervisor/KinesisSupervisorTest.java | 150 +++--- .../druid/indexing/overlord/TaskQueue.java | 451 ++++++++++-------- .../supervisor/SeekableStreamSupervisor.java | 2 +- .../indexing/overlord/TaskQueueTest.java | 15 +- .../SeekableStreamSupervisorStateTest.java | 32 +- 6 files changed, 429 insertions(+), 362 deletions(-) diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java index 3eb0d2a8df83..5ee72f9e18bf 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java @@ -122,6 +122,7 @@ import java.time.Instant; import java.time.temporal.ChronoUnit; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.HashMap; @@ -131,6 +132,8 @@ import java.util.TreeMap; import java.util.concurrent.Executor; import java.util.concurrent.ScheduledExecutorService; +import java.util.function.Function; +import java.util.stream.Collectors; @RunWith(Parameterized.class) public class KafkaSupervisorTest extends EasyMockSupport @@ -396,7 +399,7 @@ public SeekableStreamIndexTaskClient build( EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskMaster.getSupervisorManager()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -543,7 +546,7 @@ public void testNoInitialState() throws Exception Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -619,7 +622,7 @@ public void testSkipOffsetGaps() throws Exception Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -643,7 +646,7 @@ public void testMultiTask() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -722,7 +725,7 @@ public void testReplicas() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -801,7 +804,7 @@ public void testLateMessageRejectionPeriod() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -840,7 +843,7 @@ public void testEarlyMessageRejectionPeriod() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -882,7 +885,7 @@ public void testLatestOffset() throws Exception Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -934,7 +937,7 @@ public void testPartitionIdsUpdates() throws Exception Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -960,7 +963,7 @@ public void testAlwaysUsesEarliestOffsetForNewlyDiscoveredPartitions() throws Ex Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -1000,7 +1003,7 @@ public void testAlwaysUsesEarliestOffsetForNewlyDiscoveredPartitions() throws Ex addMoreEvents(9, 6); EasyMock.reset(taskQueue, taskStorage); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); Capture tmp = Capture.newInstance(); EasyMock.expect(taskQueue.add(EasyMock.capture(tmp))).andReturn(true); EasyMock.replay(taskStorage, taskQueue); @@ -1008,7 +1011,7 @@ public void testAlwaysUsesEarliestOffsetForNewlyDiscoveredPartitions() throws Ex verifyAll(); EasyMock.reset(taskQueue, taskStorage); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); Capture newcaptured = Capture.newInstance(); EasyMock.expect(taskQueue.add(EasyMock.capture(newcaptured))).andReturn(true); EasyMock.replay(taskStorage, taskQueue); @@ -1057,7 +1060,7 @@ public void testDatasourceMetadataSingleTopicToSingleTopicMatch() throws Excepti Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( new SeekableStreamStartSequenceNumbers<>( @@ -1115,7 +1118,7 @@ public void testDatasourceMetadataSingleTopicToMultiTopicMatch() throws Exceptio Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(List.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); ImmutableMap.Builder partitionSequenceNumberMap = ImmutableMap.builder(); // these should match partitionSequenceNumberMap.putAll(singlePartitionMap(topic, 0, 10L, 1, 20L, 2, 30L)); @@ -1173,7 +1176,7 @@ public void testDatasourceMetadataSingleTopicToMultiTopicNotMatch() throws Excep Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(List.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); ImmutableMap.Builder partitionSequenceNumberMap = ImmutableMap.builder(); // these should not match partitionSequenceNumberMap.putAll(singlePartitionMap("notMatch", 0, 10L, 1, 20L, 2, 30L)); @@ -1229,7 +1232,7 @@ public void testDatasourceMetadataMultiTopicToSingleTopic() throws Exception Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(List.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); ImmutableMap.Builder partitionSequenceNumberMap = ImmutableMap.builder(); // these should match partitionSequenceNumberMap.putAll(multiTopicPartitionMap(topic, 0, 10L, 1, 20L, 2, 30L)); @@ -1292,7 +1295,7 @@ public void testDatasourceMetadataMultiTopicToMultiTopic() throws Exception Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(List.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); ImmutableMap.Builder partitionSequenceNumberMap = ImmutableMap.builder(); // these should match partitionSequenceNumberMap.putAll(multiTopicPartitionMap(topic, 0, 10L, 1, 20L, 2, 30L)); @@ -1384,7 +1387,7 @@ public void testDontKillTasksWithMismatchedType() throws Exception ImmutableMap.of() ); - List existingTasks = ImmutableList.of(id2); + Map existingTasks = Map.of(id2.getId(), id2); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); @@ -1487,7 +1490,7 @@ public void testKillBadPartitionAssignment() throws Exception tuningConfig ); - List existingTasks = ImmutableList.of(id1, id2, id3, id4, id5); + Map existingTasks = toMap(id1, id2, id3, id4, id5); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); @@ -1550,7 +1553,7 @@ public void testRequeueTaskWhenFailed() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) .anyTimes(); @@ -1588,7 +1591,7 @@ public void testRequeueTaskWhenFailed() throws Exception // are all still running EasyMock.reset(taskStorage); EasyMock.reset(taskQueue); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -1607,7 +1610,7 @@ public void testRequeueTaskWhenFailed() throws Exception KafkaIndexTask iHaveFailed = (KafkaIndexTask) tasks.get(3); EasyMock.reset(taskStorage); EasyMock.reset(taskQueue); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(imStillAlive).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(imStillAlive)).anyTimes(); for (Task task : imStillAlive) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -1653,7 +1656,7 @@ public void testRequeueAdoptedTaskWhenFailed() throws Exception supervisor.getTuningConfig() ); - List existingTasks = ImmutableList.of(id1); + final Map existingTasks = Map.of(id1.getId(), id1); Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); @@ -1690,7 +1693,7 @@ public void testRequeueAdoptedTaskWhenFailed() throws Exception // test that a task failing causes a new task to be re-queued with the same parameters String runningTaskId = captured.getValue().getId(); Capture aNewTaskCapture = Capture.newInstance(); - KafkaIndexTask iHaveFailed = (KafkaIndexTask) existingTasks.get(0); + KafkaIndexTask iHaveFailed = (KafkaIndexTask) existingTasks.get(id1.getId()); EasyMock.reset(taskStorage); EasyMock.reset(taskQueue); EasyMock.reset(taskClient); @@ -1704,7 +1707,7 @@ public void testRequeueAdoptedTaskWhenFailed() throws Exception .times(1); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) - .andReturn(ImmutableList.of(captured.getValue())) + .andReturn(Map.of(captured.getValue().getId(), captured.getValue())) .anyTimes(); EasyMock.expect(taskStorage.getStatus(iHaveFailed.getId())) .andReturn(Optional.of(TaskStatus.failure(iHaveFailed.getId(), "Dummy task status failure err message"))); @@ -1748,7 +1751,7 @@ public void testQueueNextTasksOnSuccess() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) .anyTimes(); @@ -1793,7 +1796,7 @@ public void testQueueNextTasksOnSuccess() throws Exception .andReturn(Futures.immediateFuture(checkpoints2)) .times(2); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -1816,7 +1819,7 @@ public void testQueueNextTasksOnSuccess() throws Exception EasyMock.reset(taskStorage); EasyMock.reset(taskQueue); EasyMock.reset(taskClient); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(imStillRunning).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(imStillRunning)).anyTimes(); for (Task task : imStillRunning) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -1854,7 +1857,7 @@ public void testBeginPublishAndQueueNextTasks() throws Exception EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(EasyMock.anyString())).andReturn(TaskLocation.unknown()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -1876,7 +1879,7 @@ public void testBeginPublishAndQueueNextTasks() throws Exception EasyMock.reset(taskStorage, taskRunner, taskClient, taskQueue); captured = Capture.newInstance(CaptureType.ALL); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -1981,7 +1984,9 @@ public void testDiscoverExistingPublishingTask() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(task)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) + .andReturn(Map.of(task.getId(), task)) + .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(task)).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( @@ -2119,7 +2124,9 @@ public void testDiscoverExistingPublishingTaskWithDifferentPartitionAllocation() EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(task)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) + .andReturn(Map.of(task.getId(), task)) + .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(task)).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( @@ -2276,7 +2283,7 @@ public void testDiscoverExistingPublishingAndReadingTask() throws Exception EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) - .andReturn(ImmutableList.of(id1, id2)) + .andReturn(toMap(id1, id2)) .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); @@ -2398,7 +2405,7 @@ public void testReportWhenMultipleActiveTasks() throws Exception supervisor.getTuningConfig() ); - List existingTasks = ImmutableList.of(id1, id2); + Map existingTasks = toMap(id1, id2); Collection workItems = new ArrayList<>(); workItems.add(new TestTaskRunnerWorkItem(id1, null, location1)); @@ -2552,7 +2559,7 @@ public void testSupervisorIsIdleIfStreamInactive() throws Exception supervisor.getTuningConfig() ); - List existingTasks = ImmutableList.of(id1, id2); + Map existingTasks = toMap(id1, id2); Collection workItems = new ArrayList<>(); workItems.add(new TestTaskRunnerWorkItem(id1, null, location1)); @@ -2674,7 +2681,7 @@ public void testSupervisorIsIdleIfStreamInactiveWhenNoActiveTasks() throws Excep EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( @@ -2732,7 +2739,7 @@ public void testSupervisorNotIdleIfStreamInactiveWhenSuspended() throws Exceptio EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( @@ -2792,7 +2799,7 @@ public void testSupervisorIsIdleIfStreamInactiveWhenSuspended() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( @@ -2888,7 +2895,7 @@ public void testSupervisorIsIdleIfStreamInactiveWhenNoActiveTasksAndFewPendingTa supervisor.getTuningConfig() ); - List existingTasks = ImmutableList.of(id1, id2); + Map existingTasks = toMap(id1, id2); Collection workItems = new ArrayList<>(); workItems.add(new TestTaskRunnerWorkItem(id1, null, location1)); @@ -3008,7 +3015,7 @@ public void testKillUnresponsiveTasksWhileGettingStartTime() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -3037,7 +3044,7 @@ public void testKillUnresponsiveTasksWhileGettingStartTime() throws Exception .andReturn(Futures.immediateFuture(checkpoints2)) .times(2); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -3068,7 +3075,7 @@ public void testKillUnresponsiveTasksWhilePausing() throws Exception EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(EasyMock.anyString())).andReturn(TaskLocation.unknown()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -3102,7 +3109,7 @@ public void testKillUnresponsiveTasksWhilePausing() throws Exception .times(2); captured = Capture.newInstance(CaptureType.ALL); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -3166,7 +3173,7 @@ public void testKillUnresponsiveTasksWhileSettingEndOffsets() throws Exception EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(EasyMock.anyString())).andReturn(TaskLocation.unknown()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -3200,7 +3207,7 @@ public void testKillUnresponsiveTasksWhileSettingEndOffsets() throws Exception .times(2); captured = Capture.newInstance(CaptureType.ALL); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -3355,7 +3362,7 @@ public void testStopGracefully() throws Exception EasyMock.expect(taskRunner.getTaskLocation(id2.getId())).andReturn(location2).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(id3.getId())).andReturn(TaskLocation.unknown()).anyTimes(); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) - .andReturn(ImmutableList.of(id1, id2, id3)) + .andReturn(toMap(id1, id2, id3)) .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); @@ -3420,7 +3427,7 @@ public void testResetNoTasks() EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(EasyMock.anyString())).andReturn(TaskLocation.unknown()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)) .andReturn(null).anyTimes(); @@ -3449,7 +3456,7 @@ public void testResetDataSourceMetadata() throws Exception EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(EasyMock.anyString())).andReturn(TaskLocation.unknown()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)) .andReturn(null).anyTimes(); @@ -3514,7 +3521,7 @@ public void testResetNoDataSourceMetadata() EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(EasyMock.anyString())).andReturn(TaskLocation.unknown()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)) .andReturn(null).anyTimes(); @@ -3550,7 +3557,7 @@ public void testGetOffsetFromStorageForPartitionWithResetOffsetAutomatically() t EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); EasyMock.reset(indexerMetadataStorageCoordinator); @@ -3655,7 +3662,7 @@ public void testResetRunningTasks() throws Exception EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) - .andReturn(ImmutableList.of(id1, id2, id3)) + .andReturn(toMap(id1, id2, id3)) .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); @@ -3768,7 +3775,7 @@ public void testNoDataIngestionTasks() EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) - .andReturn(ImmutableList.of(id1, id2, id3)) + .andReturn(toMap(id1, id2, id3)) .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); @@ -3892,7 +3899,7 @@ public void testCheckpointForInactiveTaskGroup() EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) - .andReturn(ImmutableList.of(id1, id2, id3)) + .andReturn(toMap(id1, id2, id3)) .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); @@ -4013,7 +4020,7 @@ public void testCheckpointForUnknownTaskGroup() EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) - .andReturn(ImmutableList.of(id1, id2, id3)) + .andReturn(toMap(id1, id2, id3)) .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); @@ -4067,7 +4074,7 @@ public void testSuspendedNoRunningTasks() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -4164,7 +4171,7 @@ public void testSuspendedRunningTasks() throws Exception EasyMock.expect(taskRunner.getTaskLocation(id2.getId())).andReturn(location2).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(id3.getId())).andReturn(TaskLocation.unknown()).anyTimes(); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) - .andReturn(ImmutableList.of(id1, id2, id3)) + .andReturn(toMap(id1, id2, id3)) .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); @@ -4222,7 +4229,7 @@ public void testResetSuspended() EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(EasyMock.anyString())).andReturn(TaskLocation.unknown()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)) .andReturn(null).anyTimes(); taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); @@ -4260,7 +4267,7 @@ public void testFailedInitializationAndRecovery() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -4286,7 +4293,7 @@ public void testFailedInitializationAndRecovery() throws Exception Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -4515,7 +4522,7 @@ public void testDoNotKillCompatibleTasks() supervisor.getTuningConfig() ); - List existingTasks = ImmutableList.of(task); + Map existingTasks = Map.of(task.getId(), task); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); @@ -4588,7 +4595,7 @@ public void testKillIncompatibleTasks() supervisor.getTuningConfig() ); - List existingTasks = ImmutableList.of(task); + Map existingTasks = Map.of(task.getId(), task); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); @@ -4986,7 +4993,7 @@ public void testResumeAllActivelyReadingTasks() throws Exception supervisor.getTuningConfig() ); - List tasks = ImmutableList.of( + Map tasks = toMap( readingTask, publishingTask, pausedTask, @@ -5739,6 +5746,16 @@ private static ImmutableMap multiTopicPartitionMap( ); } + private static Map toMap(List tasks) + { + return tasks.stream().collect(Collectors.toMap(Task::getId, Function.identity())); + } + + private static Map toMap(Task... tasks) + { + return Arrays.stream(tasks).collect(Collectors.toMap(Task::getId, Function.identity())); + } + private static class TestTaskRunnerWorkItem extends TaskRunnerWorkItem { private final String taskType; diff --git a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTest.java b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTest.java index 193ad3a34828..d7b7aaa6b4a5 100644 --- a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTest.java +++ b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTest.java @@ -21,7 +21,6 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Optional; -import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.util.concurrent.Futures; @@ -113,12 +112,14 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executor; import java.util.concurrent.ScheduledExecutorService; +import java.util.function.Function; +import java.util.stream.Collectors; public class KinesisSupervisorTest extends EasyMockSupport { private static final ObjectMapper OBJECT_MAPPER = TestHelper.makeJsonMapper(); private static final InputFormat INPUT_FORMAT = new JsonInputFormat( - new JSONPathSpec(true, ImmutableList.of()), + new JSONPathSpec(true, List.of()), ImmutableMap.of(), false, false, @@ -244,7 +245,7 @@ public void testNoInitialState() throws Exception Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata( null @@ -340,7 +341,7 @@ public void testNoInitialStateWithAutoScaleOut() throws Exception Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata( null @@ -415,7 +416,7 @@ public void testNoInitialStateWithAutoScaleIn() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock .expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)) .andReturn(new KinesisDataSourceMetadata(null)) @@ -587,7 +588,7 @@ public void testMultiTask() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock .expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)) .andReturn(new KinesisDataSourceMetadata(null)) @@ -646,7 +647,7 @@ public void testReplicas() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata( null @@ -723,7 +724,7 @@ public void testLateMessageRejectionPeriod() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata( null @@ -775,7 +776,7 @@ public void testEarlyMessageRejectionPeriod() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata( null @@ -837,7 +838,7 @@ public void testDatasourceMetadata() throws Exception Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata( new SeekableStreamStartSequenceNumbers<>( @@ -900,7 +901,7 @@ public void testBadMetadataOffsets() throws Exception EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.absent()).anyTimes(); - EasyMock.expect(taskStorage.getActiveTasksByDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasksByDatasource(DATASOURCE)).andReturn(List.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata( new SeekableStreamStartSequenceNumbers<>( @@ -951,12 +952,12 @@ public void testDontKillTasksWithMismatchedType() throws Exception ImmutableMap.of() ); - List existingTasks = ImmutableList.of(id2); + List existingTasks = List.of(id2); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(existingTasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(existingTasks)).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.NOT_STARTED)) .anyTimes(); @@ -1045,12 +1046,12 @@ public void testKillBadPartitionAssignment() throws Exception null ); - List existingTasks = ImmutableList.of(id1, id2, id3, id4); + List existingTasks = List.of(id1, id2, id3, id4); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(existingTasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(existingTasks)).anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); @@ -1121,7 +1122,7 @@ public void testRequeueTaskWhenFailed() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.NOT_STARTED)) .anyTimes(); @@ -1168,7 +1169,7 @@ public void testRequeueTaskWhenFailed() throws Exception // are all still running EasyMock.reset(taskStorage); EasyMock.reset(taskQueue); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -1187,7 +1188,7 @@ public void testRequeueTaskWhenFailed() throws Exception KinesisIndexTask iHaveFailed = (KinesisIndexTask) tasks.get(3); EasyMock.reset(taskStorage); EasyMock.reset(taskQueue); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(imStillAlive).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(imStillAlive)).anyTimes(); for (Task task : imStillAlive) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -1254,13 +1255,13 @@ public void testRequeueAdoptedTaskWhenFailed() throws Exception maxi ); - List existingTasks = ImmutableList.of(id1); + List existingTasks = List.of(id1); Capture captured = Capture.newInstance(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(existingTasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(existingTasks)).anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); EasyMock.expect(taskClient.getStatusAsync("id1")) @@ -1311,7 +1312,9 @@ public void testRequeueAdoptedTaskWhenFailed() throws Exception .times(1); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(captured.getValue())).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) + .andReturn(toMap(List.of(captured.getValue()))) + .anyTimes(); EasyMock.expect(taskStorage.getStatus(iHaveFailed.getId())) .andReturn(Optional.of(TaskStatus.failure(iHaveFailed.getId(), "Dummy task status failure err message"))); EasyMock.expect(taskStorage.getStatus(runningTaskId)) @@ -1371,7 +1374,7 @@ public void testQueueNextTasksOnSuccess() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.NOT_STARTED)) .anyTimes(); @@ -1425,7 +1428,7 @@ public void testQueueNextTasksOnSuccess() throws Exception .times(2); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -1448,7 +1451,7 @@ public void testQueueNextTasksOnSuccess() throws Exception EasyMock.reset(taskStorage); EasyMock.reset(taskQueue); EasyMock.reset(taskClient); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(imStillRunning).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(imStillRunning)).anyTimes(); for (Task task : imStillRunning) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -1500,7 +1503,7 @@ public void testBeginPublishAndQueueNextTasks() throws Exception EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(EasyMock.anyString())).andReturn(TaskLocation.unknown()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata(null) ).anyTimes(); @@ -1520,7 +1523,7 @@ public void testBeginPublishAndQueueNextTasks() throws Exception EasyMock.reset(taskStorage, taskRunner, taskClient, taskQueue); final Capture secondTasks = Capture.newInstance(CaptureType.ALL); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -1660,7 +1663,7 @@ public void testDiscoverExistingPublishingTask() throws Exception EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(EasyMock.anyString())).andReturn(location).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(task)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of(task.getId(), task)).anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(task)).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( @@ -1821,7 +1824,7 @@ public void testDiscoverExistingPublishingTaskWithDifferentPartitionAllocation() EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(task)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of(task.getId(), task)).anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(task)).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( @@ -1994,7 +1997,7 @@ public void testDiscoverExistingPublishingAndReadingTask() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(id1, id2)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(List.of(id1, id2))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); @@ -2124,7 +2127,7 @@ public void testKillUnresponsiveTasksWhileGettingStartTime() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata( null @@ -2162,7 +2165,7 @@ public void testKillUnresponsiveTasksWhileGettingStartTime() throws Exception .times(2); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -2207,7 +2210,7 @@ public void testKillUnresponsiveTasksWhilePausing() throws Exception EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(EasyMock.anyString())).andReturn(TaskLocation.unknown()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata( null @@ -2244,7 +2247,7 @@ public void testKillUnresponsiveTasksWhilePausing() throws Exception .times(2); captured = Capture.newInstance(CaptureType.ALL); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -2317,7 +2320,7 @@ public void testKillUnresponsiveTasksWhileSettingEndOffsets() throws Exception EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(EasyMock.anyString())).andReturn(TaskLocation.unknown()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata( null @@ -2356,7 +2359,7 @@ public void testKillUnresponsiveTasksWhileSettingEndOffsets() throws Exception .times(2); captured = Capture.newInstance(CaptureType.ALL); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -2542,7 +2545,9 @@ public void testStopGracefully() throws Exception EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(id1.getId())).andReturn(location1).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(id2.getId())).andReturn(location2).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) + .andReturn(toMap(List.of(id1, id2, id3))) + .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); @@ -2631,7 +2636,7 @@ public void testResetNoTasks() EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); replayAll(); @@ -2660,7 +2665,7 @@ public void testResetDataSourceMetadata() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); replayAll(); @@ -2739,7 +2744,7 @@ public void testResetNoDataSourceMetadata() EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); replayAll(); @@ -2797,7 +2802,7 @@ public void testGetOffsetFromStorageForPartitionWithResetOffsetAutomatically() t .andReturn(false) .anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); @@ -2960,7 +2965,9 @@ public void testResetRunningTasks() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) + .andReturn(toMap(List.of(id1, id2, id3))) + .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); @@ -3120,7 +3127,9 @@ public void testNoDataIngestionTasks() throws Exception workItems.add(new TestTaskRunnerWorkItem(id2, null, TaskLocation.create(id2.getId(), 8100, 8100))); workItems.add(new TestTaskRunnerWorkItem(id3, null, TaskLocation.create(id3.getId(), 8100, 8100))); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) + .andReturn(toMap(List.of(id1, id2, id3))) + .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); @@ -3276,7 +3285,9 @@ public void testCheckpointForInactiveTaskGroup() throws InterruptedException EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) + .andReturn(toMap(List.of(id1, id2, id3))) + .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); @@ -3426,7 +3437,9 @@ public void testCheckpointForUnknownTaskGroup() EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) + .andReturn(Map.of(id1.getId(), id1, id2.getId(), id2, id3.getId(), id3)) + .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); @@ -3480,7 +3493,7 @@ public void testSuspendedNoRunningTasks() .anyTimes(); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KinesisDataSourceMetadata( null @@ -3602,7 +3615,9 @@ public void testSuspendedRunningTasks() throws Exception EasyMock.expect(taskRunner.getTaskLocation(id1.getId())).andReturn(location1).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(id2.getId())).andReturn(location2).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(id3.getId())).andReturn(TaskLocation.unknown()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) + .andReturn(toMap(List.of(id1, id2, id3))) + .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); @@ -3680,7 +3695,7 @@ public void testResetSuspended() EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); replayAll(); @@ -3703,7 +3718,7 @@ public void testGetCurrentTotalStats() { supervisor = getTestableSupervisor(1, 2, true, "PT1H", null, null, false); supervisor.setPartitionIdsForTests( - ImmutableList.of(SHARD_ID0, SHARD_ID1) + List.of(SHARD_ID0, SHARD_ID1) ); supervisor.addTaskGroupToActivelyReadingTaskGroup( supervisor.getTaskGroupIdForPartition(SHARD_ID0), @@ -3801,12 +3816,12 @@ public void testDoNotKillCompatibleTasks() null ); - List existingTasks = ImmutableList.of(task); + List existingTasks = List.of(task); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(existingTasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(existingTasks)).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); EasyMock.expect(taskStorage.getTask("id2")).andReturn(Optional.of(task)).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) @@ -3896,12 +3911,12 @@ public void testKillIncompatibleTasks() throws InterruptedException null ); - List existingTasks = ImmutableList.of(task); + List existingTasks = List.of(task); EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(existingTasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(existingTasks)).anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(task)).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) @@ -4239,7 +4254,7 @@ private List testShardSplitPhaseOne() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.NOT_STARTED)) .anyTimes(); @@ -4282,7 +4297,7 @@ private List testShardSplitPhaseOne() throws Exception .andReturn(Futures.immediateFuture(checkpoints1)) .times(1); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) @@ -4350,7 +4365,7 @@ private List testShardSplitPhaseTwo(List phaseOneTasks) throws Excep EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.NOT_STARTED)) .anyTimes(); @@ -4396,7 +4411,7 @@ private List testShardSplitPhaseTwo(List phaseOneTasks) throws Excep .times(1); List postSplitTasks = postSplitCaptured.getValues(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(postSplitTasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(postSplitTasks)).anyTimes(); for (Task task : postSplitTasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -4530,7 +4545,7 @@ private void testShardSplitPhaseThree(List phaseTwoTasks) throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.NOT_STARTED)) .anyTimes(); @@ -4583,7 +4598,7 @@ private void testShardSplitPhaseThree(List phaseTwoTasks) throws Exception .times(1); List postSplitTasks = postSplitCaptured.getValues(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(postSplitTasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(postSplitTasks)).anyTimes(); for (Task task : postSplitTasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -4694,7 +4709,7 @@ private List testShardMergePhaseOne() throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.NOT_STARTED)) .anyTimes(); @@ -4746,7 +4761,7 @@ private List testShardMergePhaseOne() throws Exception .andReturn(Futures.immediateFuture(checkpoints1)) .times(1); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(tasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(tasks)).anyTimes(); for (Task task : tasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -4816,7 +4831,7 @@ private List testShardMergePhaseTwo(List phaseOneTasks) throws Excep EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.NOT_STARTED)) .anyTimes(); @@ -4864,7 +4879,7 @@ private List testShardMergePhaseTwo(List phaseOneTasks) throws Excep .times(1); List postMergeTasks = postMergeCaptured.getValues(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(postMergeTasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(postMergeTasks)).anyTimes(); for (Task task : postMergeTasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -4975,7 +4990,7 @@ private void testShardMergePhaseThree(List phaseTwoTasks) throws Exception EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.NOT_STARTED)) .anyTimes(); @@ -5017,7 +5032,7 @@ private void testShardMergePhaseThree(List phaseTwoTasks) throws Exception .times(1); List postSplitTasks = postSplitCaptured.getValues(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(postSplitTasks).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(toMap(postSplitTasks)).anyTimes(); for (Task task : postSplitTasks) { EasyMock.expect(taskStorage.getStatus(task.getId())) .andReturn(Optional.of(TaskStatus.running(task.getId()))) @@ -5503,7 +5518,7 @@ private static DataSchema getDataSchema(String dataSource) new UniformGranularitySpec( Granularities.HOUR, Granularities.NONE, - ImmutableList.of() + List.of() ) ) .build(); @@ -5590,6 +5605,11 @@ private KinesisIndexTask createKinesisIndexTask( ); } + private static Map toMap(List tasks) + { + return tasks.stream().collect(Collectors.toMap(Task::getId, Function.identity())); + } + private static class TestTaskRunnerWorkItem extends TaskRunnerWorkItem { private final String taskType; diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java index 3bbafb749017..45e7ebb10f01 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java @@ -24,8 +24,8 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Optional; import com.google.common.base.Preconditions; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.Sets; +import com.google.common.collect.MapDifference; +import com.google.common.collect.Maps; import com.google.common.util.concurrent.FutureCallback; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; @@ -51,6 +51,7 @@ import org.apache.druid.indexing.overlord.config.DefaultTaskConfig; import org.apache.druid.indexing.overlord.config.TaskLockConfig; import org.apache.druid.indexing.overlord.config.TaskQueueConfig; +import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.concurrent.Execs; @@ -64,11 +65,11 @@ import org.apache.druid.metadata.PasswordProviderRedactionMixIn; import org.apache.druid.server.coordinator.stats.CoordinatorRunStats; import org.apache.druid.utils.CollectionUtils; +import org.joda.time.DateTime; import java.util.Collection; import java.util.HashMap; import java.util.HashSet; -import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -79,9 +80,13 @@ import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; -import java.util.concurrent.locks.ReentrantLock; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.function.Consumer; +import java.util.function.Function; import java.util.stream.Collectors; /** @@ -101,20 +106,11 @@ public class TaskQueue // 60 MB warning threshold since 64 MB is the default max_allowed_packet size in MySQL 8+ private static final long TASK_SIZE_WARNING_THRESHOLD = 1024 * 1024 * 60; - // Task ID -> Task, for tasks that are active in some way (submitted, running, or finished and to-be-cleaned-up). - @GuardedBy("giant") - private final LinkedHashMap tasks = new LinkedHashMap<>(); - - // Task ID -> Future from the TaskRunner - @GuardedBy("giant") - private final Map> taskFutures = new HashMap<>(); - /** - * Tasks that have recently completed and are being cleaned up. These tasks - * should not be relaunched by task management. + * Tasks being managed by this queue. This includes tasks that are currently + * active i.e. submitted, running or recently completed (pending clean up). */ - @GuardedBy("giant") - private final Set recentlyCompletedTasks = new HashSet<>(); + private final ConcurrentHashMap activeTasks = new ConcurrentHashMap<>(); private final TaskLockConfig lockConfig; private final TaskQueueConfig config; @@ -127,7 +123,13 @@ public class TaskQueue private final ObjectMapper passwordRedactingMapper; private final TaskContextEnricher taskContextEnricher; - private final ReentrantLock giant = new ReentrantLock(true); + /** + * Lock used to ensure that {@link #start} and {@link #stop} do not run while + * queue management is in progress. {@link #start} and {@link #stop} methods + * acquire a WRITE lock whereas all other critical sections acquire a READ lock. + */ + private final ReentrantReadWriteLock giant = new ReentrantReadWriteLock(true); + @SuppressWarnings("MismatchedQueryAndUpdateOfCollection") private final BlockingQueue managementMayBeNecessary = new ArrayBlockingQueue<>(8); private final ExecutorService managerExec = Executors.newSingleThreadExecutor( @@ -203,7 +205,7 @@ void setActive() @LifecycleStart public void start() { - giant.lock(); + giant.writeLock().lock(); try { Preconditions.checkState(!active, "queue must be stopped"); @@ -270,7 +272,7 @@ public void start() } } finally { - giant.unlock(); + giant.writeLock().unlock(); } } @@ -280,18 +282,17 @@ public void start() @LifecycleStop public void stop() { - giant.lock(); + giant.writeLock().lock(); try { - tasks.clear(); - taskFutures.clear(); + activeTasks.clear(); active = false; managerExec.shutdownNow(); storageSyncExec.shutdownNow(); requestManagement(); } finally { - giant.unlock(); + giant.writeLock().unlock(); } } @@ -306,7 +307,7 @@ public boolean isActive() * Other callers (such as notifyStatus) should trigger activity on the * TaskQueue thread by requesting management here. */ - void requestManagement() + private void requestManagement() { // use a BlockingQueue since the offer/poll/wait behaviour is simple // and very easy to reason about @@ -328,7 +329,7 @@ void requestManagement() * @throws InterruptedException */ @SuppressFBWarnings(value = "RV_RETURN_VALUE_IGNORED", justification = "using queue as notification mechanism, result has no value") - void awaitManagementNanos(long nanos) throws InterruptedException + private void awaitManagementNanos(long nanos) throws InterruptedException { // mitigate a busy loop, it can get pretty busy when there are a lot of start/stops try { @@ -354,11 +355,11 @@ private void manage() throws InterruptedException log.info("Beginning management in %s.", config.getStartDelay()); Thread.sleep(config.getStartDelay().getMillis()); - // Ignore return value- we'll get the IDs and futures from getKnownTasks later. + // Ignore return value - we'll get the IDs and futures from getKnownTasks later. taskRunner.restore(); while (active) { - manageInternal(); + manageQueuedTasks(); // awaitNanos because management may become necessary without this condition signalling, // due to e.g. tasks becoming ready when other folks mess with the TaskLockbox. @@ -367,58 +368,62 @@ private void manage() throws InterruptedException } @VisibleForTesting - void manageInternal() + void manageQueuedTasks() { - Set knownTaskIds = new HashSet<>(); - Map> runnerTaskFutures = new HashMap<>(); - - giant.lock(); - + giant.readLock().lock(); try { - manageInternalCritical(knownTaskIds, runnerTaskFutures); + startPendingTasksOnRunner(); } finally { - giant.unlock(); + giant.readLock().unlock(); } - - manageInternalPostCritical(knownTaskIds, runnerTaskFutures); } /** - * Management loop critical section tasks. - * - * @param knownTaskIds will be modified - filled with known task IDs - * @param runnerTaskFutures will be modified - filled with futures related to getting the running tasks + * Starts pending tasks on the {@link #taskRunner}. + * Cleans up tasks present on the {@link #taskRunner} but not present in + * {@link #activeTasks} anymore. */ @GuardedBy("giant") - private void manageInternalCritical( - final Set knownTaskIds, - final Map> runnerTaskFutures - ) + private void startPendingTasksOnRunner() { - // Task futures available from the taskRunner - for (final TaskRunnerWorkItem workItem : taskRunner.getKnownTasks()) { - if (!recentlyCompletedTasks.contains(workItem.getTaskId())) { - // Don't do anything with tasks that have recently finished; notifyStatus will handle it. - runnerTaskFutures.put(workItem.getTaskId(), workItem.getResult()); - } + final Set unknownTaskIds = new HashSet<>(); + + // Get all task futures from the taskRunner and shutdown tasks not in queue anymore + final Map> runnerTaskFutures = new HashMap<>(); + for (final TaskRunnerWorkItem workItem : List.copyOf(taskRunner.getKnownTasks())) { + final String taskId = workItem.getTaskId(); + updateTaskEntry(taskId, entry -> { + if (entry == null) { + unknownTaskIds.add(taskId); + shutdownUnknownTaskOnRunner(taskId); + } else { + runnerTaskFutures.put(taskId, workItem.getResult()); + } + }); } + log.info("Cleaned up [%,d] tasks on task runner with IDs[%s].", unknownTaskIds.size(), unknownTaskIds); + // Attain futures for all active tasks (assuming they are ready to run). // Copy tasks list, as notifyStatus may modify it. - for (final Task task : ImmutableList.copyOf(tasks.values())) { - if (recentlyCompletedTasks.contains(task.getId())) { - // Don't do anything with tasks that have recently finished; notifyStatus will handle it. - continue; - } + for (final String queuedTaskId : List.copyOf(activeTasks.keySet())) { + updateTaskEntry( + queuedTaskId, + entry -> startPendingTaskOnRunner(entry, runnerTaskFutures.get(queuedTaskId)) + ); + } + } - knownTaskIds.add(task.getId()); + @GuardedBy("giant") + private void startPendingTaskOnRunner(TaskEntry entry, ListenableFuture runnerTaskFuture) + { + // Don't do anything with tasks that have recently finished; notifyStatus will handle it. + if (entry != null && !entry.isComplete) { + final Task task = entry.task; - if (!taskFutures.containsKey(task.getId())) { - final ListenableFuture runnerTaskFuture; - if (runnerTaskFutures.containsKey(task.getId())) { - runnerTaskFuture = runnerTaskFutures.get(task.getId()); - } else { + if (entry.future == null) { + if (runnerTaskFuture == null) { // Task should be running, so run it. final boolean taskIsReady; try { @@ -437,7 +442,7 @@ private void manageInternalCritical( } TaskStatus taskStatus = TaskStatus.failure(task.getId(), errorMessage); notifyStatus(task, taskStatus, taskStatus.getErrorMsg()); - continue; + return; } if (taskIsReady) { log.info("Asking taskRunner to run task[%s]", task.getId()); @@ -446,11 +451,11 @@ private void manageInternalCritical( // Task.isReady() can internally lock intervals or segments. // We should release them if the task is not ready. taskLockbox.unlockAll(task); - continue; + return; } } attachCallbacks(task, runnerTaskFuture); - taskFutures.put(task.getId(), runnerTaskFuture); + entry.future = runnerTaskFuture; } else if (isTaskPending(task)) { // if the taskFutures contain this task and this task is pending, also let the taskRunner // to run it to guarantee it will be assigned to run @@ -460,32 +465,13 @@ private void manageInternalCritical( } } - @VisibleForTesting - private void manageInternalPostCritical( - final Set knownTaskIds, - final Map> runnerTaskFutures - ) + private void shutdownUnknownTaskOnRunner(String taskId) { - // Kill tasks that shouldn't be running - final Set tasksToKill = Sets.difference(runnerTaskFutures.keySet(), knownTaskIds); - if (!tasksToKill.isEmpty()) { - log.info("Asking taskRunner to clean up [%,d] tasks.", tasksToKill.size()); - - // On large installations running several thousands of tasks, - // concatenating the list of known task ids can be compupationally expensive. - final boolean logKnownTaskIds = log.isDebugEnabled(); - final String reason = logKnownTaskIds - ? StringUtils.format("Task is not in knownTaskIds[%s]", knownTaskIds) - : "Task is not in knownTaskIds"; - - for (final String taskId : tasksToKill) { - try { - taskRunner.shutdown(taskId, reason); - } - catch (Exception e) { - log.warn(e, "TaskRunner failed to clean up task[%s].", taskId); - } - } + try { + taskRunner.shutdown(taskId, "Task is not present in queue anymore."); + } + catch (Exception e) { + log.warn(e, "TaskRunner failed to clean up task[%s].", taskId); } } @@ -525,55 +511,89 @@ public boolean add(final Task task) taskContextEnricher.enrichContext(task); - giant.lock(); + giant.readLock().lock(); try { Preconditions.checkState(active, "Queue is not active!"); Preconditions.checkNotNull(task, "task"); - if (tasks.size() >= config.getMaxSize()) { + if (activeTasks.size() >= config.getMaxSize()) { throw DruidException.forPersona(DruidException.Persona.ADMIN) .ofCategory(DruidException.Category.CAPACITY_EXCEEDED) .build( "Task queue already contains [%d] tasks." + " Retry later or increase 'druid.indexer.queue.maxSize'[%d].", - tasks.size(), config.getMaxSize() + activeTasks.size(), config.getMaxSize() ); } // If this throws with any sort of exception, including TaskExistsException, we don't want to // insert the task into our queue. So don't catch it. + final DateTime insertTime = DateTimes.nowUtc(); taskStorage.insert(task, TaskStatus.running(task.getId())); - addTaskInternal(task); + addTaskInternal(task, insertTime); requestManagement(); return true; } finally { - giant.unlock(); + giant.readLock().unlock(); } } @GuardedBy("giant") - private void addTaskInternal(final Task task) + private void addTaskInternal(final Task task, final DateTime updateTime) { - final Task existingTask = tasks.putIfAbsent(task.getId(), task); + final AtomicBoolean added = new AtomicBoolean(false); + final TaskEntry entry = activeTasks.compute( + task.getId(), + (taskId, prevEntry) -> { + if (prevEntry == null) { + // TODO: What if this task was just removed from the queue? + // and the latest poll didn't have this update? + // we might end up re-queueing the task + // Maybe do not remove recently completed tasks except from sync thread i.e. every minute + + added.set(true); + return new TaskEntry(task); + } else if (prevEntry.lastUpdatedTime.isBefore(updateTime)) { + prevEntry.lastUpdatedTime = updateTime; + } + + return prevEntry; + } + ); - if (existingTask == null) { + if (added.get()) { taskLockbox.add(task); - } else if (!existingTask.equals(task)) { + } else if (!entry.task.equals(task)) { throw new ISE("Cannot add task ID [%s] with same ID as task that has already been added", task.getId()); } } /** - * Removes a task from {@link #tasks} and {@link #taskLockbox}, if it exists. Returns whether the task was - * removed or not. + * Removes a task from {@link #activeTasks} and {@link #taskLockbox}, if it exists. + * + * @return true if the task was removed successfully */ @GuardedBy("giant") - private boolean removeTaskInternal(final String taskId) + private boolean removeTaskInternal(final String taskId, final DateTime deleteTime) { - final Task task = tasks.remove(taskId); - if (task != null) { - taskLockbox.remove(task); + final AtomicReference removedTask = new AtomicReference<>(); + + activeTasks.compute( + taskId, + (id, prevEntry) -> { + // Remove the task only if it doesn't have a more recent update + if (prevEntry != null && prevEntry.lastUpdatedTime.isBefore(deleteTime)) { + removedTask.set(prevEntry.task); + } + + // Remove this taskId from activeTasks by mapping it to null + return null; + } + ); + + if (removedTask.get() != null) { + taskLockbox.remove(removedTask.get()); return true; } else { return false; @@ -589,38 +609,36 @@ private boolean removeTaskInternal(final String taskId) */ public void shutdown(final String taskId, String reasonFormat, Object... args) { - giant.lock(); - + giant.readLock().lock(); try { - final Task task = tasks.get(Preconditions.checkNotNull(taskId, "taskId")); - if (task != null) { - notifyStatus(task, TaskStatus.failure(taskId, StringUtils.format(reasonFormat, args)), reasonFormat, args); + final TaskEntry entry = activeTasks.get(Preconditions.checkNotNull(taskId, "taskId")); + if (entry != null) { + notifyStatus(entry.task, TaskStatus.failure(taskId, StringUtils.format(reasonFormat, args)), reasonFormat, args); } } finally { - giant.unlock(); + giant.readLock().unlock(); } } /** * Shuts down a task, but records the task status as a success, unike {@link #shutdown(String, String, Object...)} * - * @param taskId task to shutdown + * @param taskId task to shut down * @param reasonFormat A format string indicating the shutdown reason * @param args arguments for reasonFormat */ public void shutdownWithSuccess(final String taskId, String reasonFormat, Object... args) { - giant.lock(); - + giant.readLock().lock(); try { - final Task task = tasks.get(Preconditions.checkNotNull(taskId, "taskId")); - if (task != null) { - notifyStatus(task, TaskStatus.success(taskId), reasonFormat, args); + final TaskEntry entry = activeTasks.get(Preconditions.checkNotNull(taskId, "taskId")); + if (entry != null) { + notifyStatus(entry.task, TaskStatus.success(taskId), reasonFormat, args); } } finally { - giant.unlock(); + giant.readLock().unlock(); } } @@ -628,13 +646,12 @@ public void shutdownWithSuccess(final String taskId, String reasonFormat, Object * Notifies this queue that the given task has an updated status. If this update * is valid and task is now complete, the following operations are performed: *
    - *
  • Add task to {@link #recentlyCompletedTasks} to prevent re-launching them
    • + *
  • Mark task as recently completed to prevent re-launching them
    • *
  • Persist new status in the metadata storage to safeguard against crashes * and leader re-elections
    • *
  • Request {@link #taskRunner} to shutdown task (synchronously)
    • *
  • Remove all locks for task from metadata storage
    • - *
  • Remove task entry from {@link #tasks} and {@link #taskFutures}
    • - *
  • Remove task from {@link #recentlyCompletedTasks}
    • + *
  • Remove task entry from {@link #activeTasks}
    • *
  • Request task management
    • *
*

@@ -646,6 +663,7 @@ public void shutdownWithSuccess(final String taskId, String reasonFormat, Object * @throws IllegalArgumentException if the task ID does not match the status ID * @throws IllegalStateException if this queue is currently shut down */ + @GuardedBy("giant") private void notifyStatus(final Task task, final TaskStatus taskStatus, String reasonFormat, Object... args) { Preconditions.checkNotNull(task, "task"); @@ -663,14 +681,13 @@ private void notifyStatus(final Task task, final TaskStatus taskStatus, String r return; } - // Critical section: add this task to recentlyCompletedTasks, so it isn't managed while being cleaned up. - giant.lock(); - try { - recentlyCompletedTasks.add(task.getId()); - } - finally { - giant.unlock(); - } + // Mark this task as complete, so it isn't managed while being cleaned up. + updateTaskEntry(task.getId(), entry -> { + // Don't do anything if the task is not present in activeTasks + if (entry != null) { + entry.isComplete = true; + } + }); final TaskLocation taskLocation = taskRunner.getTaskLocation(task.getId()); @@ -705,20 +722,15 @@ private void notifyStatus(final Task task, final TaskStatus taskStatus, String r log.warn(e, "TaskRunner failed to cleanup task after completion: %s", task.getId()); } - // Critical section: remove this task from all of our tracking data structures. - giant.lock(); + giant.readLock().lock(); try { - if (removeTaskInternal(task.getId())) { - taskFutures.remove(task.getId()); - } else { + if (!removeTaskInternal(task.getId(), DateTimes.nowUtc())) { log.warn("Unknown task[%s] completed", task.getId()); } - - recentlyCompletedTasks.remove(task.getId()); requestManagement(); } finally { - giant.unlock(); + giant.readLock().unlock(); } } @@ -776,7 +788,13 @@ private void handleStatus(final TaskStatus status) return; } - notifyStatus(task, status, "notified status change from task"); + giant.readLock().lock(); + try { + notifyStatus(task, status, "notified status change from task"); + } + finally { + giant.readLock().unlock(); + } // Emit event and log, if the task is done if (status.isComplete()) { @@ -818,31 +836,28 @@ private void handleStatus(final TaskStatus status) */ private void syncFromStorage() { - giant.lock(); + giant.readLock().lock(); + final DateTime syncStartTime = DateTimes.nowUtc(); try { if (active) { final Map newTasks = toTaskIDMap(taskStorage.getActiveTasks()); final int tasksSynced = newTasks.size(); - final Map oldTasks = new HashMap<>(tasks); + final Map oldTasks = CollectionUtils.mapValues(activeTasks, entry -> entry.task); - // Calculate differences on IDs instead of Task Objects. - Set commonIds = Sets.newHashSet(Sets.intersection(newTasks.keySet(), oldTasks.keySet())); - for (String taskID : commonIds) { - newTasks.remove(taskID); - oldTasks.remove(taskID); - } - Collection addedTasks = newTasks.values(); - Collection removedTasks = oldTasks.values(); + // Identify the tasks that have been added or removed from the storage + final MapDifference mapDifference = Maps.difference(oldTasks, newTasks); + final Collection addedTasks = mapDifference.entriesOnlyOnRight().values(); + final Collection removedTasks = mapDifference.entriesOnlyOnLeft().values(); - // Clean up removed Tasks + // Remove tasks not present in metadata store if their lastUpdatedTime is before syncStartTime for (Task task : removedTasks) { - removeTaskInternal(task.getId()); + removeTaskInternal(task.getId(), syncStartTime); } - // Add newly Added tasks to the queue + // Add new tasks present in metadata store if their lastUpdatedTime is before syncStartTime for (Task task : addedTasks) { - addTaskInternal(task); + addTaskInternal(task, syncStartTime); } log.info( @@ -861,7 +876,7 @@ private void syncFromStorage() throw new RuntimeException(e); } finally { - giant.unlock(); + giant.readLock().unlock(); } } @@ -876,9 +891,14 @@ private static Map toTaskIDMap(List taskList) private Map getDeltaValues(Map total, Map prev) { - return total.entrySet() - .stream() - .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue() - prev.getOrDefault(e.getKey(), 0L))); + final Map deltaValues = new HashMap<>(); + total.forEach( + (dataSource, totalCount) -> deltaValues.put( + dataSource, + totalCount - prev.getOrDefault(dataSource, 0L) + ) + ); + return deltaValues; } public Map getSuccessfulTaskCount() @@ -901,15 +921,11 @@ public Map getFailedTaskCount() } } - Map getCurrentTaskDatasources() + private Map getCurrentTaskDatasources() { - giant.lock(); - try { - return tasks.values().stream().collect(Collectors.toMap(Task::getId, Task::getDataSource)); - } - finally { - giant.unlock(); - } + return activeTasks.values().stream() + .map(entry -> entry.task) + .collect(Collectors.toMap(Task::getId, Task::getDataSource)); } public Map getRunningTaskCount() @@ -943,14 +959,10 @@ public Map getWaitingTaskCount() .map(TaskRunnerWorkItem::getTaskId) .collect(Collectors.toSet()); - giant.lock(); - try { - return tasks.values().stream().filter(task -> !runnerKnownTaskIds.contains(task.getId())) - .collect(Collectors.toMap(Task::getDataSource, task -> 1L, Long::sum)); - } - finally { - giant.unlock(); - } + return activeTasks.values().stream() + .map(entry -> entry.task) + .filter(task -> !runnerKnownTaskIds.contains(task.getId())) + .collect(Collectors.toMap(Task::getDataSource, task -> 1L, Long::sum)); } /** @@ -988,18 +1000,16 @@ public CoordinatorRunStats getQueueStats() */ public Optional getActiveTask(String id) { - Task task; - giant.lock(); - try { - task = tasks.get(id); - } - finally { - giant.unlock(); + final TaskEntry entry = activeTasks.get(id); + if (entry == null) { + return Optional.absent(); } + + Task task = entry.task; if (task != null) { try { // Write and read the value using a mapper with password redaction mixin. - task = passwordRedactingMapper.readValue(passwordRedactingMapper.writeValueAsString(task), Task.class); + task = passwordRedactingMapper.readValue(passwordRedactingMapper.writeValueAsString(entry.task), Task.class); } catch (JsonProcessingException e) { log.error(e, "Failed to serialize or deserialize task with id [%s].", task.getId()); @@ -1017,43 +1027,24 @@ public Optional getActiveTask(String id) */ public List getTasks() { - giant.lock(); - try { - return List.copyOf(tasks.values()); - } - finally { - giant.unlock(); - } - } - - /** - * IDs of tasks that have recently completed and are being cleaned up. - */ - private Set getRecentlyCompletedTaskIds() - { - giant.lock(); - try { - return Set.copyOf(recentlyCompletedTasks); - } - finally { - giant.unlock(); - } + return activeTasks.values().stream().map(entry -> entry.task).collect(Collectors.toList()); } /** * Returns the list of currently active tasks for the given datasource. */ - public List getActiveTasksForDatasource(String datasource) + public Map getActiveTasksForDatasource(String datasource) { - // This method is called very frequently by streaming supervisors. - // To reduce contention, perform filtration outside the lock - final List allTasks = getTasks(); - final Set completedTaskIds = getRecentlyCompletedTaskIds(); - - return allTasks.stream().filter( - task -> task.getDataSource().equals(datasource) - && !completedTaskIds.contains(task.getId()) - ).collect(Collectors.toList()); + // TODO: ensure that iteration over a concurrent hash map is performant + // We may even maintain a separate data structure if that helps this API + // I think we should maintain a separate DS after all + return activeTasks.values().stream() + .filter( + entry -> !entry.isComplete + && entry.task.getDataSource().equals(datasource) + ) + .map(entry -> entry.task) + .collect(Collectors.toMap(Task::getId, Function.identity())); } private void validateTaskPayload(Task task) @@ -1085,4 +1076,42 @@ private void validateTaskPayload(Task task) ); } } + + /** + * Performs a thread-safe update on the task entry for the given task ID. + */ + private void updateTaskEntry(String taskId, Consumer updateOperation) + { + giant.readLock().lock(); + try { + activeTasks.compute( + taskId, + (id, existingEntry) -> { + updateOperation.accept(existingEntry); + return existingEntry; + } + ); + } + finally { + giant.readLock().unlock(); + } + } + + /** + * Represents an entry in this task queue. + */ + private static class TaskEntry + { + private final Task task; + + private DateTime lastUpdatedTime; + private ListenableFuture future = null; + private boolean isComplete = false; + + TaskEntry(Task task) + { + this.task = task; + this.lastUpdatedTime = DateTimes.nowUtc(); + } + } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java index daaf2529cc49..8e0eddf13f9f 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -4297,7 +4297,7 @@ private Map getActiveTaskMap() final Optional taskQueue = taskMaster.getTaskQueue(); final List tasks; if (taskQueue.isPresent()) { - tasks = taskQueue.get().getActiveTasksForDatasource(dataSource); + return taskQueue.get().getActiveTasksForDatasource(dataSource); } else { tasks = taskStorage.getActiveTasksByDatasource(dataSource); } diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java index 58ef76eb7d6f..391e539dbdb4 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java @@ -100,6 +100,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; public class TaskQueueTest extends IngestionTestBase @@ -142,7 +143,7 @@ public Map getContext() } @Test - public void testManageInternalReleaseLockWhenTaskIsNotReady() throws Exception + public void testManageQueuedTasksReleaseLockWhenTaskIsNotReady() throws Exception { // task1 emulates a case when there is a task that was issued before task2 and acquired locks conflicting // to task2. @@ -153,14 +154,14 @@ public void testManageInternalReleaseLockWhenTaskIsNotReady() throws Exception final TestTask task2 = new TestTask("t2", Intervals.of("2021-01-31/P1M")); taskQueue.add(task2); - taskQueue.manageInternal(); + taskQueue.manageQueuedTasks(); Assert.assertFalse(task2.isDone()); Assert.assertTrue(getLockbox().findLocksForTask(task2).isEmpty()); // task3 can run because task2 is still blocked by task1. final TestTask task3 = new TestTask("t3", Intervals.of("2021-02-01/P1M")); taskQueue.add(task3); - taskQueue.manageInternal(); + taskQueue.manageQueuedTasks(); Assert.assertFalse(task2.isDone()); Assert.assertTrue(task3.isDone()); Assert.assertTrue(getLockbox().findLocksForTask(task2).isEmpty()); @@ -170,7 +171,7 @@ public void testManageInternalReleaseLockWhenTaskIsNotReady() throws Exception taskQueue.shutdown(task3.getId(), "Emulating shutdown of task3"); // Now task2 should run. - taskQueue.manageInternal(); + taskQueue.manageQueuedTasks(); Assert.assertTrue(task2.isDone()); } @@ -377,7 +378,7 @@ public boolean isReady(TaskActionClient taskActionClient) } }; taskQueue.add(task); - taskQueue.manageInternal(); + taskQueue.manageQueuedTasks(); Optional statusOptional = getTaskStorage().getStatus(task.getId()); Assert.assertTrue(statusOptional.isPresent()); @@ -418,7 +419,7 @@ public void testKilledTasksEmitRuntimeMetricWithHttpRemote() throws InterruptedE taskQueue.setActive(); final Task task = new TestTask("t1", Intervals.of("2021-01-01/P1D")); taskQueue.add(task); - taskQueue.manageInternal(); + taskQueue.manageQueuedTasks(); // Announce the task and wait for it to start running final String taskId = task.getId(); @@ -437,7 +438,7 @@ public void testKilledTasksEmitRuntimeMetricWithHttpRemote() throws InterruptedE TaskAnnouncement.create(task, TaskStatus.failure(taskId, "shutdown"), taskLocation), workerHolder ); - taskQueue.manageInternal(); + taskQueue.manageQueuedTasks(); Thread.sleep(100); // Verify that metrics are emitted on receiving announcement diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java index 245aa5aba3df..e8dcdef91567 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java @@ -201,7 +201,7 @@ public void testRunning() { EasyMock.expect(spec.isSuspended()).andReturn(false).anyTimes(); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); replayAll(); @@ -243,7 +243,7 @@ public void testRunningStreamGetSequenceNumberReturnsNull() EasyMock.expect(recordSupplier.getLatestSequenceNumber(EasyMock.anyObject())).andReturn(null).anyTimes(); EasyMock.expect(spec.isSuspended()).andReturn(false).anyTimes(); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); replayAll(); @@ -469,7 +469,7 @@ public void testConnectingToStreamFail() EasyMock.expect(recordSupplier.getPartitionIds(STREAM)) .andThrow(new StreamException(new IllegalStateException(EXCEPTION_MSG))) .anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); replayAll(); @@ -529,7 +529,7 @@ public void testConnectingToStreamFailRecoveryFailRecovery() .andThrow(new StreamException(new IllegalStateException())) .times(3); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).times(3); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); EasyMock.expect(taskRunner.getRunningTasks()).andReturn(ImmutableList.of()).anyTimes(); @@ -597,7 +597,7 @@ public void testDiscoveringInitialTasksFailRecoveryFail() EasyMock.expect(spec.isSuspended()).andReturn(false).anyTimes(); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andThrow(new IllegalStateException(EXCEPTION_MSG)).times(3); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).times(6); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).times(6); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andThrow(new IllegalStateException(EXCEPTION_MSG)).times(3); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); @@ -704,7 +704,7 @@ public Duration getEmissionDuration() EasyMock.expect(spec.getType()).andReturn("test").anyTimes(); EasyMock.expect(spec.getSupervisorStateManagerConfig()).andReturn(supervisorConfig).anyTimes(); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); replayAll(); @@ -811,7 +811,7 @@ public Duration getEmissionDuration() EasyMock.expect(recordSupplier.isOffsetAvailable(EasyMock.anyObject(), EasyMock.anyObject())) .andReturn(true) .anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(List.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); replayAll(); @@ -849,7 +849,7 @@ public void testCreatingTasksFailRecoveryFail() { EasyMock.expect(spec.isSuspended()).andReturn(false).anyTimes(); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andThrow(new IllegalStateException(EXCEPTION_MSG)).times(3); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).times(3); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andThrow(new IllegalStateException(EXCEPTION_MSG)).times(3); @@ -927,7 +927,7 @@ public void testSuspended() { EasyMock.expect(spec.isSuspended()).andReturn(true).anyTimes(); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); replayAll(); @@ -966,7 +966,7 @@ public void testStopping() { EasyMock.expect(spec.isSuspended()).andReturn(false).anyTimes(); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); taskRunner.unregisterListener("testSupervisorId"); @@ -1007,7 +1007,7 @@ public void testStopGracefully() throws Exception { EasyMock.expect(spec.isSuspended()).andReturn(false).anyTimes(); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(List.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); taskRunner.unregisterListener("testSupervisorId"); @@ -1029,7 +1029,7 @@ public void testStoppingGracefully() { EasyMock.expect(spec.isSuspended()).andReturn(false).anyTimes(); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); taskRunner.unregisterListener("testSupervisorId"); @@ -1171,7 +1171,7 @@ public Duration getEmissionDuration() EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) - .andReturn(ImmutableList.of(id1, id2, id3)) + .andReturn(Map.of(id1.getId(), id1, id2.getId(), id2, id3.getId(), id3)) .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); @@ -1400,7 +1400,7 @@ public Duration getEmissionDuration() EasyMock.expect(taskRunner.getTaskLocation(id2.getId())).andReturn(location2).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(id3.getId())).andReturn(location3).anyTimes(); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) - .andReturn(ImmutableList.of(id1, id2, id3)) + .andReturn(Map.of(id1.getId(), id1, id2.getId(), id2, id3.getId(), id3)) .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); @@ -1583,7 +1583,7 @@ public Duration getEmissionDuration() EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); EasyMock.expect(taskRunner.getTaskLocation(id1.getId())).andReturn(location1).anyTimes(); EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)) - .andReturn(List.of(id1)) + .andReturn(Map.of(id1.getId(), id1)) .anyTimes(); EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); @@ -2580,7 +2580,7 @@ public Duration getEmissionDuration() EasyMock.expect(spec.getContextValue(DruidMetrics.TAGS)).andReturn(METRIC_TAGS).anyTimes(); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); - EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); replayAll(); From 866f14e45474e3939a57732434231d46c7ff4f5d Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Mon, 24 Mar 2025 11:09:57 +0530 Subject: [PATCH 02/10] Handle race conditions with sync --- .../druid/indexing/overlord/TaskQueue.java | 152 +++++++++--------- .../indexing/overlord/TaskQueueScaleTest.java | 3 +- .../indexing/overlord/TaskQueueTest.java | 1 - 3 files changed, 76 insertions(+), 80 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java index 45e7ebb10f01..8d351185c120 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java @@ -126,7 +126,7 @@ public class TaskQueue /** * Lock used to ensure that {@link #start} and {@link #stop} do not run while * queue management is in progress. {@link #start} and {@link #stop} methods - * acquire a WRITE lock whereas all other critical sections acquire a READ lock. + * acquire a WRITE lock whereas all other operations acquire a READ lock. */ private final ReentrantReadWriteLock giant = new ReentrantReadWriteLock(true); @@ -303,7 +303,7 @@ public boolean isActive() /** * Request management from the management thread. Non-blocking. - * + *

* Other callers (such as notifyStatus) should trigger activity on the * TaskQueue thread by requesting management here. */ @@ -322,26 +322,21 @@ private void requestManagement() /** * Await for an event to manage. - * + *

* This should only be called from the management thread to wait for activity. - * - * @param nanos - * @throws InterruptedException */ @SuppressFBWarnings(value = "RV_RETURN_VALUE_IGNORED", justification = "using queue as notification mechanism, result has no value") - private void awaitManagementNanos(long nanos) throws InterruptedException + private void awaitManagement() throws InterruptedException { // mitigate a busy loop, it can get pretty busy when there are a lot of start/stops - try { - Thread.sleep(MIN_WAIT_TIME_MS); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } + Thread.sleep(MIN_WAIT_TIME_MS); // wait for an item, if an item arrives (or is already available), complete immediately // (does not actually matter what the item is) - managementMayBeNecessary.poll(nanos - (TimeUnit.MILLISECONDS.toNanos(MIN_WAIT_TIME_MS)), TimeUnit.NANOSECONDS); + managementMayBeNecessary.poll( + TaskQueue.MANAGEMENT_WAIT_TIMEOUT_NANOS - TimeUnit.MILLISECONDS.toNanos(MIN_WAIT_TIME_MS), + TimeUnit.NANOSECONDS + ); // there may have been multiple requests, clear them all managementMayBeNecessary.clear(); @@ -352,7 +347,7 @@ private void awaitManagementNanos(long nanos) throws InterruptedException */ private void manage() throws InterruptedException { - log.info("Beginning management in %s.", config.getStartDelay()); + log.info("Beginning management in [%s].", config.getStartDelay()); Thread.sleep(config.getStartDelay().getMillis()); // Ignore return value - we'll get the IDs and futures from getKnownTasks later. @@ -361,9 +356,9 @@ private void manage() throws InterruptedException while (active) { manageQueuedTasks(); - // awaitNanos because management may become necessary without this condition signalling, - // due to e.g. tasks becoming ready when other folks mess with the TaskLockbox. - awaitManagementNanos(MANAGEMENT_WAIT_TIMEOUT_NANOS); + // wait because management may become necessary without this condition signalling, + // e.g., due to tasks becoming ready when other folks mess with the TaskLockbox. + awaitManagement(); } } @@ -441,7 +436,7 @@ private void startPendingTaskOnRunner(TaskEntry entry, ListenableFuture { if (prevEntry == null) { - // TODO: What if this task was just removed from the queue? - // and the latest poll didn't have this update? - // we might end up re-queueing the task - // Maybe do not remove recently completed tasks except from sync thread i.e. every minute - added.set(true); return new TaskEntry(task); } else if (prevEntry.lastUpdatedTime.isBefore(updateTime)) { @@ -570,12 +560,16 @@ private void addTaskInternal(final Task task, final DateTime updateTime) } /** - * Removes a task from {@link #activeTasks} and {@link #taskLockbox}, if it exists. - * - * @return true if the task was removed successfully + * Removes a task from {@link #activeTasks} and {@link #taskLockbox}, if required. + *

+ * This method must be called only from {@link #syncFromStorage()} to avoid + * race conditions. For example, if a task finishes and is removed from + * {@link #activeTasks} while a sync is in progress, the polled results from + * the DB would still have the task as RUNNING and {@link #syncFromStorage()} + * might add it back to the queue, thus causing a duplicate run of the task. */ @GuardedBy("giant") - private boolean removeTaskInternal(final String taskId, final DateTime deleteTime) + private void removeTaskInternal(final String taskId, final DateTime deleteTime) { final AtomicReference removedTask = new AtomicReference<>(); @@ -593,10 +587,7 @@ private boolean removeTaskInternal(final String taskId, final DateTime deleteTim ); if (removedTask.get() != null) { - taskLockbox.remove(removedTask.get()); - return true; - } else { - return false; + removeTaskLock(removedTask.get()); } } @@ -609,12 +600,18 @@ private boolean removeTaskInternal(final String taskId, final DateTime deleteTim */ public void shutdown(final String taskId, String reasonFormat, Object... args) { + Preconditions.checkNotNull(taskId, "taskId"); giant.readLock().lock(); try { - final TaskEntry entry = activeTasks.get(Preconditions.checkNotNull(taskId, "taskId")); - if (entry != null) { - notifyStatus(entry.task, TaskStatus.failure(taskId, StringUtils.format(reasonFormat, args)), reasonFormat, args); - } + updateTaskEntry( + taskId, + entry -> notifyStatus( + entry, + TaskStatus.failure(taskId, StringUtils.format(reasonFormat, args)), + reasonFormat, + args + ) + ); } finally { giant.readLock().unlock(); @@ -630,12 +627,13 @@ public void shutdown(final String taskId, String reasonFormat, Object... args) */ public void shutdownWithSuccess(final String taskId, String reasonFormat, Object... args) { + Preconditions.checkNotNull(taskId, "taskId"); giant.readLock().lock(); try { - final TaskEntry entry = activeTasks.get(Preconditions.checkNotNull(taskId, "taskId")); - if (entry != null) { - notifyStatus(entry.task, TaskStatus.success(taskId), reasonFormat, args); - } + updateTaskEntry( + taskId, + entry -> notifyStatus(entry, TaskStatus.success(taskId), reasonFormat, args) + ); } finally { giant.readLock().unlock(); @@ -646,14 +644,15 @@ public void shutdownWithSuccess(final String taskId, String reasonFormat, Object * Notifies this queue that the given task has an updated status. If this update * is valid and task is now complete, the following operations are performed: *

    - *
  • Mark task as recently completed to prevent re-launching them
    • + *
  • Mark task as completed to prevent re-launching it
    • *
  • Persist new status in the metadata storage to safeguard against crashes * and leader re-elections
    • *
  • Request {@link #taskRunner} to shutdown task (synchronously)
    • *
  • Remove all locks for task from metadata storage
    • - *
  • Remove task entry from {@link #activeTasks}
    • *
  • Request task management
    • *
+ * This method does not remove the task from {@link #activeTasks} to avoid + * race conditions with {@link #syncFromStorage()}. *

* Since this operation involves DB updates and synchronous remote calls, it * must be invoked on a dedicated executor so that task runner and worker sync @@ -662,10 +661,16 @@ public void shutdownWithSuccess(final String taskId, String reasonFormat, Object * @throws NullPointerException if task or status is null * @throws IllegalArgumentException if the task ID does not match the status ID * @throws IllegalStateException if this queue is currently shut down + * @see #removeTaskInternal */ - @GuardedBy("giant") - private void notifyStatus(final Task task, final TaskStatus taskStatus, String reasonFormat, Object... args) + private void notifyStatus(final TaskEntry entry, final TaskStatus taskStatus, String reasonFormat, Object... args) { + // Don't do anything if the task has no entry in activeTasks + if (entry == null) { + return; + } + + final Task task = entry.task; Preconditions.checkNotNull(task, "task"); Preconditions.checkNotNull(taskStatus, "status"); Preconditions.checkState(active, "Queue is not active!"); @@ -682,12 +687,7 @@ private void notifyStatus(final Task task, final TaskStatus taskStatus, String r } // Mark this task as complete, so it isn't managed while being cleaned up. - updateTaskEntry(task.getId(), entry -> { - // Don't do anything if the task is not present in activeTasks - if (entry != null) { - entry.isComplete = true; - } - }); + entry.isComplete = true; final TaskLocation taskLocation = taskRunner.getTaskLocation(task.getId()); @@ -722,16 +722,8 @@ private void notifyStatus(final Task task, final TaskStatus taskStatus, String r log.warn(e, "TaskRunner failed to cleanup task after completion: %s", task.getId()); } - giant.readLock().lock(); - try { - if (!removeTaskInternal(task.getId(), DateTimes.nowUtc())) { - log.warn("Unknown task[%s] completed", task.getId()); - } - requestManagement(); - } - finally { - giant.readLock().unlock(); - } + removeTaskLock(task); + requestManagement(); } /** @@ -788,13 +780,10 @@ private void handleStatus(final TaskStatus status) return; } - giant.readLock().lock(); - try { - notifyStatus(task, status, "notified status change from task"); - } - finally { - giant.readLock().unlock(); - } + updateTaskEntry( + task.getId(), + entry -> notifyStatus(entry, status, "notified status change from task") + ); // Emit event and log, if the task is done if (status.isComplete()) { @@ -1035,16 +1024,14 @@ public List getTasks() */ public Map getActiveTasksForDatasource(String datasource) { - // TODO: ensure that iteration over a concurrent hash map is performant - // We may even maintain a separate data structure if that helps this API - // I think we should maintain a separate DS after all - return activeTasks.values().stream() - .filter( - entry -> !entry.isComplete - && entry.task.getDataSource().equals(datasource) - ) - .map(entry -> entry.task) - .collect(Collectors.toMap(Task::getId, Function.identity())); + return activeTasks.values().stream().filter( + entry -> !entry.isComplete + && entry.task.getDataSource().equals(datasource) + ).map( + entry -> entry.task + ).collect( + Collectors.toMap(Task::getId, Function.identity()) + ); } private void validateTaskPayload(Task task) @@ -1078,7 +1065,8 @@ private void validateTaskPayload(Task task) } /** - * Performs a thread-safe update on the task entry for the given task ID. + * Performs a thread-safe update on the task entry for the given task ID, + * and sets the {@link TaskEntry#lastUpdatedTime} to now. */ private void updateTaskEntry(String taskId, Consumer updateOperation) { @@ -1088,6 +1076,9 @@ private void updateTaskEntry(String taskId, Consumer updateOperation) taskId, (id, existingEntry) -> { updateOperation.accept(existingEntry); + if (existingEntry != null) { + existingEntry.lastUpdatedTime = DateTimes.nowUtc(); + } return existingEntry; } ); @@ -1097,6 +1088,11 @@ private void updateTaskEntry(String taskId, Consumer updateOperation) } } + private void removeTaskLock(Task task) + { + taskLockbox.remove(task); + } + /** * Represents an entry in this task queue. */ diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueScaleTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueScaleTest.java index 80026f137269..2307c4f662ca 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueScaleTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueScaleTest.java @@ -143,6 +143,7 @@ public RetType submit(TaskAction taskAction) new NoopTaskContextEnricher() ); + storageCoordinator.start(); taskQueue.start(); closer.register(taskQueue::stop); } @@ -322,7 +323,7 @@ public void shutdown(String taskid, String reason) } if (!existingTask.getResult().isDone()) { exec.schedule(() -> { - existingTask.setResult(TaskStatus.failure("taskId", "stopped")); + existingTask.setResult(TaskStatus.failure(existingTask.getTaskId(), "stopped")); synchronized (knownTasks) { knownTasks.remove(taskid); } diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java index 391e539dbdb4..5858f7274f16 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java @@ -100,7 +100,6 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; public class TaskQueueTest extends IngestionTestBase From 18b62fa33206d1b7a29487054499c669e1f5bb68 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Fri, 28 Mar 2025 10:19:39 +0530 Subject: [PATCH 03/10] Rename giant to startStopLock --- .../druid/indexing/overlord/TaskQueue.java | 42 +++++++++---------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java index 8d351185c120..a8bee7555223 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java @@ -128,7 +128,7 @@ public class TaskQueue * queue management is in progress. {@link #start} and {@link #stop} methods * acquire a WRITE lock whereas all other operations acquire a READ lock. */ - private final ReentrantReadWriteLock giant = new ReentrantReadWriteLock(true); + private final ReentrantReadWriteLock startStopLock = new ReentrantReadWriteLock(true); @SuppressWarnings("MismatchedQueryAndUpdateOfCollection") private final BlockingQueue managementMayBeNecessary = new ArrayBlockingQueue<>(8); @@ -205,7 +205,7 @@ void setActive() @LifecycleStart public void start() { - giant.writeLock().lock(); + startStopLock.writeLock().lock(); try { Preconditions.checkState(!active, "queue must be stopped"); @@ -272,7 +272,7 @@ public void start() } } finally { - giant.writeLock().unlock(); + startStopLock.writeLock().unlock(); } } @@ -282,7 +282,7 @@ public void start() @LifecycleStop public void stop() { - giant.writeLock().lock(); + startStopLock.writeLock().lock(); try { activeTasks.clear(); @@ -292,7 +292,7 @@ public void stop() requestManagement(); } finally { - giant.writeLock().unlock(); + startStopLock.writeLock().unlock(); } } @@ -365,12 +365,12 @@ private void manage() throws InterruptedException @VisibleForTesting void manageQueuedTasks() { - giant.readLock().lock(); + startStopLock.readLock().lock(); try { startPendingTasksOnRunner(); } finally { - giant.readLock().unlock(); + startStopLock.readLock().unlock(); } } @@ -380,7 +380,7 @@ void manageQueuedTasks() * Cleans up tasks present on the {@link #taskRunner} but not present in * {@link #activeTasks} anymore. */ - @GuardedBy("giant") + @GuardedBy("startStopLock") private void startPendingTasksOnRunner() { final Set unknownTaskIds = new HashSet<>(); @@ -410,7 +410,7 @@ private void startPendingTasksOnRunner() } } - @GuardedBy("giant") + @GuardedBy("startStopLock") private void startPendingTaskOnRunner(TaskEntry entry, ListenableFuture runnerTaskFuture) { // Don't do anything with tasks that have recently finished; notifyStatus will handle it. @@ -506,7 +506,7 @@ public boolean add(final Task task) taskContextEnricher.enrichContext(task); - giant.readLock().lock(); + startStopLock.readLock().lock(); try { Preconditions.checkState(active, "Queue is not active!"); @@ -530,11 +530,11 @@ public boolean add(final Task task) return true; } finally { - giant.readLock().unlock(); + startStopLock.readLock().unlock(); } } - @GuardedBy("giant") + @GuardedBy("startStopLock") private void addTaskInternal(final Task task, final DateTime updateTime) { final AtomicBoolean added = new AtomicBoolean(false); @@ -568,7 +568,7 @@ private void addTaskInternal(final Task task, final DateTime updateTime) * the DB would still have the task as RUNNING and {@link #syncFromStorage()} * might add it back to the queue, thus causing a duplicate run of the task. */ - @GuardedBy("giant") + @GuardedBy("startStopLock") private void removeTaskInternal(final String taskId, final DateTime deleteTime) { final AtomicReference removedTask = new AtomicReference<>(); @@ -601,7 +601,7 @@ private void removeTaskInternal(final String taskId, final DateTime deleteTime) public void shutdown(final String taskId, String reasonFormat, Object... args) { Preconditions.checkNotNull(taskId, "taskId"); - giant.readLock().lock(); + startStopLock.readLock().lock(); try { updateTaskEntry( taskId, @@ -614,7 +614,7 @@ public void shutdown(final String taskId, String reasonFormat, Object... args) ); } finally { - giant.readLock().unlock(); + startStopLock.readLock().unlock(); } } @@ -628,7 +628,7 @@ public void shutdown(final String taskId, String reasonFormat, Object... args) public void shutdownWithSuccess(final String taskId, String reasonFormat, Object... args) { Preconditions.checkNotNull(taskId, "taskId"); - giant.readLock().lock(); + startStopLock.readLock().lock(); try { updateTaskEntry( taskId, @@ -636,7 +636,7 @@ public void shutdownWithSuccess(final String taskId, String reasonFormat, Object ); } finally { - giant.readLock().unlock(); + startStopLock.readLock().unlock(); } } @@ -825,7 +825,7 @@ private void handleStatus(final TaskStatus status) */ private void syncFromStorage() { - giant.readLock().lock(); + startStopLock.readLock().lock(); final DateTime syncStartTime = DateTimes.nowUtc(); try { @@ -865,7 +865,7 @@ private void syncFromStorage() throw new RuntimeException(e); } finally { - giant.readLock().unlock(); + startStopLock.readLock().unlock(); } } @@ -1070,7 +1070,7 @@ private void validateTaskPayload(Task task) */ private void updateTaskEntry(String taskId, Consumer updateOperation) { - giant.readLock().lock(); + startStopLock.readLock().lock(); try { activeTasks.compute( taskId, @@ -1084,7 +1084,7 @@ private void updateTaskEntry(String taskId, Consumer updateOperation) ); } finally { - giant.readLock().unlock(); + startStopLock.readLock().unlock(); } } From 536beaa4370ad35a9f633145dd25e8de5804ee7a Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Wed, 14 May 2025 09:03:52 +0530 Subject: [PATCH 04/10] Add CollectionUtils.toMap --- .../druid/indexing/overlord/TaskQueue.java | 13 ++++++------- .../supervisor/SeekableStreamSupervisor.java | 19 +++++++++---------- .../apache/druid/utils/CollectionUtils.java | 10 ++++++++++ .../druid/utils/CollectionUtilsTest.java | 16 ++++++++++++++++ 4 files changed, 41 insertions(+), 17 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java index a8bee7555223..b7a301cacaf6 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java @@ -830,9 +830,10 @@ private void syncFromStorage() try { if (active) { - final Map newTasks = toTaskIDMap(taskStorage.getActiveTasks()); - final int tasksSynced = newTasks.size(); - final Map oldTasks = CollectionUtils.mapValues(activeTasks, entry -> entry.task); + final Map newTasks = + CollectionUtils.toMap(taskStorage.getActiveTasks(), Task::getId, Function.identity()); + final Map oldTasks = + CollectionUtils.mapValues(activeTasks, entry -> entry.task); // Identify the tasks that have been added or removed from the storage final MapDifference mapDifference = Maps.difference(oldTasks, newTasks); @@ -850,10 +851,8 @@ private void syncFromStorage() } log.info( - "Synced %d tasks from storage (%d tasks added, %d tasks removed).", - tasksSynced, - addedTasks.size(), - removedTasks.size() + "Synced [%d] tasks from storage (%d tasks added, %d tasks removed).", + newTasks.size(), addedTasks.size(), removedTasks.size() ); requestManagement(); } else { diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java index 8e0eddf13f9f..5df8be280a4f 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -98,6 +98,7 @@ import org.apache.druid.segment.incremental.ParseExceptionReport; import org.apache.druid.segment.incremental.RowIngestionMetersFactory; import org.apache.druid.segment.indexing.DataSchema; +import org.apache.druid.utils.CollectionUtils; import org.joda.time.DateTime; import org.joda.time.Duration; @@ -4287,24 +4288,22 @@ public void setPartitionIdsForTests( } /** - * Get all active tasks from metadata storage + * Get all active tasks either from memory or metadata store, if not cached. * * @return map from taskId to Task */ private Map getActiveTaskMap() { - final ImmutableMap.Builder activeTaskMap = ImmutableMap.builder(); final Optional taskQueue = taskMaster.getTaskQueue(); - final List tasks; if (taskQueue.isPresent()) { return taskQueue.get().getActiveTasksForDatasource(dataSource); } else { - tasks = taskStorage.getActiveTasksByDatasource(dataSource); - } - for (Task task : tasks) { - activeTaskMap.put(task.getId(), task); + return CollectionUtils.toMap( + taskStorage.getActiveTasksByDatasource(dataSource), + Task::getId, + task -> task + ); } - return activeTaskMap.build(); } /** @@ -4312,7 +4311,7 @@ private Map getActiveTaskMap() * * @return specific instance of Kafka/Kinesis IOConfig */ - protected abstract SeekableStreamIndexTaskIOConfig createTaskIoConfig( + protected abstract SeekableStreamIndexTaskIOConfig createTaskIoConfig( int groupId, Map startPartitions, Map endPartitions, @@ -4335,7 +4334,7 @@ protected abstract List> sequenceOffsets, - SeekableStreamIndexTaskIOConfig taskIoConfig, + SeekableStreamIndexTaskIOConfig taskIoConfig, SeekableStreamIndexTaskTuningConfig taskTuningConfig, RowIngestionMetersFactory rowIngestionMetersFactory ) throws JsonProcessingException; diff --git a/processing/src/main/java/org/apache/druid/utils/CollectionUtils.java b/processing/src/main/java/org/apache/druid/utils/CollectionUtils.java index 1309d85f4e12..46a7ff39dbac 100644 --- a/processing/src/main/java/org/apache/druid/utils/CollectionUtils.java +++ b/processing/src/main/java/org/apache/druid/utils/CollectionUtils.java @@ -39,6 +39,7 @@ import java.util.TreeSet; import java.util.function.Function; import java.util.function.Supplier; +import java.util.stream.Collectors; import java.util.stream.Stream; public final class CollectionUtils @@ -124,6 +125,15 @@ public static Map mapKeys(Map map, Function keyMa return result; } + /** + * Creates an immutable map by mapping each entry in the given collection to + * a key and a value. + */ + public static Map toMap(Collection collection, Function keyMapper, Function valueMapper) + { + return collection.stream().collect(Collectors.toMap(keyMapper, valueMapper)); + } + /** * Returns a LinkedHashMap with an appropriate size based on the callers expectedSize. This methods functionality * mirrors that of com.google.common.collect.Maps#newLinkedHashMapWithExpectedSize in Guava 19+. Thus, this method diff --git a/processing/src/test/java/org/apache/druid/utils/CollectionUtilsTest.java b/processing/src/test/java/org/apache/druid/utils/CollectionUtilsTest.java index d84a56bb613e..ed22b8bcbb0c 100644 --- a/processing/src/test/java/org/apache/druid/utils/CollectionUtilsTest.java +++ b/processing/src/test/java/org/apache/druid/utils/CollectionUtilsTest.java @@ -29,6 +29,9 @@ import org.junit.internal.matchers.ThrowableMessageMatcher; import java.util.Collections; +import java.util.List; +import java.util.Locale; +import java.util.Map; import java.util.Set; public class CollectionUtilsTest @@ -95,4 +98,17 @@ public void testGetOnlyElement_two() ); MatcherAssert.assertThat(e, ThrowableMessageMatcher.hasMessage(CoreMatchers.equalTo("oops"))); } + + @Test + public void test_toMap() + { + Assert.assertEquals( + Map.of("a", "A", "b", "B"), + CollectionUtils.toMap( + List.of("a", "b"), + entry -> entry, + entry -> entry.toUpperCase(Locale.ROOT) + ) + ); + } } From 154b08c4f1d9cf690f2130a46a9109760afc818f Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Thu, 15 May 2025 19:54:34 +0530 Subject: [PATCH 05/10] Add TaskQueueConcurrencyTest --- .../druid/indexing/overlord/TaskQueue.java | 40 +- .../supervisor/SeekableStreamSupervisor.java | 4 +- .../overlord/TaskQueueConcurrencyTest.java | 420 ++++++++++++++++++ .../indexing/overlord/TaskQueueTest.java | 11 +- 4 files changed, 454 insertions(+), 21 deletions(-) create mode 100644 indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java index b7a301cacaf6..123ae87bacd6 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java @@ -538,9 +538,9 @@ public boolean add(final Task task) private void addTaskInternal(final Task task, final DateTime updateTime) { final AtomicBoolean added = new AtomicBoolean(false); - final TaskEntry entry = activeTasks.compute( + final TaskEntry entry = addOrUpdateTaskEntry( task.getId(), - (taskId, prevEntry) -> { + prevEntry -> { if (prevEntry == null) { added.set(true); return new TaskEntry(task); @@ -573,9 +573,9 @@ private void removeTaskInternal(final String taskId, final DateTime deleteTime) { final AtomicReference removedTask = new AtomicReference<>(); - activeTasks.compute( + addOrUpdateTaskEntry( taskId, - (id, prevEntry) -> { + prevEntry -> { // Remove the task only if it doesn't have a more recent update if (prevEntry != null && prevEntry.lastUpdatedTime.isBefore(deleteTime)) { removedTask.set(prevEntry.task); @@ -823,7 +823,8 @@ private void handleStatus(final TaskStatus status) * Resync the contents of this task queue with our storage facility. Useful to make sure our in-memory state * corresponds to the storage facility even if the latter is manually modified. */ - private void syncFromStorage() + @VisibleForTesting + void syncFromStorage() { startStopLock.readLock().lock(); final DateTime syncStartTime = DateTimes.nowUtc(); @@ -1068,18 +1069,29 @@ private void validateTaskPayload(Task task) * and sets the {@link TaskEntry#lastUpdatedTime} to now. */ private void updateTaskEntry(String taskId, Consumer updateOperation) + { + addOrUpdateTaskEntry( + taskId, + existingEntry -> { + updateOperation.accept(existingEntry); + if (existingEntry != null) { + existingEntry.lastUpdatedTime = DateTimes.nowUtc(); + } + return existingEntry; + } + ); + } + + /** + * Performs a thread-safe upsert operation on the task entry for the given task ID. + */ + TaskEntry addOrUpdateTaskEntry(String taskId, Function updateOperation) { startStopLock.readLock().lock(); try { - activeTasks.compute( + return activeTasks.compute( taskId, - (id, existingEntry) -> { - updateOperation.accept(existingEntry); - if (existingEntry != null) { - existingEntry.lastUpdatedTime = DateTimes.nowUtc(); - } - return existingEntry; - } + (id, existingEntry) -> updateOperation.apply(existingEntry) ); } finally { @@ -1095,7 +1107,7 @@ private void removeTaskLock(Task task) /** * Represents an entry in this task queue. */ - private static class TaskEntry + static class TaskEntry { private final Task task; diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java index 8e227e2b441b..7b89f6e9f26f 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -4333,7 +4333,7 @@ private Map getActiveTaskMap() * * @return specific instance of Kafka/Kinesis IOConfig */ - protected abstract SeekableStreamIndexTaskIOConfig createTaskIoConfig( + protected abstract SeekableStreamIndexTaskIOConfig createTaskIoConfig( int groupId, Map startPartitions, Map endPartitions, @@ -4356,7 +4356,7 @@ protected abstract List> sequenceOffsets, - SeekableStreamIndexTaskIOConfig taskIoConfig, + SeekableStreamIndexTaskIOConfig taskIoConfig, SeekableStreamIndexTaskTuningConfig taskTuningConfig, RowIngestionMetersFactory rowIngestionMetersFactory ) throws JsonProcessingException; diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java new file mode 100644 index 000000000000..09aeeadf3c78 --- /dev/null +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java @@ -0,0 +1,420 @@ +/* + * 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.indexing.overlord; + +import com.google.common.base.Optional; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.indexing.common.task.IngestionTestBase; +import org.apache.druid.indexing.common.task.NoopTask; +import org.apache.druid.indexing.common.task.NoopTaskContextEnricher; +import org.apache.druid.indexing.common.task.Task; +import org.apache.druid.indexing.overlord.config.DefaultTaskConfig; +import org.apache.druid.indexing.overlord.config.TaskLockConfig; +import org.apache.druid.indexing.overlord.config.TaskQueueConfig; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.concurrent.Execs; +import org.apache.druid.segment.TestDataSource; +import org.apache.druid.server.metrics.NoopServiceEmitter; +import org.junit.Assert; +import org.junit.Test; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; + +/** + * Unit tests to verify that operations performed on {@link TaskQueue} for a + * given task ID are thread-safe, so that the {@link TaskQueue} always has a + * consistent view of the status of a task. + */ +public class TaskQueueConcurrencyTest extends IngestionTestBase +{ + private static final String THREAD_NAME_FORMAT = "TaskQueueConcurrencyTest-%s"; + + private TaskQueue taskQueue; + + private int threadId; + private Map threadToCriticalUpdate; + + @Override + public void setUpIngestionTestBase() throws IOException + { + super.setUpIngestionTestBase(); + + threadId = 0; + threadToCriticalUpdate = new HashMap<>(); + + taskQueue = new TaskQueue( + new TaskLockConfig(), + new TaskQueueConfig(3, null, null, null, null, null), + new DefaultTaskConfig(), + getTaskStorage(), + new TaskQueueTest.SimpleTaskRunner(NoopServiceEmitter.instance()) + { + @Override + public ListenableFuture run(Task task) + { + return Futures.immediateFuture(TaskStatus.success(task.getId())); + } + }, + createActionClientFactory(), + getLockbox(), + NoopServiceEmitter.instance(), + getObjectMapper(), + new NoopTaskContextEnricher() + ) + { + @Override + TaskEntry addOrUpdateTaskEntry(String taskId, Function updateOperation) + { + // Use a modified update for this thread + final String threadName = Thread.currentThread().getName(); + final CriticalUpdate criticalUpdate = threadToCriticalUpdate.get(threadName); + + return criticalUpdate == null + ? super.addOrUpdateTaskEntry(taskId, updateOperation) + : super.addOrUpdateTaskEntry( + taskId, + existing -> criticalUpdate.apply(existing, updateOperation) + ); + } + }; + + taskQueue.setActive(); + } + + @Test + public void test_add_blocks_syncFromStorage_forSameTaskId() + { + final String taskId = "t2"; + final Task task = createTask(taskId); + + ActionVerifier.verifyThat( + update(() -> taskQueue.add(task)) + .withEndState( + () -> Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(taskId)) + ) + ).blocks( + update(() -> taskQueue.syncFromStorage()) + .withEndState( + () -> Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(taskId)) + ) + ); + } + + @Test + public void test_syncFromStorage_blocks_add_forSameTaskId() + { + final String taskId = "t2"; + final Task task = createTask(taskId); + + // Add the task to queue and storage + taskQueue.add(task); + Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(taskId)); + Assert.assertEquals(Optional.of(task), getTaskStorage().getTask(taskId)); + + // Mark the task as completed and remove it from storage but not queue + taskQueue.shutdown(taskId, "test"); + getTaskStorage().removeTasksOlderThan(DateTimes.nowUtc().plusDays(1).getMillis()); + + ActionVerifier.verifyThat( + update( + () -> taskQueue.syncFromStorage() + ).withEndState( + () -> Assert.assertEquals(Optional.absent(), taskQueue.getActiveTask(taskId)) + ) + ).blocks( + update( + () -> taskQueue.add(task) + ).withEndState( + () -> Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(taskId)) + ) + ); + } + + @Test(timeout = 20_000L) + public void test_manageQueuedTasks_blocks_shutdown_forSameTaskId() + { + final String taskId = "t2"; + final Task task = createTask(taskId); + + taskQueue.add(task); + Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(taskId)); + + ActionVerifier.verifyThat( + update( + () -> taskQueue.manageQueuedTasks() + ).withEndState( + () -> Assert.assertEquals( + Optional.of(TaskStatus.running(taskId)), + taskQueue.getTaskStatus(taskId) + ) + ) + ).blocks( + update( + () -> taskQueue.shutdown(taskId, "test") + ).withEndState( + () -> Assert.assertEquals( + Optional.of(TaskStatus.failure(taskId, "test")), + taskQueue.getTaskStatus(taskId) + ) + ) + ); + } + + @Test(timeout = 20_000L) + public void test_add_blocks_shutdownWithSuccess_forSameTaskId() + { + final String taskId = "t2"; + final Task task = createTask(taskId); + + ActionVerifier.verifyThat( + update( + () -> taskQueue.add(task) + ).withEndState( + () -> Assert.assertEquals( + Optional.of(task), + taskQueue.getActiveTask(taskId) + ) + ) + ).blocks( + update( + () -> taskQueue.shutdownWithSuccess(task.getId(), "test") + ).withEndState( + () -> Assert.assertEquals( + Optional.of(TaskStatus.success(taskId)), + taskQueue.getTaskStatus(taskId) + ) + ) + ); + } + + @Test(timeout = 20_000L) + public void test_add_blocks_shutdown_forSameTaskId() + { + final String taskId = "t1"; + final Task task = createTask(taskId); + + ActionVerifier.verifyThat( + update( + () -> taskQueue.add(task) + ).withEndState( + () -> Assert.assertEquals( + Optional.of(task), + taskQueue.getActiveTask(taskId) + ) + ) + ).blocks( + update( + () -> taskQueue.shutdown(task.getId(), "test") + ).withEndState( + () -> Assert.assertEquals( + Optional.of(TaskStatus.failure(taskId, "test")), + taskQueue.getTaskStatus(taskId) + ) + ) + ); + } + + private UpdateAction update(Action action) + { + return new UpdateAction(action); + } + + private static Task createTask(String id) + { + return new NoopTask(id, id, TestDataSource.WIKI, 0L, 0L, Map.of()); + } + + private static void waitFor(CountDownLatch latch) + { + try { + boolean done = latch.await(5, TimeUnit.SECONDS); + Assert.assertTrue(StringUtils.format("Latch[%s] is still blocked", latch), done); + } + catch (Exception e) { + throw new RuntimeException(e); + } + } + + @FunctionalInterface + private interface Action + { + void perform(); + } + + /** + * Wrapper around a critical update action on a task in the {@link TaskQueue}. + * This class is used to identify the current state of the update and pause or + * resume it to force certain race conditions. + */ + private static class CriticalUpdate + { + final CountDownLatch hasStarted = new NamedLatch("hasStarted"); + final CountDownLatch resume = new NamedLatch("resume"); + + final CountDownLatch isReadyToFinish = new NamedLatch("isReadyToFinish"); + final CountDownLatch finish = new NamedLatch("finish"); + + TaskQueue.TaskEntry apply( + TaskQueue.TaskEntry existingEntry, Function updateOperation + ) + { + hasStarted.countDown(); + waitFor(resume); + + // Actual update operation + final TaskQueue.TaskEntry updatedEntry = updateOperation.apply(existingEntry); + + isReadyToFinish.countDown(); + waitFor(finish); + + return updatedEntry; + } + } + + /** + * CountDownLatch with a name to easily identify bugs. + */ + private static class NamedLatch extends CountDownLatch + { + final String name; + + private NamedLatch(String name) + { + super(1); + this.name = name; + } + + @Override + public String toString() + { + return name; + } + } + + /** + * An update action with a critical part and a verification step. + */ + private class UpdateAction + { + final CountDownLatch finished = new NamedLatch("finished"); + final CriticalUpdate critical = new CriticalUpdate(); + + final Action action; + Action verifyState; + + UpdateAction(Action action) + { + this.action = action; + threadToCriticalUpdate.put(StringUtils.format(THREAD_NAME_FORMAT, threadId++), critical); + } + + UpdateAction withEndState(Action verifyAction) + { + this.verifyState = verifyAction; + return this; + } + + void perform() + { + try { + action.perform(); + finished.countDown(); + } + catch (Throwable t) { + t.printStackTrace(System.out); + } + } + + void waitToFinishAndVerify() + { + waitFor(finished); + verifyState.perform(); + } + } + + /** + * Verifies thread-safety between two actions. + */ + private static class ActionVerifier + { + UpdateAction update1; + + static ActionVerifier verifyThat(UpdateAction update1) + { + ActionVerifier verifier = new ActionVerifier(); + verifier.update1 = update1; + return verifier; + } + + /** + * Verifies that the critical update section of {@code operation1} completely + * blocks the critical update in {@code operation2}. + */ + void blocks(UpdateAction update2) + { + final ExecutorService executor = Execs.multiThreaded(2, "TaskQueueConcurrencyTest-%s"); + + // Start update 1 and wait for it to enter critical section + executor.submit(update1::perform); + waitFor(update1.critical.hasStarted); + + executor.submit(update2::perform); + + // Wait for some time and verify that update 2 critical has not started yet + try { + Thread.sleep(1000); + } + catch (Exception e) { + throw new RuntimeException(e); + } + Assert.assertEquals(1, update2.critical.hasStarted.getCount()); + + update1.critical.resume.countDown(); + + // Wait for update 1 critical to reach finish + // and verify that update 2 critical has not started yet + waitFor(update1.critical.isReadyToFinish); + Assert.assertEquals(1, update2.critical.hasStarted.getCount()); + + // Finish update 1 and verify that update 2 critical has now started + update1.critical.finish.countDown(); + waitFor(update2.critical.hasStarted); + + update1.waitToFinishAndVerify(); + + update2.critical.resume.countDown(); + waitFor(update2.critical.isReadyToFinish); + update2.critical.finish.countDown(); + + update2.waitToFinishAndVerify(); + + executor.shutdownNow(); + } + } +} diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java index 5858f7274f16..d5408135ed37 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java @@ -76,6 +76,7 @@ import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.granularity.Granularities; import org.apache.druid.java.util.common.granularity.Granularity; +import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.java.util.http.client.HttpClient; import org.apache.druid.java.util.metrics.StubServiceEmitter; import org.apache.druid.metadata.DefaultPasswordProvider; @@ -131,7 +132,7 @@ public Map getContext() } }, getTaskStorage(), - new SimpleTaskRunner(), + new SimpleTaskRunner(serviceEmitter), actionClientFactory, getLockbox(), serviceEmitter, @@ -232,7 +233,7 @@ public Map getContext() } }, getTaskStorage(), - new SimpleTaskRunner(), + new SimpleTaskRunner(serviceEmitter), actionClientFactory, getLockbox(), serviceEmitter, @@ -701,14 +702,14 @@ public boolean isDone() } } - private class SimpleTaskRunner extends SingleTaskBackgroundRunner + static class SimpleTaskRunner extends SingleTaskBackgroundRunner { - SimpleTaskRunner() + SimpleTaskRunner(ServiceEmitter emitter) { super( EasyMock.createMock(TaskToolboxFactory.class), null, - serviceEmitter, + emitter, new DruidNode("overlord", "localhost", false, 8091, null, true, false), null ); From c58ea8e57a88bf8a2124696e67ddb8eb69ad517e Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Thu, 15 May 2025 23:31:51 +0530 Subject: [PATCH 06/10] Fix checkstyle --- .../overlord/TaskQueueConcurrencyTest.java | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java index 09aeeadf3c78..1765bc51da49 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java @@ -242,7 +242,10 @@ public void test_add_blocks_shutdown_forSameTaskId() private UpdateAction update(Action action) { - return new UpdateAction(action); + final UpdateAction updateAction = new UpdateAction(action); + threadToCriticalUpdate.put(StringUtils.format(THREAD_NAME_FORMAT, threadId++), updateAction.critical); + + return updateAction; } private static Task createTask(String id) @@ -281,7 +284,8 @@ private static class CriticalUpdate final CountDownLatch finish = new NamedLatch("finish"); TaskQueue.TaskEntry apply( - TaskQueue.TaskEntry existingEntry, Function updateOperation + TaskQueue.TaskEntry existingEntry, + Function updateOperation ) { hasStarted.countDown(); @@ -320,23 +324,22 @@ public String toString() /** * An update action with a critical part and a verification step. */ - private class UpdateAction + private static class UpdateAction { final CountDownLatch finished = new NamedLatch("finished"); final CriticalUpdate critical = new CriticalUpdate(); final Action action; - Action verifyState; + Action verifyAction; UpdateAction(Action action) { this.action = action; - threadToCriticalUpdate.put(StringUtils.format(THREAD_NAME_FORMAT, threadId++), critical); } UpdateAction withEndState(Action verifyAction) { - this.verifyState = verifyAction; + this.verifyAction = verifyAction; return this; } @@ -354,7 +357,7 @@ void perform() void waitToFinishAndVerify() { waitFor(finished); - verifyState.perform(); + verifyAction.perform(); } } From 1391166c6deaadf88840df9c0eba28014cc461f6 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Fri, 16 May 2025 12:29:27 +0530 Subject: [PATCH 07/10] Simplify tests --- .../overlord/TaskQueueConcurrencyTest.java | 55 ++++++++++--------- 1 file changed, 30 insertions(+), 25 deletions(-) diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java index 1765bc51da49..dc5e2d906b7b 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java @@ -31,6 +31,7 @@ import org.apache.druid.indexing.overlord.config.TaskLockConfig; import org.apache.druid.indexing.overlord.config.TaskQueueConfig; import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.concurrent.Execs; import org.apache.druid.segment.TestDataSource; @@ -45,6 +46,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; import java.util.function.Function; +import java.util.function.Supplier; /** * Unit tests to verify that operations performed on {@link TaskQueue} for a @@ -99,7 +101,7 @@ TaskEntry addOrUpdateTaskEntry(String taskId, Function upd ? super.addOrUpdateTaskEntry(taskId, updateOperation) : super.addOrUpdateTaskEntry( taskId, - existing -> criticalUpdate.apply(existing, updateOperation) + existing -> criticalUpdate.perform(() -> updateOperation.apply(existing)) ); } }; @@ -114,15 +116,17 @@ public void test_add_blocks_syncFromStorage_forSameTaskId() final Task task = createTask(taskId); ActionVerifier.verifyThat( - update(() -> taskQueue.add(task)) - .withEndState( - () -> Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(taskId)) - ) + update( + () -> taskQueue.add(task) + ).withEndState( + () -> Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(taskId)) + ) ).blocks( - update(() -> taskQueue.syncFromStorage()) - .withEndState( - () -> Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(taskId)) - ) + update( + () -> taskQueue.syncFromStorage() + ).withEndState( + () -> Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(taskId)) + ) ); } @@ -277,27 +281,28 @@ private interface Action */ private static class CriticalUpdate { - final CountDownLatch hasStarted = new NamedLatch("hasStarted"); - final CountDownLatch resume = new NamedLatch("resume"); + final CountDownLatch isReadyToStart = new NamedLatch("hasStarted"); + final CountDownLatch start = new NamedLatch("resume"); final CountDownLatch isReadyToFinish = new NamedLatch("isReadyToFinish"); final CountDownLatch finish = new NamedLatch("finish"); - TaskQueue.TaskEntry apply( - TaskQueue.TaskEntry existingEntry, - Function updateOperation - ) + synchronized V perform(Supplier updateComputation) { - hasStarted.countDown(); - waitFor(resume); + if (isReadyToStart.getCount() == 0) { + throw new ISE("Update has already run on another thread"); + } + + isReadyToStart.countDown(); + waitFor(start); // Actual update operation - final TaskQueue.TaskEntry updatedEntry = updateOperation.apply(existingEntry); + final V updatedValue = updateComputation.get(); isReadyToFinish.countDown(); waitFor(finish); - return updatedEntry; + return updatedValue; } } @@ -385,7 +390,7 @@ void blocks(UpdateAction update2) // Start update 1 and wait for it to enter critical section executor.submit(update1::perform); - waitFor(update1.critical.hasStarted); + waitFor(update1.critical.isReadyToStart); executor.submit(update2::perform); @@ -396,22 +401,22 @@ void blocks(UpdateAction update2) catch (Exception e) { throw new RuntimeException(e); } - Assert.assertEquals(1, update2.critical.hasStarted.getCount()); + Assert.assertEquals(1, update2.critical.isReadyToStart.getCount()); - update1.critical.resume.countDown(); + update1.critical.start.countDown(); // Wait for update 1 critical to reach finish // and verify that update 2 critical has not started yet waitFor(update1.critical.isReadyToFinish); - Assert.assertEquals(1, update2.critical.hasStarted.getCount()); + Assert.assertEquals(1, update2.critical.isReadyToStart.getCount()); // Finish update 1 and verify that update 2 critical has now started update1.critical.finish.countDown(); - waitFor(update2.critical.hasStarted); + waitFor(update2.critical.isReadyToStart); update1.waitToFinishAndVerify(); - update2.critical.resume.countDown(); + update2.critical.start.countDown(); waitFor(update2.critical.isReadyToFinish); update2.critical.finish.countDown(); From 9aed4e8d5d58b75e46235ea1e6f660ad2a14d659 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Fri, 16 May 2025 14:33:31 +0530 Subject: [PATCH 08/10] Minor cleanup --- .../overlord/TaskQueueConcurrencyTest.java | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java index dc5e2d906b7b..d68273087c15 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java @@ -45,6 +45,7 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; +import java.util.function.BiConsumer; import java.util.function.Function; import java.util.function.Supplier; @@ -55,11 +56,8 @@ */ public class TaskQueueConcurrencyTest extends IngestionTestBase { - private static final String THREAD_NAME_FORMAT = "TaskQueueConcurrencyTest-%s"; - private TaskQueue taskQueue; - private int threadId; private Map threadToCriticalUpdate; @Override @@ -67,7 +65,6 @@ public void setUpIngestionTestBase() throws IOException { super.setUpIngestionTestBase(); - threadId = 0; threadToCriticalUpdate = new HashMap<>(); taskQueue = new TaskQueue( @@ -93,7 +90,7 @@ public ListenableFuture run(Task task) @Override TaskEntry addOrUpdateTaskEntry(String taskId, Function updateOperation) { - // Use a modified update for this thread + // Use a modified update so that we can track and control the progress final String threadName = Thread.currentThread().getName(); final CriticalUpdate criticalUpdate = threadToCriticalUpdate.get(threadName); @@ -246,10 +243,7 @@ public void test_add_blocks_shutdown_forSameTaskId() private UpdateAction update(Action action) { - final UpdateAction updateAction = new UpdateAction(action); - threadToCriticalUpdate.put(StringUtils.format(THREAD_NAME_FORMAT, threadId++), updateAction.critical); - - return updateAction; + return new UpdateAction(action, (k, v) -> threadToCriticalUpdate.put(k, v)); } private static Task createTask(String id) @@ -276,8 +270,8 @@ private interface Action /** * Wrapper around a critical update action on a task in the {@link TaskQueue}. - * This class is used to identify the current state of the update and pause or - * resume it to force certain race conditions. + * This class contains latches to track and control the progress of the update + * and verify behaviour in race conditions. */ private static class CriticalUpdate { @@ -334,12 +328,15 @@ private static class UpdateAction final CountDownLatch finished = new NamedLatch("finished"); final CriticalUpdate critical = new CriticalUpdate(); + final BiConsumer mapThreadNameToUpdate; + final Action action; Action verifyAction; - UpdateAction(Action action) + UpdateAction(Action action, BiConsumer mapThreadNameToUpdate) { this.action = action; + this.mapThreadNameToUpdate = mapThreadNameToUpdate; } UpdateAction withEndState(Action verifyAction) @@ -350,6 +347,9 @@ UpdateAction withEndState(Action verifyAction) void perform() { + // Map the thread name to the CriticalUpdate so that it can be invoked later in this thread + mapThreadNameToUpdate.accept(Thread.currentThread().getName(), critical); + try { action.perform(); finished.countDown(); From d8158d10b8a472588a8e8a18cc88d9791c669c66 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Sun, 18 May 2025 09:53:46 +0530 Subject: [PATCH 09/10] Add more tests with start/stop --- .../druid/indexing/overlord/TaskQueue.java | 13 +- .../overlord/TaskQueueConcurrencyTest.java | 111 +++++++++++++++--- .../indexing/overlord/TaskQueueTest.java | 8 +- 3 files changed, 105 insertions(+), 27 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java index 123ae87bacd6..5a26cb773393 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java @@ -194,9 +194,14 @@ public TaskQueue( } @VisibleForTesting - void setActive() + void setActive(boolean active) { - this.active = true; + startStopLock.writeLock().lock(); + try { + this.active = active; + } finally { + startStopLock.writeLock().unlock(); + } } /** @@ -209,7 +214,7 @@ public void start() try { Preconditions.checkState(!active, "queue must be stopped"); - active = true; + setActive(true); syncFromStorage(); // Mark these tasks as failed as they could not reacuire the lock // Clean up needs to happen after tasks have been synced from storage @@ -285,8 +290,8 @@ public void stop() startStopLock.writeLock().lock(); try { + setActive(false); activeTasks.clear(); - active = false; managerExec.shutdownNow(); storageSyncExec.shutdownNow(); requestManagement(); diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java index d68273087c15..bcd75090f007 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueConcurrencyTest.java @@ -41,11 +41,11 @@ import java.io.IOException; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; -import java.util.function.BiConsumer; import java.util.function.Function; import java.util.function.Supplier; @@ -58,14 +58,14 @@ public class TaskQueueConcurrencyTest extends IngestionTestBase { private TaskQueue taskQueue; - private Map threadToCriticalUpdate; + private Map threadToUpdateAction; @Override public void setUpIngestionTestBase() throws IOException { super.setUpIngestionTestBase(); - threadToCriticalUpdate = new HashMap<>(); + threadToUpdateAction = new HashMap<>(); taskQueue = new TaskQueue( new TaskLockConfig(), @@ -90,25 +90,87 @@ public ListenableFuture run(Task task) @Override TaskEntry addOrUpdateTaskEntry(String taskId, Function updateOperation) { - // Use a modified update so that we can track and control the progress + // Override this critical action so that we can track and control progress final String threadName = Thread.currentThread().getName(); - final CriticalUpdate criticalUpdate = threadToCriticalUpdate.get(threadName); + final UpdateAction updateAction = threadToUpdateAction.remove(threadName); - return criticalUpdate == null + return updateAction == null ? super.addOrUpdateTaskEntry(taskId, updateOperation) : super.addOrUpdateTaskEntry( taskId, - existing -> criticalUpdate.perform(() -> updateOperation.apply(existing)) + existing -> updateAction.critical.perform(() -> updateOperation.apply(existing)) ); } + + @Override + void setActive(boolean active) + { + // Override this critical action so that we can track and control progress + final String threadName = Thread.currentThread().getName(); + final UpdateAction updateAction = threadToUpdateAction.remove(threadName); + + if (updateAction == null) { + super.setActive(active); + } else { + updateAction.critical.perform(() -> { + super.setActive(active); + return 0; + }); + } + } }; + } - taskQueue.setActive(); + @Test + public void test_start_blocks_add_forAnyTaskId() + { + // Add task1 to storage and mark it as running + final Task task1 = createTask("t1"); + getTaskStorage().insert(task1, TaskStatus.running(task1.getId())); + + final Task task2 = createTask("t2"); + + ActionVerifier.verifyThat( + update( + () -> taskQueue.start() + ).withEndState( + () -> Assert.assertEquals(List.of(task1), taskQueue.getTasks()) + ) + ).blocks( + update( + () -> taskQueue.add(task2) + ).withEndState( + () -> Assert.assertEquals(List.of(task1, task2), taskQueue.getTasks()) + ) + ); + } + + @Test + public void test_add_blocks_stop() + { + taskQueue.setActive(true); + + final Task task = createTask("t1"); + ActionVerifier.verifyThat( + update( + () -> taskQueue.add(task) + ).withEndState( + () -> Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(task.getId())) + ) + ).blocks( + update( + () -> taskQueue.stop() + ).withEndState( + () -> Assert.assertEquals(Optional.absent(), taskQueue.getActiveTask(task.getId())) + ) + ); } @Test public void test_add_blocks_syncFromStorage_forSameTaskId() { + taskQueue.setActive(true); + final String taskId = "t2"; final Task task = createTask(taskId); @@ -134,6 +196,7 @@ public void test_syncFromStorage_blocks_add_forSameTaskId() final Task task = createTask(taskId); // Add the task to queue and storage + taskQueue.setActive(true); taskQueue.add(task); Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(taskId)); Assert.assertEquals(Optional.of(task), getTaskStorage().getTask(taskId)); @@ -163,6 +226,7 @@ public void test_manageQueuedTasks_blocks_shutdown_forSameTaskId() final String taskId = "t2"; final Task task = createTask(taskId); + taskQueue.setActive(true); taskQueue.add(task); Assert.assertEquals(Optional.of(task), taskQueue.getActiveTask(taskId)); @@ -190,6 +254,8 @@ public void test_manageQueuedTasks_blocks_shutdown_forSameTaskId() @Test(timeout = 20_000L) public void test_add_blocks_shutdownWithSuccess_forSameTaskId() { + taskQueue.setActive(true); + final String taskId = "t2"; final Task task = createTask(taskId); @@ -217,6 +283,8 @@ public void test_add_blocks_shutdownWithSuccess_forSameTaskId() @Test(timeout = 20_000L) public void test_add_blocks_shutdown_forSameTaskId() { + taskQueue.setActive(true); + final String taskId = "t1"; final Task task = createTask(taskId); @@ -243,7 +311,7 @@ public void test_add_blocks_shutdown_forSameTaskId() private UpdateAction update(Action action) { - return new UpdateAction(action, (k, v) -> threadToCriticalUpdate.put(k, v)); + return new UpdateAction(action, threadToUpdateAction::put); } private static Task createTask(String id) @@ -268,15 +336,21 @@ private interface Action void perform(); } + @FunctionalInterface + private interface UpdateStartNotifier + { + void onUpdateStart(String threadName, UpdateAction action); + } + /** - * Wrapper around a critical update action on a task in the {@link TaskQueue}. + * Wrapper around the critical part of an update action. * This class contains latches to track and control the progress of the update * and verify behaviour in race conditions. */ private static class CriticalUpdate { - final CountDownLatch isReadyToStart = new NamedLatch("hasStarted"); - final CountDownLatch start = new NamedLatch("resume"); + final CountDownLatch isReadyToStart = new NamedLatch("isReadyToStart"); + final CountDownLatch start = new NamedLatch("start"); final CountDownLatch isReadyToFinish = new NamedLatch("isReadyToFinish"); final CountDownLatch finish = new NamedLatch("finish"); @@ -284,7 +358,7 @@ private static class CriticalUpdate synchronized V perform(Supplier updateComputation) { if (isReadyToStart.getCount() == 0) { - throw new ISE("Update has already run on another thread"); + throw new ISE("Critical update has already run on another thread"); } isReadyToStart.countDown(); @@ -328,15 +402,15 @@ private static class UpdateAction final CountDownLatch finished = new NamedLatch("finished"); final CriticalUpdate critical = new CriticalUpdate(); - final BiConsumer mapThreadNameToUpdate; - final Action action; + final UpdateStartNotifier startNotifier; + Action verifyAction; - UpdateAction(Action action, BiConsumer mapThreadNameToUpdate) + UpdateAction(Action action, UpdateStartNotifier startNotifier) { this.action = action; - this.mapThreadNameToUpdate = mapThreadNameToUpdate; + this.startNotifier = startNotifier; } UpdateAction withEndState(Action verifyAction) @@ -347,8 +421,7 @@ UpdateAction withEndState(Action verifyAction) void perform() { - // Map the thread name to the CriticalUpdate so that it can be invoked later in this thread - mapThreadNameToUpdate.accept(Thread.currentThread().getName(), critical); + startNotifier.onUpdateStart(Thread.currentThread().getName(), this); try { action.perform(); diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java index d5408135ed37..040d54fea622 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskQueueTest.java @@ -139,7 +139,7 @@ public Map getContext() getObjectMapper(), new NoopTaskContextEnricher() ); - taskQueue.setActive(); + taskQueue.setActive(true); } @Test @@ -240,7 +240,7 @@ public Map getContext() getObjectMapper(), new NoopTaskContextEnricher() ); - maxPayloadTaskQueue.setActive(); + maxPayloadTaskQueue.setActive(true); // 1 MB is not too large char[] context = new char[1024 * 1024]; @@ -416,7 +416,7 @@ public void testKilledTasksEmitRuntimeMetricWithHttpRemote() throws InterruptedE getObjectMapper(), new NoopTaskContextEnricher() ); - taskQueue.setActive(); + taskQueue.setActive(true); final Task task = new TestTask("t1", Intervals.of("2021-01-01/P1D")); taskQueue.add(task); taskQueue.manageQueuedTasks(); @@ -504,7 +504,7 @@ public void testGetTaskStatus() getObjectMapper(), new NoopTaskContextEnricher() ); - taskQueue.setActive(); + taskQueue.setActive(true); Assert.assertEquals(TaskStatus.running(newTask), taskQueue.getTaskStatus(newTask).get()); Assert.assertEquals(TaskStatus.running(waitingTask), taskQueue.getTaskStatus(waitingTask).get()); From 6c4dcf4478535e9161f85fe2bbe9cbc77d0988b2 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Sun, 18 May 2025 10:30:20 +0530 Subject: [PATCH 10/10] Fix checkstyle --- .../java/org/apache/druid/indexing/overlord/TaskQueue.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java index 5a26cb773393..260b34330d09 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java @@ -199,7 +199,8 @@ void setActive(boolean active) startStopLock.writeLock().lock(); try { this.active = active; - } finally { + } + finally { startStopLock.writeLock().unlock(); } }