feat: 11347: introduce ReconnectHalfMillionNodesBench (#11487)

Signed-off-by: Anthony Petrov <anthony@swirldslabs.com>
Signed-off-by: Timo Brandstätter <timo@swirldslabs.com>

build-logic/project-plugins/src/main/kotlin/com.hedera.hashgraph.benchmark-conventions.gradle.kts

@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+import me.champeau.jmh.JMHTask
+
 plugins { id("me.champeau.jmh") }
 
 jmh {
@@ -23,6 +25,12 @@ jmh {
 
 tasks.jmh { outputs.upToDateWhen { false } }
 
+tasks.withType<JMHTask>().configureEach {
+    group = "jmh"
+    jarArchive = tasks.jmhJar.flatMap { it.archiveFile }
+    jvm = javaToolchains.launcherFor(java.toolchain).map { it.executablePath }.get().asFile.path
+}
+
 tasks.jmhJar { manifest { attributes(mapOf("Multi-Release" to true)) } }
 
 configurations {

platform-sdk/swirlds-virtualmap/build.gradle.kts

@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+import me.champeau.jmh.JMHTask
+
 plugins {
     id("com.hedera.hashgraph.sdk.conventions")
     id("com.hedera.hashgraph.platform-maven-publish")
@@ -25,7 +27,12 @@ mainModuleInfo { annotationProcessor("com.swirlds.config.processor") }
 
 jmhModuleInfo {
     requires("com.swirlds.common")
+    requires("com.swirlds.common.test.fixtures")
+    requires("com.swirlds.config.api")
+    requires("com.swirlds.config.extensions.test.fixtures")
+    requires("com.swirlds.virtualmap.test.fixtures")
     requires("jmh.core")
+    requires("org.junit.jupiter.api")
 }
 
 testModuleInfo {
@@ -46,3 +53,13 @@ hammerModuleInfo {
     requires("org.junit.jupiter.api")
     runtimeOnly("com.swirlds.config.impl")
 }
+
+tasks.register<JMHTask>("jmhReconnect") {
+    includes.set(listOf("Reconnect.*"))
+    jvmArgs.set(listOf("-Xmx16g"))
+    fork.set(1)
+    warmupIterations.set(2)
+    iterations.set(5)
+
+    resultsFile.convention(layout.buildDirectory.file("results/jmh/results-reconnect.txt"))
+}

platform-sdk/swirlds-virtualmap/src/jmh/java/com/swirlds/virtualmap/benchmark/reconnect/ReconnectHalfMillionNodesBench.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright (C) 2024 Hedera Hashgraph, LLC
+ *
+ * Licensed 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 com.swirlds.virtualmap.benchmark.reconnect;
+
+import com.swirlds.common.constructable.ConstructableRegistryException;
+import com.swirlds.virtualmap.test.fixtures.TestKey;
+import com.swirlds.virtualmap.test.fixtures.TestValue;
+import java.io.FileNotFoundException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+
+@BenchmarkMode(Mode.AverageTime)
+@State(Scope.Thread)
+public class ReconnectHalfMillionNodesBench extends VirtualMapReconnectBenchBase {
+
+    private static final Map<TestKey, TestValue> testTeacherMap = new HashMap<>();
+    private static final Map<TestKey, TestValue> testLearnerMap = new HashMap<>();
+
+    static {
+        try {
+            VirtualMapReconnectBenchBase.startup();
+        } catch (ConstructableRegistryException e) {
+            throw new RuntimeException(e);
+        } catch (FileNotFoundException e) {
+            throw new RuntimeException(e);
+        }
+
+        // Create a state to be reused in every run
+        StateBuilder.buildState(new Random(9823452658L), 500_000, 0.15, 0.15, testTeacherMap::put, testLearnerMap::put);
+    }
+
+    @Setup(Level.Invocation)
+    @Override
+    public void setupEach() {
+        super.setupEach();
+
+        testTeacherMap.entrySet().forEach(e -> teacherMap.put(e.getKey(), e.getValue()));
+        testLearnerMap.entrySet().forEach(e -> learnerMap.put(e.getKey(), e.getValue()));
+    }
+
+    @Benchmark
+    public void reconnectHalfMillionNodes() throws Exception {
+        super.reconnect();
+    }
+}

platform-sdk/swirlds-virtualmap/src/jmh/java/com/swirlds/virtualmap/benchmark/reconnect/StateBuilder.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright (C) 2024 Hedera Hashgraph, LLC
+ *
+ * Licensed 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 com.swirlds.virtualmap.benchmark.reconnect;
+
+import com.swirlds.common.test.fixtures.RandomUtils;
+import com.swirlds.virtualmap.test.fixtures.TestKey;
+import com.swirlds.virtualmap.test.fixtures.TestValue;
+import java.util.Random;
+import java.util.function.BiConsumer;
+import java.util.stream.LongStream;
+
+/**
+ * A utility class to help build random states.
+ */
+public class StateBuilder {
+
+    /** Return {@code true} with the given probability. */
+    private static boolean isRandomOutcome(final Random random, final double probability) {
+        return random.nextDouble(1.) < probability;
+    }
+
+    /**
+     * Build a random state and pass it to the provided teacher and learner populators.
+     *
+     * @param random a Random instance
+     * @param size the number of nodes in the teacher state.
+     *          The learner will have the same number of nodes or less.
+     * @param learnerMissingProbability the probability of a key to be missing in the learner state
+     * @param learnerDifferentProbability the probability of a node under a given key in the learner state
+     *          to have a value that is different from the value under the same key in the teacher state.
+     * @param teacherPopulator a BiConsumer that persists the teacher state (Map::put or similar)
+     * @param learnerPopulator a BiConsumer that persists the learner state (Map::put or similar)
+     */
+    public static void buildState(
+            final Random random,
+            final long size,
+            final double learnerMissingProbability,
+            final double learnerDifferentProbability,
+            final BiConsumer<TestKey, TestValue> teacherPopulator,
+            final BiConsumer<TestKey, TestValue> learnerPopulator) {
+        LongStream.range(1, size).forEach(i -> {
+            final TestKey key = new TestKey(i);
+
+            final TestValue teacherValue = new TestValue(RandomUtils.randomString(random, random.nextInt(1, 64)));
+            teacherPopulator.accept(key, teacherValue);
+
+            if (isRandomOutcome(random, learnerMissingProbability)) {
+                return;
+            }
+
+            final TestValue learnerValue;
+            if (isRandomOutcome(random, learnerDifferentProbability)) {
+                learnerValue = new TestValue(RandomUtils.randomString(random, random.nextInt(1, 64)));
+            } else {
+                learnerValue = teacherValue;
+            }
+            learnerPopulator.accept(key, learnerValue);
+        });
+    }
+}

platform-sdk/swirlds-virtualmap/src/jmh/java/com/swirlds/virtualmap/benchmark/reconnect/VirtualMapReconnectBenchBase.java

@@ -0,0 +1,118 @@
+/*
+ * Copyright (C) 2021-2024 Hedera Hashgraph, LLC
+ *
+ * Licensed 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 com.swirlds.virtualmap.benchmark.reconnect;
+
+import static com.swirlds.common.test.fixtures.io.ResourceLoader.loadLog4jContext;
+
+import com.swirlds.common.constructable.ClassConstructorPair;
+import com.swirlds.common.constructable.ConstructableRegistry;
+import com.swirlds.common.constructable.ConstructableRegistryException;
+import com.swirlds.common.merkle.MerkleInternal;
+import com.swirlds.common.merkle.synchronization.config.ReconnectConfig;
+import com.swirlds.common.merkle.synchronization.config.ReconnectConfig_;
+import com.swirlds.common.merkle.synchronization.internal.QueryResponse;
+import com.swirlds.common.test.fixtures.merkle.dummy.DummyMerkleInternal;
+import com.swirlds.common.test.fixtures.merkle.dummy.DummyMerkleLeaf;
+import com.swirlds.common.test.fixtures.merkle.util.MerkleTestUtils;
+import com.swirlds.config.api.Configuration;
+import com.swirlds.config.extensions.test.fixtures.TestConfigBuilder;
+import com.swirlds.virtualmap.VirtualMap;
+import com.swirlds.virtualmap.datasource.VirtualDataSourceBuilder;
+import com.swirlds.virtualmap.datasource.VirtualLeafRecord;
+import com.swirlds.virtualmap.internal.merkle.VirtualMapState;
+import com.swirlds.virtualmap.internal.merkle.VirtualRootNode;
+import com.swirlds.virtualmap.internal.pipeline.VirtualRoot;
+import com.swirlds.virtualmap.test.fixtures.InMemoryBuilder;
+import com.swirlds.virtualmap.test.fixtures.TestKey;
+import com.swirlds.virtualmap.test.fixtures.TestValue;
+import java.io.FileNotFoundException;
+import org.junit.jupiter.api.Assertions;
+
+/**
+ * The code is partially borrowed from VirtualMapReconnectTestBase.java in swirlds-virtualmap/src/test/.
+ * Ideally, it belongs to a shared test fixture, but I was unable to find a way to resolve dependencies
+ * between projects and modules, so I created this copy here and removed a few static definitions that
+ * are irrelevant to JMH benchmarks. In the future, this JMH-specific copy may in fact diverge
+ * from the unit test base class if/when we implement performance testing-related features here
+ * (e.g. artificial latencies etc.)
+ */
+public abstract class VirtualMapReconnectBenchBase {
+
+    protected VirtualMap<TestKey, TestValue> teacherMap;
+    protected VirtualMap<TestKey, TestValue> learnerMap;
+    protected VirtualDataSourceBuilder<TestKey, TestValue> teacherBuilder;
+    protected VirtualDataSourceBuilder<TestKey, TestValue> learnerBuilder;
+
+    protected final Configuration configuration = new TestConfigBuilder().getOrCreateConfig();
+    protected final ReconnectConfig reconnectConfig = configuration.getConfigData(ReconnectConfig.class);
+
+    protected VirtualDataSourceBuilder<TestKey, TestValue> createBuilder() {
+        return new InMemoryBuilder();
+    }
+
+    protected void setupEach() {
+        teacherBuilder = createBuilder();
+        learnerBuilder = createBuilder();
+        teacherMap = new VirtualMap<>("Teacher", teacherBuilder);
+        learnerMap = new VirtualMap<>("Learner", learnerBuilder);
+    }
+
+    protected static void startup() throws ConstructableRegistryException, FileNotFoundException {
+        loadLog4jContext();
+        final ConstructableRegistry registry = ConstructableRegistry.getInstance();
+        registry.registerConstructables("com.swirlds.common");
+        registry.registerConstructables("com.swirlds.virtualmap");
+        registry.registerConstructable(new ClassConstructorPair(QueryResponse.class, QueryResponse::new));
+        registry.registerConstructable(new ClassConstructorPair(DummyMerkleInternal.class, DummyMerkleInternal::new));
+        registry.registerConstructable(new ClassConstructorPair(DummyMerkleLeaf.class, DummyMerkleLeaf::new));
+        registry.registerConstructable(new ClassConstructorPair(VirtualLeafRecord.class, VirtualLeafRecord::new));
+        registry.registerConstructable(new ClassConstructorPair(VirtualMap.class, VirtualMap::new));
+        registry.registerConstructable(new ClassConstructorPair(VirtualMapState.class, VirtualMapState::new));
+        registry.registerConstructable(new ClassConstructorPair(VirtualRootNode.class, VirtualRootNode::new));
+        registry.registerConstructable(new ClassConstructorPair(TestKey.class, TestKey::new));
+        registry.registerConstructable(new ClassConstructorPair(TestValue.class, TestValue::new));
+
+        new TestConfigBuilder()
+                .withValue(ReconnectConfig_.ACTIVE, "true")
+                // This is lower than the default, helps test that is supposed to fail to finish faster.
+                .withValue(ReconnectConfig_.ASYNC_STREAM_TIMEOUT, "5000ms")
+                .getOrCreateConfig();
+    }
+
+    protected MerkleInternal createTreeForMap(VirtualMap<TestKey, TestValue> map) {
+        final var tree = MerkleTestUtils.buildLessSimpleTree();
+        tree.getChild(1).asInternal().setChild(3, map);
+        tree.reserve();
+        return tree;
+    }
+
+    protected void reconnect() throws Exception {
+        final MerkleInternal teacherTree = createTreeForMap(teacherMap);
+        final VirtualMap<TestKey, TestValue> copy = teacherMap.copy();
+        final MerkleInternal learnerTree = createTreeForMap(learnerMap);
+        try {
+            final var node = MerkleTestUtils.hashAndTestSynchronization(learnerTree, teacherTree, reconnectConfig);
+            node.release();
+            final VirtualRoot root = learnerMap.getRight();
+            Assertions.assertTrue(root.isHashed(), "Learner root node must be hashed");
+        } finally {
+            teacherTree.release();
+            learnerTree.release();
+            copy.release();
+        }
+    }
+}