Filter pending segments upgraded with transactional replace

extensions-contrib/materialized-view-maintenance/src/main/java/org/apache/druid/indexing/materializedview/MaterializedViewSupervisor.java

@@ -296,6 +296,12 @@ public LagStats computeLagStats()
     throw new UnsupportedOperationException("Compute Lag Stats not supported in MaterializedViewSupervisor");
   }
 
+  @Override
+  public Set<String> getActiveRealtimeSequencePrefixes()
+  {
+    throw new UnsupportedOperationException("Get Active sequence names is not supported in MaterializedViewSupervisor");
+  }
+
   @Override
   public int getActiveTaskGroupsCount()
   {

extensions-contrib/materialized-view-maintenance/src/test/java/org/apache/druid/indexing/materializedview/MaterializedViewSupervisorSpecTest.java

@@ -217,6 +217,13 @@ public void testMaterializedViewSupervisorSpecCreated()
         Assert.assertTrue(e instanceof UnsupportedOperationException);
       }
 
+      try {
+        supervisor.getActiveRealtimeSequencePrefixes();
+      }
+      catch (Exception e) {
+        Assert.assertTrue(e instanceof UnsupportedOperationException);
+      }
+
       Callable<Integer> noop = new Callable<Integer>() {
         @Override
         public Integer call()

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

@@ -145,8 +145,11 @@ private void tryUpgradeOverlappingPendingSegments(Task task, TaskActionToolbox t
       return;
     }
 
+    final Set<String> activeRealtimeSequencePrefixes
+        = supervisorManager.getActiveRealtimeSequencePrefixes(activeSupervisorIdWithAppendLock.get());
     Map<SegmentIdWithShardSpec, SegmentIdWithShardSpec> upgradedPendingSegments =
-        toolbox.getIndexerMetadataStorageCoordinator().upgradePendingSegmentsOverlappingWith(segments);
+        toolbox.getIndexerMetadataStorageCoordinator()
+               .upgradePendingSegmentsOverlappingWith(segments, activeRealtimeSequencePrefixes);
     log.info(
         "Upgraded [%d] pending segments for REPLACE task[%s]: [%s]",
         upgradedPendingSegments.size(), task.getId(), upgradedPendingSegments

indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java

@@ -39,6 +39,7 @@
 
 import javax.annotation.Nullable;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -110,6 +111,14 @@ public Optional<String> getActiveSupervisorIdForDatasourceWithAppendLock(String
     return Optional.absent();
   }
 
+  public Set<String> getActiveRealtimeSequencePrefixes(String activeSupervisorId)
+  {
+    if (!supervisors.containsKey(activeSupervisorId)) {
+      return Collections.emptySet();
+    }
+    return supervisors.get(activeSupervisorId).lhs.getActiveRealtimeSequencePrefixes();
+  }
+
   public Optional<SupervisorSpec> getSupervisorSpec(String id)
   {
     Pair<Supervisor, SupervisorSpec> supervisor = supervisors.get(id);

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

@@ -1093,6 +1093,28 @@ public void resetOffsets(@Nonnull DataSourceMetadata resetDataSourceMetadata)
     addNotice(new ResetOffsetsNotice(resetDataSourceMetadata));
   }
 
+  /**
+   * The base sequence name of a seekable stream task group is used as a prefix of the sequence names
+   * of pending segments published by it.
+   * This method can be used to identify the active pending segments for a datasource
+   * by checking if the sequence name begins with any of the active realtime sequence prefix returned by this method
+   * @return the set of base sequence names of both active and pending completion task gruops.
+   */
+  @Override
+  public Set<String> getActiveRealtimeSequencePrefixes()
+  {
+    final Set<String> activeBaseSequences = new HashSet<>();
+    for (TaskGroup taskGroup : activelyReadingTaskGroups.values()) {
+      activeBaseSequences.add(taskGroup.baseSequenceName);
+    }
+    for (List<TaskGroup> taskGroupList : pendingCompletionTaskGroups.values()) {
+      for (TaskGroup taskGroup : taskGroupList) {
+        activeBaseSequences.add(taskGroup.baseSequenceName);
+      }
+    }
+    return activeBaseSequences;
+  }
+
   public void registerNewVersionOfPendingSegment(
       SegmentIdWithShardSpec basePendingSegment,
       SegmentIdWithShardSpec newSegmentVersion

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

@@ -1313,6 +1313,48 @@ public void testSupervisorResetOneTaskSpecificOffsetsWithCheckpoints() throws In
     validateSupervisorStateAfterResetOffsets(supervisor, resetOffsets, 0);
   }
 
+  @Test
+  public void testGetActiveRealtimeSequencePrefixes()
+  {
+    EasyMock.expect(spec.isSuspended()).andReturn(false);
+
+    replayAll();
+
+    final TestSeekableStreamSupervisor supervisor = new TestSeekableStreamSupervisor();
+
+    // Spin off two active tasks with each task serving one partition.
+    supervisor.getIoConfig().setTaskCount(3);
+    supervisor.start();
+    supervisor.addTaskGroupToActivelyReadingTaskGroup(
+        supervisor.getTaskGroupIdForPartition("0"),
+        ImmutableMap.of("0", "5"),
+        Optional.absent(),
+        Optional.absent(),
+        ImmutableSet.of("task1"),
+        ImmutableSet.of()
+    );
+
+    supervisor.addTaskGroupToActivelyReadingTaskGroup(
+        supervisor.getTaskGroupIdForPartition("1"),
+        ImmutableMap.of("1", "6"),
+        Optional.absent(),
+        Optional.absent(),
+        ImmutableSet.of("task2"),
+        ImmutableSet.of()
+    );
+
+    supervisor.addTaskGroupToPendingCompletionTaskGroup(
+        supervisor.getTaskGroupIdForPartition("2"),
+        ImmutableMap.of("2", "100"),
+        Optional.absent(),
+        Optional.absent(),
+        ImmutableSet.of("task3"),
+        ImmutableSet.of()
+    );
+
+    Assert.assertEquals(3, supervisor.getActiveRealtimeSequencePrefixes().size());
+  }
+
   @Test
   public void testSupervisorResetSpecificOffsetsTasksWithCheckpoints() throws InterruptedException, IOException
   {

indexing-service/src/test/java/org/apache/druid/indexing/test/TestIndexerMetadataStorageCoordinator.java

@@ -238,7 +238,10 @@ public SegmentIdWithShardSpec allocatePendingSegment(
   }
 
   @Override
-  public Map<SegmentIdWithShardSpec, SegmentIdWithShardSpec> upgradePendingSegmentsOverlappingWith(Set<DataSegment> replaceSegments)
+  public Map<SegmentIdWithShardSpec, SegmentIdWithShardSpec> upgradePendingSegmentsOverlappingWith(
+      Set<DataSegment> replaceSegments,
+      Set<String> activeBaseSequenceNames
+  )
   {
     return Collections.emptyMap();
   }

server/src/main/java/org/apache/druid/indexing/overlord/IndexerMetadataStorageCoordinator.java

@@ -347,10 +347,13 @@ SegmentPublishResult commitReplaceSegments(
    * </ul>
    *
    * @param replaceSegments Segments being committed by a REPLACE task
+   * @param activeRealtimeSequencePrefixes Set of base sequence names of active and pending completion task groups
+   *                                       of the supervisor (if any) for this datasource
    * @return Map from originally allocated pending segment to its new upgraded ID.
    */
   Map<SegmentIdWithShardSpec, SegmentIdWithShardSpec> upgradePendingSegmentsOverlappingWith(
-      Set<DataSegment> replaceSegments
+      Set<DataSegment> replaceSegments,
+      Set<String> activeRealtimeSequencePrefixes
   );
 
   /**

server/src/main/java/org/apache/druid/indexing/overlord/supervisor/NoopSupervisorSpec.java

@@ -31,6 +31,7 @@
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
@@ -185,6 +186,12 @@ public int getActiveTaskGroupsCount()
       {
         return -1;
       }
+
+      @Override
+      public Set<String> getActiveRealtimeSequencePrefixes()
+      {
+        return Collections.emptySet();
+      }
     };
   }
 

server/src/main/java/org/apache/druid/indexing/overlord/supervisor/Supervisor.java

@@ -29,6 +29,7 @@
 import javax.annotation.Nullable;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 public interface Supervisor
 {
@@ -93,4 +94,9 @@ default Boolean isHealthy()
   LagStats computeLagStats();
 
   int getActiveTaskGroupsCount();
+
+  /**
+   * @return active base sequence names for reading and pending completion task groups of a seekable stream supervisor
+   */
+  Set<String> getActiveRealtimeSequencePrefixes();
 }