feat: Add new TaskLockType KILL for embedded kill tasks

benchmarks/src/test/java/org/apache/druid/benchmark/indexing/SqlSegmentsMetadataQueryBenchmark.java

@@ -25,6 +25,7 @@
 import org.apache.druid.java.util.common.parsers.CloseableIterator;
 import org.apache.druid.metadata.IndexerSqlMetadataStorageCoordinatorTestBase;
 import org.apache.druid.metadata.MetadataStorageTablesConfig;
+import org.apache.druid.metadata.SegmentsMetadataManagerConfig;
 import org.apache.druid.metadata.SqlSegmentsMetadataQuery;
 import org.apache.druid.metadata.TestDerbyConnector;
 import org.apache.druid.segment.TestDataSource;
@@ -137,7 +138,13 @@ private <T> Set<T> readAsSet(Function<SqlSegmentsMetadataQuery, CloseableIterato
 
     return derbyConnector.inReadOnlyTransaction((handle, status) -> {
       final SqlSegmentsMetadataQuery query =
-          SqlSegmentsMetadataQuery.forHandle(handle, derbyConnector, tablesConfig, TestHelper.JSON_MAPPER);
+          SqlSegmentsMetadataQuery.forHandle(
+              handle,
+              derbyConnector,
+              tablesConfig,
+              new SegmentsMetadataManagerConfig(null, null, null),
+              TestHelper.JSON_MAPPER
+          );
 
       try (CloseableIterator<T> iterator = iterableReader.apply(query)) {
         return ImmutableSet.copyOf(iterator);

docs/configuration/index.md

@@ -1040,7 +1040,9 @@ None of the configs that apply to [auto-kill performed by the Coordinator](../da
 |Property|Description|Default|
 |--------|-----------|-------|
 |`druid.manager.segments.killUnused.enabled`|Boolean flag to enable auto-kill of eligible unused segments on the Overlord. This feature can be used only when [segment metadata caching](#segment-metadata-cache) is enabled on the Overlord and MUST NOT be enabled if `druid.coordinator.kill.on` is already set to `true` on the Coordinator.|`true`|
-|`druid.manager.segments.killUnused.bufferPeriod`|Period after which a segment marked as unused becomes eligible for auto-kill on the Overlord. This config is effective only if `druid.manager.segments.killUnused.enabled` is set to `true`.|`P30D` (30 days)|
+|`druid.manager.segments.killUnused.bufferPeriod`|ISO8601 Period after which a segment marked as unused cannot be marked as used anymore and becomes eligible for auto-kill on the Overlord. This config is effective only if `druid.manager.segments.killUnused.enabled` is set to `true`.|`P30D` (30 days)|
+|`druid.manager.segments.killUnused.dutyPeriod`|ISO8601 Period defining the frequency at which unused segments should be added to the kill queue. If the queue already has some unused segments, new segments are not added. This config is effective only if `druid.manager.segments.killUnused.enabled` is set to `true`.|`PT1H` (1 hour)|
+|`druid.manager.segments.killUnused.maxSegmentsToKill`|Maximum number of unused segments that can be added to the kill queue in a single cycle. A very large value for this config may cause the metadata store to slow down while fetching unused segments to kill, whereas a very small value would cause the kill operation to be ineffective as it wouldn't be able to catch up with the number of old unused segments in the cluster. This config is effective only if `druid.manager.segments.killUnused.enabled` is set to `true`.|200,000|
 
 #### Overlord dynamic configuration
 

docs/data-management/delete.md

@@ -143,7 +143,7 @@ These embedded tasks offer several advantages over auto-kill performed by the Co
 - run on the Overlord and do not take up task slots.
 - finish faster as they save on the overhead of launching a task process.
 - kill a small number of segments per task, to ensure that locks on an interval are not held for too long.
-- skip locked intervals to avoid head-of-line blocking in kill tasks.
+- use a dedicated `KILL` lock type which allows killing of segments in the background while an ingestion proceeds for the same interval.
 - require little to no configuration.
 - can keep up with a large number of unused segments in the cluster.
 - take advantage of the segment metadata cache on the Overlord.

embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/KafkaClusterMetricsTest.java

@@ -30,6 +30,7 @@
 import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorSpec;
 import org.apache.druid.indexing.overlord.Segments;
 import org.apache.druid.java.util.common.HumanReadableBytes;
+import org.apache.druid.metadata.UnusedSegmentKillerConfig;
 import org.apache.druid.query.DruidMetrics;
 import org.apache.druid.rpc.UpdateResponse;
 import org.apache.druid.rpc.indexing.OverlordClient;
@@ -105,14 +106,16 @@ public void stop()
       }
     };
 
+    final Period killBufferPeriod = Period.millis(100).minus(UnusedSegmentKillerConfig.GRACE_PERIOD);
+
     indexer.setServerMemory(1_000_000_000L)
            .addProperty("druid.segment.handoff.pollDuration", "PT0.1s")
            .addProperty("druid.worker.capacity", "10");
     overlord.addProperty("druid.indexer.task.default.context", "{\"useConcurrentLocks\": true}")
             .addProperty("druid.manager.segments.useIncrementalCache", "ifSynced")
             .addProperty("druid.manager.segments.pollDuration", "PT0.1s")
             .addProperty("druid.manager.segments.killUnused.enabled", "true")
-            .addProperty("druid.manager.segments.killUnused.bufferPeriod", "PT0.1s")
+            .addProperty("druid.manager.segments.killUnused.bufferPeriod", killBufferPeriod.toString())
             .addProperty("druid.manager.segments.killUnused.dutyPeriod", "PT1s");
     coordinator.addProperty("druid.manager.segments.useIncrementalCache", "ifSynced");
     cluster.addExtension(KafkaIndexTaskModule.class)
@@ -203,7 +206,7 @@ public void test_ingest10kRows_ofSelfClusterMetrics_andVerifyValues()
   @MethodSource("getCompactionSupervisorTestParams")
   @ParameterizedTest(name = "engine={0}, policy={1}")
   @Timeout(120)
-  public void test_ingestClusterMetrics_withConcurrentCompactionSupervisor_andSkipKillOfUnusedSegments(
+  public void test_ingestClusterMetrics_withConcurrentCompactionSupervisor_andKillUnusedSegments(
       CompactionEngine engine,
       CompactionCandidateSearchPolicy policy
   )
@@ -303,9 +306,16 @@ public void test_ingestClusterMetrics_withConcurrentCompactionSupervisor_andSkip
         agg -> agg.hasSumAtLeast(1)
     );
 
-    // Verify that the segments are skipped since the interval is still being appended to
+    // Verify that some unused segments have been killed from metadata store
+    overlord.latchableEmitter().waitForEventAggregate(
+        event -> event.hasMetricName("segment/killed/metadataStore/count")
+                      .hasDimension(DruidMetrics.DATASOURCE, dataSource),
+        agg -> agg.hasSumAtLeast(1)
+    );
+
+    // Verify that some segments were not deleted from deep store due to being upgraded
     overlord.latchableEmitter().waitForEventAggregate(
-        event -> event.hasMetricName("segment/kill/skippedIntervals/count")
+        event -> event.hasMetricName("segment/kill/deepStorageSkipped/count")
                       .hasDimension(DruidMetrics.DATASOURCE, dataSource),
         agg -> agg.hasSumAtLeast(1)
     );

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

@@ -24,6 +24,7 @@
  * 2) An appending task may use only EXCLUSIVE, SHARED or APPEND locks
  * 3) A replacing task may use only EXCLUSIVE or REPLACE locks
  * 4) REPLACE and APPEND locks can only be used with timechunk locking
+ * 5) Only kill tasks (task type "kill") may acquire a KILL lock
  */
 public enum TaskLockType
 {
@@ -49,5 +50,12 @@ public enum TaskLockType
    * and with at most one REPLACE lock whose interval encloses that of the APPEND lock.
    * They are incompatible with all other active locks.
    */
-  APPEND
+  APPEND,
+  /**
+   * There can be at most one active KILL lock for a given interval.
+   * It can co-exist with any other active lock type (EXCLUSIVE, SHARED, REPLACE, APPEND),
+   * but not with another KILL lock whose interval overlaps.
+   * Only tasks of type "kill" may acquire a KILL lock.
+   */
+  KILL
 }

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

@@ -315,7 +315,16 @@ public TaskStatus runTask(TaskToolbox toolbox) throws Exception
 
       // 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());
+
+      final int numSegmentsDeletedFromDeepStore = segmentsToKillFromDeepStore.size();
+      emitMetric(toolbox.getEmitter(), TaskMetrics.SEGMENTS_DELETED_FROM_DEEPSTORE, numSegmentsDeletedFromDeepStore);
+      if (numSegmentsDeletedFromMetadataStore > numSegmentsDeletedFromDeepStore) {
+        emitMetric(
+            toolbox.getEmitter(),
+            TaskMetrics.SEGMENTS_SKIPPED_DEEPSTORE_KILL,
+            numSegmentsDeletedFromMetadataStore - numSegmentsDeletedFromDeepStore
+        );
+      }
 
       numBatchesProcessed++;
       totalSegmentsDeletedFromMetadataStore += numSegmentsDeletedFromMetadataStore;

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

@@ -33,5 +33,5 @@ private TaskMetrics()
 
   public static final String SEGMENTS_DELETED_FROM_METADATA_STORE = "segment/killed/metadataStore/count";
   public static final String SEGMENTS_DELETED_FROM_DEEPSTORE = "segment/killed/deepStorage/count";
-  public static final String FILES_DELETED_FROM_DEEPSTORE = "segment/killed/deepStorageFile/count";
+  public static final String SEGMENTS_SKIPPED_DEEPSTORE_KILL = "segment/kill/deepStorageSkipped/count";
 }

indexing-service/src/main/java/org/apache/druid/indexing/overlord/GlobalTaskLockbox.java

@@ -80,6 +80,11 @@ public GlobalTaskLockbox(
    * Syncs the current in-memory state with the {@link TaskStorage}.
    * This method should be called only from {@link TaskQueue#start()}.
    * If the sync fails, no other operation can be performed on this lockbox.
+   * <p>
+   * The sync does not restore a lock, if it is associated with a dummy task ID
+   * that is not persisted in the task storage (e.g. embedded kill tasks).
+   * This is okay since the sync happens only when Overlord becomes leader, at
+   * which point there wouldn't be any embedded (kill) tasks running anyway.
    *
    * @return SyncResult which needs to be processed by the caller
    */

indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskLockbox.java

@@ -39,6 +39,7 @@
 import org.apache.druid.indexing.common.task.PendingSegmentAllocatingTask;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.indexing.common.task.Tasks;
+import org.apache.druid.indexing.overlord.duty.UnusedSegmentsKiller;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.java.util.common.Pair;
@@ -571,6 +572,14 @@ private TaskLockPosse createOrFindLockPosse(LockRequest request, Task task, bool
         throw new ISE("Unable to grant LockPosse to inactive Task [%s]", task.getId());
       }
 
+      if (request.getType() == TaskLockType.KILL
+          && !UnusedSegmentsKiller.TASK_TYPE_EMBEDDED_KILL.equals(task.getType())) {
+        throw new ISE(
+            "Task[%s] of type[%s] cannot acquire a KILL lock. Only tasks of type[%s] may use KILL locks.",
+            task.getId(), task.getType(), UnusedSegmentsKiller.TASK_TYPE_EMBEDDED_KILL
+        );
+      }
+
       final TaskLockPosse posseToUse;
       final List<TaskLockPosse> foundPosses = findLockPossesOverlapsInterval(
           request.getInterval()
@@ -587,7 +596,7 @@ private TaskLockPosse createOrFindLockPosse(LockRequest request, Task task, bool
         final List<TaskLockPosse> reusablePosses = foundPosses
             .stream()
             .filter(posse -> posse.reusableFor(request))
-            .collect(Collectors.toList());
+            .toList();
 
         if (reusablePosses.isEmpty()) {
           // case 1) this task doesn't have any lock, but others do
@@ -835,8 +844,9 @@ public void revokeLock(String taskId, TaskLock lock)
         throw new ISE("Cannot revoke lock for inactive task[%s]", taskId);
       }
 
+      // Embedded kill tasks can hold locks but are not persisted in TaskStorage
       final Task task = taskStorage.getTask(taskId).orNull();
-      if (task == null) {
+      if (task == null && lock.getType() != TaskLockType.KILL) {
         throw new ISE("Cannot revoke lock for unknown task[%s]", taskId);
       }
 
@@ -963,26 +973,29 @@ public List<Interval> getLockedIntervals(List<LockFilterPolicy> lockFilterPolici
             }
 
             final int priority = lockFilter.getPriority();
-            final boolean isReplaceLock = TaskLockType.REPLACE.name().equals(
-                lockFilter.getContext().getOrDefault(
-                    Tasks.TASK_LOCK_TYPE,
-                    Tasks.DEFAULT_TASK_LOCK_TYPE
-                )
+            final TaskLockType lockType = QueryContexts.getAsEnum(
+                Tasks.TASK_LOCK_TYPE,
+                lockFilter.getContext().get(Tasks.TASK_LOCK_TYPE),
+                TaskLockType.class,
+                Tasks.DEFAULT_TASK_LOCK_TYPE
             );
-            final boolean isUsingConcurrentLocks = Boolean.TRUE.equals(
-                lockFilter.getContext().getOrDefault(
-                    Tasks.USE_CONCURRENT_LOCKS,
-                    Tasks.DEFAULT_USE_CONCURRENT_LOCKS
-                )
+            final boolean isReplaceLock = TaskLockType.REPLACE == lockType;
+            final boolean isUsingConcurrentLocks = QueryContexts.getAsBoolean(
+                Tasks.USE_CONCURRENT_LOCKS,
+                lockFilter.getContext().get(Tasks.USE_CONCURRENT_LOCKS),
+                Tasks.DEFAULT_USE_CONCURRENT_LOCKS
             );
             final boolean ignoreAppendLocks = isUsingConcurrentLocks || isReplaceLock;
+            final boolean ignoreKillLocks = lockType != TaskLockType.KILL;
 
             running.forEach(
                 (startTime, startTimeLocks) -> startTimeLocks.forEach(
                     (interval, taskLockPosses) -> taskLockPosses.forEach(
                         taskLockPosse -> {
                           if (taskLockPosse.getTaskLock().isRevoked()) {
                             // do nothing
+                          } else if (ignoreKillLocks && TaskLockType.KILL.equals(taskLockPosse.getTaskLock().getType())) {
+                            // do nothing
                           } else if (ignoreAppendLocks
                                      && TaskLockType.APPEND.equals(taskLockPosse.getTaskLock().getType())) {
                             // do nothing
@@ -1363,7 +1376,7 @@ Optional<TaskLockPosse> getOnlyTaskLockPosseContainingInterval(Task task, Interv
       final List<TaskLockPosse> filteredPosses = findLockPossesContainingInterval(interval)
           .stream()
           .filter(lockPosse -> lockPosse.containsTask(task))
-          .collect(Collectors.toList());
+          .toList();
 
       if (filteredPosses.isEmpty()) {
         throw new ISE("Cannot find any lock for task[%s] and interval[%s]", task.getId(), interval);
@@ -1422,6 +1435,8 @@ private boolean canLockCoexist(List<TaskLockPosse> conflictPosses, LockRequest r
         return canSharedLockCoexist(conflictPosses);
       case EXCLUSIVE:
         return canExclusiveLockCoexist(conflictPosses);
+      case KILL:
+        return canKillLockCoexist(conflictPosses);
       default:
         throw new UOE("Unsupported lock type: " + request.getType());
     }
@@ -1430,7 +1445,7 @@ private boolean canLockCoexist(List<TaskLockPosse> conflictPosses, LockRequest r
   /**
    * Check if an APPEND lock can coexist with a given set of conflicting posses.
    * An APPEND lock can coexist with any number of other APPEND locks
-   *    OR with at most one REPLACE lock over an interval which encloes this request.
+   *    OR with at most one REPLACE lock over an interval which encloses this request.
    * @param conflictPosses conflicting lock posses
    * @param appendRequest append lock request
    * @return true iff append lock can coexist with all its conflicting locks
@@ -1477,6 +1492,8 @@ private boolean canReplaceLockCoexist(List<TaskLockPosse> conflictPosses, LockRe
           || posse.getTaskLock().getType().equals(TaskLockType.REPLACE)) {
         return false;
       }
+      // REPLACE lock can coexist with an APPEND lock only if the append interval
+      // is fully contained within the replace interval
       if (posse.getTaskLock().getType().equals(TaskLockType.APPEND)
           && !replaceLock.getInterval().contains(posse.getTaskLock().getInterval())) {
         return false;
@@ -1487,7 +1504,8 @@ private boolean canReplaceLockCoexist(List<TaskLockPosse> conflictPosses, LockRe
 
   /**
    * Check if a SHARED lock can coexist with a given set of conflicting posses.
-   * A SHARED lock can coexist with any number of other active SHARED locks
+   * A SHARED lock can coexist with any number of other active SHARED locks or
+   * a KILL lock.
    * @param conflictPosses conflicting lock posses
    * @return true iff shared lock can coexist with all its conflicting locks
    */
@@ -1508,7 +1526,8 @@ private boolean canSharedLockCoexist(List<TaskLockPosse> conflictPosses)
 
   /**
    * Check if an EXCLUSIVE lock can coexist with a given set of conflicting posses.
-   * An EXCLUSIVE lock cannot coexist with any other overlapping active locks
+   * An EXCLUSIVE lock cannot coexist with any other overlapping active locks,
+   * except a KILL lock.
    * @param conflictPosses conflicting lock posses
    * @return true iff the exclusive lock can coexist with all its conflicting locks
    */
@@ -1518,11 +1537,36 @@ private boolean canExclusiveLockCoexist(List<TaskLockPosse> conflictPosses)
       if (posse.getTaskLock().isRevoked()) {
         continue;
       }
-      return false;
+      if (!isLockTypeKill(posse)) {
+        return false;
+      }
     }
     return true;
   }
 
+  /**
+   * Check if a KILL lock can coexist with a given set of conflicting posses.
+   * A KILL lock can coexist with any other lock type but not with another KILL lock.
+   * @param conflictPosses conflicting lock posses
+   * @return true iff the kill lock can coexist with all its conflicting locks
+   */
+  private boolean canKillLockCoexist(List<TaskLockPosse> conflictPosses)
+  {
+    for (TaskLockPosse posse : conflictPosses) {
+      if (posse.getTaskLock().isRevoked()) {
+        continue;
+      }
+      if (isLockTypeKill(posse)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  private boolean isLockTypeKill(TaskLockPosse posse)
+  {
+    return posse.getTaskLock().getType().equals(TaskLockType.KILL);
+  }
 
   /**
    * Verify if every incompatible active lock is revokable. If yes, revoke all of them.
@@ -1545,7 +1589,9 @@ private boolean revokeAllIncompatibleActiveLocksIfPossible(
     final List<TaskLockPosse> possesToRevoke = new ArrayList<>();
 
     for (TaskLockPosse posse : conflictPosses) {
-      if (posse.getTaskLock().isRevoked()) {
+      // No need to revoke an already revoked lock or a KILL lock (unless the new LockRequest is also a KILL)
+      if (posse.getTaskLock().isRevoked()
+          || (isLockTypeKill(posse) && type != TaskLockType.KILL)) {
         continue;
       }
       switch (type) {
@@ -1582,6 +1628,15 @@ private boolean revokeAllIncompatibleActiveLocksIfPossible(
             possesToRevoke.add(posse);
           }
           break;
+        case KILL:
+          // KILL locks are incompatible only with other KILL locks
+          if (isLockTypeKill(posse)) {
+            if (posse.getTaskLock().getNonNullPriority() >= priority) {
+              return false;
+            }
+            possesToRevoke.add(posse);
+          }
+          break;
         default:
           throw new UOE("Unsupported lock type: " + type);
       }
@@ -1686,6 +1741,7 @@ boolean reusableFor(LockRequest request)
           case REPLACE:
           case APPEND:
           case SHARED:
+          case KILL:
             if (request instanceof TimeChunkLockRequest) {
               return taskLock.getInterval().contains(request.getInterval())
                      && taskLock.getGroupId().equals(request.getGroupId());