Fix segment rebalance race with load queue snapshots

Author: jtuglu1

server/src/main/java/org/apache/druid/server/coordinator/duty/PrepareBalancerAndLoadQueues.java

@@ -81,34 +81,42 @@ public PrepareBalancerAndLoadQueues(
   @Override
   public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params)
   {
-    List<ImmutableDruidServer> currentServers = prepareCurrentServers();
-    taskMaster.resetPeonsForNewServers(currentServers);
-
-    final CoordinatorDynamicConfig dynamicConfig = params.getCoordinatorDynamicConfig();
-    final SegmentLoadingConfig segmentLoadingConfig
-        = SegmentLoadingConfig.create(dynamicConfig, params.getUsedSegmentCount());
-
-    final DruidCluster cluster = prepareCluster(dynamicConfig, segmentLoadingConfig, currentServers);
-    cancelLoadsOnDecommissioningServers(cluster);
-
-    final CoordinatorRunStats stats = params.getCoordinatorStats();
-    collectHistoricalStats(cluster, stats);
-    collectUsedSegmentStats(params, stats);
-    collectDebugStats(segmentLoadingConfig, stats);
-
-    final int numBalancerThreads = segmentLoadingConfig.getBalancerComputeThreads();
-    final BalancerStrategy balancerStrategy = balancerStrategyFactory.createBalancerStrategy(numBalancerThreads);
-    log.debug(
-        "Using balancer strategy[%s] with [%d] threads.",
-        balancerStrategy.getClass().getSimpleName(), numBalancerThreads
-    );
-
-    return params.buildFromExisting()
-                 .withDruidCluster(cluster)
-                 .withBalancerStrategy(balancerStrategy)
-                 .withSegmentLoadingConfig(segmentLoadingConfig)
-                 .withSegmentAssignerUsing(loadQueueManager)
-                 .build();
+    // Prevent callbacks from firing while balancer/load queue accounting is taking place.
+    // This ensures snapshot validity at the expense of potentially stale snapshots (in some cases).
+    taskMaster.getCallbackLock().writeLock().lock();
+    try {
+      List<ImmutableDruidServer> currentServers = prepareCurrentServers();
+      taskMaster.resetPeonsForNewServers(currentServers);
+
+      final CoordinatorDynamicConfig dynamicConfig = params.getCoordinatorDynamicConfig();
+      final SegmentLoadingConfig segmentLoadingConfig
+          = SegmentLoadingConfig.create(dynamicConfig, params.getUsedSegmentCount());
+
+      final DruidCluster cluster = prepareCluster(dynamicConfig, segmentLoadingConfig, currentServers);
+      cancelLoadsOnDecommissioningServers(cluster);
+
+      final CoordinatorRunStats stats = params.getCoordinatorStats();
+      collectHistoricalStats(cluster, stats);
+      collectUsedSegmentStats(params, stats);
+      collectDebugStats(segmentLoadingConfig, stats);
+
+      final int numBalancerThreads = segmentLoadingConfig.getBalancerComputeThreads();
+      final BalancerStrategy balancerStrategy = balancerStrategyFactory.createBalancerStrategy(numBalancerThreads);
+      log.debug(
+          "Using balancer strategy[%s] with [%d] threads.",
+          balancerStrategy.getClass().getSimpleName(), numBalancerThreads
+      );
+
+      return params.buildFromExisting()
+                   .withDruidCluster(cluster)
+                   .withBalancerStrategy(balancerStrategy)
+                   .withSegmentLoadingConfig(segmentLoadingConfig)
+                   .withSegmentAssignerUsing(loadQueueManager)
+                   .build();
+    }
+    finally {
+      taskMaster.getCallbackLock().writeLock().unlock();
+    }
   }
 
   /**

server/src/main/java/org/apache/druid/server/coordinator/loading/HttpLoadQueuePeon.java

@@ -72,6 +72,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.ReadWriteLock;
 import java.util.function.Supplier;
 
 /**
@@ -123,6 +124,7 @@ public class HttpLoadQueuePeon implements LoadQueuePeon
 
   private final AtomicBoolean mainLoopInProgress = new AtomicBoolean(false);
   private final ExecutorService callBackExecutor;
+  private final ReadWriteLock callbackSync;
   private final Supplier<SegmentLoadingMode> loadingModeSupplier;
 
   private final ObjectWriter requestBodyWriter;
@@ -135,7 +137,8 @@ public HttpLoadQueuePeon(
       HttpLoadQueuePeonConfig config,
       Supplier<SegmentLoadingMode> loadingModeSupplier,
       ScheduledExecutorService processingExecutor,
-      ExecutorService callBackExecutor
+      ExecutorService callBackExecutor,
+      ReadWriteLock callbackSync
   )
   {
     this.jsonMapper = jsonMapper;
@@ -144,6 +147,7 @@ public HttpLoadQueuePeon(
     this.config = config;
     this.processingExecutor = processingExecutor;
     this.callBackExecutor = callBackExecutor;
+    this.callbackSync = callbackSync;
 
     this.serverId = baseUrl;
     this.loadingModeSupplier = loadingModeSupplier;
@@ -637,7 +641,14 @@ private void executeCallbacks(SegmentHolder holder, boolean success)
   {
     callBackExecutor.execute(() -> {
       for (LoadPeonCallback callback : holder.getCallbacks()) {
-        callback.execute(success);
+        // Load queue peons acquire read lock to increase tput (they operate safely on underlying state).
+        callbackSync.readLock().lock();
+        try {
+          callback.execute(success);
+        }
+        finally {
+          callbackSync.readLock().unlock();
+        }
       }
     });
   }

server/src/main/java/org/apache/druid/server/coordinator/loading/LoadQueueTaskMaster.java

@@ -35,6 +35,8 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.function.Supplier;
 
 /**
@@ -47,6 +49,7 @@ public class LoadQueueTaskMaster
   private final ObjectMapper jsonMapper;
   private final ScheduledExecutorService peonExec;
   private final ExecutorService callbackExec;
+  private final ReadWriteLock callbackLock;
   private final HttpLoadQueuePeonConfig config;
   private final HttpClient httpClient;
   private final Supplier<CoordinatorDynamicConfig> coordinatorDynamicConfigSupplier;
@@ -68,6 +71,7 @@ public LoadQueueTaskMaster(
     this.jsonMapper = jsonMapper;
     this.peonExec = peonExec;
     this.callbackExec = callbackExec;
+    this.callbackLock = new ReentrantReadWriteLock();
     this.config = config;
     this.httpClient = httpClient;
     this.coordinatorDynamicConfigSupplier = coordinatorDynamicConfigSupplier;
@@ -82,7 +86,8 @@ private LoadQueuePeon createPeon(ImmutableDruidServer server)
         config,
         () -> coordinatorDynamicConfigSupplier.get().getLoadingModeForServer(server.getName()),
         peonExec,
-        callbackExec
+        callbackExec,
+        callbackLock
     );
   }
 
@@ -154,4 +159,9 @@ public boolean isHttpLoading()
   {
     return true;
   }
+
+  public ReadWriteLock getCallbackLock()
+  {
+    return callbackLock;
+  }
 }