SOLR-16257: ZkStateReader changes to avoid race condition between collectionWatches and watchedCollectionStates

solr/core/src/test/org/apache/solr/cloud/overseer/ZkStateReaderTest.java

@@ -16,28 +16,32 @@
  */
 package org.apache.solr.cloud.overseer;
 
+import java.lang.invoke.MethodHandles;
 import java.nio.file.Path;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
+import java.util.*;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.lucene.util.IOUtils;
 import org.apache.solr.SolrTestCaseJ4;
 import org.apache.solr.cloud.OverseerTest;
 import org.apache.solr.cloud.Stats;
 import org.apache.solr.cloud.ZkController;
 import org.apache.solr.cloud.ZkTestServer;
-import org.apache.solr.common.cloud.ClusterState;
-import org.apache.solr.common.cloud.DocCollection;
-import org.apache.solr.common.cloud.DocRouter;
-import org.apache.solr.common.cloud.SolrZkClient;
-import org.apache.solr.common.cloud.ZkStateReader;
+import org.apache.solr.common.cloud.*;
+import org.apache.solr.common.util.ExecutorUtil;
+import org.apache.solr.common.util.SolrNamedThreadFactory;
 import org.apache.solr.common.util.TimeSource;
 import org.apache.solr.handler.admin.ConfigSetsHandler;
 import org.apache.solr.util.TimeOut;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class ZkStateReaderTest extends SolrTestCaseJ4 {
-
+  private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
   private static final long TIMEOUT = 30;
 
   public void testExternalCollectionWatchedNotWatched() throws Exception {
@@ -196,4 +200,251 @@ public void testWatchedCollectionCreation() throws Exception {
       server.shutdown();
     }
   }
+
+  public void testForciblyRefreshAllClusterState() throws Exception {
+    Path zkDir = createTempDir("testForciblyRefreshAllClusterState");
+
+    ZkTestServer server = new ZkTestServer(zkDir);
+
+    SolrZkClient zkClient = null;
+    ZkStateReader reader = null;
+
+    try {
+      server.run();
+
+      zkClient = new SolrZkClient(server.getZkAddress(), OverseerTest.DEFAULT_CONNECTION_TIMEOUT);
+      ZkController.createClusterZkNodes(zkClient);
+
+      reader = new ZkStateReader(zkClient);
+      reader.createClusterStateWatchersAndUpdate();
+      reader.registerCore("c1"); // watching c1, so it should get non lazy reference
+      zkClient.makePath(ZkStateReader.COLLECTIONS_ZKNODE + "/c1", true);
+
+      reader.forciblyRefreshAllClusterStateSlow();
+      // Initially there should be no c1 collection.
+      assertNull(reader.getClusterState().getCollectionRef("c1"));
+
+      ZkStateWriter writer = new ZkStateWriter(reader, new Stats());
+
+      // create new collection
+      DocCollection state =
+          new DocCollection(
+              "c1",
+              new HashMap<>(),
+              Map.of(ZkStateReader.CONFIGNAME_PROP, ConfigSetsHandler.DEFAULT_CONFIGSET_NAME),
+              DocRouter.DEFAULT,
+              0);
+      ZkWriteCommand wc = new ZkWriteCommand("c1", state);
+      writer.enqueueUpdate(reader.getClusterState(), Collections.singletonList(wc), null);
+      writer.writePendingUpdates();
+
+      assertTrue(zkClient.exists(ZkStateReader.COLLECTIONS_ZKNODE + "/c1/state.json", true));
+
+      reader.forciblyRefreshAllClusterStateSlow();
+      ClusterState.CollectionRef ref = reader.getClusterState().getCollectionRef("c1");
+      assertNotNull(ref);
+      assertFalse(ref.isLazilyLoaded());
+      assertEquals(0, ref.get().getZNodeVersion());
+
+      // update the collection
+      state =
+          new DocCollection(
+              "c1",
+              new HashMap<>(),
+              Map.of(ZkStateReader.CONFIGNAME_PROP, ConfigSetsHandler.DEFAULT_CONFIGSET_NAME),
+              DocRouter.DEFAULT,
+              ref.get().getZNodeVersion());
+      wc = new ZkWriteCommand("c1", state);
+      writer.enqueueUpdate(reader.getClusterState(), Collections.singletonList(wc), null);
+      writer.writePendingUpdates();
+
+      reader.forciblyRefreshAllClusterStateSlow();
+      ref = reader.getClusterState().getCollectionRef("c1");
+      assertNotNull(ref);
+      assertFalse(ref.isLazilyLoaded());
+      assertEquals(1, ref.get().getZNodeVersion());
+
+      // delete the collection c1, add a collection c2 that is NOT watched
+      ZkWriteCommand wc1 = new ZkWriteCommand("c1", null);
+
+      zkClient.makePath(ZkStateReader.COLLECTIONS_ZKNODE + "/c2", true);
+      state =
+          new DocCollection(
+              "c2",
+              new HashMap<>(),
+              Map.of(ZkStateReader.CONFIGNAME_PROP, ConfigSetsHandler.DEFAULT_CONFIGSET_NAME),
+              DocRouter.DEFAULT,
+              0);
+      ZkWriteCommand wc2 = new ZkWriteCommand("c2", state);
+
+      writer.enqueueUpdate(reader.getClusterState(), Arrays.asList(wc1, wc2), null);
+      writer.writePendingUpdates();
+
+      reader.forciblyRefreshAllClusterStateSlow();
+      ref = reader.getClusterState().getCollectionRef("c1");
+      assertNull(ref);
+
+      ref = reader.getClusterState().getCollectionRef("c2");
+      assertNotNull(ref);
+      assert (ref.isLazilyLoaded()); // c2 should be lazily loaded as it's not watched
+      assertEquals(0, ref.get().getZNodeVersion());
+    } finally {
+      IOUtils.close(reader, zkClient);
+      server.shutdown();
+    }
+  }
+
+  public void testGetCurrentCollections() throws Exception {
+    Path zkDir = createTempDir("testGetCurrentCollections");
+
+    ZkTestServer server = new ZkTestServer(zkDir);
+
+    SolrZkClient zkClient = null;
+    ZkStateReader reader = null;
+
+    try {
+      server.run();
+
+      zkClient = new SolrZkClient(server.getZkAddress(), OverseerTest.DEFAULT_CONNECTION_TIMEOUT);
+      ZkController.createClusterZkNodes(zkClient);
+
+      reader = new ZkStateReader(zkClient);
+      reader.createClusterStateWatchersAndUpdate();
+      reader.registerCore("c1"); // listen to c1. not yet exist
+      zkClient.makePath(ZkStateReader.COLLECTIONS_ZKNODE + "/c1", true);
+      reader.forceUpdateCollection("c1");
+      Set<String> currentCollections = reader.getCurrentCollections();
+      assertEquals(0, currentCollections.size()); // no active collections yet
+
+      // now create both c1 (watched) and c2 (not watched)
+      DocCollection state1 =
+          new DocCollection(
+              "c1",
+              new HashMap<>(),
+              Map.of(ZkStateReader.CONFIGNAME_PROP, ConfigSetsHandler.DEFAULT_CONFIGSET_NAME),
+              DocRouter.DEFAULT,
+              0);
+      ZkWriteCommand wc1 = new ZkWriteCommand("c1", state1);
+      DocCollection state2 =
+          new DocCollection(
+              "c2",
+              new HashMap<>(),
+              Map.of(ZkStateReader.CONFIGNAME_PROP, ConfigSetsHandler.DEFAULT_CONFIGSET_NAME),
+              DocRouter.DEFAULT,
+              0);
+
+      // do not listen to c2
+      zkClient.makePath(ZkStateReader.COLLECTIONS_ZKNODE + "/c2", true);
+      ZkWriteCommand wc2 = new ZkWriteCommand("c2", state2);
+
+      ZkStateWriter writer = new ZkStateWriter(reader, new Stats());
+      writer.enqueueUpdate(reader.getClusterState(), Arrays.asList(wc1, wc2), null);
+      writer.writePendingUpdates();
+
+      reader.forceUpdateCollection("c1");
+      reader.forceUpdateCollection("c2");
+      currentCollections =
+          reader.getCurrentCollections(); // should detect both collections (c1 watched, c2 lazy
+      // loaded)
+      assertEquals(2, currentCollections.size());
+    } finally {
+      IOUtils.close(reader, zkClient);
+      server.shutdown();
+    }
+  }
+
+  public void testWatchRaceCondition() throws Exception {
+    final int RUN_COUNT = 10000;
+    Path zkDir = createTempDir("testWatchRaceCondition");
+
+    ZkTestServer server = new ZkTestServer(zkDir);
+
+    SolrZkClient zkClient = null;
+    ZkStateReader reader = null;
+    ExecutorService executorService =
+        ExecutorUtil.newMDCAwareSingleThreadExecutor(
+            new SolrNamedThreadFactory("zkStateReaderTest"));
+
+    try {
+      server.run();
+
+      zkClient = new SolrZkClient(server.getZkAddress(), OverseerTest.DEFAULT_CONNECTION_TIMEOUT);
+      ZkController.createClusterZkNodes(zkClient);
+
+      reader = new ZkStateReader(zkClient);
+      final ZkStateReader readerRef = reader;
+      reader.createClusterStateWatchersAndUpdate();
+      zkClient.makePath(ZkStateReader.COLLECTIONS_ZKNODE + "/c1", true);
+
+      // start another thread to constantly updating the state
+      final AtomicBoolean stopMutatingThread = new AtomicBoolean(false);
+      final ZkStateWriter writer = new ZkStateWriter(reader, new Stats());
+      final AtomicInteger updateCounts = new AtomicInteger(0);
+      final AtomicReference<Exception> updateException = new AtomicReference<>();
+      executorService.submit(
+          () -> {
+            try {
+              ClusterState clusterState = readerRef.getClusterState();
+              while (!stopMutatingThread.get()) {
+                DocCollection collection = clusterState.getCollectionOrNull("c1");
+                int currentVersion = collection != null ? collection.getZNodeVersion() : 0;
+                // create new collection
+                DocCollection state =
+                    new DocCollection(
+                        "c1",
+                        new HashMap<>(),
+                        Map.of(
+                            ZkStateReader.CONFIGNAME_PROP,
+                            ConfigSetsHandler.DEFAULT_CONFIGSET_NAME),
+                        DocRouter.DEFAULT,
+                        currentVersion);
+                ZkWriteCommand wc = new ZkWriteCommand("c1", state);
+                writer.enqueueUpdate(clusterState, Collections.singletonList(wc), null);
+                clusterState = writer.writePendingUpdates();
+              }
+            } catch (Exception e) {
+              updateException.set(e);
+            }
+            return null;
+          });
+      executorService.shutdown();
+
+      for (int i = 0; i < RUN_COUNT; i++) {
+        final CountDownLatch latch = new CountDownLatch(2);
+
+        // remove itself on 2nd trigger
+        DocCollectionWatcher dummyWatcher =
+            collection -> {
+              latch.countDown();
+              return latch.getCount() == 0;
+            };
+        reader.registerDocCollectionWatcher("c1", dummyWatcher);
+        latch.await(10, TimeUnit.SECONDS);
+        reader.removeDocCollectionWatcher("c1", dummyWatcher);
+
+        boolean refLazilyLoaded = false;
+        for (int j = 0; j < 10; j++) {
+          if (reader.getClusterState().getCollectionRef("c1").isLazilyLoaded()) {
+            refLazilyLoaded = true; // it should eventually be lazily loaded
+            break;
+          }
+          int attempt = j + 1;
+          log.info("ref is not lazily loaded yet. Attempt : {}", attempt);
+
+          TimeUnit.MILLISECONDS.sleep(100);
+        }
+        assert (refLazilyLoaded);
+      }
+
+      stopMutatingThread.set(true);
+      if (updateException.get() != null) {
+        throw (updateException.get());
+      }
+
+    } finally {
+      IOUtils.close(reader, zkClient);
+      executorService.awaitTermination(10, TimeUnit.SECONDS);
+      server.shutdown();
+    }
+  }
 }