Retry segment publish task actions without holding locks

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

@@ -76,6 +76,7 @@
 import org.apache.druid.segment.indexing.IngestionSpec;
 import org.apache.druid.segment.indexing.TuningConfig;
 import org.apache.druid.segment.realtime.appenderator.SegmentIdWithShardSpec;
+import org.apache.druid.segment.realtime.appenderator.TransactionalSegmentPublisher;
 import org.apache.druid.segment.transform.CompactionTransformSpec;
 import org.apache.druid.timeline.CompactionState;
 import org.apache.druid.timeline.DataSegment;
@@ -456,6 +457,30 @@ protected TaskAction<SegmentPublishResult> buildPublishAction(
     }
   }
 
+  protected TransactionalSegmentPublisher buildSegmentPublisher(TaskToolbox toolbox)
+  {
+    return new TransactionalSegmentPublisher()
+    {
+      @Override
+      public SegmentPublishResult publishAnnotatedSegments(
+          @Nullable Set<DataSegment> segmentsToBeOverwritten,
+          Set<DataSegment> segmentsToPublish,
+          @Nullable Object commitMetadata,
+          @Nullable SegmentSchemaMapping schemaMapping
+      ) throws IOException
+      {
+        return toolbox.getTaskActionClient().submit(
+            buildPublishAction(
+                segmentsToBeOverwritten,
+                segmentsToPublish,
+                schemaMapping,
+                getTaskLockHelper().getLockTypeToUse()
+            )
+        );
+      }
+    };
+  }
+
   protected boolean tryTimeChunkLock(TaskActionClient client, List<Interval> intervals) throws IOException
   {
     // The given intervals are first converted to align with segment granularity. This is because,

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

@@ -48,7 +48,6 @@
 import org.apache.druid.indexer.partitions.PartitionsSpec;
 import org.apache.druid.indexer.partitions.SecondaryPartitionType;
 import org.apache.druid.indexer.report.TaskReport;
-import org.apache.druid.indexing.common.TaskLockType;
 import org.apache.druid.indexing.common.TaskRealtimeMetricsMonitorBuilder;
 import org.apache.druid.indexing.common.TaskToolbox;
 import org.apache.druid.indexing.common.actions.TaskActionClient;
@@ -866,11 +865,7 @@ private TaskStatus generateAndPublishSegments(
         throw new UOE("[%s] secondary partition type is not supported", partitionsSpec.getType());
     }
 
-    final TaskLockType taskLockType = getTaskLockHelper().getLockTypeToUse();
-    final TransactionalSegmentPublisher publisher =
-        (segmentsToBeOverwritten, segmentsToPublish, commitMetadata, map) -> toolbox.getTaskActionClient().submit(
-            buildPublishAction(segmentsToBeOverwritten, segmentsToPublish, map, taskLockType)
-        );
+    final TransactionalSegmentPublisher publisher = buildSegmentPublisher(toolbox);
 
     String effectiveId = getContextValue(CompactionTask.CTX_KEY_APPENDERATOR_TRACKING_TASK_ID, null);
     if (effectiveId == null) {

indexing-service/src/main/java/org/apache/druid/indexing/common/task/batch/parallel/ParallelIndexSupervisorTask.java

@@ -45,7 +45,6 @@
 import org.apache.druid.indexer.report.IngestionStatsAndErrors;
 import org.apache.druid.indexer.report.IngestionStatsAndErrorsTaskReport;
 import org.apache.druid.indexer.report.TaskReport;
-import org.apache.druid.indexing.common.TaskLockType;
 import org.apache.druid.indexing.common.TaskToolbox;
 import org.apache.druid.indexing.common.actions.TaskActionClient;
 import org.apache.druid.indexing.common.task.AbstractBatchIndexTask;
@@ -1191,12 +1190,7 @@ private void publishSegments(
       }
     }
 
-    final TaskLockType taskLockType = getTaskLockHelper().getLockTypeToUse();
-    final TransactionalSegmentPublisher publisher =
-        (segmentsToBeOverwritten, segmentsToPublish, commitMetadata, map) -> toolbox.getTaskActionClient().submit(
-            buildPublishAction(segmentsToBeOverwritten, segmentsToPublish, map, taskLockType)
-        );
-
+    final TransactionalSegmentPublisher publisher = buildSegmentPublisher(toolbox);
     final boolean published =
         newSegments.isEmpty()
         || publisher.publishSegments(oldSegments, newSegments, annotateFunction, null, segmentSchemaMapping).isSuccess();

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

@@ -331,7 +331,7 @@ TransactionalSegmentPublisher createPublisher(
   }
 
   private class SequenceMetadataTransactionalSegmentPublisher
-      implements TransactionalSegmentPublisher
+      extends TransactionalSegmentPublisher
   {
     private final SeekableStreamIndexTaskRunner<PartitionIdType, SequenceOffsetType, ?> runner;
     private final TaskToolbox toolbox;

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

@@ -22,7 +22,6 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import org.apache.druid.error.InvalidInput;
 import org.apache.druid.indexing.common.TaskLockType;
 import org.apache.druid.indexing.common.task.NoopTask;
 import org.apache.druid.indexing.common.task.Task;
@@ -150,11 +149,9 @@ public void testFailTransactionalUpdateDataSourceMetadata() throws Exception
     );
 
     Assert.assertEquals(
-        SegmentPublishResult.fail(
-            InvalidInput.exception(
-                "The new start metadata state[ObjectMetadata{theObject=[1]}] is"
-                + " ahead of the last committed end state[null]. Try resetting the supervisor."
-            ).toString()
+        SegmentPublishResult.retryableFailure(
+            "The new start metadata state[ObjectMetadata{theObject=[1]}] is"
+            + " ahead of the last committed end state[null]. Try resetting the supervisor."
         ),
         result
     );

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

@@ -22,7 +22,7 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableSet;
+import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.metadata.PendingSegmentRecord;
 import org.apache.druid.segment.SegmentUtils;
 import org.apache.druid.timeline.DataSegment;
@@ -34,59 +34,59 @@
 import java.util.Set;
 
 /**
- * Result of an operation that attempts to publish segments. Indicates the set of segments actually published
- * and whether or not the transaction was a success.
- *
- * If "success" is false then the segments set will be empty.
- *
- * It's possible for the segments set to be empty even if "success" is true, since the segments set only
- * includes segments actually published as part of the transaction. The requested segments could have been
- * published by a different transaction (e.g. in the case of replica sets) and this one would still succeed.
+ * Result of a segment publish operation.
  */
 public class SegmentPublishResult
 {
   private final Set<DataSegment> segments;
   private final boolean success;
-  @Nullable
+  private final boolean retryable;
   private final String errorMsg;
-  @Nullable
   private final List<PendingSegmentRecord> upgradedPendingSegments;
 
   public static SegmentPublishResult ok(Set<DataSegment> segments)
   {
-    return new SegmentPublishResult(segments, true, null);
+    return new SegmentPublishResult(segments, true, false, null);
   }
 
   public static SegmentPublishResult ok(Set<DataSegment> segments, List<PendingSegmentRecord> upgradedPendingSegments)
   {
-    return new SegmentPublishResult(segments, true, null, upgradedPendingSegments);
+    return new SegmentPublishResult(segments, true, false, null, upgradedPendingSegments);
   }
 
-  public static SegmentPublishResult fail(String errorMsg)
+  public static SegmentPublishResult fail(String errorMsg, Object... args)
   {
-    return new SegmentPublishResult(ImmutableSet.of(), false, errorMsg);
+    return new SegmentPublishResult(Set.of(), false, false, StringUtils.format(errorMsg, args), null);
+  }
+
+  public static SegmentPublishResult retryableFailure(String errorMsg, Object... args)
+  {
+    return new SegmentPublishResult(Set.of(), false, true, StringUtils.format(errorMsg, args), null);
   }
 
   @JsonCreator
   private SegmentPublishResult(
       @JsonProperty("segments") Set<DataSegment> segments,
       @JsonProperty("success") boolean success,
+      @JsonProperty("retryable") boolean retryable,
       @JsonProperty("errorMsg") @Nullable String errorMsg
   )
   {
-    this(segments, success, errorMsg, null);
+    this(segments, success, retryable, errorMsg, null);
   }
 
   private SegmentPublishResult(
       Set<DataSegment> segments,
       boolean success,
-       @Nullable String errorMsg,
+      boolean retryable,
+      @Nullable String errorMsg,
       List<PendingSegmentRecord> upgradedPendingSegments
   )
   {
     this.segments = Preconditions.checkNotNull(segments, "segments");
     this.success = success;
     this.errorMsg = errorMsg;
+    this.retryable = retryable;
     this.upgradedPendingSegments = upgradedPendingSegments;
 
     if (!success) {
@@ -98,6 +98,12 @@ private SegmentPublishResult(
     }
   }
 
+  /**
+   * Set of segments published successfully.
+   *
+   * @return Empty set if the publish operation failed or if all the segments had
+   * already been published by a different transaction.
+   */
   @JsonProperty
   public Set<DataSegment> getSegments()
   {
@@ -117,6 +123,12 @@ public String getErrorMsg()
     return errorMsg;
   }
 
+  @JsonProperty
+  public boolean isRetryable()
+  {
+    return retryable;
+  }
+
   @Nullable
   public List<PendingSegmentRecord> getUpgradedPendingSegments()
   {
@@ -134,14 +146,15 @@ public boolean equals(Object o)
     }
     SegmentPublishResult that = (SegmentPublishResult) o;
     return success == that.success &&
+           retryable == that.retryable &&
            Objects.equals(segments, that.segments) &&
            Objects.equals(errorMsg, that.errorMsg);
   }
 
   @Override
   public int hashCode()
   {
-    return Objects.hash(segments, success, errorMsg);
+    return Objects.hash(segments, success, errorMsg, retryable);
   }
 
   @Override
@@ -150,6 +163,7 @@ public String toString()
     return "SegmentPublishResult{" +
            "segments=" + SegmentUtils.commaSeparatedIdentifiers(segments) +
            ", success=" + success +
+           ", retryable=" + retryable +
            ", errorMsg='" + errorMsg + '\'' +
            '}';
   }