SOLR-18298: only recover after ZooKeeper session expiry

changelog/unreleased/SOLR-18298-zk-reconnect-session-expiry.yml

@@ -0,0 +1,8 @@
+title: ZkController no longer treats every ZooKeeper reconnect as session expiration. Re-election and core re-registration now run only after Curator reports ConnectionState.LOST, restoring Solr 9 behavior during ZK rolling restarts.
+type: fixed
+authors:
+  - name: Shrey Narayan (NextBrick)
+    url: https://nextbrick.com
+links:
+  - name: SOLR-18298
+    url: https://issues.apache.org/jira/browse/SOLR-18298

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

@@ -54,6 +54,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -214,6 +215,7 @@ public String toString() {
           new SolrNamedThreadFactory("zkConnectionListenerCallback"));
   private final OnReconnect onReconnect = this::onReconnect;
   private final OnDisconnect onDisconnect = this::onDisconnect;
+  private final AtomicBoolean zkSessionExpired = new AtomicBoolean();
 
   private final String zkServerAddress; // example: 127.0.0.1:54062/solr
 
@@ -402,6 +404,10 @@ public ZkController(
   }
 
   private void onDisconnect(boolean sessionExpired) {
+    if (!sessionExpired) {
+      return;
+    }
+    zkSessionExpired.set(true);
     try {
       overseer.close();
     } catch (Exception e) {
@@ -431,6 +437,9 @@ private <T> T loadPluginOrDefault(
   }
 
   private void onReconnect() {
+    if (!zkSessionExpired.compareAndSet(true, false)) {
+      return;
+    }
     // on reconnect, reload cloud info
     log.info("ZooKeeper session re-connected ... refreshing core states after session expiration.");
     clearZkCollectionTerms();

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

@@ -36,6 +36,9 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import org.apache.curator.CuratorZookeeperClient;
+import org.apache.curator.test.InstanceSpec;
+import org.apache.curator.test.TestingCluster;
 import org.apache.solr.SolrTestCaseJ4;
 import org.apache.solr.client.api.util.SolrVersion;
 import org.apache.solr.client.solrj.jetty.HttpJettySolrClient;
@@ -50,6 +53,7 @@
 import org.apache.solr.common.cloud.ZkStateReader;
 import org.apache.solr.common.params.CollectionParams;
 import org.apache.solr.common.util.ExecutorUtil;
+import org.apache.solr.common.util.RetryUtil;
 import org.apache.solr.common.util.SolrNamedThreadFactory;
 import org.apache.solr.common.util.Utils;
 import org.apache.solr.core.CloudConfig;
@@ -766,6 +770,53 @@ public void testOverseerEnabledClusterPropertyTrue() throws Exception {
     }
   }
 
+  @Test
+  public void testReconnectRecoveryRequiresSessionExpiration() throws Exception {
+    try (TestingCluster zkCluster = new TestingCluster(3)) {
+      zkCluster.start();
+      CoreContainer cc = getCoreContainer();
+      try {
+        CloudConfig cloudConfig = new CloudConfig.CloudConfigBuilder("127.0.0.1", 8983).build();
+        try (ZkController zkController =
+            new ZkController(cc, zkCluster.getConnectString(), TIMEOUT, cloudConfig)) {
+          AtomicInteger recoveries = new AtomicInteger();
+          zkController.addOnReconnectListener(recoveries::incrementAndGet);
+          CuratorZookeeperClient curatorClient =
+              zkController.getZkClient().getCuratorFramework().getZookeeperClient();
+
+          InstanceSpec connected = zkCluster.findConnectionInstance(curatorClient.getZooKeeper());
+          assertNotNull(connected);
+          zkCluster.killServer(connected);
+          RetryUtil.retryUntil(
+              "Solr did not connect to another ZooKeeper server",
+              30,
+              200,
+              TimeUnit.MILLISECONDS,
+              () -> zkCluster.findConnectionInstance(curatorClient.getZooKeeper()),
+              current -> current != null && !current.equals(connected));
+          assertEquals(
+              "A transient ZooKeeper reconnect must not trigger session-expiration recovery",
+              0,
+              recoveries.get());
+
+          curatorClient.getZooKeeper().getTestable().injectSessionExpiration();
+          RetryUtil.retryUntil(
+              "A reconnect after session expiration did not trigger recovery",
+              30,
+              200,
+              TimeUnit.MILLISECONDS,
+              () -> recoveries.get() == 1);
+          assertEquals(1, recoveries.get());
+        }
+      } finally {
+        cc.shutdown();
+      }
+    } finally {
+      // TestingCluster closes its quorum asynchronously; allow its worker threads to terminate.
+      Thread.sleep(3000);
+    }
+  }
+
   private CoreContainer getCoreContainer() {
     return new MockCoreContainer();
   }

solr/solrj-zookeeper/src/java/org/apache/solr/common/cloud/ZkStateReader.java

@@ -390,6 +390,9 @@ private StatefulCollectionWatch(StateWatcher associatedWatcher) {
   private final SolrZkClient zkClient;
 
   private final boolean closeClient;
+  private final AtomicBoolean zkSessionExpired = new AtomicBoolean();
+  private final OnDisconnect onDisconnect = this::onDisconnect;
+  private final OnReconnect onReconnect = this::onReconnect;
 
   private volatile boolean closed = false;
 
@@ -424,29 +427,34 @@ public ZkStateReader(
             .withConnTimeOut(zkClientConnectTimeout, TimeUnit.MILLISECONDS)
             .withUseDefaultCredsAndACLs(canUseZkACLs)
             .build();
-    this.zkClient
-        .getCuratorFramework()
-        .getConnectionStateListenable()
-        .addListener(
-            (OnReconnect)
-                () -> {
-                  // on reconnect, reload cloud info
-                  try {
-                    this.createClusterStateWatchersAndUpdate();
-                  } catch (InterruptedException e) {
-                    // Restore the interrupted status
-                    Thread.currentThread().interrupt();
-                    log.warn("Interrupted", e);
-                  } catch (Throwable e) {
-                    log.error("An error has occurred while updating the cluster state", e);
-                  }
-                });
+    this.zkClient.getCuratorFramework().getConnectionStateListenable().addListener(onReconnect);
+    this.zkClient.getCuratorFramework().getConnectionStateListenable().addListener(onDisconnect);
     this.closeClient = true;
     this.securityNodeWatcher = null;
     collectionPropertiesZkStateReader = new CollectionPropertiesZkStateReader(this);
     assert ObjectReleaseTracker.track(this);
   }
 
+  private void onDisconnect(boolean sessionExpired) {
+    if (sessionExpired) {
+      zkSessionExpired.set(true);
+    }
+  }
+
+  private void onReconnect() {
+    if (!zkSessionExpired.compareAndSet(true, false)) {
+      return;
+    }
+    try {
+      createClusterStateWatchersAndUpdate();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      log.warn("Interrupted", e);
+    } catch (Throwable e) {
+      log.error("An error has occurred while updating the cluster state", e);
+    }
+  }
+
   /**
    * Forcibly refresh cluster state from ZK. Do this only to avoid race conditions because it's
    * expensive.

solr/solrj-zookeeper/src/test/org/apache/solr/common/cloud/TestOnReconnectSessionExpiry.java

@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.solr.common.cloud;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.curator.framework.state.ConnectionState;
+import org.apache.solr.SolrTestCase;
+import org.junit.Test;
+
+/** Verifies the shared Curator listener adapters retain their general-purpose behavior. */
+public class TestOnReconnectSessionExpiry extends SolrTestCase {
+
+  @Test
+  public void testReconnectFiresForEveryReconnectedEvent() {
+    AtomicInteger reconnects = new AtomicInteger();
+    OnReconnect listener = () -> reconnects.incrementAndGet();
+
+    listener.stateChanged(null, ConnectionState.SUSPENDED);
+    listener.stateChanged(null, ConnectionState.RECONNECTED);
+    listener.stateChanged(null, ConnectionState.LOST);
+    listener.stateChanged(null, ConnectionState.RECONNECTED);
+
+    listener.stateChanged(null, ConnectionState.RECONNECTED);
+    assertEquals(3, reconnects.get());
+  }
+
+  @Test
+  public void testDisconnectDistinguishesSuspensionFromSessionLoss() {
+    AtomicInteger suspensions = new AtomicInteger();
+    AtomicInteger expirations = new AtomicInteger();
+    OnDisconnect listener =
+        sessionExpired -> {
+          if (sessionExpired) {
+            expirations.incrementAndGet();
+          } else {
+            suspensions.incrementAndGet();
+          }
+        };
+
+    listener.stateChanged(null, ConnectionState.SUSPENDED);
+    listener.stateChanged(null, ConnectionState.LOST);
+
+    assertEquals(1, suspensions.get());
+    assertEquals(1, expirations.get());
+  }
+}