From f0188bb4ca73c7e8c279db72a9400958c2946731 Mon Sep 17 00:00:00 2001 From: ohad Date: Fri, 31 Jan 2025 11:15:11 -0500 Subject: [PATCH] Implement global (within factory) state for alternating connections. --- .../MultiServerConnectionFactory.java | 30 ++++++- .../MultiServerConnectionFactoryTest.java | 81 +++++++++++++++++++ 2 files changed, 108 insertions(+), 3 deletions(-) create mode 100644 yaml-tests/src/test/java/MultiServerConnectionFactoryTest.java diff --git a/yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/MultiServerConnectionFactory.java b/yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/MultiServerConnectionFactory.java index 02cc5bf17b..8d8ba9c60d 100644 --- a/yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/MultiServerConnectionFactory.java +++ b/yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/MultiServerConnectionFactory.java @@ -36,6 +36,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -59,11 +60,12 @@ public enum ConnectionSelectionPolicy { DEFAULT, ALTERNATE } @Nonnull private final ConnectionSelectionPolicy connectionSelectionPolicy; - private final int initialConnection; @Nonnull private final YamlRunner.YamlConnectionFactory defaultFactory; @Nonnull private final List alternateFactories; + @Nonnull + private final AtomicInteger currentConnection; public MultiServerConnectionFactory(@Nonnull final YamlRunner.YamlConnectionFactory defaultFactory, @Nonnull final List alternateFactories) { @@ -75,18 +77,18 @@ public MultiServerConnectionFactory(@Nonnull final ConnectionSelectionPolicy con @Nonnull final YamlRunner.YamlConnectionFactory defaultFactory, @Nonnull final List alternateFactories) { this.connectionSelectionPolicy = connectionSelectionPolicy; - this.initialConnection = initialConnection; this.defaultFactory = defaultFactory; this.alternateFactories = alternateFactories; this.versionsUnderTest = Stream.concat(Stream.of(defaultFactory), alternateFactories.stream()) .flatMap(yamlConnectionFactory -> yamlConnectionFactory.getVersionsUnderTest().stream()) .collect(Collectors.toSet()); + this.currentConnection = new AtomicInteger(initialConnection); } @Override public RelationalConnection getNewConnection(@Nonnull URI connectPath) throws SQLException { - return new MultiServerRelationalConnection(connectionSelectionPolicy, initialConnection, defaultFactory.getNewConnection(connectPath), alternateConnections(connectPath)); + return new MultiServerRelationalConnection(connectionSelectionPolicy, getNextConnectionNumber(), defaultFactory.getNewConnection(connectPath), alternateConnections(connectPath)); } @Override @@ -99,6 +101,10 @@ public boolean isMultiServer() { return true; } + public int getCurrentConnection() { + return currentConnection.get(); + } + @Nonnull private List alternateConnections(URI connectPath) { return alternateFactories.stream().map(factory -> { @@ -110,6 +116,24 @@ private List alternateConnections(URI connectPath) { }).collect(Collectors.toList()); } + /** + * Increment and return the next connection's initialConnection number. + * This allows us to better distribute the connections positions as connections are created per query in the tests + * and thus connections with a single query should increment their position between connection creation or else they + * will always execute with the same initial connection. + * @return the next initial connection number to use + */ + private int getNextConnectionNumber() { + switch (connectionSelectionPolicy) { + case DEFAULT: + return DEFAULT_CONNECTION; + case ALTERNATE: + return currentConnection.addAndGet(1) % (1 + alternateFactories.size()); + default: + throw new IllegalArgumentException("Unknown policy"); + } + } + /** * A connection that wraps around multiple connections. */ diff --git a/yaml-tests/src/test/java/MultiServerConnectionFactoryTest.java b/yaml-tests/src/test/java/MultiServerConnectionFactoryTest.java new file mode 100644 index 0000000000..dd89c4669b --- /dev/null +++ b/yaml-tests/src/test/java/MultiServerConnectionFactoryTest.java @@ -0,0 +1,81 @@ +/* + * MultiServerFactoryTest.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * 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. + */ + +import com.apple.foundationdb.relational.api.RelationalConnection; +import com.apple.foundationdb.relational.yamltests.MultiServerConnectionFactory; +import com.apple.foundationdb.relational.yamltests.YamlRunner; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; + +import javax.annotation.Nonnull; +import java.net.URI; +import java.sql.SQLException; +import java.util.List; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +public class MultiServerConnectionFactoryTest { + @ParameterizedTest + @CsvSource({"0", "1", "2", "3", "4"}) + void testDefaultPolicy(String initialConnection) throws SQLException { + int ic = Integer.parseInt(initialConnection); + MultiServerConnectionFactory classUnderTest = new MultiServerConnectionFactory( + MultiServerConnectionFactory.ConnectionSelectionPolicy.DEFAULT, + ic, + dummyConnectionFactory(), + List.of(dummyConnectionFactory())); + assertEquals(ic, classUnderTest.getCurrentConnection()); + classUnderTest.getNewConnection(URI.create("Blah")); + assertEquals(ic, classUnderTest.getCurrentConnection()); + classUnderTest.getNewConnection(URI.create("Blah")); + assertEquals(ic, classUnderTest.getCurrentConnection()); + } + + @ParameterizedTest + @CsvSource({"0", "1", "2", "3", "4"}) + void testAlternatePolicy(String initialConnection) throws SQLException { + int ic = Integer.parseInt(initialConnection); + MultiServerConnectionFactory classUnderTest = new MultiServerConnectionFactory( + MultiServerConnectionFactory.ConnectionSelectionPolicy.ALTERNATE, + ic, + dummyConnectionFactory(), + List.of(dummyConnectionFactory())); + assertEquals(ic, classUnderTest.getCurrentConnection()); + classUnderTest.getNewConnection(URI.create("Blah")); + assertEquals(ic + 1, classUnderTest.getCurrentConnection()); + classUnderTest.getNewConnection(URI.create("Blah")); + assertEquals(ic + 2, classUnderTest.getCurrentConnection()); + } + + YamlRunner.YamlConnectionFactory dummyConnectionFactory() { + return new YamlRunner.YamlConnectionFactory() { + @Override + public RelationalConnection getNewConnection(@Nonnull URI connectPath) throws SQLException { + return null; + } + + @Override + public Set getVersionsUnderTest() { + return Set.of("0.0.0.0"); + } + }; + } +}