perf: Improve efficiency of embedded kill tasks

docs/operations/metrics.md

@@ -418,6 +418,7 @@ These metrics are emitted only if [auto-kill of unused segments](../data-managem
 |`segment/killed/metadataStore/count`|Number of segments permanently deleted from the metadata store.|`taskId`, `groupId`, `taskType`(=`kill`), `dataSource`|
 |`segment/killed/deepStorage/count`|Number of segments permanently deleted from the deep storage.|`taskId`, `groupId`, `taskType`(=`kill`), `dataSource`|
 |`segment/kill/unusedIntervals/count`|Number of intervals containing unused segments for a given datasource.|`dataSource`|
+|`segment/kill/eligibleSegments/count`|Number of unused segments in an interval that are eligible for kill.|`dataSource`, `interval`|
 |`segment/kill/skippedIntervals/count`|Number of intervals that were skipped for kill due to being already locked by another task.|`taskId`, `groupId`, `taskType`(=`kill`), `dataSource`|
 |`segment/kill/queueReset/time`|Time taken in milliseconds to reset the kill queue.||
 |`segment/kill/queueProcess/time`|Time taken in milliseconds to fully process the kill queue.||

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

@@ -39,8 +39,10 @@
 
 /**
  * Permanently deletes unused segments from the metadata store.
+ * {@link #perform(Task, TaskActionToolbox)} returns the number of segments
+ * deleted from the metadata store, but older versions return a null value.
  */
-public class SegmentNukeAction implements TaskAction<Void>
+public class SegmentNukeAction implements TaskAction<Integer>
 {
   private static final Logger log = new Logger(SegmentNukeAction.class);
 
@@ -61,19 +63,20 @@ public Set<DataSegment> getSegments()
   }
 
   @Override
-  public TypeReference<Void> getReturnTypeReference()
+  public TypeReference<Integer> getReturnTypeReference()
   {
     return new TypeReference<>() {};
   }
 
   @Override
-  public Void perform(Task task, TaskActionToolbox toolbox)
+  public Integer perform(Task task, TaskActionToolbox toolbox)
   {
     TaskLocks.checkLockCoversSegments(task, toolbox.getTaskLockbox(), segments);
 
+    final int numDeletedSegments;
     try {
       final Set<Interval> intervals = segments.stream().map(DataSegment::getInterval).collect(Collectors.toSet());
-      int numDeletedSegments = toolbox.getTaskLockbox().doInCriticalSection(
+      numDeletedSegments = toolbox.getTaskLockbox().doInCriticalSection(
           task,
           intervals,
           CriticalAction.<Integer>builder().onValidLocks(
@@ -106,7 +109,7 @@ public Void perform(Task task, TaskActionToolbox toolbox)
       toolbox.getEmitter().emit(metricBuilder.setMetric("segment/nuked/bytes", segment.getSize()));
     }
 
-    return null;
+    return numDeletedSegments;
   }
 
   @Override

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

@@ -33,6 +33,7 @@
 import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.metadata.ReplaceTaskLock;
+import org.apache.druid.query.QueryContexts;
 import org.apache.druid.timeline.DataSegment;
 import org.joda.time.DateTime;
 import org.joda.time.Interval;
@@ -170,6 +171,61 @@ public static TaskLockType determineLockTypeForAppend(
     }
   }
 
+  /**
+   * Checks if an append task with the given context should use concurrent locks.
+   */
+  public static boolean shouldUseConcurrentLocksForAppend(Map<String, Object> context, boolean defaultValue)
+  {
+    return shouldUseConcurrentLocks(context, defaultValue, TaskLockType.APPEND);
+  }
+
+  /**
+   * Checks if a replace task with the given context should use concurrent locks.
+   */
+  public static boolean shouldUseConcurrentLocksForReplace(Map<String, Object> context, boolean defaultValue)
+  {
+    return shouldUseConcurrentLocks(context, defaultValue, TaskLockType.REPLACE);
+  }
+
+  /**
+   * Checks if an append or replace task with the given context should use
+   * concurrent locks.
+   *
+   * @param context                 Context used by the task.
+   * @param defaultValue            Default behaviour in case the context does not
+   *                                explicitly specify any value for {@code useConcurrentLocks}
+   *                                or {@code taskLockType}.
+   * @param validConcurrentLockType Expected concurrent lock type (APPEND or REPLACE)
+   *                                for the given task.
+   */
+  private static boolean shouldUseConcurrentLocks(
+      Map<String, Object> context,
+      boolean defaultValue,
+      TaskLockType validConcurrentLockType
+  )
+  {
+    if (context == null) {
+      return defaultValue;
+    }
+    Boolean useConcurrentLocks = QueryContexts.getAsBoolean(
+        Tasks.USE_CONCURRENT_LOCKS,
+        context.get(Tasks.USE_CONCURRENT_LOCKS)
+    );
+    if (useConcurrentLocks != null) {
+      return useConcurrentLocks;
+    }
+    TaskLockType taskLockType = QueryContexts.getAsEnum(
+        Tasks.TASK_LOCK_TYPE,
+        context.get(Tasks.TASK_LOCK_TYPE),
+        TaskLockType.class
+    );
+    if (taskLockType != null) {
+      return taskLockType == validConcurrentLockType;
+    }
+
+    return defaultValue;
+  }
+
   /**
    * Finds locks of type {@link TaskLockType#REPLACE} for each of the given segments
    * that have an interval completely covering the interval of the respective segments.

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

@@ -47,6 +47,7 @@
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.QueryContexts;
 import org.apache.druid.server.coordination.BroadcastDatasourceLoadingSpec;
 import org.apache.druid.server.http.DataSegmentPlus;
 import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
@@ -87,10 +88,16 @@
  * <li> Filter the set of unreferenced segments using load specs from the set of used segments. </li>
  * <li> Kill the filtered set of segments from deep storage. </li>
  * </ol>
- * Note: When {@link Tasks#USE_CONCURRENT_LOCKS} is true, keep a large buffer
- * period before killing segments after they have been marked as unused.
- * Otherwise, there may be a potential data loss if a concurrent APPEND job
- * upgrades one of the segments that are being killed.
+ * <p>
+ * <b>Use concurrent locks:</b> A kill task using REPLACE locks (enabled by
+ * {@link Tasks#USE_CONCURRENT_LOCKS}) is affected by other tasks in the following
+ * manner:
+ * <ul>
+ * <li>Not affected by tasks using EXCLUSIVE, SHARED or another REPLACE lock
+ * since they are mutually exclusive with this REPLACE lock.</li>
+ * <li>Not affected by APPEND tasks since they can only upgrade used segments
+ * and do not modify unused segments.</li>
+ * </ul>
  */
 public class KillUnusedSegmentsTask extends AbstractFixedIntervalTask
 {
@@ -210,6 +217,7 @@ public TaskStatus runTask(TaskToolbox toolbox) throws Exception
   {
     // Track stats for reporting
     int numSegmentsKilled = 0;
+    int totalSegmentsDeletedFromMetadataStore = 0;
     int numBatchesProcessed = 0;
 
     // List unused segments
@@ -294,14 +302,23 @@ public TaskStatus runTask(TaskToolbox toolbox) throws Exception
       }
 
       // 3. Nuke all eligible unused segments
-      taskActionClient.submit(new SegmentNukeAction(unusedSegments));
-      emitMetric(toolbox.getEmitter(), TaskMetrics.SEGMENTS_DELETED_FROM_METADATA_STORE, unusedIdToSegmentPlus.size());
+      final Number nukeResult = taskActionClient.submit(new SegmentNukeAction(unusedSegments));
+
+      // Older versions of the Overlord return a null value from the SegmentNukeAction
+      final int numSegmentsDeletedFromMetadataStore =
+          nukeResult == null ? unusedSegments.size() : nukeResult.intValue();
+      emitMetric(
+          toolbox.getEmitter(),
+          TaskMetrics.SEGMENTS_DELETED_FROM_METADATA_STORE,
+          numSegmentsDeletedFromMetadataStore
+      );
 
       // 4. Delete deep store files only for segments which do not share load specs with other segments
       toolbox.getDataSegmentKiller().kill(segmentsToKillFromDeepStore);
       emitMetric(toolbox.getEmitter(), TaskMetrics.SEGMENTS_DELETED_FROM_DEEPSTORE, segmentsToKillFromDeepStore.size());
 
       numBatchesProcessed++;
+      totalSegmentsDeletedFromMetadataStore += numSegmentsDeletedFromMetadataStore;
       numSegmentsKilled += segmentsToKillFromDeepStore.size();
 
       logInfo("Processed [%d] batches for kill task[%s].", numBatchesProcessed, getId());
@@ -312,8 +329,10 @@ public TaskStatus runTask(TaskToolbox toolbox) throws Exception
     final String taskId = getId();
     logInfo(
         "Finished kill task[%s] for dataSource[%s] and interval[%s]."
-        + " Deleted total [%d] unused segments in [%d] batches.",
-        taskId, getDataSource(), getInterval(), numSegmentsKilled, numBatchesProcessed
+        + " Deleted [%d] unused segments from metadata store and files for"
+        + " [%d] segments from deep store in [%d] batches.",
+        taskId, getDataSource(), getInterval(),
+        totalSegmentsDeletedFromMetadataStore, numSegmentsKilled, numBatchesProcessed
     );
 
     final KillTaskReport.Stats stats =
@@ -512,13 +531,13 @@ public BroadcastDatasourceLoadingSpec getBroadcastDatasourceLoadingSpec()
   public boolean isReady(TaskActionClient taskActionClient) throws Exception
   {
     final boolean useConcurrentLocks = Boolean.TRUE.equals(
-        getContextValue(
+        QueryContexts.getAsBoolean(
             Tasks.USE_CONCURRENT_LOCKS,
-            Tasks.DEFAULT_USE_CONCURRENT_LOCKS
+            getContextValue(Tasks.USE_CONCURRENT_LOCKS)
         )
     );
 
-    TaskLockType actualLockType = determineLockType(useConcurrentLocks);
+    final TaskLockType actualLockType = determineLockType(useConcurrentLocks);
 
     final TaskLock lock = taskActionClient.submit(
         new TimeChunkLockTryAcquireAction(
@@ -539,9 +558,13 @@ private TaskLockType determineLockType(boolean useConcurrentLocks)
     if (useConcurrentLocks) {
       actualLockType = TaskLockType.REPLACE;
     } else {
-      actualLockType = getContextValue(Tasks.TASK_LOCK_TYPE, TaskLockType.EXCLUSIVE);
+      actualLockType = QueryContexts.getAsEnum(
+          Tasks.TASK_LOCK_TYPE,
+          getContextValue(Tasks.TASK_LOCK_TYPE),
+          TaskLockType.class,
+          TaskLockType.EXCLUSIVE
+      );
     }
     return actualLockType;
   }
-
 }