feat: targetServerType=preferPrimary connection parameter (#2483)

mitya555 · web-flow · commit 8444ed6a8594 · 2022-05-02T09:36:06.000-04:00
* preferPrimary: documentation, code/logic, tests

* preferPrimary: fixed code style

* preferPrimary: new tests are in alphabetical order

* preferPrimary: simplify MultiHostChooser.candidateIterator(), fix the optimization in there
diff --git a/README.md b/README.md
@@ -128,7 +128,7 @@ In addition to the standard connection parameters the driver supports a number o
 | disableColumnSanitiser        | Boolean | false   | Enable optimization that disables column name sanitiser |
 | assumeMinServerVersion        | String  | null    | Assume the server is at least that version |
 | currentSchema                 | String  | null    | Specify the schema (or several schema separated by commas) to be set in the search-path |
-| targetServerType              | String  | any     | Specifies what kind of server to connect, possible values: any, master, slave (deprecated), secondary, preferSlave (deprecated), preferSecondary |
+| targetServerType              | String  | any     | Specifies what kind of server to connect, possible values: any, master, slave (deprecated), secondary, preferSlave (deprecated), preferSecondary, preferPrimary |
 | hostRecheckSeconds            | Integer | 10      | Specifies period (seconds) after which the host status is checked again in case it has changed |
 | loadBalanceHosts              | Boolean | false   | If disabled hosts are connected in the given order. If enabled hosts are chosen randomly from the set of suitable candidates |
 | socketFactory                 | String  | null    | Specify a socket factory for socket creation |
diff --git a/docs/documentation/head/connect.md b/docs/documentation/head/connect.md
@@ -475,10 +475,12 @@ Connection conn = DriverManager.getConnection(url);
 * **targetServerType** = String
 
 	Allows opening connections to only servers with required state, 
-	the allowed values are any, primary, master, slave, secondary, preferSlave and preferSecondary. 
+	the allowed values are any, primary, master, slave, secondary, preferSlave, preferSecondary and preferPrimary. 
 	The primary/secondary distinction is currently done by observing if the server allows writes. 
 	The value preferSecondary tries to connect to secondary if any are available, 
 	otherwise allows falls back to connecting also to primary.
+	The value preferPrimary tries to connect to primary if it is available, 
+	otherwise allows falls back to connecting to secondaries available.
 	- *N.B.* the words master and slave are being deprecated. We will silently accept them, but primary
 	and secondary are encouraged.
 
diff --git a/pgjdbc/src/main/java/org/postgresql/PGProperty.java b/pgjdbc/src/main/java/org/postgresql/PGProperty.java
@@ -693,7 +693,7 @@ public enum PGProperty {
       "any",
       "Specifies what kind of server to connect",
       false,
-      new String [] {"any", "primary", "master", "slave", "secondary",  "preferSlave", "preferSecondary"}),
+      new String [] {"any", "primary", "master", "slave", "secondary",  "preferSlave", "preferSecondary", "preferPrimary"}),
 
   /**
    * Enable or disable TCP keep-alive. The default is {@code false}.
diff --git a/pgjdbc/src/main/java/org/postgresql/hostchooser/HostRequirement.java b/pgjdbc/src/main/java/org/postgresql/hostchooser/HostRequirement.java
@@ -40,6 +40,11 @@ public boolean allowConnectingTo(@Nullable HostStatus status) {
     public boolean allowConnectingTo(@Nullable HostStatus status) {
       return status != HostStatus.ConnectFail;
     }
+  },
+  preferPrimary {
+    public boolean allowConnectingTo(@Nullable HostStatus status) {
+      return status != HostStatus.ConnectFail;
+    }
   };
 
   public abstract boolean allowConnectingTo(@Nullable HostStatus status);
diff --git a/pgjdbc/src/main/java/org/postgresql/hostchooser/MultiHostChooser.java b/pgjdbc/src/main/java/org/postgresql/hostchooser/MultiHostChooser.java
@@ -14,7 +14,6 @@
 import java.util.AbstractList;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Properties;
@@ -47,41 +46,42 @@ public Iterator<CandidateHost> iterator() {
       // In case all the candidate hosts are unavailable or do not match, try all the hosts just in case
       List<HostSpec> allHosts = Arrays.asList(hostSpecs);
       if (loadBalance) {
-        allHosts = new ArrayList<HostSpec>(allHosts);
-        Collections.shuffle(allHosts);
+        allHosts = new ArrayList<>(allHosts);
+        shuffle(allHosts);
       }
       res = withReqStatus(targetServerType, allHosts).iterator();
     }
     return res;
   }
 
   private Iterator<CandidateHost> candidateIterator() {
-    if (targetServerType != HostRequirement.preferSecondary) {
+    if (   targetServerType != HostRequirement.preferSecondary
+        && targetServerType != HostRequirement.preferPrimary   ) {
       return getCandidateHosts(targetServerType).iterator();
     }
 
+    HostRequirement preferredServerType =
+        targetServerType == HostRequirement.preferSecondary
+          ? HostRequirement.secondary
+          : HostRequirement.primary;
+
     // preferSecondary tries to find secondary hosts first
     // Note: sort does not work here since there are "unknown" hosts,
     // and that "unknown" might turn out to be master, so we should discard that
     // if other secondaries exist
-    List<CandidateHost> secondaries = getCandidateHosts(HostRequirement.secondary);
+    // Same logic as the above works for preferPrimary if we replace "secondary"
+    // with "primary" and vice versa
+    List<CandidateHost> preferred = getCandidateHosts(preferredServerType);
     List<CandidateHost> any = getCandidateHosts(HostRequirement.any);
 
-    if (secondaries.isEmpty()) {
-      return any.iterator();
-    }
-
-    if (any.isEmpty()) {
-      return secondaries.iterator();
-    }
-
-    if (secondaries.get(secondaries.size() - 1).equals(any.get(0))) {
-      // When the last secondary's hostspec is the same as the first in "any" list, there's no need
-      // to attempt to connect it as "secondary"
+    if (  !preferred.isEmpty() && !any.isEmpty()
+        && preferred.get(preferred.size() - 1).hostSpec.equals(any.get(0).hostSpec)) {
+      // When the last preferred host's hostspec is the same as the first in "any" list, there's no need
+      // to attempt to connect it as "preferred"
       // Note: this is only an optimization
-      secondaries = rtrim(1, secondaries);
+      preferred = rtrim(1, preferred);
     }
-    return append(secondaries, any).iterator();
+    return append(preferred, any).iterator();
   }
 
   private List<CandidateHost> getCandidateHosts(HostRequirement hostRequirement) {
diff --git a/pgjdbc/src/test/java/org/postgresql/test/hostchooser/MultiHostsConnectionTest.java b/pgjdbc/src/test/java/org/postgresql/test/hostchooser/MultiHostsConnectionTest.java
@@ -14,6 +14,7 @@
 import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeTrue;
 import static org.postgresql.hostchooser.HostRequirement.any;
+import static org.postgresql.hostchooser.HostRequirement.preferPrimary;
 import static org.postgresql.hostchooser.HostRequirement.preferSecondary;
 import static org.postgresql.hostchooser.HostRequirement.primary;
 import static org.postgresql.hostchooser.HostRequirement.secondary;
@@ -237,6 +238,27 @@ public void testConnectToMaster() throws SQLException {
     assertGlobalState(secondary1, "Secondary"); // was unknown, so tried to connect in order
   }
 
+  @Test
+  public void testConnectToPrimaryFirst() throws SQLException {
+    getConnection(preferPrimary, true, fake1, primary1, secondary1);
+    assertRemote(primaryIp);
+    assertGlobalState(fake1, "ConnectFail");
+    assertGlobalState(primary1, "Primary");
+    assertGlobalState(secondary1, null);
+
+    getConnection(primary, false, fake1, secondary1, primary1);
+    assertRemote(primaryIp);
+    assertGlobalState(fake1, "ConnectFail");
+    assertGlobalState(primary1, "Primary");
+    assertGlobalState(secondary1, "Secondary"); // tried as it was unknown
+
+    getConnection(preferPrimary, true, fake1, secondary1, primary1);
+    assertRemote(primaryIp);
+    assertGlobalState(fake1, "ConnectFail");
+    assertGlobalState(primary1, "Primary");
+    assertGlobalState(secondary1, "Secondary");
+  }
+
   @Test
   public void testConnectToPrimaryWithReadonlyTransactionMode() throws SQLException {
     con = TestUtil.openPrivilegedDB();
@@ -322,6 +344,47 @@ public void testLoadBalancing() throws SQLException {
     assertTrue("Never tried to connect to fake node", fake1FoundTried);
   }
 
+  @Test
+  public void testLoadBalancing_preferPrimary() throws SQLException {
+    Set<String> connectedHosts = new HashSet<String>();
+    Set<HostSpec> tryConnectedHosts = new HashSet<HostSpec>();
+    for (int i = 0; i < 20; ++i) {
+      getConnection(preferPrimary, true, true, fake1, secondary1, secondary2, primary1);
+      connectedHosts.add(getRemoteHostSpec());
+      tryConnectedHosts.addAll(hostStatusMap.keySet());
+      if (tryConnectedHosts.size() == 4) {
+        break;
+      }
+    }
+
+    assertRemote(primaryIp);
+    assertEquals("Connected to hosts other than primary", new HashSet<String>(asList(primaryIp)),
+        connectedHosts);
+    assertEquals("Never tried to connect to fake node", 4, tryConnectedHosts.size());
+
+    getConnection(preferPrimary, false, true, fake1, secondary1, primary1);
+    assertRemote(primaryIp);
+
+    // connect to secondaries when there's no primary - with load balancing
+    connectedHosts.clear();
+    for (int i = 0; i < 20; ++i) {
+      getConnection(preferPrimary, false, true, fake1, secondary1, secondary2);
+      connectedHosts.add(getRemoteHostSpec());
+      if (connectedHosts.size() == 2) {
+        break;
+      }
+    }
+    assertEquals("Never connected to all secondary hosts", new HashSet<String>(asList(secondaryIP, secondaryIP2)),
+        connectedHosts);
+
+    // connect to secondary when there's no primary
+    getConnection(preferPrimary, true, true, fake1, secondary1);
+    assertRemote(secondaryIP);
+
+    getConnection(preferPrimary, false, true, fake1, secondary1);
+    assertRemote(secondaryIP);
+  }
+
   @Test
   public void testLoadBalancing_preferSecondary() throws SQLException {
     Set<String> connectedHosts = new HashSet<String>();
