Swith to a new connection if possible when asked to do so

Try to avoid reusing the same connection we want to avoid using,
only using it when there is just one source

4 files changed, 49 insertions, 15 deletions

diff --git a/config/src/main/java/com/yahoo/config/subscription/impl/MockConnection.java b/config/src/main/java/com/yahoo/config/subscription/impl/MockConnection.java
index 58eed7f9e78..3a06aa02dc0 100644
--- a/config/src/main/java/com/yahoo/config/subscription/impl/MockConnection.java
+++ b/config/src/main/java/com/yahoo/config/subscription/impl/MockConnection.java
@@ -7,6 +7,7 @@ import com.yahoo.jrt.Supervisor;
 import com.yahoo.vespa.config.ConfigPayload;
 import com.yahoo.vespa.config.Connection;
 import com.yahoo.vespa.config.ConnectionPool;
+import com.yahoo.vespa.config.JRTConnection;
 import com.yahoo.vespa.config.protocol.JRTServerConfigRequestV3;
 import com.yahoo.vespa.config.protocol.Payload;
 import com.yahoo.vespa.config.util.ConfigUtils;
@@ -87,9 +88,7 @@ public class MockConnection implements ConnectionPool, com.yahoo.vespa.config.Co
     }
 
     @Override
-    public Connection setNewCurrentConnection() {
-        return this;
-    }
+    public Connection switchConnection() { return this; }
 
     @Override
     public int getSize() {
diff --git a/config/src/main/java/com/yahoo/vespa/config/ConnectionPool.java b/config/src/main/java/com/yahoo/vespa/config/ConnectionPool.java
index 5a6f8a8848b..949cc16ce68 100644
--- a/config/src/main/java/com/yahoo/vespa/config/ConnectionPool.java
+++ b/config/src/main/java/com/yahoo/vespa/config/ConnectionPool.java
@@ -10,11 +10,32 @@ public interface ConnectionPool {
 
     void close();
 
+    /**
+     * Sets the supplied Connection to have an error, implementations are expected to call
+     * {@link #switchConnection()} after setting state for the supplied Connection.
+     *
+     */
     void setError(Connection connection, int i);
 
     Connection getCurrent();
 
-    Connection setNewCurrentConnection();
+    /**
+     * Switches to another JRTConnection instance by randomly choosing
+     * from the available sources, disregarding the current connection if there is
+     * more than one source. Returns the resulting Connection. See also {@link #setError(Connection, int)}
+     *
+     * @return a JRTConnection
+     */
+    Connection switchConnection();
+
+    /**
+     * Sets the current JRTConnection instance by randomly choosing
+     * from the available sources and returns the result.
+     *
+     * @return a JRTConnection
+     */
+    @Deprecated
+    default Connection setNewCurrentConnection() { return switchConnection(); };
 
     int getSize();
 
diff --git a/config/src/main/java/com/yahoo/vespa/config/JRTConnectionPool.java b/config/src/main/java/com/yahoo/vespa/config/JRTConnectionPool.java
index 2c1d00f295e..da77abbb648 100644
--- a/config/src/main/java/com/yahoo/vespa/config/JRTConnectionPool.java
+++ b/config/src/main/java/com/yahoo/vespa/config/JRTConnectionPool.java
@@ -13,6 +13,7 @@ import java.util.concurrent.ThreadLocalRandom;
 import java.util.logging.Logger;
 
 import static java.util.logging.Level.FINE;
+import static java.util.logging.Level.INFO;
 
 /**
  * A pool of JRT connections to a config source (either a config server or a config proxy).
@@ -53,7 +54,7 @@ public class JRTConnectionPool implements ConnectionPool {
                 connections.put(address, new JRTConnection(address, supervisor));
             }
         }
-        setNewCurrentConnection();
+        initialize();
     }
 
     /**
@@ -65,14 +66,27 @@ public class JRTConnectionPool implements ConnectionPool {
         return currentConnection;
     }
 
-    /**
-     * Returns and set the current JRTConnection instance by randomly choosing
-     * from the available sources (this means that you might end up using
-     * the same connection).
-     *
-     * @return a JRTConnection
-     */
+    @Override
+    public synchronized JRTConnection switchConnection() {
+        List<JRTConnection> sources = getSources();
+        if (sources.size() > 1) {
+            Connection previousConnection = currentConnection;
+            List<JRTConnection> sourcesWithoutCurrent = new ArrayList<>();
+            sources.stream()
+                    .filter(source -> ! source.equals(currentConnection))
+                    .forEach(sourcesWithoutCurrent::add);
+            currentConnection = sourcesWithoutCurrent.get(ThreadLocalRandom.current().nextInt(0, sourcesWithoutCurrent.size()));
+            log.log(INFO, () -> "Switching from " + previousConnection + " to " + currentConnection);
+        }
+        return currentConnection;
+    }
+
+    @Deprecated
     public synchronized JRTConnection setNewCurrentConnection() {
+        return initialize();
+    }
+
+    public synchronized JRTConnection initialize() {
         List<JRTConnection> sources = getSources();
         currentConnection = sources.get(ThreadLocalRandom.current().nextInt(0, sources.size()));
         log.log(FINE, () -> "Choosing new connection: " + currentConnection);
@@ -94,7 +108,7 @@ public class JRTConnectionPool implements ConnectionPool {
     @Override
     public void setError(Connection connection, int errorCode) {
         connection.setError(errorCode);
-        setNewCurrentConnection();
+        switchConnection();
     }
 
     public JRTConnectionPool updateSources(List<String> addresses) {
diff --git a/config/src/test/java/com/yahoo/vespa/config/JRTConnectionPoolTest.java b/config/src/test/java/com/yahoo/vespa/config/JRTConnectionPoolTest.java
index cc46301e869..56f40daf0fd 100644
--- a/config/src/test/java/com/yahoo/vespa/config/JRTConnectionPoolTest.java
+++ b/config/src/test/java/com/yahoo/vespa/config/JRTConnectionPoolTest.java
@@ -29,7 +29,7 @@ public class JRTConnectionPoolTest {
 
         Map<String, Integer> sourceOccurrences = new HashMap<>();
         for (int i = 0; i < 1000; i++) {
-            final String address = sourcePool.setNewCurrentConnection().getAddress();
+            final String address = sourcePool.switchConnection().getAddress();
             if (sourceOccurrences.containsKey(address)) {
                 sourceOccurrences.put(address, sourceOccurrences.get(address) + 1);
             } else {
@@ -57,7 +57,7 @@ public class JRTConnectionPoolTest {
 
         int count = 1000;
         for (int i = 0; i < count; i++) {
-            String address = sourcePool.setNewCurrentConnection().getAddress();
+            String address = sourcePool.switchConnection().getAddress();
             if (timesUsed.containsKey(address)) {
                 int times = timesUsed.get(address);
                 timesUsed.put(address, times + 1);