Add delay before the peon drops the segments after publishing them (#…

…15373)

Currently in the realtime ingestion (Kafka/Kinesis) case, after publishing the segments, upon acknowledgement from the coordinator that the segments are already placed in some historicals, the peon would unannounce the segments (basically saying the segments are not in this peon anymore to the whole cluster) and drop the segments from cache and sink timeline in one shot.

The in transit queries from the brokers that still thinks the segments are in the peon can get a NullPointer exception when the peon is unsetting the hydrants in the sinks.

The fix would let the peon to wait for a configurable delay period before dropping segments, remove segments from cache etc after the peon unannounce the segments.

This delayed approach is similar to how the historicals handle segments moving out.

indexing-service/src/main/java/org/apache/druid/indexing/common/TaskToolbox.java

@@ -55,6 +55,7 @@
 import org.apache.druid.segment.loading.DataSegmentMover;
 import org.apache.druid.segment.loading.DataSegmentPusher;
 import org.apache.druid.segment.loading.SegmentCacheManager;
+import org.apache.druid.segment.loading.SegmentLoaderConfig;
 import org.apache.druid.segment.realtime.appenderator.AppenderatorsManager;
 import org.apache.druid.segment.realtime.appenderator.UnifiedIndexerAppenderatorsManager;
 import org.apache.druid.segment.realtime.firehose.ChatHandlerProvider;
@@ -78,6 +79,7 @@
  */
 public class TaskToolbox
 {
+  private final SegmentLoaderConfig segmentLoaderConfig;
   private final TaskConfig config;
   private final DruidNode taskExecutorNode;
   private final TaskActionClient taskActionClient;
@@ -130,6 +132,7 @@ public class TaskToolbox
   private final String attemptId;
 
   public TaskToolbox(
+      SegmentLoaderConfig segmentLoaderConfig,
       TaskConfig config,
       DruidNode taskExecutorNode,
       TaskActionClient taskActionClient,
@@ -171,6 +174,7 @@ public TaskToolbox(
       String attemptId
   )
   {
+    this.segmentLoaderConfig = segmentLoaderConfig;
     this.config = config;
     this.taskExecutorNode = taskExecutorNode;
     this.taskActionClient = taskActionClient;
@@ -213,6 +217,11 @@ public TaskToolbox(
     this.attemptId = attemptId;
   }
 
+  public SegmentLoaderConfig getSegmentLoaderConfig()
+  {
+    return segmentLoaderConfig;
+  }
+
   public TaskConfig getConfig()
   {
     return config;
@@ -504,6 +513,7 @@ public static RuntimeInfo createAdjustedRuntimeInfo(
 
   public static class Builder
   {
+    private SegmentLoaderConfig segmentLoaderConfig;
     private TaskConfig config;
     private DruidNode taskExecutorNode;
     private TaskActionClient taskActionClient;
@@ -550,6 +560,7 @@ public Builder()
 
     public Builder(TaskToolbox other)
     {
+      this.segmentLoaderConfig = other.segmentLoaderConfig;
       this.config = other.config;
       this.taskExecutorNode = other.taskExecutorNode;
       this.taskActionClient = other.taskActionClient;
@@ -589,6 +600,12 @@ public Builder(TaskToolbox other)
       this.shuffleClient = other.shuffleClient;
     }
 
+    public Builder config(final SegmentLoaderConfig segmentLoaderConfig)
+    {
+      this.segmentLoaderConfig = segmentLoaderConfig;
+      return this;
+    }
+
     public Builder config(final TaskConfig config)
     {
       this.config = config;
@@ -826,6 +843,7 @@ public Builder attemptId(final String attemptId)
     public TaskToolbox build()
     {
       return new TaskToolbox(
+          segmentLoaderConfig,
           config,
           taskExecutorNode,
           taskActionClient,

indexing-service/src/main/java/org/apache/druid/indexing/common/TaskToolboxFactory.java

@@ -56,6 +56,7 @@
 import org.apache.druid.segment.loading.DataSegmentKiller;
 import org.apache.druid.segment.loading.DataSegmentMover;
 import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.SegmentLoaderConfig;
 import org.apache.druid.segment.realtime.appenderator.AppenderatorsManager;
 import org.apache.druid.segment.realtime.firehose.ChatHandlerProvider;
 import org.apache.druid.server.DruidNode;
@@ -72,6 +73,7 @@
  */
 public class TaskToolboxFactory
 {
+  private final SegmentLoaderConfig segmentLoaderConfig;
   private final TaskConfig config;
   private final DruidNode taskExecutorNode;
   private final TaskActionClientFactory taskActionClientFactory;
@@ -115,6 +117,7 @@ public class TaskToolboxFactory
 
   @Inject
   public TaskToolboxFactory(
+      SegmentLoaderConfig segmentLoadConfig,
       TaskConfig config,
       @Parent DruidNode taskExecutorNode,
       TaskActionClientFactory taskActionClientFactory,
@@ -155,6 +158,7 @@ public TaskToolboxFactory(
       @AttemptId String attemptId
   )
   {
+    this.segmentLoaderConfig = segmentLoadConfig;
     this.config = config;
     this.taskExecutorNode = taskExecutorNode;
     this.taskActionClientFactory = taskActionClientFactory;
@@ -210,6 +214,7 @@ public TaskToolbox build(TaskConfig config, Task task)
     final File taskWorkDir = config.getTaskWorkDir(task.getId());
     return new TaskToolbox.Builder()
         .config(config)
+        .config(segmentLoaderConfig)
         .taskExecutorNode(taskExecutorNode)
         .taskActionClient(taskActionClientFactory.create(task))
         .emitter(emitter)

indexing-service/src/main/java/org/apache/druid/indexing/common/task/AppenderatorDriverRealtimeIndexTask.java

@@ -775,6 +775,7 @@ private Appenderator newAppenderator(
   )
   {
     return toolbox.getAppenderatorsManager().createRealtimeAppenderatorForTask(
+        null,
         getId(),
         dataSchema,
         tuningConfig.withBasePersistDirectory(toolbox.getPersistDir()),

indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTask.java

@@ -187,6 +187,7 @@ public Appenderator newAppenderator(
   )
   {
     return toolbox.getAppenderatorsManager().createRealtimeAppenderatorForTask(
+        toolbox.getSegmentLoaderConfig(),
         getId(),
         dataSchema,
         tuningConfig.withBasePersistDirectory(toolbox.getPersistDir()),

indexing-service/src/test/java/org/apache/druid/indexing/common/TaskToolboxTest.java

@@ -48,6 +48,7 @@
 import org.apache.druid.segment.loading.DataSegmentKiller;
 import org.apache.druid.segment.loading.DataSegmentMover;
 import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.SegmentLoaderConfig;
 import org.apache.druid.segment.loading.SegmentLocalCacheManager;
 import org.apache.druid.segment.realtime.appenderator.AppenderatorsManager;
 import org.apache.druid.segment.realtime.appenderator.UnifiedIndexerAppenderatorsManager;
@@ -94,6 +95,7 @@ public class TaskToolboxTest
   private IndexIO mockIndexIO = EasyMock.createMock(IndexIO.class);
   private Cache mockCache = EasyMock.createMock(Cache.class);
   private CacheConfig mockCacheConfig = EasyMock.createMock(CacheConfig.class);
+  private SegmentLoaderConfig segmentLoaderConfig = EasyMock.createMock(SegmentLoaderConfig.class);
 
   @Rule
   public TemporaryFolder temporaryFolder = new TemporaryFolder();
@@ -115,6 +117,7 @@ public void setUp() throws IOException
         .build();
 
     taskToolbox = new TaskToolboxFactory(
+        segmentLoaderConfig,
         taskConfig,
         new DruidNode("druid/middlemanager", "localhost", false, 8091, null, true, false),
         mockTaskActionClientFactory,
@@ -162,6 +165,12 @@ public void testGetDataSegmentArchiver()
     Assert.assertEquals(mockDataSegmentArchiver, taskToolbox.build(task).getDataSegmentArchiver());
   }
 
+  @Test
+  public void testGetSegmentLoaderConfig()
+  {
+    Assert.assertEquals(segmentLoaderConfig, taskToolbox.build(task).getSegmentLoaderConfig());
+  }
+
   @Test
   public void testGetSegmentAnnouncer()
   {

indexing-service/src/test/java/org/apache/druid/indexing/common/task/AppenderatorDriverRealtimeIndexTaskTest.java

@@ -1606,6 +1606,7 @@ public void close()
     };
     final TestUtils testUtils = new TestUtils();
     taskToolboxFactory = new TaskToolboxFactory(
+        null,
         taskConfig,
         new DruidNode("druid/middlemanager", "localhost", false, 8091, null, true, false),
         taskActionClientFactory,

indexing-service/src/test/java/org/apache/druid/indexing/common/task/RealtimeIndexTaskTest.java

@@ -979,6 +979,7 @@ public void close()
     };
     final TestUtils testUtils = new TestUtils();
     final TaskToolboxFactory toolboxFactory = new TaskToolboxFactory(
+        null,
         taskConfig,
         null, // taskExecutorNode
         taskActionClientFactory,

indexing-service/src/test/java/org/apache/druid/indexing/common/task/TestAppenderatorsManager.java

@@ -36,6 +36,7 @@
 import org.apache.druid.segment.indexing.DataSchema;
 import org.apache.druid.segment.join.JoinableFactory;
 import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.SegmentLoaderConfig;
 import org.apache.druid.segment.realtime.FireDepartmentMetrics;
 import org.apache.druid.segment.realtime.appenderator.Appenderator;
 import org.apache.druid.segment.realtime.appenderator.AppenderatorConfig;
@@ -50,6 +51,7 @@ public class TestAppenderatorsManager implements AppenderatorsManager
 
   @Override
   public Appenderator createRealtimeAppenderatorForTask(
+      SegmentLoaderConfig segmentLoaderConfig,
       String taskId,
       DataSchema schema,
       AppenderatorConfig config,
@@ -72,6 +74,7 @@ public Appenderator createRealtimeAppenderatorForTask(
   )
   {
     realtimeAppenderator = Appenderators.createRealtime(
+        segmentLoaderConfig,
         taskId,
         schema,
         config,

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

@@ -97,6 +97,7 @@ public void setup() throws IOException
     final ServiceEmitter emitter = new NoopServiceEmitter();
     EmittingLogger.registerEmitter(emitter);
     final TaskToolboxFactory toolboxFactory = new TaskToolboxFactory(
+        null,
         taskConfig,
         null,
         EasyMock.createMock(TaskActionClientFactory.class),

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

@@ -619,6 +619,7 @@ private TaskToolboxFactory setUpTaskToolboxFactory(
         .build();
 
     return new TaskToolboxFactory(
+        null,
         taskConfig,
         new DruidNode("druid/middlemanager", "localhost", false, 8091, null, true, false),
         tac,

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

@@ -77,6 +77,7 @@ public TestTaskToolboxFactory(
   )
   {
     super(
+        null,
         bob.config,
         bob.taskExecutorNode,
         bob.taskActionClientFactory,

indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskTestBase.java

@@ -662,6 +662,7 @@ public void close()
     final DataSegmentPusher dataSegmentPusher = new LocalDataSegmentPusher(dataSegmentPusherConfig);
 
     toolboxFactory = new TaskToolboxFactory(
+        null,
         taskConfig,
         null, // taskExecutorNode
         taskActionClientFactory,

indexing-service/src/test/java/org/apache/druid/indexing/worker/WorkerTaskManagerTest.java

@@ -128,6 +128,7 @@ private WorkerTaskManager createWorkerTaskManager()
         jsonMapper,
         new TestTaskRunner(
             new TaskToolboxFactory(
+                null,
                 taskConfig,
                 null,
                 taskActionClientFactory,

indexing-service/src/test/java/org/apache/druid/indexing/worker/WorkerTaskMonitorTest.java

@@ -170,6 +170,7 @@ private WorkerTaskMonitor createTaskMonitor()
         jsonMapper,
         new SingleTaskBackgroundRunner(
             new TaskToolboxFactory(
+                null,
                 taskConfig,
                 null,
                 taskActionClientFactory,

server/src/main/java/org/apache/druid/segment/realtime/appenderator/Appenderators.java

@@ -35,6 +35,7 @@
 import org.apache.druid.segment.join.JoinableFactory;
 import org.apache.druid.segment.join.JoinableFactoryWrapper;
 import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.SegmentLoaderConfig;
 import org.apache.druid.segment.realtime.FireDepartmentMetrics;
 import org.apache.druid.server.coordination.DataSegmentAnnouncer;
 import org.apache.druid.server.coordination.NoopDataSegmentAnnouncer;
@@ -43,6 +44,7 @@
 public class Appenderators
 {
   public static Appenderator createRealtime(
+      SegmentLoaderConfig segmentLoaderConfig,
       String id,
       DataSchema schema,
       AppenderatorConfig config,
@@ -65,6 +67,7 @@ public static Appenderator createRealtime(
   )
   {
     return new StreamAppenderator(
+        segmentLoaderConfig,
         id,
         schema,
         config,

server/src/main/java/org/apache/druid/segment/realtime/appenderator/AppenderatorsManager.java

@@ -36,6 +36,7 @@
 import org.apache.druid.segment.indexing.DataSchema;
 import org.apache.druid.segment.join.JoinableFactory;
 import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.SegmentLoaderConfig;
 import org.apache.druid.segment.realtime.FireDepartmentMetrics;
 import org.apache.druid.server.coordination.DataSegmentAnnouncer;
 import org.joda.time.Interval;
@@ -64,6 +65,7 @@ public interface AppenderatorsManager
    * used for query processing.
    */
   Appenderator createRealtimeAppenderatorForTask(
+      SegmentLoaderConfig segmentLoaderConfig,
       String taskId,
       DataSchema schema,
       AppenderatorConfig config,

server/src/main/java/org/apache/druid/segment/realtime/appenderator/DefaultRealtimeAppenderatorFactory.java

@@ -96,6 +96,7 @@ public Appenderator build(
   {
     final RowIngestionMeters rowIngestionMeters = new NoopRowIngestionMeters();
     return Appenderators.createRealtime(
+        null,
         schema.getDataSource(),
         schema,
         config.withBasePersistDirectory(

server/src/main/java/org/apache/druid/segment/realtime/appenderator/DummyForInjectionAppenderatorsManager.java

@@ -37,6 +37,7 @@
 import org.apache.druid.segment.indexing.DataSchema;
 import org.apache.druid.segment.join.JoinableFactory;
 import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.SegmentLoaderConfig;
 import org.apache.druid.segment.realtime.FireDepartmentMetrics;
 import org.apache.druid.server.coordination.DataSegmentAnnouncer;
 import org.joda.time.Interval;
@@ -55,6 +56,7 @@ public class DummyForInjectionAppenderatorsManager implements AppenderatorsManag
 
   @Override
   public Appenderator createRealtimeAppenderatorForTask(
+      SegmentLoaderConfig segmentLoaderConfig,
       String taskId,
       DataSchema schema,
       AppenderatorConfig config,

server/src/main/java/org/apache/druid/segment/realtime/appenderator/PeonAppenderatorsManager.java

@@ -37,6 +37,7 @@
 import org.apache.druid.segment.indexing.DataSchema;
 import org.apache.druid.segment.join.JoinableFactory;
 import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.SegmentLoaderConfig;
 import org.apache.druid.segment.realtime.FireDepartmentMetrics;
 import org.apache.druid.server.coordination.DataSegmentAnnouncer;
 import org.joda.time.Interval;
@@ -61,6 +62,7 @@ public class PeonAppenderatorsManager implements AppenderatorsManager
 
   @Override
   public Appenderator createRealtimeAppenderatorForTask(
+      SegmentLoaderConfig segmentLoaderConfig,
       String taskId,
       DataSchema schema,
       AppenderatorConfig config,
@@ -88,6 +90,7 @@ public Appenderator createRealtimeAppenderatorForTask(
       throw new ISE("A batch appenderator was already created for this peon's task.");
     } else {
       realtimeAppenderator = Appenderators.createRealtime(
+          segmentLoaderConfig,
           taskId,
           schema,
           config,