Maintain TaskLockbox at datasource level for higher concurrency

extensions-core/testing-tools/src/main/java/org/apache/druid/guice/ClusterTestingModule.java

@@ -31,7 +31,7 @@
 import org.apache.druid.guice.annotations.Self;
 import org.apache.druid.indexing.common.actions.RemoteTaskActionClientFactory;
 import org.apache.druid.indexing.common.task.Task;
-import org.apache.druid.indexing.overlord.TaskLockbox;
+import org.apache.druid.indexing.overlord.GlobalTaskLockbox;
 import org.apache.druid.initialization.DruidModule;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.rpc.indexing.OverlordClient;
@@ -106,7 +106,7 @@ private void bindDependenciesForClusterTestingMode(Binder binder)
     } else if (roles.contains(NodeRole.OVERLORD)) {
       // If this is the Overlord, bind a faulty storage coordinator
       log.warn("Running Overlord in cluster testing mode.");
-      binder.bind(TaskLockbox.class)
+      binder.bind(GlobalTaskLockbox.class)
             .to(FaultyTaskLockbox.class)
             .in(LazySingleton.class);
     }

extensions-core/testing-tools/src/main/java/org/apache/druid/testing/cluster/overlord/FaultyTaskLockbox.java

@@ -22,6 +22,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.inject.Inject;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.indexing.overlord.GlobalTaskLockbox;
 import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator;
 import org.apache.druid.indexing.overlord.TaskLockbox;
 import org.apache.druid.indexing.overlord.TaskStorage;
@@ -34,7 +35,7 @@
  * <li>Skip cleanup of pending segments</li>
  * </ul>
  */
-public class FaultyTaskLockbox extends TaskLockbox
+public class FaultyTaskLockbox extends GlobalTaskLockbox
 {
   private static final Logger log = new Logger(FaultyTaskLockbox.class);
 

extensions-core/testing-tools/src/test/java/org/apache/druid/guice/ClusterTestingModuleTest.java

@@ -40,7 +40,7 @@
 import org.apache.druid.indexing.common.task.batch.parallel.ParallelIndexSupervisorTask;
 import org.apache.druid.indexing.common.task.batch.parallel.ParallelIndexTuningConfig;
 import org.apache.druid.indexing.input.DruidInputSource;
-import org.apache.druid.indexing.overlord.TaskLockbox;
+import org.apache.druid.indexing.overlord.GlobalTaskLockbox;
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.query.filter.TrueDimFilter;
 import org.apache.druid.rpc.indexing.OverlordClient;
@@ -244,7 +244,7 @@ public void test_overlordService_hasFaultyStorageCoordinator_ifTestingIsEnabled(
 
       final Injector overlordInjector = overlord.makeInjector(Set.of(NodeRole.OVERLORD));
 
-      TaskLockbox taskLockbox = overlordInjector.getInstance(TaskLockbox.class);
+      GlobalTaskLockbox taskLockbox = overlordInjector.getInstance(GlobalTaskLockbox.class);
       Assert.assertTrue(taskLockbox instanceof FaultyTaskLockbox);
     }
     finally {

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

@@ -24,8 +24,8 @@
 import org.apache.druid.indexing.common.LockGranularity;
 import org.apache.druid.indexing.common.task.IndexTaskUtils;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.indexing.overlord.GlobalTaskLockbox;
 import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator;
-import org.apache.druid.indexing.overlord.TaskLockbox;
 import org.apache.druid.indexing.overlord.config.TaskLockConfig;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.Intervals;
@@ -74,7 +74,7 @@ public class SegmentAllocationQueue
 
   private final long maxWaitTimeMillis;
 
-  private final TaskLockbox taskLockbox;
+  private final GlobalTaskLockbox taskLockbox;
   private final IndexerMetadataStorageCoordinator metadataStorage;
   private final AtomicBoolean isLeader = new AtomicBoolean(false);
   private final ServiceEmitter emitter;
@@ -91,7 +91,7 @@ public class SegmentAllocationQueue
 
   @Inject
   public SegmentAllocationQueue(
-      TaskLockbox taskLockbox,
+      GlobalTaskLockbox taskLockbox,
       TaskLockConfig taskLockConfig,
       IndexerMetadataStorageCoordinator metadataStorage,
       ServiceEmitter emitter,

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

@@ -23,8 +23,8 @@
 import com.google.common.base.Optional;
 import com.google.inject.Inject;
 import org.apache.druid.guice.annotations.Json;
+import org.apache.druid.indexing.overlord.GlobalTaskLockbox;
 import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator;
-import org.apache.druid.indexing.overlord.TaskLockbox;
 import org.apache.druid.indexing.overlord.TaskRunner;
 import org.apache.druid.indexing.overlord.TaskRunnerFactory;
 import org.apache.druid.indexing.overlord.TaskStorage;
@@ -33,7 +33,7 @@
 
 public class TaskActionToolbox
 {
-  private final TaskLockbox taskLockbox;
+  private final GlobalTaskLockbox taskLockbox;
   private final TaskStorage taskStorage;
   private final SegmentAllocationQueue segmentAllocationQueue;
   private final IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator;
@@ -44,7 +44,7 @@ public class TaskActionToolbox
 
   @Inject
   public TaskActionToolbox(
-      TaskLockbox taskLockbox,
+      GlobalTaskLockbox taskLockbox,
       TaskStorage taskStorage,
       IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator,
       SegmentAllocationQueue segmentAllocationQueue,
@@ -63,7 +63,7 @@ public TaskActionToolbox(
   }
 
   public TaskActionToolbox(
-      TaskLockbox taskLockbox,
+      GlobalTaskLockbox taskLockbox,
       TaskStorage taskStorage,
       IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator,
       ServiceEmitter emitter,
@@ -82,7 +82,7 @@ public TaskActionToolbox(
     );
   }
 
-  public TaskLockbox getTaskLockbox()
+  public GlobalTaskLockbox getTaskLockbox()
   {
     return taskLockbox;
   }

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

@@ -28,7 +28,7 @@
 import org.apache.druid.indexing.common.task.AbstractBatchIndexTask;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.indexing.common.task.Tasks;
-import org.apache.druid.indexing.overlord.TaskLockbox;
+import org.apache.druid.indexing.overlord.GlobalTaskLockbox;
 import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.metadata.ReplaceTaskLock;
@@ -52,7 +52,7 @@ public class TaskLocks
 {
   static void checkLockCoversSegments(
       final Task task,
-      final TaskLockbox taskLockbox,
+      final GlobalTaskLockbox taskLockbox,
       final Collection<DataSegment> segments
   )
   {
@@ -69,7 +69,7 @@ static void checkLockCoversSegments(
   @VisibleForTesting
   static boolean isLockCoversSegments(
       final Task task,
-      final TaskLockbox taskLockbox,
+      final GlobalTaskLockbox taskLockbox,
       final Collection<DataSegment> segments
   )
   {
@@ -176,7 +176,7 @@ public static TaskLockType determineLockTypeForAppend(
    */
   public static Map<DataSegment, ReplaceTaskLock> findReplaceLocksCoveringSegments(
       final String datasource,
-      final TaskLockbox taskLockbox,
+      final GlobalTaskLockbox taskLockbox,
       final Set<DataSegment> segments
   )
   {
@@ -206,7 +206,7 @@ public static Map<DataSegment, ReplaceTaskLock> findReplaceLocksCoveringSegments
 
   public static List<TaskLock> findLocksForSegments(
       final Task task,
-      final TaskLockbox taskLockbox,
+      final GlobalTaskLockbox taskLockbox,
       final Collection<DataSegment> segments
   )
   {
@@ -245,7 +245,7 @@ public static List<TaskLock> findLocksForSegments(
     return found;
   }
 
-  private static NavigableMap<DateTime, List<TaskLock>> getTaskLockMap(TaskLockbox taskLockbox, Task task)
+  private static NavigableMap<DateTime, List<TaskLock>> getTaskLockMap(GlobalTaskLockbox taskLockbox, Task task)
   {
     final List<TaskLock> taskLocks = taskLockbox.findLocksForTask(task);
     final NavigableMap<DateTime, List<TaskLock>> taskLockMap = new TreeMap<>();

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

@@ -27,13 +27,13 @@
  * This class represents a critical action must be done while the task's lock is guaranteed to not be revoked in the
  * middle of the action.
  *
- * Implementations must not change the lock state by calling {@link TaskLockbox#lock)}, {@link TaskLockbox#tryLock)},
- * or {@link TaskLockbox#unlock(Task, Interval)}.
+ * Implementations must not change the lock state by calling {@link GlobalTaskLockbox#lock)}, {@link GlobalTaskLockbox#tryLock)},
+ * or {@link GlobalTaskLockbox#unlock(Task, Interval)}.
  *
- * Also, implementations should be finished as soon as possible because all methods in {@link TaskLockbox} are blocked
+ * Also, implementations should be finished as soon as possible because all methods in {@link GlobalTaskLockbox} are blocked
  * until this action is finished.
  *
- * @see TaskLockbox#doInCriticalSection
+ * @see GlobalTaskLockbox#doInCriticalSection
  */
 public class CriticalAction<T>
 {

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

@@ -80,7 +80,7 @@ public DruidOverlord(
       final TaskLockConfig taskLockConfig,
       final TaskQueueConfig taskQueueConfig,
       final DefaultTaskConfig defaultTaskConfig,
-      final TaskLockbox taskLockbox,
+      final GlobalTaskLockbox taskLockbox,
       final TaskStorage taskStorage,
       final TaskActionClientFactory taskActionClientFactory,
       @Self final DruidNode selfNode,