Add ability to limit the number of segments killed in kill task

docs/configuration/index.md

@@ -934,6 +934,7 @@ A sample Coordinator dynamic config JSON object is shown below:
   "replicantLifetime": 15,
   "replicationThrottleLimit": 10,
   "killDataSourceWhitelist": ["wikipedia", "testDatasource"],
+  "killTaskSlotRatio": 0.10,
   "decommissioningNodes": ["localhost:8182", "localhost:8282"],
   "decommissioningMaxPercentOfMaxSegmentsToMove": 70,
   "pauseCoordination": false,
@@ -955,6 +956,7 @@ Issuing a GET request at the same URL will return the spec that is currently in
 |`replicationThrottleLimit`|The maximum number of segment replicas that can be assigned to a historical tier in a single Coordinator run. This property prevents historicals from becoming overwhelmed when loading extra replicas of segments that are already available in the cluster.|500|
 |`balancerComputeThreads`|Thread pool size for computing moving cost of segments during segment balancing. Consider increasing this if you have a lot of segments and moving segments begins to stall.|1|
 |`killDataSourceWhitelist`|List of specific data sources for which kill tasks are sent if property `druid.coordinator.kill.on` is true. This can be a list of comma-separated data source names or a JSON array.|none|
+|`killTaskSlotRatio`| Ratio of total available task slots, including autoscaling if applicable that will be allowed for kill tasks. This limit only applies for kill tasks that are spawned automatically by the coordinator's auto kill duty, which is enabled when `druid.coordinator.kill.on` is true.| none - no limit                |
 |`killPendingSegmentsSkipList`|List of data sources for which pendingSegments are _NOT_ cleaned up if property `druid.coordinator.kill.pendingSegments.on` is true. This can be a list of comma-separated data sources or a JSON array.|none|
 |`maxSegmentsInNodeLoadingQueue`|The maximum number of segments allowed in the load queue of any given server. Use this parameter to load segments faster if, for example, the cluster contains slow-loading nodes or if there are too many segments to be replicated to a particular node (when faster loading is preferred to better segments distribution). The optimal value depends on the loading speed of segments, acceptable replication time and number of nodes. |500|
 |`useRoundRobinSegmentAssignment`|Boolean flag for whether segments should be assigned to historicals in a round robin fashion. When disabled, segment assignment is done using the chosen balancer strategy. When enabled, this can speed up segment assignments leaving balancing to move the segments to their optimal locations (based on the balancer strategy) lazily. |true|

docs/data-management/delete.md

@@ -95,9 +95,15 @@ The available grammar is:
     "id": <task_id>,
     "dataSource": <task_datasource>,
     "interval" : <all_unused_segments_in_this_interval_will_die!>,
-    "context": <task context>
+    "context": <task context>,
+    "maxSegmentsToKill": <the maximum number of segents to delete>
 }
 ```
 
 **WARNING:** The `kill` task permanently removes all information about the affected segments from the metadata store and
 deep storage. This operation cannot be undone.
+
+Note: If `maxSegmentsToKill` is not specified, all matched segments are deleted. If `maxSegmentsToKill` is less than
+the number of matching segments found, then only that number of matching segments will be deleted, but all matching
+segments will still be marked unused, if specified to be. In this case, any remaining unused segments can be deleted
+with a subsequent kill task issued, or via [Automated unused segment deletion](../operations/clean-metadata-store.md#segment-records-and-segments-in-deep-storage-kill-task)

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

@@ -23,6 +23,7 @@
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableSet;
 import org.apache.druid.client.indexing.ClientKillUnusedSegmentsTaskQuery;
 import org.apache.druid.indexer.TaskStatus;
@@ -41,6 +42,8 @@
 import org.joda.time.Interval;
 
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -58,17 +61,20 @@
  */
 public class KillUnusedSegmentsTask extends AbstractFixedIntervalTask
 {
+  public static final String TYPE = "kill";
   private static final Logger LOG = new Logger(KillUnusedSegmentsTask.class);
 
   private final boolean markAsUnused;
+  @Nullable private final Integer maxSegmentsToKill;
 
   @JsonCreator
   public KillUnusedSegmentsTask(
       @JsonProperty("id") String id,
       @JsonProperty("dataSource") String dataSource,
       @JsonProperty("interval") Interval interval,
       @JsonProperty("context") Map<String, Object> context,
-      @JsonProperty("markAsUnused") Boolean markAsUnused
+      @JsonProperty("markAsUnused") Boolean markAsUnused,
+      @JsonProperty("maxSegmentsToKill") @Nullable Integer maxSegmentsToKill
   )
   {
     super(
@@ -77,7 +83,12 @@ public KillUnusedSegmentsTask(
         interval,
         context
     );
+    if (null != maxSegmentsToKill) {
+      Preconditions.checkArgument(maxSegmentsToKill > 0, "maxSegmentsToKill must be > 0");
+    }
     this.markAsUnused = markAsUnused != null && markAsUnused;
+    this.maxSegmentsToKill = maxSegmentsToKill;
+
   }
 
   @JsonProperty
@@ -87,6 +98,13 @@ public boolean isMarkAsUnused()
     return markAsUnused;
   }
 
+  @Nullable
+  @JsonProperty
+  public Integer getMaxSegmentsToKill()
+  {
+    return maxSegmentsToKill;
+  }
+
   @Override
   public String getType()
   {
@@ -114,7 +132,7 @@ public TaskStatus runTask(TaskToolbox toolbox) throws Exception
     }
 
     // List unused segments
-    final List<DataSegment> unusedSegments = toolbox
+    List<DataSegment> unusedSegments = toolbox
         .getTaskActionClient()
         .submit(new RetrieveUnusedSegmentsAction(getDataSource(), getInterval()));
 
@@ -128,6 +146,9 @@ public TaskStatus runTask(TaskToolbox toolbox) throws Exception
     }
 
     // Kill segments
+    unusedSegments = maxSegmentsToKill == null
+        ? unusedSegments
+        : unusedSegments.subList(0, Math.min(maxSegmentsToKill, unusedSegments.size()));
     toolbox.getTaskActionClient().submit(new SegmentNukeAction(new HashSet<>(unusedSegments)));
     for (DataSegment segment : unusedSegments) {
       toolbox.getDataSegmentKiller().kill(segment);

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

@@ -61,7 +61,7 @@
  */
 @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type")
 @JsonSubTypes(value = {
-    @Type(name = "kill", value = KillUnusedSegmentsTask.class),
+    @Type(name = KillUnusedSegmentsTask.TYPE, value = KillUnusedSegmentsTask.class),
     @Type(name = "move", value = MoveTask.class),
     @Type(name = "archive", value = ArchiveTask.class),
     @Type(name = "restore", value = RestoreTask.class),

indexing-service/src/test/java/org/apache/druid/indexing/common/task/ClientKillUnusedSegmentsTaskQuerySerdeTest.java

@@ -51,14 +51,16 @@ public void testClientKillUnusedSegmentsTaskQueryToKillUnusedSegmentsTask() thro
         "killTaskId",
         "datasource",
         Intervals.of("2020-01-01/P1D"),
-        true
+        true,
+        5
     );
     final byte[] json = objectMapper.writeValueAsBytes(taskQuery);
     final KillUnusedSegmentsTask fromJson = (KillUnusedSegmentsTask) objectMapper.readValue(json, Task.class);
     Assert.assertEquals(taskQuery.getId(), fromJson.getId());
     Assert.assertEquals(taskQuery.getDataSource(), fromJson.getDataSource());
     Assert.assertEquals(taskQuery.getInterval(), fromJson.getInterval());
     Assert.assertEquals(taskQuery.getMarkAsUnused(), fromJson.isMarkAsUnused());
+    Assert.assertEquals(taskQuery.getMaxSegmentsToKill(), fromJson.getMaxSegmentsToKill());
   }
 
   @Test
@@ -69,7 +71,8 @@ public void testKillUnusedSegmentsTaskToClientKillUnusedSegmentsTaskQuery() thro
         "datasource",
         Intervals.of("2020-01-01/P1D"),
         null,
-        true
+        true,
+        null
     );
     final byte[] json = objectMapper.writeValueAsBytes(task);
     final ClientKillUnusedSegmentsTaskQuery taskQuery = (ClientKillUnusedSegmentsTaskQuery) objectMapper.readValue(
@@ -80,5 +83,6 @@ public void testKillUnusedSegmentsTaskToClientKillUnusedSegmentsTaskQuery() thro
     Assert.assertEquals(task.getDataSource(), taskQuery.getDataSource());
     Assert.assertEquals(task.getInterval(), taskQuery.getInterval());
     Assert.assertEquals(task.isMarkAsUnused(), taskQuery.getMarkAsUnused());
+    Assert.assertNull(task.getMaxSegmentsToKill());
   }
 }

indexing-service/src/test/java/org/apache/druid/indexing/common/task/KillUnusedSegmentsTaskTest.java

@@ -79,7 +79,8 @@ public void testKill() throws Exception
             DATA_SOURCE,
             Intervals.of("2019-03-01/2019-04-01"),
             null,
-            false
+            false,
+            null
         );
 
     Assert.assertEquals(TaskState.SUCCESS, taskRunner.run(task).get().getStatusCode());
@@ -124,7 +125,8 @@ public void testKillWithMarkUnused() throws Exception
             DATA_SOURCE,
             Intervals.of("2019-03-01/2019-04-01"),
             null,
-            true
+            true,
+            null
         );
 
     Assert.assertEquals(TaskState.SUCCESS, taskRunner.run(task).get().getStatusCode());
@@ -151,7 +153,8 @@ public void testGetInputSourceResources()
             DATA_SOURCE,
             Intervals.of("2019-03-01/2019-04-01"),
             null,
-            true
+            true,
+            null
         );
     Assert.assertTrue(task.getInputSourceResources().isEmpty());
   }

indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskLifecycleTest.java

@@ -947,7 +947,8 @@ public DataSegment apply(String input)
             "test_kill_task",
             Intervals.of("2011-04-01/P4D"),
             null,
-            false
+            false,
+            null
         );
 
     final TaskStatus status = runTask(killUnusedSegmentsTask);
@@ -967,6 +968,103 @@ public DataSegment apply(String input)
     }
   }
 
+  @Test
+  public void testKillUnusedSegmentsTaskWithMaxSegmentsToKill() throws Exception
+  {
+    final File tmpSegmentDir = temporaryFolder.newFolder();
+
+    List<DataSegment> expectedUnusedSegments = Lists.transform(
+        ImmutableList.of(
+            "2011-04-01/2011-04-02",
+            "2011-04-02/2011-04-03",
+            "2011-04-04/2011-04-05"
+        ), new Function<String, DataSegment>()
+        {
+          @Override
+          public DataSegment apply(String input)
+          {
+            final Interval interval = Intervals.of(input);
+            try {
+              return DataSegment.builder()
+                  .dataSource("test_kill_task")
+                  .interval(interval)
+                  .loadSpec(
+                      ImmutableMap.of(
+                          "type",
+                          "local",
+                          "path",
+                          tmpSegmentDir.getCanonicalPath()
+                          + "/druid/localStorage/wikipedia/"
+                          + interval.getStart()
+                          + "-"
+                          + interval.getEnd()
+                          + "/"
+                          + "2011-04-6T16:52:46.119-05:00"
+                          + "/0/index.zip"
+                      )
+                  )
+                  .version("2011-04-6T16:52:46.119-05:00")
+                  .dimensions(ImmutableList.of())
+                  .metrics(ImmutableList.of())
+                  .shardSpec(NoneShardSpec.instance())
+                  .binaryVersion(9)
+                  .size(0)
+                  .build();
+            }
+            catch (IOException e) {
+              throw new ISE(e, "Error creating segments");
+            }
+          }
+        }
+    );
+
+    mdc.setUnusedSegments(expectedUnusedSegments);
+
+    // manually create local segments files
+    List<File> segmentFiles = new ArrayList<>();
+    for (DataSegment segment : mdc.retrieveUnusedSegmentsForInterval("test_kill_task", Intervals.of("2011-04-01/P4D"))) {
+      File file = new File((String) segment.getLoadSpec().get("path"));
+      FileUtils.mkdirp(file.getParentFile());
+      Files.write(file.toPath(), ByteArrays.EMPTY_ARRAY);
+      segmentFiles.add(file);
+    }
+
+    final int maxSegmentsToKill = 2;
+    final Task killUnusedSegmentsTask =
+        new KillUnusedSegmentsTask(
+            null,
+            "test_kill_task",
+            Intervals.of("2011-04-01/P4D"),
+            null,
+            false,
+            maxSegmentsToKill
+        );
+
+    final TaskStatus status = runTask(killUnusedSegmentsTask);
+    Assert.assertEquals(taskLocation, status.getLocation());
+    Assert.assertEquals("merged statusCode", TaskState.SUCCESS, status.getStatusCode());
+    Assert.assertEquals("num segments published", 0, mdc.getPublished().size());
+    Assert.assertEquals("num segments nuked", maxSegmentsToKill, mdc.getNuked().size());
+    Assert.assertTrue(
+        "expected unused segments get killed",
+        expectedUnusedSegments.containsAll(mdc.getNuked())
+    );
+
+    int expectedNumOfSegmentsRemaining = segmentFiles.size() - maxSegmentsToKill;
+    int actualNumOfSegmentsRemaining = 0;
+    for (File file : segmentFiles) {
+      if (file.exists()) {
+        actualNumOfSegmentsRemaining++;
+      }
+    }
+
+    Assert.assertEquals(
+        "Expected of segments deleted did not match expectations",
+        expectedNumOfSegmentsRemaining,
+        actualNumOfSegmentsRemaining
+    );
+  }
+
   @Test
   public void testRealtimeishTask() throws Exception
   {

server/src/main/java/org/apache/druid/client/indexing/ClientKillUnusedSegmentsTaskQuery.java

@@ -24,6 +24,8 @@
 import com.google.common.base.Preconditions;
 import org.joda.time.Interval;
 
+import javax.annotation.Nullable;
+
 import java.util.Objects;
 
 /**
@@ -39,19 +41,25 @@ public class ClientKillUnusedSegmentsTaskQuery implements ClientTaskQuery
   private final String dataSource;
   private final Interval interval;
   private final Boolean markAsUnused;
+  @Nullable private final Integer maxSegmentsToKill;
 
   @JsonCreator
   public ClientKillUnusedSegmentsTaskQuery(
       @JsonProperty("id") String id,
       @JsonProperty("dataSource") String dataSource,
       @JsonProperty("interval") Interval interval,
-      @JsonProperty("markAsUnused") Boolean markAsUnused
+      @JsonProperty("markAsUnused") Boolean markAsUnused,
+      @JsonProperty("maxSegmentsToKill") Integer maxSegmentsToKill
   )
   {
     this.id = Preconditions.checkNotNull(id, "id");
     this.dataSource = dataSource;
     this.interval = interval;
     this.markAsUnused = markAsUnused;
+    if (null != maxSegmentsToKill) {
+      Preconditions.checkArgument(maxSegmentsToKill > 0, "maxSegmentsToKill must be > 0");
+    }
+    this.maxSegmentsToKill = maxSegmentsToKill;
   }
 
   @JsonProperty
@@ -87,6 +95,14 @@ public Boolean getMarkAsUnused()
     return markAsUnused;
   }
 
+  @JsonProperty
+  @Nullable
+  public Integer getMaxSegmentsToKill()
+  {
+    return maxSegmentsToKill;
+  }
+
+
   @Override
   public boolean equals(Object o)
   {
@@ -100,12 +116,13 @@ public boolean equals(Object o)
     return Objects.equals(id, that.id)
            && Objects.equals(dataSource, that.dataSource)
            && Objects.equals(interval, that.interval)
-           && Objects.equals(markAsUnused, that.markAsUnused);
+           && Objects.equals(markAsUnused, that.markAsUnused)
+           && Objects.equals(maxSegmentsToKill, that.maxSegmentsToKill);
   }
 
   @Override
   public int hashCode()
   {
-    return Objects.hash(id, dataSource, interval, markAsUnused);
+    return Objects.hash(id, dataSource, interval, markAsUnused, maxSegmentsToKill);
   }
 }