Implement global (within factory) state for alternating connections.

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<YamlRunner.YamlConnectionFactory> alternateFactories;
+    @Nonnull
+    private final AtomicInteger currentConnection;
 
     public MultiServerConnectionFactory(@Nonnull final YamlRunner.YamlConnectionFactory defaultFactory,
                                         @Nonnull final List<YamlRunner.YamlConnectionFactory> alternateFactories) {
@@ -75,18 +77,18 @@ public MultiServerConnectionFactory(@Nonnull final ConnectionSelectionPolicy con
                                         @Nonnull final YamlRunner.YamlConnectionFactory defaultFactory,
                                         @Nonnull final List<YamlRunner.YamlConnectionFactory> 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<RelationalConnection> alternateConnections(URI connectPath) {
         return alternateFactories.stream().map(factory -> {
@@ -110,6 +116,24 @@ private List<RelationalConnection> 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.
      */

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<String> getVersionsUnderTest() {
+                return Set.of("0.0.0.0");
+            }
+        };
+    }
+}