Fix sequence change race condition

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

@@ -804,26 +804,33 @@ public void onFailure(Throwable t)
       // We need to copy sequences here, because the success callback in publishAndRegisterHandoff removes items from
       // the sequence list. If a publish finishes before we finish iterating through the sequence list, we can
       // end up skipping some sequences.
-      List<SequenceMetadata<PartitionIdType, SequenceOffsetType>> sequencesSnapshot = new ArrayList<>(sequences);
-      for (int i = 0; i < sequencesSnapshot.size(); i++) {
-        final SequenceMetadata<PartitionIdType, SequenceOffsetType> sequenceMetadata = sequencesSnapshot.get(i);
-        if (!publishingSequences.contains(sequenceMetadata.getSequenceName())
-            && !publishedSequences.contains(sequenceMetadata.getSequenceName())) {
-          final boolean isLast = i == (sequencesSnapshot.size() - 1);
-          if (isLast) {
-            // Shorten endOffsets of the last sequence to match currOffsets.
-            sequenceMetadata.setEndOffsets(currOffsets);
-          }
+      List<SequenceMetadata<PartitionIdType, SequenceOffsetType>> waitForPublishingSequences = new ArrayList<>();
+
+      synchronized (publishingSequences) {
+        List<SequenceMetadata<PartitionIdType, SequenceOffsetType>> sequencesSnapshot = new ArrayList<>(sequences);
+        for (int i = 0; i < sequencesSnapshot.size(); i++) {
+          final SequenceMetadata<PartitionIdType, SequenceOffsetType> sequenceMetadata = sequencesSnapshot.get(i);
+          if (!publishingSequences.contains(sequenceMetadata.getSequenceName())) {
+            final boolean isLast = i == (sequencesSnapshot.size() - 1);
+            if (isLast) {
+              // Shorten endOffsets of the last sequence to match currOffsets.
+              sequenceMetadata.setEndOffsets(currOffsets);
+            }
 
-          // Update assignments of the sequence, which should clear them. (This will be checked later, when the
-          // Committer is built.)
-          sequenceMetadata.updateAssignments(currOffsets, this::isMoreToReadAfterReadingRecord);
-          publishingSequences.add(sequenceMetadata.getSequenceName());
-          // persist already done in finally, so directly add to publishQueue
-          publishAndRegisterHandoff(sequenceMetadata);
+            // Update assignments of the sequence, which should clear them. (This will be checked later, when the
+            // Committer is built.)
+            sequenceMetadata.updateAssignments(currOffsets, this::isMoreToReadAfterReadingRecord);
+            publishingSequences.add(sequenceMetadata.getSequenceName());
+            waitForPublishingSequences.add(sequenceMetadata);
+          }
         }
       }
 
+      for (SequenceMetadata<PartitionIdType, SequenceOffsetType> sequenceMetadata : waitForPublishingSequences) {
+        // persist already done in finally, so directly add to publishQueue
+        publishAndRegisterHandoff(sequenceMetadata);
+      }
+
       if (backgroundThreadException != null) {
         throw new RuntimeException(backgroundThreadException);
       }
@@ -1011,9 +1018,10 @@ public void onSuccess(SegmentsAndCommitMetadata publishedSegmentsAndCommitMetada
             );
             log.infoSegments(publishedSegmentsAndCommitMetadata.getSegments(), "Published segments");
 
-            publishedSequences.add(sequenceMetadata.getSequenceName());
-            sequences.remove(sequenceMetadata);
-            publishingSequences.remove(sequenceMetadata.getSequenceName());
+            synchronized (publishingSequences) {
+              sequences.remove(sequenceMetadata);
+              publishingSequences.remove(sequenceMetadata.getSequenceName());
+            }
 
             try {
               persistSequences();
@@ -1158,27 +1166,34 @@ private Map<String, Object> getTaskCompletionRowStats()
   private void maybePersistAndPublishSequences(Supplier<Committer> committerSupplier)
       throws InterruptedException
   {
-    for (SequenceMetadata<PartitionIdType, SequenceOffsetType> sequenceMetadata : sequences) {
-      sequenceMetadata.updateAssignments(currOffsets, this::isMoreToReadBeforeReadingRecord);
-      if (!sequenceMetadata.isOpen()
-          && !publishingSequences.contains(sequenceMetadata.getSequenceName())
-          && !publishedSequences.contains(sequenceMetadata.getSequenceName())) {
-        publishingSequences.add(sequenceMetadata.getSequenceName());
-        try {
-          final Object result = driver.persist(committerSupplier.get());
-          log.debug(
-              "Persist completed with metadata [%s], adding sequence [%s] to publish queue.",
-              result,
-              sequenceMetadata.getSequenceName()
-          );
-          publishAndRegisterHandoff(sequenceMetadata);
-        }
-        catch (InterruptedException e) {
-          log.warn("Interrupted while persisting metadata for sequence [%s].", sequenceMetadata.getSequenceName());
-          throw e;
+    List<SequenceMetadata<PartitionIdType, SequenceOffsetType>> waitForPublishingSequences = new ArrayList<>();
+
+    synchronized (publishingSequences) {
+      List<SequenceMetadata<PartitionIdType, SequenceOffsetType>> sequencesSnapshot = new ArrayList<>(sequences);
+      for (SequenceMetadata<PartitionIdType, SequenceOffsetType> sequenceMetadata : sequencesSnapshot) {
+        sequenceMetadata.updateAssignments(currOffsets, this::isMoreToReadBeforeReadingRecord);
+        if (!sequenceMetadata.isOpen() && !publishingSequences.contains(sequenceMetadata.getSequenceName())) {
+          publishingSequences.add(sequenceMetadata.getSequenceName());
+          waitForPublishingSequences.add(sequenceMetadata);
         }
       }
     }
+
+    for (SequenceMetadata<PartitionIdType, SequenceOffsetType> sequenceMetadata : waitForPublishingSequences) {
+      try {
+        final Object result = driver.persist(committerSupplier.get());
+        log.debug(
+                "Persist completed with metadata [%s], adding sequence [%s] to publish queue.",
+                result,
+                sequenceMetadata.getSequenceName()
+        );
+        publishAndRegisterHandoff(sequenceMetadata);
+      }
+      catch (InterruptedException e) {
+        log.warn("Interrupted while persisting metadata for sequence [%s].", sequenceMetadata.getSequenceName());
+        throw e;
+      }
+    }
   }
 
   private Set<StreamPartition<PartitionIdType>> assignPartitions(