SOLR-17971: Fix REBALANCELEADERS wait when the Overseer is disabled

changelog/unreleased/solr-17971-rebalance-leaders-distributed.yml

@@ -0,0 +1,7 @@
+title: "REBALANCELEADERS no longer waits until timeout or leaves async ids behind when the Overseer is disabled."
+type: fixed
+authors:
+  - name: Amrit Sarkar
+links:
+  - name: SOLR-17971
+    url: https://issues.apache.org/jira/browse/SOLR-17971

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

@@ -286,6 +286,11 @@ public boolean deleteSingleAsyncId(String asyncId) throws Exception {
         && trackedAsyncTasks.remove(asyncId);
   }
 
+  @VisibleForTesting
+  public Collection<String> getTrackedAsyncIds() throws Exception {
+    return trackedAsyncTasks.keys();
+  }
+
   /**
    * Deletes all async id's for completed or failed async jobs. Does not touch id's for submitted or
    * running jobs.

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

@@ -2510,6 +2510,24 @@ public boolean clearAsyncId(String asyncId) throws KeeperException {
     }
   }
 
+  /**
+   * Removes the stored response of a completed or failed async operation and releases its asyncId,
+   * in both Overseer and distributed modes. An asyncId with no stored response is left claimed,
+   * since its task may still be running or pending.
+   *
+   * @return True if a stored response was found and removed.
+   */
+  public boolean removeStoredAsyncResponse(String asyncId) throws Exception {
+    if (distributedCommandRunner.isPresent()) {
+      return distributedCommandRunner.get().deleteSingleAsyncId(asyncId);
+    }
+    if (overseerCompletedMap.remove(asyncId) || overseerFailureMap.remove(asyncId)) {
+      clearAsyncId(asyncId);
+      return true;
+    }
+    return false;
+  }
+
   public Overseer getOverseer() {
     return overseer;
   }

solr/core/src/java/org/apache/solr/cloud/api/collections/DistributedCollectionConfigSetCommandRunner.java

@@ -24,8 +24,10 @@
 import static org.apache.solr.common.cloud.ZkStateReader.SHARD_ID_PROP;
 import static org.apache.solr.common.params.CommonParams.NAME;
 
+import com.google.common.annotations.VisibleForTesting;
 import java.io.IOException;
 import java.lang.invoke.MethodHandles;
+import java.util.Collection;
 import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutorService;
@@ -152,6 +154,12 @@ public boolean deleteSingleAsyncId(String asyncId) throws Exception {
     return asyncTaskTracker.deleteSingleAsyncId(asyncId);
   }
 
+  /** See {@link DistributedApiAsyncTracker#getTrackedAsyncIds()} */
+  @VisibleForTesting
+  public Collection<String> getTrackedAsyncIds() throws Exception {
+    return asyncTaskTracker.getTrackedAsyncIds();
+  }
+
   /** See {@link DistributedApiAsyncTracker#deleteAllAsyncIds()} */
   public void deleteAllAsyncIds() throws Exception {
     asyncTaskTracker.deleteAllAsyncIds();

solr/core/src/java/org/apache/solr/handler/admin/RebalanceLeaders.java

@@ -122,7 +122,7 @@ class RebalanceLeaders {
     coreContainer = collectionsHandler.getCoreContainer();
   }
 
-  void execute() throws KeeperException, InterruptedException {
+  void execute() throws Exception {
     DocCollection dc = checkParams();
 
     int max = req.getParams().getInt(MAX_AT_ONCE_PROP, Integer.MAX_VALUE);
@@ -471,8 +471,7 @@ private void rejoinElectionQueue(
   // mean failure!
   //
 
-  private boolean waitAsyncRequests(final int maxWaitSecs, Boolean waitForAll)
-      throws KeeperException, InterruptedException {
+  private boolean waitAsyncRequests(final int maxWaitSecs, Boolean waitForAll) throws Exception {
 
     if (asyncRequests.size() == 0) {
       return true;
@@ -482,15 +481,7 @@ private boolean waitAsyncRequests(final int maxWaitSecs, Boolean waitForAll)
       Iterator<String> iter = asyncRequests.iterator();
       boolean foundChange = false;
       while (iter.hasNext()) {
-        String asyncId = iter.next();
-        if (coreContainer.getZkController().getOverseerFailureMap().contains(asyncId)) {
-          coreContainer.getZkController().getOverseerFailureMap().remove(asyncId);
-          coreContainer.getZkController().clearAsyncId(asyncId);
-          iter.remove();
-          foundChange = true;
-        } else if (coreContainer.getZkController().getOverseerCompletedMap().contains(asyncId)) {
-          coreContainer.getZkController().getOverseerCompletedMap().remove(asyncId);
-          coreContainer.getZkController().clearAsyncId(asyncId);
+        if (coreContainer.getZkController().removeStoredAsyncResponse(iter.next())) {
           iter.remove();
           foundChange = true;
         }

solr/core/src/java/org/apache/solr/handler/admin/api/ClusterCommands.java

@@ -129,20 +129,7 @@ public DeleteClusterCommandStatusResponse deleteClusterCommandStatus(String requ
 
     final DeleteClusterCommandStatusResponse response =
         instantiateJerseyResponse(DeleteClusterCommandStatusResponse.class);
-    final ZkController zkController = coreContainer.getZkController();
-
-    if (zkController.getDistributedCommandRunner().isEmpty()) {
-      if (zkController.getOverseerCompletedMap().remove(requestId)) {
-        zkController.clearAsyncId(requestId);
-        response.status = "successfully removed stored response for [" + requestId + "]";
-      } else if (zkController.getOverseerFailureMap().remove(requestId)) {
-        zkController.clearAsyncId(requestId);
-        response.status = "successfully removed stored response for [" + requestId + "]";
-      } else {
-        // Don't call zkController.clearAsyncId for this, since it could be a running/pending task
-        response.status = "[" + requestId + "] not found in stored responses";
-      }
-    } else if (zkController.getDistributedCommandRunner().get().deleteSingleAsyncId(requestId)) {
+    if (coreContainer.getZkController().removeStoredAsyncResponse(requestId)) {
       response.status = "successfully removed stored response for [" + requestId + "]";
     } else {
       response.status = "[" + requestId + "] not found in stored responses";

solr/core/src/test/org/apache/solr/cloud/TestRebalanceLeaders.java

@@ -24,6 +24,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.TreeSet;
@@ -36,6 +37,7 @@
 import org.apache.solr.client.solrj.request.GenericSolrRequest;
 import org.apache.solr.client.solrj.response.CollectionAdminResponse;
 import org.apache.solr.client.solrj.response.SimpleSolrResponse;
+import org.apache.solr.cloud.api.collections.DistributedCollectionConfigSetCommandRunner;
 import org.apache.solr.common.cloud.DocCollection;
 import org.apache.solr.common.cloud.Replica;
 import org.apache.solr.common.cloud.Slice;
@@ -157,6 +159,54 @@ public void testRebalanceLeaders() throws Exception {
     checkPreferredsAreLeaders();
   }
 
+  @Test
+  public void testRebalanceLeadersClearsAsyncIds() throws Exception {
+    DocCollection docCollection =
+        cluster.getSolrClient().getClusterState().getCollection(COLLECTION_NAME);
+    for (Slice slice : docCollection.getSlices()) {
+      Replica nonLeader =
+          slice.getReplicas().stream()
+              .filter(rep -> rep.getName().equals(slice.getLeader().getName()) == false)
+              .findFirst()
+              .orElseThrow();
+      setProp(slice, nonLeader, "preferredLeader");
+    }
+
+    ZkController zkController =
+        cluster.getRandomJetty(random()).getCoreContainer().getZkController();
+    Set<String> idsBefore = asyncIds(zkController);
+
+    int maxWaitSeconds = 60;
+    long start = System.nanoTime();
+    CollectionAdminResponse resp =
+        CollectionAdminRequest.rebalanceLeaders(COLLECTION_NAME)
+            .setMaxWaitSeconds(maxWaitSeconds)
+            .process(cluster.getSolrClient());
+    long elapsedSeconds = TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() - start);
+    assertTrue(
+        resp.getResponse().toString(),
+        resp.getResponse().get("Summary").toString().contains("Success"));
+    assertNotNull("No leader change was requested", resp.getResponse().get("successes"));
+    assertTrue(
+        "Waited " + elapsedSeconds + "s, hit maxWaitSeconds", elapsedSeconds < maxWaitSeconds);
+
+    Set<String> leftover = asyncIds(zkController);
+    leftover.removeAll(idsBefore);
+    assertTrue("Async ids not cleared: " + leftover, leftover.isEmpty());
+  }
+
+  private Set<String> asyncIds(ZkController zkController) throws Exception {
+    Optional<DistributedCollectionConfigSetCommandRunner> runner =
+        zkController.getDistributedCommandRunner();
+    if (runner.isPresent()) {
+      return new HashSet<>(runner.get().getTrackedAsyncIds());
+    }
+    Set<String> ids = new HashSet<>(zkController.getOverseerCompletedMap().keys());
+    ids.addAll(zkController.getOverseerFailureMap().keys());
+    ids.addAll(Overseer.getAsyncIdsMap(zkController.getZkClient()).keys());
+    return ids;
+  }
+
   // Ensure that the property is set on only one replica per slice when changing a unique property
   // on an individual
   // replica.