Merge pull request #29546 from vespa-engine/vekterli/use-fake-zk-database-for-subset-of-cc-testsv8.268.18

Use fake ZooKeeper database implementation for subset of CC tests

4 files changed, 159 insertions, 8 deletions

diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/DistributionBitCountTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/DistributionBitCountTest.java
index 0e48a9bbc45..1e13ed4eef5 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/DistributionBitCountTest.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/DistributionBitCountTest.java
@@ -14,9 +14,12 @@ import java.util.List;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
-@ExtendWith(CleanupZookeeperLogsOnSuccess.class)
 public class DistributionBitCountTest extends FleetControllerTest {
 
+    DistributionBitCountTest() {
+        useRealZooKeeperInTest(false);
+    }
+
     private FleetControllerOptions setUpSystem() throws Exception {
         List<ConfiguredNode> configuredNodes = new ArrayList<>();
         for (int i = 0 ; i < 10; i++) {
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/FakeZooKeeperDatabase.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/FakeZooKeeperDatabase.java
new file mode 100644
index 00000000000..057d137650f
--- /dev/null
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/FakeZooKeeperDatabase.java
@@ -0,0 +1,136 @@
+// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.clustercontroller.core;
+
+import com.yahoo.vdslib.state.Node;
+import com.yahoo.vdslib.state.NodeState;
+import com.yahoo.vespa.clustercontroller.core.database.Database;
+import com.yahoo.vespa.clustercontroller.core.database.DatabaseFactory;
+
+import java.util.Map;
+import java.util.TreeMap;
+
+/**
+ * Memory-backed fake DB implementation that tries to mirror the semantics of the
+ * (synchronous) ZooKeeper DB implementation. By itself this fake acts as if a quorum
+ * with a _single_, local ZK instance has been configured. This DB instance cannot be
+ * used across multiple cluster controller instances.
+ *
+ * Threading note: we expect all invocations on this instance to happen from the
+ * main cluster controller thread (i.e. "as-if" single threaded), but we wrap everything
+ * in a mutex to stay on the safe side since this isn't explicitly documented as
+ * part of the API,
+ */
+public class FakeZooKeeperDatabase extends Database {
+
+    public static class Factory implements DatabaseFactory {
+        private final FleetControllerContext context;
+        public Factory(FleetControllerContext context) {
+            this.context = context;
+        }
+        @Override
+        public Database create(Params params) {
+            return new FakeZooKeeperDatabase(context, params.listener);
+        }
+    }
+
+    private final FleetControllerContext context;
+    private final Database.DatabaseListener listener;
+
+    private final Object mutex = new Object();
+    private boolean closed = false;
+    private Integer persistedLatestStateVersion = null;
+    private Map<Integer, Integer> persistedLeaderVotes = new TreeMap<>();
+    private Map<Node, NodeState> persistedWantedStates = new TreeMap<>();
+    private Map<Node, Long> persistedStartTimestamps = new TreeMap<>();
+    private ClusterStateBundle persistedBundle = ClusterStateBundle.ofBaselineOnly(AnnotatedClusterState.emptyState());
+
+    public FakeZooKeeperDatabase(FleetControllerContext context, DatabaseListener listener) {
+        this.context = context;
+        this.listener = listener;
+    }
+
+    @Override
+    public void close() {
+        synchronized (mutex) {
+            closed = true;
+        }
+    }
+
+    @Override
+    public boolean isClosed() {
+        synchronized (mutex) {
+            return closed;
+        }
+    }
+
+    @Override
+    public boolean storeMasterVote(int voteForNode) {
+        Map<Integer, Integer> voteState;
+        synchronized (mutex) {
+            persistedLeaderVotes.put(context.id().index(), voteForNode);
+            voteState = Map.copyOf(persistedLeaderVotes);
+        }
+        listener.handleMasterData(voteState);
+        return true;
+    }
+
+    @Override
+    public boolean storeLatestSystemStateVersion(int version) {
+        synchronized (mutex) {
+            persistedLatestStateVersion = version;
+            return true;
+        }
+    }
+
+    @Override
+    public Integer retrieveLatestSystemStateVersion() {
+        synchronized (mutex) {
+            return persistedLatestStateVersion;
+        }
+    }
+
+    @Override
+    public boolean storeWantedStates(Map<Node, NodeState> states) {
+        synchronized (mutex) {
+            persistedWantedStates = Map.copyOf(states);
+        }
+        return true;
+    }
+
+    @Override
+    public Map<Node, NodeState> retrieveWantedStates() {
+        synchronized (mutex) {
+            return Map.copyOf(persistedWantedStates);
+        }
+    }
+
+    @Override
+    public boolean storeStartTimestamps(Map<Node, Long> timestamps) {
+        synchronized (mutex) {
+            persistedStartTimestamps = Map.copyOf(timestamps);
+            return true;
+        }
+    }
+
+    @Override
+    public Map<Node, Long> retrieveStartTimestamps() {
+        synchronized (mutex) {
+            return Map.copyOf(persistedStartTimestamps);
+        }
+    }
+
+    @Override
+    public boolean storeLastPublishedStateBundle(ClusterStateBundle stateBundle) {
+        synchronized (mutex) {
+            persistedBundle = stateBundle;
+            return true;
+        }
+    }
+
+    @Override
+    public ClusterStateBundle retrieveLastPublishedStateBundle() {
+        synchronized (mutex) {
+            return persistedBundle;
+        }
+    }
+}
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/FleetControllerTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/FleetControllerTest.java
index 400f4a14d24..ee0506070b5 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/FleetControllerTest.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/FleetControllerTest.java
@@ -16,7 +16,6 @@ import com.yahoo.vdslib.state.NodeState;
 import com.yahoo.vdslib.state.NodeType;
 import com.yahoo.vdslib.state.State;
 import com.yahoo.vespa.clustercontroller.core.database.DatabaseHandler;
-import com.yahoo.vespa.clustercontroller.core.database.ZooKeeperDatabaseFactory;
 import com.yahoo.vespa.clustercontroller.core.rpc.RPCCommunicator;
 import com.yahoo.vespa.clustercontroller.core.rpc.RpcServer;
 import com.yahoo.vespa.clustercontroller.core.rpc.SlobrokClient;
@@ -49,13 +48,14 @@ import static org.junit.jupiter.api.Assertions.fail;
  */
 public abstract class FleetControllerTest implements Waiter {
 
-    private static final Logger log = Logger.getLogger(FleetControllerTest.class.getName());
+    protected static final Logger log = Logger.getLogger(FleetControllerTest.class.getName());
     private static final int DEFAULT_NODE_COUNT = 10;
 
     private final Duration timeout = Duration.ofSeconds(30);
     protected Slobrok slobrok;
 
     protected FleetControllerOptions options;
+    private boolean useRealZooKeeperInTest = true;
     ZooKeeperTestServer zooKeeperServer;
     protected final List<FleetController> fleetControllers = new ArrayList<>();
     protected List<DummyVdsNode> nodes = new ArrayList<>();
@@ -73,6 +73,10 @@ public abstract class FleetControllerTest implements Waiter {
         LogSetup.initVespaLogging("fleetcontroller");
     }
 
+    protected void useRealZooKeeperInTest(boolean useRealZk) {
+        this.useRealZooKeeperInTest = useRealZk;
+    }
+
     protected static FleetControllerOptions.Builder defaultOptions() {
         return defaultOptions(IntStream.range(0, DEFAULT_NODE_COUNT)
                                       .mapToObj(i -> new ConfiguredNode(i, false))
@@ -121,7 +125,7 @@ public abstract class FleetControllerTest implements Waiter {
         var log = new EventLog(timer, metricUpdater);
         var cluster = new ContentCluster(options.clusterName(), options.nodes(), options.storageDistribution());
         var stateGatherer = new NodeStateGatherer(timer, timer, log);
-        var database = new DatabaseHandler(context, new ZooKeeperDatabaseFactory(context), timer, options.zooKeeperServerAddress(), timer);
+        var database = new DatabaseHandler(context, options.dbFactoryFn().apply(context), timer, options.zooKeeperServerAddress(), timer);
         // Setting this <1000 ms causes ECONNREFUSED on socket trying to connect to ZK server, in ZooKeeper,
         // after creating a new ZooKeeper (session).  This causes ~10s extra time to connect after connection loss.
         // Reasons unknown.  Larger values like the default 10_000 causes that much additional running time for some tests.
@@ -139,7 +143,13 @@ public abstract class FleetControllerTest implements Waiter {
     }
 
     protected FleetControllerOptions setUpFleetController(Timer timer, FleetControllerOptions.Builder builder) throws Exception {
-        setUpZooKeeperServer(builder);
+        // TODO consolidate CC setup in tests; currently partial duplication of
+        //  setup/init code across test subclasses.
+        if (useRealZooKeeperInTest) {
+            setUpZooKeeperServer(builder);
+        } else {
+            builder.setDbFactoryFn(FakeZooKeeperDatabase.Factory::new);
+        }
         builder.setSlobrokConnectionSpecs(getSlobrokConnectionSpecs(slobrok));
         options = builder.build();
         startFleetController(timer);
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/StateChangeTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/StateChangeTest.java
index 01987cf1f4c..9e9206526e5 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/StateChangeTest.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/StateChangeTest.java
@@ -9,7 +9,6 @@ import com.yahoo.vdslib.state.NodeType;
 import com.yahoo.vdslib.state.State;
 import com.yahoo.vespa.clustercontroller.core.testutils.StateWaiter;
 import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.ExtendWith;
 
 import java.time.Duration;
 import java.util.ArrayList;
@@ -23,7 +22,6 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
-@ExtendWith(CleanupZookeeperLogsOnSuccess.class)
 public class StateChangeTest extends FleetControllerTest {
 
     private final FakeTimer timer = new FakeTimer();
@@ -31,6 +29,10 @@ public class StateChangeTest extends FleetControllerTest {
     private FleetController ctrl;
     private DummyCommunicator communicator;
 
+    StateChangeTest() {
+        useRealZooKeeperInTest(false);
+    }
+
     private void initialize(FleetControllerOptions.Builder builder) throws Exception {
         List<Node> nodes = new ArrayList<>();
         for (int i = 0; i < builder.nodes().size(); ++i) {
@@ -38,7 +40,7 @@ public class StateChangeTest extends FleetControllerTest {
             nodes.add(new Node(NodeType.DISTRIBUTOR, i));
         }
 
-        setUpZooKeeperServer(builder);
+        builder.setDbFactoryFn(FakeZooKeeperDatabase.Factory::new);
         communicator = new DummyCommunicator(nodes, timer);
         boolean start = false;
         FleetControllerOptions options = builder.build();