Associate pending segments with the tasks that requested them

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/MSQControllerTask.java

@@ -40,6 +40,7 @@
 import org.apache.druid.indexing.common.actions.TimeChunkLockTryAcquireAction;
 import org.apache.druid.indexing.common.config.TaskConfig;
 import org.apache.druid.indexing.common.task.AbstractTask;
+import org.apache.druid.indexing.common.task.PendingSegmentAllocatingTask;
 import org.apache.druid.indexing.common.task.Tasks;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.StringUtils;
@@ -71,7 +72,7 @@
 import java.util.Set;
 
 @JsonTypeName(MSQControllerTask.TYPE)
-public class MSQControllerTask extends AbstractTask implements ClientTaskQuery
+public class MSQControllerTask extends AbstractTask implements ClientTaskQuery, PendingSegmentAllocatingTask
 {
   public static final String TYPE = "query_controller";
   public static final String DUMMY_DATASOURCE_FOR_SELECT = "__query_select";
@@ -157,6 +158,12 @@ public Set<ResourceAction> getInputSourceResources()
     return ImmutableSet.of();
   }
 
+  @Override
+  public String getTaskAllocatorId()
+  {
+    return getId();
+  }
+
   @JsonProperty("spec")
   public MSQSpec getQuerySpec()
   {

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/MSQWorkerTask.java

@@ -32,6 +32,7 @@
 import org.apache.druid.indexing.common.actions.TaskActionClient;
 import org.apache.druid.indexing.common.config.TaskConfig;
 import org.apache.druid.indexing.common.task.AbstractTask;
+import org.apache.druid.indexing.common.task.PendingSegmentAllocatingTask;
 import org.apache.druid.indexing.common.task.Tasks;
 import org.apache.druid.msq.exec.MSQTasks;
 import org.apache.druid.msq.exec.Worker;
@@ -45,7 +46,7 @@
 import java.util.Set;
 
 @JsonTypeName(MSQWorkerTask.TYPE)
-public class MSQWorkerTask extends AbstractTask
+public class MSQWorkerTask extends AbstractTask implements PendingSegmentAllocatingTask
 {
   public static final String TYPE = "query_worker";
 
@@ -125,6 +126,12 @@ public Set<ResourceAction> getInputSourceResources()
     return ImmutableSet.of();
   }
 
+  @Override
+  public String getTaskAllocatorId()
+  {
+    return getControllerTaskId();
+  }
+
 
   @Override
   public boolean isReady(final TaskActionClient taskActionClient)

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/MSQControllerTaskTest.java

@@ -35,7 +35,7 @@
 
 public class MSQControllerTaskTest
 {
-  MSQSpec MSQ_SPEC = MSQSpec
+  private final MSQSpec MSQ_SPEC = MSQSpec
       .builder()
       .destination(new DataSourceMSQDestination(
           "target",
@@ -59,15 +59,33 @@ public class MSQControllerTaskTest
   @Test
   public void testGetInputSourceResources()
   {
-    MSQControllerTask msqWorkerTask = new MSQControllerTask(
+    MSQControllerTask controllerTask = new MSQControllerTask(
         null,
         MSQ_SPEC,
         null,
         null,
         null,
         null,
         null,
-        null);
-    Assert.assertTrue(msqWorkerTask.getInputSourceResources().isEmpty());
+        null
+    );
+    Assert.assertTrue(controllerTask.getInputSourceResources().isEmpty());
+  }
+
+  @Test
+  public void testGetTaskAllocatorId()
+  {
+    final String taskId = "taskId";
+    MSQControllerTask controllerTask = new MSQControllerTask(
+        taskId,
+        MSQ_SPEC,
+        null,
+        null,
+        null,
+        null,
+        null,
+        null
+    );
+    Assert.assertEquals(taskId, controllerTask.getTaskAllocatorId());
   }
 }

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/MSQWorkerTaskTest.java

@@ -47,7 +47,6 @@ public class MSQWorkerTaskTest
   @Test
   public void testEquals()
   {
-    Assert.assertNotEquals(msqWorkerTask, 0);
     Assert.assertEquals(msqWorkerTask, msqWorkerTask);
     Assert.assertEquals(
         msqWorkerTask,
@@ -110,4 +109,11 @@ public void testGetInputSourceResources()
     Assert.assertTrue(msqWorkerTask.getInputSourceResources().isEmpty());
   }
 
+  @Test
+  public void testGetTaskAllocatorId()
+  {
+    MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId, dataSource, workerNumber, context, retryCount);
+    Assert.assertEquals(controllerTaskId, msqWorkerTask.getTaskAllocatorId());
+  }
+
 }

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

@@ -23,8 +23,10 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.google.common.base.Preconditions;
+import org.apache.druid.error.DruidException;
 import org.apache.druid.indexing.common.LockGranularity;
 import org.apache.druid.indexing.common.TaskLockType;
+import org.apache.druid.indexing.common.task.PendingSegmentAllocatingTask;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator;
 import org.apache.druid.indexing.overlord.LockRequestForNewSegment;
@@ -210,6 +212,12 @@ public SegmentIdWithShardSpec perform(
       final TaskActionToolbox toolbox
   )
   {
+    if (!(task instanceof PendingSegmentAllocatingTask)) {
+      throw DruidException.defensive(
+          "Task[%s] of type[%s] cannot allocate segments as it does not implement PendingSegmentAllocatingTask.",
+          task.getId(), task.getType()
+      );
+    }
     int attempt = 0;
     while (true) {
       attempt++;

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

@@ -22,10 +22,12 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.core.type.TypeReference;
+import org.apache.druid.error.DruidException;
 import org.apache.druid.error.InvalidInput;
 import org.apache.druid.indexing.common.TaskLock;
 import org.apache.druid.indexing.common.TaskLockType;
 import org.apache.druid.indexing.common.task.IndexTaskUtils;
+import org.apache.druid.indexing.common.task.PendingSegmentAllocatingTask;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.indexing.overlord.CriticalAction;
 import org.apache.druid.indexing.overlord.DataSourceMetadata;
@@ -41,8 +43,20 @@
 import java.util.stream.Collectors;
 
 /**
+ *
  * Append segments to metadata storage. The segment versions must all be less than or equal to a lock held by
  * your task for the segment intervals.
+ *
+ * <pre>
+ * Pseudo code (for a single interval):
+ * For an append lock held over an interval:
+ *     transaction {
+ *       commit input segments contained within interval
+ *       if there is an active replace lock over the interval:
+ *         add an entry for the inputSegment corresponding to the replace lock's task in the upgradeSegments table
+ *       fetch pending segments with parent contained within the input segments, and commit them
+ *     }
+ * </pre>
  */
 public class SegmentTransactionalAppendAction implements TaskAction<SegmentPublishResult>
 {
@@ -114,6 +128,13 @@ public TypeReference<SegmentPublishResult> getReturnTypeReference()
   @Override
   public SegmentPublishResult perform(Task task, TaskActionToolbox toolbox)
   {
+    if (!(task instanceof PendingSegmentAllocatingTask)) {
+      throw DruidException.defensive(
+          "Task[%s] of type[%s] cannot append segments as it does not implement PendingSegmentAllocatingTask.",
+          task.getId(),
+          task.getType()
+      );
+    }
     // Verify that all the locks are of expected type
     final List<TaskLock> locks = toolbox.getTaskLockbox().findLocksForTask(task);
     for (TaskLock lock : locks) {
@@ -132,17 +153,20 @@ public SegmentPublishResult perform(Task task, TaskActionToolbox toolbox)
         = TaskLocks.findReplaceLocksCoveringSegments(datasource, toolbox.getTaskLockbox(), segments);
 
     final CriticalAction.Action<SegmentPublishResult> publishAction;
+    final String taskAllocatorId = ((PendingSegmentAllocatingTask) task).getTaskAllocatorId();
     if (startMetadata == null) {
       publishAction = () -> toolbox.getIndexerMetadataStorageCoordinator().commitAppendSegments(
           segments,
-          segmentToReplaceLock
+          segmentToReplaceLock,
+          taskAllocatorId
       );
     } else {
       publishAction = () -> toolbox.getIndexerMetadataStorageCoordinator().commitAppendSegmentsAndMetadata(
           segments,
           segmentToReplaceLock,
           startMetadata,
-          endMetadata
+          endMetadata,
+          taskAllocatorId
       );
     }
 

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

@@ -30,18 +30,35 @@
 import org.apache.druid.indexing.overlord.SegmentPublishResult;
 import org.apache.druid.indexing.overlord.supervisor.SupervisorManager;
 import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.metadata.PendingSegmentRecord;
 import org.apache.druid.metadata.ReplaceTaskLock;
 import org.apache.druid.segment.SegmentUtils;
 import org.apache.druid.segment.realtime.appenderator.SegmentIdWithShardSpec;
 import org.apache.druid.timeline.DataSegment;
 
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
 /**
  * Replace segments in metadata storage. The segment versions must all be less than or equal to a lock held by
  * your task for the segment intervals.
+ *
+ * <pre>
+ *  Pseudo code (for a single interval)
+ *- For a replace lock held over an interval:
+ *     transaction {
+ *       commit input segments contained within interval
+ *       upgrade ids in the upgradeSegments table corresponding to this task to the replace lock's version and commit them
+ *       fetch payload, task_allocator_id for pending segments
+ *       upgrade each such pending segment to the replace lock's version with the corresponding root segment
+ *     }
+ * For every pending segment with version == replace lock version:
+ *    Fetch payload, group_id or the pending segment and relay them to the supervisor
+ *    The supervisor relays the payloads to all the tasks with the corresponding group_id to serve realtime queries
+ * </pre>
  */
 public class SegmentTransactionalReplaceAction implements TaskAction<SegmentPublishResult>
 {
@@ -123,7 +140,7 @@ public SegmentPublishResult perform(Task task, TaskActionToolbox toolbox)
     // failure to upgrade pending segments does not affect success of the commit
     if (publishResult.isSuccess() && toolbox.getSupervisorManager() != null) {
       try {
-        tryUpgradeOverlappingPendingSegments(task, toolbox);
+        registerUpgradedPendingSegmentsOnSupervisor(task, toolbox);
       }
       catch (Exception e) {
         log.error(e, "Error while upgrading pending segments for task[%s]", task.getId());
@@ -134,34 +151,55 @@ public SegmentPublishResult perform(Task task, TaskActionToolbox toolbox)
   }
 
   /**
-   * Tries to upgrade any pending segments that overlap with the committed segments.
+   * Registers upgraded pending segments on the active supervisor, if any
    */
-  private void tryUpgradeOverlappingPendingSegments(Task task, TaskActionToolbox toolbox)
+  private void registerUpgradedPendingSegmentsOnSupervisor(Task task, TaskActionToolbox toolbox)
   {
     final SupervisorManager supervisorManager = toolbox.getSupervisorManager();
     final Optional<String> activeSupervisorIdWithAppendLock =
         supervisorManager.getActiveSupervisorIdForDatasourceWithAppendLock(task.getDataSource());
+
     if (!activeSupervisorIdWithAppendLock.isPresent()) {
       return;
     }
 
-    final Set<String> activeRealtimeSequencePrefixes
-        = supervisorManager.getActiveRealtimeSequencePrefixes(activeSupervisorIdWithAppendLock.get());
-    Map<SegmentIdWithShardSpec, SegmentIdWithShardSpec> upgradedPendingSegments =
-        toolbox.getIndexerMetadataStorageCoordinator()
-               .upgradePendingSegmentsOverlappingWith(segments, activeRealtimeSequencePrefixes);
-    log.info(
-        "Upgraded [%d] pending segments for REPLACE task[%s]: [%s]",
-        upgradedPendingSegments.size(), task.getId(), upgradedPendingSegments
-    );
+    final Set<ReplaceTaskLock> replaceLocksForTask = toolbox
+        .getTaskLockbox()
+        .getAllReplaceLocksForDatasource(task.getDataSource())
+        .stream()
+        .filter(lock -> task.getId().equals(lock.getSupervisorTaskId()))
+        .collect(Collectors.toSet());
+
 
-    upgradedPendingSegments.forEach(
-        (oldId, newId) -> toolbox.getSupervisorManager()
-                                 .registerNewVersionOfPendingSegmentOnSupervisor(
-                                     activeSupervisorIdWithAppendLock.get(),
-                                     oldId,
-                                     newId
-                                 )
+    Set<PendingSegmentRecord> pendingSegments = new HashSet<>();
+    for (ReplaceTaskLock replaceLock : replaceLocksForTask) {
+      pendingSegments.addAll(
+          toolbox.getIndexerMetadataStorageCoordinator()
+                 .getPendingSegments(task.getDataSource(), replaceLock.getInterval())
+      );
+    }
+    Map<String, SegmentIdWithShardSpec> idToPendingSegment = new HashMap<>();
+    pendingSegments.forEach(pendingSegment -> idToPendingSegment.put(
+        pendingSegment.getId().asSegmentId().toString(),
+        pendingSegment.getId()
+    ));
+    Map<SegmentIdWithShardSpec, SegmentIdWithShardSpec> segmentToParent = new HashMap<>();
+    pendingSegments.forEach(pendingSegment -> {
+      if (pendingSegment.getUpgradedFromSegmentId() != null
+          && !pendingSegment.getUpgradedFromSegmentId().equals(pendingSegment.getId().asSegmentId().toString())) {
+        segmentToParent.put(
+            pendingSegment.getId(),
+            idToPendingSegment.get(pendingSegment.getUpgradedFromSegmentId())
+        );
+      }
+    });
+
+    segmentToParent.forEach(
+        (newId, oldId) -> supervisorManager.registerNewVersionOfPendingSegmentOnSupervisor(
+            activeSupervisorIdWithAppendLock.get(),
+            oldId,
+            newId
+        )
     );
   }
 

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

@@ -119,7 +119,8 @@
 import java.util.concurrent.TimeoutException;
 
 @Deprecated
-public class AppenderatorDriverRealtimeIndexTask extends AbstractTask implements ChatHandler
+public class AppenderatorDriverRealtimeIndexTask extends AbstractTask
+    implements ChatHandler, PendingSegmentAllocatingTask
 {
   private static final String CTX_KEY_LOOKUP_TIER = "lookupTier";
 
@@ -259,6 +260,12 @@ public boolean isReady(TaskActionClient taskActionClient)
     return true;
   }
 
+  @Override
+  public String getTaskAllocatorId()
+  {
+    return getGroupId();
+  }
+
   @Override
   public TaskStatus runTask(final TaskToolbox toolbox)
   {

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

@@ -127,7 +127,7 @@
  * serialization fields of this class must correspond to those of {@link
  * ClientCompactionTaskQuery}.
  */
-public class CompactionTask extends AbstractBatchIndexTask
+public class CompactionTask extends AbstractBatchIndexTask implements PendingSegmentAllocatingTask
 {
   private static final Logger log = new Logger(CompactionTask.class);
   private static final Clock UTC_CLOCK = Clock.systemUTC();
@@ -400,6 +400,12 @@ public String getType()
     return TYPE;
   }
 
+  @Override
+  public String getTaskAllocatorId()
+  {
+    return getGroupId();
+  }
+
   @Nonnull
   @JsonIgnore
   @Override