Add support for concurrent batch Append and Replace

indexing-hadoop/src/main/java/org/apache/druid/indexer/updater/MetadataStorageUpdaterJobSpec.java

@@ -90,6 +90,7 @@ public MetadataStorageTablesConfig getMetadataStorageTablesConfig()
         null,
         null,
         null,
+        null,
         segmentTable,
         null,
         null,

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

@@ -30,9 +30,6 @@
 import java.util.Set;
 
 /**
- * Insert segments into metadata storage. The segment versions must all be less than or equal to a lock held by
- * your task for the segment intervals.
- * <p/>
  * Word of warning: Very large "segments" sets can cause oversized audit log entries, which is bad because it means
  * that the task cannot actually complete. Callers should avoid this by avoiding inserting too many segments in the
  * same action.

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

@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.indexing.common.actions;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.type.TypeReference;
+import org.apache.druid.indexing.common.task.IndexTaskUtils;
+import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.indexing.overlord.CriticalAction;
+import org.apache.druid.indexing.overlord.SegmentPublishResult;
+import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
+import org.apache.druid.metadata.ReplaceTaskLock;
+import org.apache.druid.segment.SegmentUtils;
+import org.apache.druid.timeline.DataSegment;
+
+import java.util.Map;
+import java.util.Set;
+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.
+ */
+public class SegmentTransactionalAppendAction implements TaskAction<SegmentPublishResult>
+{
+  private final Set<DataSegment> segments;
+
+  public static SegmentTransactionalAppendAction create(Set<DataSegment> segments)
+  {
+    return new SegmentTransactionalAppendAction(segments);
+  }
+
+  @JsonCreator
+  private SegmentTransactionalAppendAction(
+      @JsonProperty("segments") Set<DataSegment> segments
+  )
+  {
+    this.segments = segments;
+  }
+
+  @JsonProperty
+  public Set<DataSegment> getSegments()
+  {
+    return segments;
+  }
+
+  @Override
+  public TypeReference<SegmentPublishResult> getReturnTypeReference()
+  {
+    return new TypeReference<SegmentPublishResult>()
+    {
+    };
+  }
+
+  /**
+   * Performs some sanity checks and publishes the given segments.
+   */
+  @Override
+  public SegmentPublishResult perform(Task task, TaskActionToolbox toolbox)
+  {
+    TaskLocks.checkLockCoversSegments(task, toolbox.getTaskLockbox(), segments);
+
+    final String datasource = task.getDataSource();
+    final Map<DataSegment, ReplaceTaskLock> segmentToReplaceLock
+        = TaskLocks.findReplaceLocksCoveringSegments(datasource, toolbox.getTaskLockbox(), segments);
+
+    final SegmentPublishResult retVal;
+    try {
+      retVal = toolbox.getTaskLockbox().doInCriticalSection(
+          task,
+          segments.stream().map(DataSegment::getInterval).collect(Collectors.toSet()),
+          CriticalAction.<SegmentPublishResult>builder()
+              .onValidLocks(
+                  () -> toolbox.getIndexerMetadataStorageCoordinator().commitAppendSegments(
+                      segments,
+                      segmentToReplaceLock
+                  )
+              )
+              .onInvalidLocks(
+                  () -> SegmentPublishResult.fail(
+                      "Invalid task locks. Maybe they are revoked by a higher priority task."
+                      + " Please check the overlord log for details."
+                  )
+              )
+              .build()
+      );
+    }
+    catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+
+    // Emit metrics
+    final ServiceMetricEvent.Builder metricBuilder = new ServiceMetricEvent.Builder();
+    IndexTaskUtils.setTaskDimensions(metricBuilder, task);
+
+    if (retVal.isSuccess()) {
+      toolbox.getEmitter().emit(metricBuilder.setMetric("segment/txn/success", 1));
+      for (DataSegment segment : retVal.getSegments()) {
+        IndexTaskUtils.setSegmentDimensions(metricBuilder, segment);
+        toolbox.getEmitter().emit(metricBuilder.setMetric("segment/added/bytes", segment.getSize()));
+      }
+    } else {
+      toolbox.getEmitter().emit(metricBuilder.setMetric("segment/txn/failure", 1));
+    }
+
+    return retVal;
+  }
+
+  @Override
+  public boolean isAudited()
+  {
+    return true;
+  }
+
+  @Override
+  public String toString()
+  {
+    return "SegmentTransactionalAppendAction{" +
+           "segments=" + SegmentUtils.commaSeparatedIdentifiers(segments) +
+           '}';
+  }
+}

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

@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.indexing.common.actions;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.google.common.collect.ImmutableSet;
+import org.apache.druid.indexing.common.task.IndexTaskUtils;
+import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.indexing.overlord.CriticalAction;
+import org.apache.druid.indexing.overlord.SegmentPublishResult;
+import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
+import org.apache.druid.metadata.ReplaceTaskLock;
+import org.apache.druid.query.DruidMetrics;
+import org.apache.druid.segment.SegmentUtils;
+import org.apache.druid.timeline.DataSegment;
+
+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.
+ */
+public class SegmentTransactionalReplaceAction implements TaskAction<SegmentPublishResult>
+{
+  /**
+   * Set of segments to be inserted into metadata storage
+   */
+  private final Set<DataSegment> segments;
+
+  public static SegmentTransactionalReplaceAction create(
+      Set<DataSegment> segmentsToPublish
+  )
+  {
+    return new SegmentTransactionalReplaceAction(segmentsToPublish);
+  }
+
+  @JsonCreator
+  private SegmentTransactionalReplaceAction(
+      @JsonProperty("segments") Set<DataSegment> segments
+  )
+  {
+    this.segments = ImmutableSet.copyOf(segments);
+  }
+
+  @JsonProperty
+  public Set<DataSegment> getSegments()
+  {
+    return segments;
+  }
+
+  @Override
+  public TypeReference<SegmentPublishResult> getReturnTypeReference()
+  {
+    return new TypeReference<SegmentPublishResult>()
+    {
+    };
+  }
+
+  /**
+   * Performs some sanity checks and publishes the given segments.
+   */
+  @Override
+  public SegmentPublishResult perform(Task task, TaskActionToolbox toolbox)
+  {
+    TaskLocks.checkLockCoversSegments(task, toolbox.getTaskLockbox(), segments);
+
+    // Find the active replace locks held only by this task
+    final Set<ReplaceTaskLock> replaceLocksForTask
+        = toolbox.getTaskLockbox().findReplaceLocksForTask(task);
+
+    final SegmentPublishResult retVal;
+    try {
+      retVal = toolbox.getTaskLockbox().doInCriticalSection(
+          task,
+          segments.stream().map(DataSegment::getInterval).collect(Collectors.toSet()),
+          CriticalAction.<SegmentPublishResult>builder()
+              .onValidLocks(
+                  () -> toolbox.getIndexerMetadataStorageCoordinator()
+                               .commitReplaceSegments(segments, replaceLocksForTask)
+              )
+              .onInvalidLocks(
+                  () -> SegmentPublishResult.fail(
+                      "Invalid task locks. Maybe they are revoked by a higher priority task."
+                      + " Please check the overlord log for details."
+                  )
+              )
+              .build()
+      );
+    }
+    catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+
+    // Emit metrics
+    final ServiceMetricEvent.Builder metricBuilder = new ServiceMetricEvent.Builder();
+    IndexTaskUtils.setTaskDimensions(metricBuilder, task);
+
+    if (retVal.isSuccess()) {
+      toolbox.getEmitter().emit(metricBuilder.setMetric("segment/txn/success", 1));
+
+      for (DataSegment segment : retVal.getSegments()) {
+        final String partitionType = segment.getShardSpec() == null ? null : segment.getShardSpec().getType();
+        metricBuilder.setDimension(DruidMetrics.PARTITIONING_TYPE, partitionType);
+        metricBuilder.setDimension(DruidMetrics.INTERVAL, segment.getInterval().toString());
+        toolbox.getEmitter().emit(metricBuilder.setMetric("segment/added/bytes", segment.getSize()));
+      }
+    } else {
+      toolbox.getEmitter().emit(metricBuilder.setMetric("segment/txn/failure", 1));
+    }
+
+    return retVal;
+  }
+
+  @Override
+  public boolean isAudited()
+  {
+    return true;
+  }
+
+  @Override
+  public String toString()
+  {
+    return "SegmentTransactionalReplaceAction{" +
+           "segments=" + SegmentUtils.commaSeparatedIdentifiers(segments) +
+           '}';
+  }
+}

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

@@ -36,6 +36,8 @@
     @JsonSubTypes.Type(name = "lockRelease", value = LockReleaseAction.class),
     @JsonSubTypes.Type(name = "segmentInsertion", value = SegmentInsertAction.class),
     @JsonSubTypes.Type(name = "segmentTransactionalInsert", value = SegmentTransactionalInsertAction.class),
+    @JsonSubTypes.Type(name = "segmentTransactionalAppend", value = SegmentTransactionalAppendAction.class),
+    @JsonSubTypes.Type(name = "segmentTransactionalReplace", value = SegmentTransactionalReplaceAction.class),
     // Type name doesn't correspond to the name of the class for backward compatibility.
     @JsonSubTypes.Type(name = "segmentListUsed", value = RetrieveUsedSegmentsAction.class),
     // Type name doesn't correspond to the name of the class for backward compatibility.