apoc.util.TestcontainersCausalCluster Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of apoc-test-utils Show documentation
Show all versions of apoc-test-utils Show documentation
Test utils for Neo4j Procedures
/*
* Copyright (c) "Neo4j"
* Neo4j Sweden AB [http://neo4j.com]
*
* This file is part of Neo4j.
*
* 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 apoc.util;
import static apoc.util.TestContainerUtil.password;
import static java.util.stream.Collectors.joining;
import static java.util.stream.Collectors.toList;
import java.net.URI;
import java.time.Duration;
import java.util.AbstractMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.function.IntFunction;
import java.util.stream.IntStream;
import java.util.stream.Stream;
import org.neo4j.driver.AuthTokens;
import org.neo4j.driver.Driver;
import org.neo4j.driver.GraphDatabase;
import org.neo4j.driver.Session;
import org.testcontainers.containers.Network;
import org.testcontainers.containers.SocatContainer;
/*
* Thanks to Michael Simons that inspired this.
*
* https://github.com/michael-simons/junit-jupiter-causal-cluster-testcontainer-extension/blob/master/src/main/java/org/neo4j/junit/jupiter/causal_cluster/CausalCluster.java
*/
public class TestcontainersCausalCluster {
private static int MINUTES_TO_WAIT = 5;
private static final int DEFAULT_BOLT_PORT = 7687;
public enum ClusterInstanceType {
CORE(DEFAULT_BOLT_PORT, "PRIMARY"),
READ_REPLICA(DEFAULT_BOLT_PORT + 1000, "SECONDARY");
private final int port;
private final String mode;
ClusterInstanceType(int port, String mode) {
this.port = port;
this.mode = mode;
}
}
private static Stream> iterateMembers(
int numOfMembers, ClusterInstanceType instanceType) {
final IntFunction generateInstanceName = i -> String.format("neo4j-%s-%d", instanceType.toString(), i);
return IntStream.rangeClosed(1, numOfMembers)
.mapToObj(i -> new AbstractMap.SimpleEntry<>(i - 1, generateInstanceName.apply(i)));
}
public static TestcontainersCausalCluster create(
List apocPackages,
int numberOfCoreMembers,
int numberOfReadReplica,
Duration timeout,
Map neo4jConfig,
Map envSettings) {
if (numberOfCoreMembers < 3) {
throw new IllegalArgumentException("numberOfCoreMembers must be >= 3");
}
if (numberOfReadReplica < 0) {
throw new IllegalArgumentException("numberOfReadReplica must be >= 0");
}
// Setup a naming strategy and the initial discovery members
final String initialDiscoveryMembers = iterateMembers(numberOfCoreMembers, ClusterInstanceType.CORE)
.map(n -> String.format("%s:5000", n.getValue()))
.collect(joining(","));
// Prepare one shared network for those containers
Network network = Network.newNetwork();
// Prepare proxy as sidecar
final SocatContainer proxy = new SocatContainer().withNetwork(network);
iterateMembers(numberOfCoreMembers, ClusterInstanceType.CORE)
.forEach(member -> proxy.withTarget(
ClusterInstanceType.CORE.port + member.getKey(), member.getValue(), DEFAULT_BOLT_PORT));
iterateMembers(numberOfReadReplica, ClusterInstanceType.READ_REPLICA)
.forEach(member -> proxy.withTarget(
ClusterInstanceType.READ_REPLICA.port + member.getKey(), member.getValue(), DEFAULT_BOLT_PORT));
proxy.start();
// Build the core/read_replica
List members = iterateMembers(numberOfCoreMembers, ClusterInstanceType.CORE)
.map(member -> createInstance(
apocPackages,
member.getValue(),
ClusterInstanceType.CORE,
network,
initialDiscoveryMembers,
neo4jConfig,
envSettings)
// Allocate the user database neo4j in every instance
// This is because the containers wait for /db/neo4j/cluster/available to return 200
// but by default from 5.x onwards not every database is allocated in every instance
// so the endpoint would return 404 and we would not complete container startup
.withNeo4jConfig("initial.dbms.default_primaries_count", Integer.toString(numberOfCoreMembers))
.withNeo4jConfig("server.default_advertised_address", member.getValue())
.withNeo4jConfig(
"server.bolt.advertised_address",
String.format(
"%s:%d",
proxy.getContainerIpAddress(),
proxy.getMappedPort(ClusterInstanceType.CORE.port + member.getKey()))))
.collect(toList());
members.addAll(iterateMembers(numberOfReadReplica, ClusterInstanceType.READ_REPLICA)
.map(member -> createInstance(
apocPackages,
member.getValue(),
ClusterInstanceType.READ_REPLICA,
network,
initialDiscoveryMembers,
neo4jConfig,
envSettings)
.withNeo4jConfig("dbms.default_advertised_address", member.getValue())
.withNeo4jConfig(
"server.bolt.advertised_address",
String.format(
"%s:%d",
proxy.getContainerIpAddress(),
proxy.getMappedPort(ClusterInstanceType.READ_REPLICA.port + member.getKey()))))
.collect(toList()));
// Start all of them in parallel
final CountDownLatch latch = new CountDownLatch(numberOfCoreMembers + numberOfReadReplica);
members.forEach(instance -> CompletableFuture.runAsync(() -> {
instance.start();
latch.countDown();
}));
try {
latch.await(MINUTES_TO_WAIT, TimeUnit.MINUTES);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
return new TestcontainersCausalCluster(members, proxy);
}
private static Neo4jContainerExtension createInstance(
List apocPackages,
String name,
ClusterInstanceType instanceType,
Network network,
String initialDiscoveryMembers,
Map neo4jConfig,
Map envSettings) {
Neo4jContainerExtension container = TestContainerUtil.createEnterpriseDB(
apocPackages, !TestUtil.isRunningInCI())
.withLabel("memberType", instanceType.toString())
.withNetwork(network)
.withNetworkAliases(name)
.withCreateContainerCmdModifier(cmd -> cmd.withHostName(name))
.withNeo4jConfig("initial.server.mode_constraint", instanceType.mode)
.withNeo4jConfig("server.default_listen_address", "0.0.0.0")
.withNeo4jConfig("dbms.cluster.raft.leader_transfer.balancing_strategy", "NO_BALANCING")
.withNeo4jConfig("dbms.cluster.discovery.endpoints", initialDiscoveryMembers)
.withStartupTimeout(Duration.ofMinutes(MINUTES_TO_WAIT));
if (withRoutingEnabled(envSettings)) {
container
.withEnv("NEO4J_server_routing_listen__address", "0.0.0.0:7618")
.withEnv("NEO4J_dbms_routing_default__router", "SERVER")
.withEnv("NEO4J_server_routing_advertised__address", name + ":7618");
} else {
container.withoutDriver();
}
neo4jConfig.forEach((conf, value) -> container.withNeo4jConfig(conf, String.valueOf(value)));
container.withEnv(envSettings);
return container;
}
private static boolean withRoutingEnabled(Map envSettings) {
return "true".equals(envSettings.get("NEO4J_dbms_routing_enabled"));
}
private final List clusterMembers;
private final SocatContainer sidecar;
private Driver driver;
private Session session;
public TestcontainersCausalCluster(List clusterMembers, SocatContainer sidecars) {
this.clusterMembers = clusterMembers;
this.sidecar = sidecars;
this.driver = GraphDatabase.driver(getURI(), AuthTokens.basic("neo4j", password));
this.session = driver.session();
}
public List getClusterMembers() {
return clusterMembers;
}
public Driver getDriver() {
return driver;
}
public Session getSession() {
return session;
}
public URI getURI() {
return Optional.of(this.sidecar)
.map(instance -> String.format(
"neo4j://%s:%d", instance.getContainerIpAddress(), instance.getMappedPort(DEFAULT_BOLT_PORT)))
.map(URI::create)
.orElseThrow(() -> new IllegalStateException("No sidecar as entrypoint into the cluster available."));
}
public void close() {
getSession().close();
getDriver().close();
sidecar.stop();
clusterMembers.forEach(Neo4jContainerExtension::stop);
}
}