SOLR-17754 Fix race condition in overseer main loop. (#3350)

This fixes the overseer main loop so we never submit more than 100 concurrent tasks to the thread pool. Instead of manually tracking when a task is complete, we check the status using a standard Future.

Author: psalagnac

solr/CHANGES.txt

@@ -56,7 +56,7 @@ Optimizations
 
 Bug Fixes
 ---------------------
-* SOLR-17772: Tests for CLI examples were failing on Windows due to a legacy bug uncovered by fix in SOLR-7962. 
+* SOLR-17772: Tests for CLI examples were failing on Windows due to a legacy bug uncovered by fix in SOLR-7962.
   Additionally achieves simplification of CLI tests (Rahul Goswami via Eric Pugh)
 
 Deprecation Removals
@@ -195,7 +195,7 @@ Other Changes
 New Features
 ---------------------
 * SOLR-17582: The CLUSTERSTATUS API will now stream each collection's status to the response,
-  fetching and computing it on the fly.  To avoid a backwards compatibilty concern, this won't work
+  fetching and computing it on the fly.  To avoid a backwards compatibility concern, this won't work
   for wt=javabin.  (Matthew Biscocho, David Smiley)
 
 * SOLR-17626: Add RawTFSimilarityFactory class. (Christine Poerschke)
@@ -296,9 +296,12 @@ Bug Fixes
 * SOLR-17758: The NumFieldLimitingUpdateRequestProcessor's "warnOnly" mode has been fixed, and now processes documents even
   when the limit has been exceeded. (Jason Gerlowski, Rahul Goswami)
 
-* SOLR-7962: Passing additional arguments to solr.cmd using "--jvm-opts" (formerly "-a") in conjunction with "-e" (examples like 'techproducts') 
+* SOLR-7962: Passing additional arguments to solr.cmd using "--jvm-opts" (formerly "-a") in conjunction with "-e" (examples like 'techproducts')
   wouldn't reflect on Windows (Rahul Goswami via Eric Pugh)
 
+* SOLR-17754: Fix rare bug in overseer main loop in case of high load, that may cause the overseer be fully stuck until
+  server restart. (Pierre Salagnac)
+
 Dependency Upgrades
 ---------------------
 * SOLR-17471: Upgrade Lucene to 9.12.1. (Pierre Salagnac, Christine Poerschke)

solr/core/src/java/org/apache/solr/cloud/OverseerTaskProcessor.java

@@ -24,13 +24,13 @@
 import java.lang.invoke.MethodHandles;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.Iterator;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
 import java.util.concurrent.SynchronousQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Predicate;
@@ -70,35 +70,27 @@ public class OverseerTaskProcessor implements Runnable, Closeable {
 
   private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
 
-  private OverseerTaskQueue workQueue;
-  private DistributedMap runningMap;
-  private DistributedMap completedMap;
-  private DistributedMap failureMap;
+  private final OverseerTaskQueue workQueue;
+  private final DistributedMap runningMap;
+  private final DistributedMap completedMap;
+  private final DistributedMap failureMap;
 
   /**
-   * Set that maintains a list of all the tasks that are running. This is keyed on zk id of the
-   * task.
+   * All the tasks that have been submitted to the runner thread pool. This is keyed on zk id of the
+   * task. It may contain tasks that have completed execution, have been entered into the
+   * completed/failed map in zk but not deleted from the work-queue as that is a batched operation.
+   * This collection is accessed only in the main loop and does not need to be synchronized.
    */
-  private final Set<String> runningTasks;
+  private final Map<String, Future<?>> runningTasks;
 
-  /** List of completed tasks. This is used to clean up workQueue in zk. */
-  private final ConcurrentHashMap<String, QueueEvent> completedTasks;
+  private final String myId;
 
-  private volatile String myId;
-
-  private volatile ZkStateReader zkStateReader;
+  private final ZkStateReader zkStateReader;
 
   private boolean isClosed;
 
-  private volatile Stats stats;
-  private SolrMetricsContext overseerTaskProcessorMetricsContext;
-
-  /**
-   * Set of tasks that have been picked up for processing but not cleaned up from zk work-queue. It
-   * may contain tasks that have completed execution, have been entered into the completed/failed
-   * map in zk but not deleted from the work-queue as that is a batched operation.
-   */
-  private final Set<String> runningZKTasks;
+  private final Stats stats;
+  private final SolrMetricsContext overseerTaskProcessorMetricsContext;
 
   /**
    * This map may contain tasks which are read from work queue but could not be executed because
@@ -116,7 +108,7 @@ public class OverseerTaskProcessor implements Runnable, Closeable {
       new Predicate<>() {
         @Override
         public boolean test(String s) {
-          return runningTasks.contains(s) || blockedTasks.containsKey(s);
+          return runningTasks.containsKey(s) || blockedTasks.containsKey(s);
         }
 
         @Override
@@ -129,9 +121,9 @@ public String toString() {
 
   protected OverseerMessageHandlerSelector selector;
 
-  private OverseerNodePrioritizer prioritizer;
+  private final OverseerNodePrioritizer prioritizer;
 
-  private String thisNode;
+  private final String thisNode;
 
   public OverseerTaskProcessor(
       ZkStateReader zkStateReader,
@@ -153,9 +145,7 @@ public OverseerTaskProcessor(
     this.runningMap = runningMap;
     this.completedMap = completedMap;
     this.failureMap = failureMap;
-    this.runningZKTasks = ConcurrentHashMap.newKeySet();
-    this.runningTasks = ConcurrentHashMap.newKeySet();
-    this.completedTasks = new ConcurrentHashMap<>();
+    this.runningTasks = new HashMap<>();
     thisNode = MDCLoggingContext.getNodeName();
 
     overseerTaskProcessorMetricsContext = solrMetricsContext.getChildContext(this);
@@ -235,13 +225,7 @@ public void run() {
             continue; // not a no, not a yes, try asking again
           }
 
-          if (log.isDebugEnabled()) {
-            log.debug(
-                "Cleaning up work-queue. #Running tasks: {} #Completed tasks: {}",
-                runningTasks.size(),
-                completedTasks.size());
-          }
-          cleanUpWorkQueue();
+          cleanUpRunningTasks();
 
           printTrackingMaps();
 
@@ -254,7 +238,9 @@ public void run() {
             waited = true;
           }
 
-          if (waited) cleanUpWorkQueue();
+          if (waited) {
+            cleanUpRunningTasks();
+          }
 
           ArrayList<QueueEvent> heads = new ArrayList<>(blockedTasks.size() + MAX_PARALLEL_TASKS);
           heads.addAll(blockedTasks.values());
@@ -309,17 +295,16 @@ public void run() {
           // on the OverseerCollectionMessageHandler
           batchSessionId++;
 
-          boolean tooManyTasks = false;
-          for (QueueEvent head : heads) {
-            if (!tooManyTasks) {
-              tooManyTasks = runningTasks.size() >= MAX_PARALLEL_TASKS;
-            }
-            if (tooManyTasks) {
+          for (int i = 0; i < heads.size(); i++) {
+            QueueEvent head = heads.get(i);
+
+            if (runningTasks.size() >= MAX_PARALLEL_TASKS) {
               // Too many tasks are running, just shove the rest into the "blocked" queue.
-              blockedTasks.put(head.getId(), head);
-              continue;
+              heads.subList(i, heads.size()).forEach(h -> blockedTasks.put(h.getId(), h));
+              break;
             }
-            if (runningZKTasks.contains(head.getId())) continue;
+
+            if (runningTasks.containsKey(head.getId())) continue;
             final ZkNodeProps message = ZkNodeProps.load(head.getBytes());
             final String asyncId = message.getStr(ASYNC);
             if (hasLeftOverItems) {
@@ -329,14 +314,14 @@ public void run() {
                 log.debug(
                     "Found already processed task in workQueue, cleaning up. AsyncId [{}]",
                     asyncId);
-                workQueue.remove(head);
+                workQueue.remove(head, false);
                 continue;
               }
             }
             String operation = message.getStr(Overseer.QUEUE_OPERATION);
             if (operation == null) {
               log.error("Msg does not have required {} : {}", Overseer.QUEUE_OPERATION, message);
-              workQueue.remove(head);
+              workQueue.remove(head, asyncId == null);
               continue;
             }
             OverseerMessageHandler messageHandler = selector.selectOverseerMessageHandler(message);
@@ -349,7 +334,11 @@ public void run() {
               continue;
             }
             try {
-              markTaskAsRunning(head, asyncId);
+              if (asyncId != null) {
+                // Store the async ID in ZK distributed map before trying to execute the task
+                runningMap.put(asyncId, null);
+              }
+
               if (log.isDebugEnabled()) {
                 log.debug("Marked task [{}] as running", head.getId());
               }
@@ -372,8 +361,9 @@ public void run() {
                   head.getId(),
                   message);
             }
-            Runner runner = new Runner(messageHandler, message, operation, head, lock);
-            tpe.execute(runner);
+            Runner runner = createRunner(messageHandler, message, operation, head, lock);
+            Future<?> future = tpe.submit(runner);
+            runningTasks.put(head.getId(), future);
           }
 
         } catch (KeeperException e) {
@@ -405,14 +395,12 @@ public void run() {
     }
   }
 
-  private void cleanUpWorkQueue() throws KeeperException, InterruptedException {
-    Iterator<Map.Entry<String, QueueEvent>> it = completedTasks.entrySet().iterator();
-    while (it.hasNext()) {
-      Map.Entry<String, QueueEvent> entry = it.next();
-      workQueue.remove(entry.getValue());
-      runningZKTasks.remove(entry.getKey());
-      it.remove();
+  /** Remove all entries from {@link #runningTasks} that are done. */
+  private void cleanUpRunningTasks() {
+    if (log.isDebugEnabled()) {
+      log.debug("Cleaning up running tasks. #Running tasks: {}", runningTasks.size());
     }
+    runningTasks.entrySet().removeIf(e -> e.getValue().isDone());
   }
 
   @Override
@@ -429,29 +417,19 @@ public void close() {
 
   public static List<String> getSortedOverseerNodeNames(SolrZkClient zk)
       throws KeeperException, InterruptedException {
-    List<String> children = null;
-    try {
-      children = zk.getChildren(Overseer.OVERSEER_ELECT + LeaderElector.ELECTION_NODE, null, true);
-    } catch (Exception e) {
-      log.warn("error ", e);
-      return new ArrayList<>();
-    }
-    LeaderElector.sortSeqs(children);
+    List<String> children =
+        getSortedElectionNodes(zk, Overseer.OVERSEER_ELECT + LeaderElector.ELECTION_NODE);
+
     ArrayList<String> nodeNames = new ArrayList<>(children.size());
     for (String c : children) nodeNames.add(LeaderElector.getNodeName(c));
     return nodeNames;
   }
 
   public static List<String> getSortedElectionNodes(SolrZkClient zk, String path)
       throws KeeperException, InterruptedException {
-    List<String> children = null;
-    try {
-      children = zk.getChildren(path, null, true);
-      LeaderElector.sortSeqs(children);
-      return children;
-    } catch (Exception e) {
-      throw e;
-    }
+    List<String> children = zk.getChildren(path, null, true);
+    LeaderElector.sortSeqs(children);
+    return children;
   }
 
   public static String getLeaderNode(SolrZkClient zkClient)
@@ -518,12 +496,17 @@ public boolean isClosed() {
     return isClosed;
   }
 
-  private void markTaskAsRunning(QueueEvent head, String asyncId)
-      throws KeeperException, InterruptedException {
-    runningZKTasks.add(head.getId());
-    runningTasks.add(head.getId());
-
-    if (asyncId != null) runningMap.put(asyncId, null);
+  /**
+   * Create a runner instance to execute a single task. This method mosty exists to provide an
+   * extension point for tests.
+   */
+  protected Runner createRunner(
+      OverseerMessageHandler messageHandler,
+      ZkNodeProps message,
+      String operation,
+      QueueEvent head,
+      OverseerMessageHandler.Lock lock) {
+    return new Runner(messageHandler, message, operation, head, lock);
   }
 
   protected class Runner implements Runnable {
@@ -589,7 +572,7 @@ public void run() {
           }
         }
 
-        markTaskComplete(head.getId(), asyncId);
+        markTaskComplete(asyncId);
         if (log.isDebugEnabled()) {
           log.debug("Marked task [{}] as completed.", head.getId());
         }
@@ -609,15 +592,15 @@ public void run() {
         log.error("KeeperException", e);
       } catch (InterruptedException e) {
         // Reset task from tracking data structures so that it can be retried.
-        resetTaskWithException(messageHandler, head.getId(), asyncId, taskKey, message);
+        resetTaskWithException(head.getId(), asyncId, taskKey);
         log.warn("Resetting task {} as the thread was interrupted.", head.getId());
         Thread.currentThread().interrupt();
       } finally {
         lock.unlock();
         if (!success) {
           // Reset task from tracking data structures so that it can be retried.
           try {
-            resetTaskWithException(messageHandler, head.getId(), asyncId, taskKey, message);
+            resetTaskWithException(head.getId(), asyncId, taskKey);
           } catch (IllegalStateException ignore) {
 
           }
@@ -628,26 +611,26 @@ public void run() {
       }
     }
 
-    private void markTaskComplete(String id, String asyncId)
-        throws KeeperException, InterruptedException {
-      completedTasks.put(id, head);
-      runningTasks.remove(id);
+    /**
+     * Invoked by the runner thread once the task is successfully completed (can be with error
+     * status). We remove the task from the work queue.
+     */
+    private void markTaskComplete(String asyncId) throws KeeperException, InterruptedException {
 
       if (asyncId != null) {
         if (!runningMap.remove(asyncId)) {
           log.warn("Could not find and remove async call [{}] from the running map.", asyncId);
         }
       }
 
-      workQueue.remove(head);
+      workQueue.remove(head, asyncId == null);
     }
 
-    private void resetTaskWithException(
-        OverseerMessageHandler messageHandler,
-        String id,
-        String asyncId,
-        String taskKey,
-        ZkNodeProps message) {
+    /**
+     * Reset the task so it will be retried by a future iteration of the main loop. We remove the
+     * async ID from the running map, but we keep the task in the work queue.
+     */
+    private void resetTaskWithException(String id, String asyncId, String taskKey) {
       log.warn("Resetting task: {}, requestid: {}, taskKey: {}", id, asyncId, taskKey);
       try {
         if (asyncId != null) {
@@ -656,7 +639,6 @@ private void resetTaskWithException(
           }
         }
 
-        runningTasks.remove(id);
       } catch (KeeperException e) {
         log.error("KeeperException", e);
       } catch (InterruptedException e) {
@@ -682,10 +664,8 @@ private boolean isSuccessful() {
 
   private void printTrackingMaps() {
     if (log.isDebugEnabled()) {
-      log.debug("RunningTasks: {}", runningTasks);
+      log.debug("RunningTasks: {}", runningTasks.keySet());
       log.debug("BlockedTasks: {}", blockedTasks.keySet()); // nowarn
-      log.debug("CompletedTasks: {}", completedTasks.keySet()); // nowarn
-      log.debug("RunningZKTasks: {}", runningZKTasks); // nowarn
     }
   }