A Replacing task must read segments created before it acquired its lock

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

@@ -40,6 +40,7 @@
 import org.apache.druid.indexing.overlord.supervisor.autoscaler.LagStats;
 import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.java.util.common.IAE;
+import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.java.util.common.JodaUtils;
 import org.apache.druid.java.util.common.Pair;
 import org.apache.druid.java.util.common.StringUtils;
@@ -364,7 +365,7 @@ Pair<SortedMap<Interval, String>, Map<Interval, List<DataSegment>>> checkSegment
     // Pair<interval -> max(created_date), interval -> list<DataSegment>>
     Pair<Map<Interval, String>, Map<Interval, List<DataSegment>>> baseSegmentsSnapshot =
         getMaxCreateDateAndBaseSegments(
-            metadataStorageCoordinator.retrieveUsedSegmentsAndCreatedDates(spec.getBaseDataSource())
+            metadataStorageCoordinator.retrieveUsedSegmentsAndCreatedDates(spec.getBaseDataSource(), Intervals.ETERNITY)
         );
     // baseSegments are used to create HadoopIndexTask
     Map<Interval, List<DataSegment>> baseSegments = baseSegmentsSnapshot.rhs;

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

@@ -0,0 +1,171 @@
+/*
+ * 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.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.type.TypeReference;
+import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.indexing.common.task.batch.parallel.AbstractBatchSubtask;
+import org.apache.druid.indexing.overlord.Segments;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.Pair;
+import org.apache.druid.metadata.ReplaceTaskLock;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.Partitions;
+import org.apache.druid.timeline.SegmentTimeline;
+import org.joda.time.Interval;
+
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * This TaskAction returns a collection of segments which have data within the specified interval and are marked as
+ * used, and have been created before a REPLACE lock, if any, was acquired.
+ *
+ * The order of segments within the returned collection is unspecified, but each segment is guaranteed to appear in
+ * the collection only once.
+ *
+ * @implNote This action doesn't produce a {@link Set} because it's implemented via {@link
+ * org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator#retrieveUsedSegmentsForIntervals} which returns
+ * a collection. Producing a {@link Set} would require an unnecessary copy of segments collection.
+ */
+public class RetrieveSegmentsToReplaceAction implements TaskAction<Collection<DataSegment>>
+{
+  @JsonIgnore
+  private final String dataSource;
+
+  @JsonIgnore
+  private final Interval interval;
+
+  @JsonCreator
+  public RetrieveSegmentsToReplaceAction(
+      @JsonProperty("dataSource") String dataSource,
+      @JsonProperty("interval") Interval interval
+  )
+  {
+    this.dataSource = dataSource;
+    this.interval = interval;
+  }
+
+  @JsonProperty
+  public String getDataSource()
+  {
+    return dataSource;
+  }
+
+  @JsonProperty
+  public Interval getInterval()
+  {
+    return interval;
+  }
+
+  @Override
+  public TypeReference<Collection<DataSegment>> getReturnTypeReference()
+  {
+    return new TypeReference<Collection<DataSegment>>() {};
+  }
+
+  @Override
+  public Collection<DataSegment> perform(Task task, TaskActionToolbox toolbox)
+  {
+    if (!task.getDataSource().equals(dataSource)) {
+      return toolbox.getIndexerMetadataStorageCoordinator()
+                    .retrieveUsedSegmentsForInterval(dataSource, interval, Segments.ONLY_VISIBLE);
+    }
+
+    final String supervisorId;
+    if (task instanceof AbstractBatchSubtask) {
+      supervisorId = ((AbstractBatchSubtask) task).getSupervisorTaskId();
+    } else {
+      supervisorId = task.getId();
+    }
+
+    final Set<ReplaceTaskLock> replaceLocksForTask = toolbox
+        .getTaskLockbox()
+        .getAllReplaceLocksForDatasource(task.getDataSource())
+        .stream()
+        .filter(lock -> supervisorId.equals(lock.getSupervisorTaskId()))
+        .collect(Collectors.toSet());
+
+    Collection<Pair<DataSegment, String>> segmentsAndCreatedDates =
+        toolbox.getIndexerMetadataStorageCoordinator().retrieveUsedSegmentsAndCreatedDates(dataSource, interval);
+
+    Set<DataSegment> allSegmentsToBeReplaced = new HashSet<>();
+    segmentsAndCreatedDates.forEach(segmentAndCreatedDate -> allSegmentsToBeReplaced.add(segmentAndCreatedDate.lhs));
+    for (Pair<DataSegment, String> segmentAndCreatedDate : segmentsAndCreatedDates) {
+      final DataSegment segment = segmentAndCreatedDate.lhs;
+      final String createdDate = segmentAndCreatedDate.rhs;
+      for (ReplaceTaskLock replaceLock : replaceLocksForTask) {
+        if (replaceLock.getInterval().contains(segment.getInterval())
+            && replaceLock.getVersion().compareTo(createdDate) < 0) {
+          // If a REPLACE lock covers a segment but has a version less than the segment's created date, remove it
+          allSegmentsToBeReplaced.remove(segment);
+        }
+      }
+    }
+
+    return SegmentTimeline.forSegments(allSegmentsToBeReplaced)
+                          .findNonOvershadowedObjectsInInterval(Intervals.ETERNITY, Partitions.ONLY_COMPLETE);
+  }
+
+  @Override
+  public boolean isAudited()
+  {
+    return false;
+  }
+
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+
+    RetrieveSegmentsToReplaceAction that = (RetrieveSegmentsToReplaceAction) o;
+
+    if (!dataSource.equals(that.dataSource)) {
+      return false;
+    }
+    return interval.equals(that.interval);
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return Objects.hash(dataSource, interval);
+  }
+
+  @Override
+  public String toString()
+  {
+    return getClass().getSimpleName() + "{" +
+           "dataSource='" + dataSource + '\'' +
+           ", interval=" + interval +
+           '}';
+  }
+}

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

@@ -40,6 +40,7 @@
     @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),
+    @JsonSubTypes.Type(name = "segmentListLocked", value = RetrieveSegmentsToReplaceAction.class),
     // Type name doesn't correspond to the name of the class for backward compatibility.
     @JsonSubTypes.Type(name = "segmentListUnused", value = RetrieveUnusedSegmentsAction.class),
     @JsonSubTypes.Type(name = "markSegmentsAsUnused", value = MarkSegmentsAsUnusedAction.class),

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

@@ -19,6 +19,7 @@
 
 package org.apache.druid.indexing.common.task.batch.parallel;
 
+import com.fasterxml.jackson.annotation.JsonProperty;
 import org.apache.druid.indexing.common.task.AbstractBatchIndexTask;
 import org.apache.druid.indexing.common.task.TaskResource;
 
@@ -29,21 +30,34 @@
 public abstract class AbstractBatchSubtask extends AbstractBatchIndexTask
 {
 
+  private final String supervisorTaskId;
+
   protected AbstractBatchSubtask(
       String id,
       @Nullable String groupId,
       @Nullable TaskResource taskResource,
       String dataSource,
       @Nullable Map<String, Object> context,
-      @Nonnull IngestionMode ingestionMode
+      @Nonnull IngestionMode ingestionMode,
+      @Nonnull String supervisorTaskId
   )
   {
     super(id, groupId, taskResource, dataSource, context, -1, ingestionMode);
+    this.supervisorTaskId = supervisorTaskId;
   }
 
   /**
    * Returns the ID of {@link SubTaskSpec} that is assigned to this subtask.
    * This ID is used to identify duplicate work of retry tasks for the same spec.
    */
   public abstract String getSubtaskSpecId();
+
+  /**
+   * @return the id of the subtask's supervisor
+   */
+  @JsonProperty
+  public String getSupervisorTaskId()
+  {
+    return supervisorTaskId;
+  }
 }

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

@@ -69,7 +69,6 @@ public class PartialDimensionCardinalityTask extends PerfectRollupWorkerTask
 
   private final int numAttempts;
   private final ParallelIndexIngestionSpec ingestionSchema;
-  private final String supervisorTaskId;
   private final String subtaskSpecId;
 
   private final ObjectMapper jsonMapper;
@@ -95,7 +94,8 @@ public class PartialDimensionCardinalityTask extends PerfectRollupWorkerTask
         taskResource,
         ingestionSchema.getDataSchema(),
         ingestionSchema.getTuningConfig(),
-        context
+        context,
+        supervisorTaskId
     );
 
     Preconditions.checkArgument(
@@ -107,7 +107,6 @@ public class PartialDimensionCardinalityTask extends PerfectRollupWorkerTask
     this.subtaskSpecId = subtaskSpecId;
     this.numAttempts = numAttempts;
     this.ingestionSchema = ingestionSchema;
-    this.supervisorTaskId = supervisorTaskId;
     this.jsonMapper = jsonMapper;
   }
 
@@ -123,12 +122,6 @@ private ParallelIndexIngestionSpec getIngestionSchema()
     return ingestionSchema;
   }
 
-  @JsonProperty
-  private String getSupervisorTaskId()
-  {
-    return supervisorTaskId;
-  }
-
   @JsonProperty
   @Override
   public String getSubtaskSpecId()
@@ -163,7 +156,7 @@ public boolean isReady(TaskActionClient taskActionClient) throws Exception
   {
     if (!getIngestionSchema().getDataSchema().getGranularitySpec().inputIntervals().isEmpty()) {
       return tryTimeChunkLock(
-          new SurrogateTaskActionClient(supervisorTaskId, taskActionClient),
+          new SurrogateTaskActionClient(getSupervisorTaskId(), taskActionClient),
           getIngestionSchema().getDataSchema().getGranularitySpec().inputIntervals()
       );
     } else {
@@ -274,7 +267,7 @@ private void sendReport(TaskToolbox toolbox, DimensionCardinalityReport report)
   {
     final ParallelIndexSupervisorTaskClient taskClient =
         toolbox.getSupervisorTaskClientProvider().build(
-            supervisorTaskId,
+            getSupervisorTaskId(),
             ingestionSchema.getTuningConfig().getChatHandlerTimeout(),
             ingestionSchema.getTuningConfig().getChatHandlerNumRetries()
         );

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

@@ -82,7 +82,6 @@ public class PartialDimensionDistributionTask extends PerfectRollupWorkerTask
 
   private final int numAttempts;
   private final ParallelIndexIngestionSpec ingestionSchema;
-  private final String supervisorTaskId;
   private final String subtaskSpecId;
 
   // For testing
@@ -136,7 +135,8 @@ public class PartialDimensionDistributionTask extends PerfectRollupWorkerTask
         taskResource,
         ingestionSchema.getDataSchema(),
         ingestionSchema.getTuningConfig(),
-        context
+        context,
+        supervisorTaskId
     );
 
     Preconditions.checkArgument(
@@ -148,7 +148,6 @@ public class PartialDimensionDistributionTask extends PerfectRollupWorkerTask
     this.subtaskSpecId = subtaskSpecId;
     this.numAttempts = numAttempts;
     this.ingestionSchema = ingestionSchema;
-    this.supervisorTaskId = supervisorTaskId;
     this.dedupInputRowFilterSupplier = dedupRowDimValueFilterSupplier;
   }
 
@@ -164,12 +163,6 @@ private ParallelIndexIngestionSpec getIngestionSchema()
     return ingestionSchema;
   }
 
-  @JsonProperty
-  private String getSupervisorTaskId()
-  {
-    return supervisorTaskId;
-  }
-
   @JsonProperty
   @Override
   public String getSubtaskSpecId()
@@ -204,7 +197,7 @@ public boolean isReady(TaskActionClient taskActionClient) throws Exception
   {
     if (!getIngestionSchema().getDataSchema().getGranularitySpec().inputIntervals().isEmpty()) {
       return tryTimeChunkLock(
-          new SurrogateTaskActionClient(supervisorTaskId, taskActionClient),
+          new SurrogateTaskActionClient(getSupervisorTaskId(), taskActionClient),
           getIngestionSchema().getDataSchema().getGranularitySpec().inputIntervals()
       );
     } else {
@@ -326,7 +319,7 @@ private Map<Interval, StringDistribution> determineDistribution(
   private void sendReport(TaskToolbox toolbox, DimensionDistributionReport report)
   {
     final ParallelIndexSupervisorTaskClient taskClient = toolbox.getSupervisorTaskClientProvider().build(
-        supervisorTaskId,
+        getSupervisorTaskId(),
         ingestionSchema.getTuningConfig().getChatHandlerTimeout(),
         ingestionSchema.getTuningConfig().getChatHandlerNumRetries()
     );