chore(spanner): add package protected setter to enable mux for RW (#3327)

* chore(spanner): add package protected setter to enable mux for RW

* Update google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java

Co-authored-by: Knut Olav Løite <[email protected]>

---------

Co-authored-by: Knut Olav Løite <[email protected]>

Author: harshachinta

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java

@@ -73,6 +73,8 @@ public class SessionPoolOptions {
 
   private final boolean useMultiplexedSession;
 
+  private final boolean useMultiplexedSessionForRW;
+
   // TODO: Change to use java.time.Duration.
   private final Duration multiplexedSessionMaintenanceDuration;
 
@@ -108,6 +110,13 @@ private SessionPoolOptions(Builder builder) {
         (useMultiplexedSessionFromEnvVariable != null)
             ? useMultiplexedSessionFromEnvVariable
             : builder.useMultiplexedSession;
+    // useMultiplexedSessionForRW priority => Environment var > private setter > client default
+    Boolean useMultiplexedSessionForRWFromEnvVariable =
+        getUseMultiplexedSessionForRWFromEnvVariable();
+    this.useMultiplexedSessionForRW =
+        (useMultiplexedSessionForRWFromEnvVariable != null)
+            ? useMultiplexedSessionForRWFromEnvVariable
+            : builder.useMultiplexedSessionForRW;
     this.multiplexedSessionMaintenanceDuration = builder.multiplexedSessionMaintenanceDuration;
   }
 
@@ -144,6 +153,7 @@ public boolean equals(Object o) {
             this.inactiveTransactionRemovalOptions, other.inactiveTransactionRemovalOptions)
         && Objects.equals(this.poolMaintainerClock, other.poolMaintainerClock)
         && Objects.equals(this.useMultiplexedSession, other.useMultiplexedSession)
+        && Objects.equals(this.useMultiplexedSessionForRW, other.useMultiplexedSessionForRW)
         && Objects.equals(
             this.multiplexedSessionMaintenanceDuration,
             other.multiplexedSessionMaintenanceDuration);
@@ -174,6 +184,7 @@ public int hashCode() {
         this.inactiveTransactionRemovalOptions,
         this.poolMaintainerClock,
         this.useMultiplexedSession,
+        this.useMultiplexedSessionForRW,
         this.multiplexedSessionMaintenanceDuration);
   }
 
@@ -307,6 +318,14 @@ public boolean getUseMultiplexedSession() {
     return useMultiplexedSession;
   }
 
+  @VisibleForTesting
+  @InternalApi
+  public boolean getUseMultiplexedSessionForRW() {
+    // Multiplexed sessions for R/W are enabled only if both global multiplexed sessions and
+    // read-write multiplexed session flags are set to true.
+    return getUseMultiplexedSession() && useMultiplexedSessionForRW;
+  }
+
   private static Boolean getUseMultiplexedSessionFromEnvVariable() {
     String useMultiplexedSessionFromEnvVariable =
         System.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS");
@@ -323,6 +342,12 @@ private static Boolean getUseMultiplexedSessionFromEnvVariable() {
     return null;
   }
 
+  private static Boolean getUseMultiplexedSessionForRWFromEnvVariable() {
+    // Checks the value of env, GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW
+    // This returns null until RW is supported.
+    return null;
+  }
+
   Duration getMultiplexedSessionMaintenanceDuration() {
     return multiplexedSessionMaintenanceDuration;
   }
@@ -529,6 +554,12 @@ public static class Builder {
     // Set useMultiplexedSession to true to make multiplexed session the default.
     private boolean useMultiplexedSession = false;
 
+    // This field controls the default behavior of session management for RW operations in Java
+    // client.
+    // Set useMultiplexedSessionForRW to true to make multiplexed session for RW operations the
+    // default.
+    private boolean useMultiplexedSessionForRW = false;
+
     private Duration multiplexedSessionMaintenanceDuration = Duration.ofDays(7);
     private Clock poolMaintainerClock = Clock.INSTANCE;
 
@@ -570,6 +601,7 @@ private Builder(SessionPoolOptions options) {
       this.randomizePositionQPSThreshold = options.randomizePositionQPSThreshold;
       this.inactiveTransactionRemovalOptions = options.inactiveTransactionRemovalOptions;
       this.useMultiplexedSession = options.useMultiplexedSession;
+      this.useMultiplexedSessionForRW = options.useMultiplexedSessionForRW;
       this.multiplexedSessionMaintenanceDuration = options.multiplexedSessionMaintenanceDuration;
       this.poolMaintainerClock = options.poolMaintainerClock;
     }
@@ -757,6 +789,15 @@ Builder setUseMultiplexedSession(boolean useMultiplexedSession) {
       return this;
     }
 
+    /**
+     * Sets whether the client should use multiplexed session for R/W operations or not. This method
+     * is intentionally package-private and intended for internal use.
+     */
+    Builder setUseMultiplexedSessionForRW(boolean useMultiplexedSessionForRW) {
+      this.useMultiplexedSessionForRW = useMultiplexedSessionForRW;
+      return this;
+    }
+
     @VisibleForTesting
     Builder setMultiplexedSessionMaintenanceDuration(
         Duration multiplexedSessionMaintenanceDuration) {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java

@@ -277,6 +277,9 @@ public DatabaseClient getDatabaseClient(DatabaseId db) {
 
         boolean useMultiplexedSession =
             getOptions().getSessionPoolOptions().getUseMultiplexedSession();
+        boolean useMultiplexedSessionForRW =
+            getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
+
         MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient =
             useMultiplexedSession
                 ? new MultiplexedSessionDatabaseClient(SpannerImpl.this.getSessionClient(db))

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolOptionsTest.java

@@ -299,6 +299,52 @@ public void testUseMultiplexedSession() {
             .getUseMultiplexedSession());
   }
 
+  @Test
+  public void testUseMultiplexedSessionForRW() {
+    // skip these tests since this configuration can have dual behaviour in different test-runners
+    assumeFalse(SessionPoolOptions.newBuilder().build().getUseMultiplexedSession());
+    assumeFalse(SessionPoolOptions.newBuilder().build().getUseMultiplexedSessionForRW());
+
+    // Verify default client behavior for multiplexed sessions in R/W transactions
+    assertEquals(false, SessionPoolOptions.newBuilder().build().getUseMultiplexedSessionForRW());
+
+    // Client will use multiplexed sessions for R/W transactions if both the fields are set to true.
+    assertEquals(
+        true,
+        SessionPoolOptions.newBuilder()
+            .setUseMultiplexedSession(true)
+            .setUseMultiplexedSessionForRW(true)
+            .build()
+            .getUseMultiplexedSessionForRW());
+    // Client will not use multiplexed sessions for R/W transactions, since one of the field is set
+    // to false.
+    assertEquals(
+        false,
+        SessionPoolOptions.newBuilder()
+            .setUseMultiplexedSession(true)
+            .setUseMultiplexedSessionForRW(false)
+            .build()
+            .getUseMultiplexedSessionForRW());
+    // Client will not use multiplexed sessions for R/W transactions, since one of the field is set
+    // to false.
+    assertEquals(
+        false,
+        SessionPoolOptions.newBuilder()
+            .setUseMultiplexedSession(false)
+            .setUseMultiplexedSessionForRW(true)
+            .build()
+            .getUseMultiplexedSessionForRW());
+    // Client will not use multiplexed sessions for R/W transactions, since both the fields are set
+    // to false.
+    assertEquals(
+        false,
+        SessionPoolOptions.newBuilder()
+            .setUseMultiplexedSession(false)
+            .setUseMultiplexedSessionForRW(false)
+            .build()
+            .getUseMultiplexedSessionForRW());
+  }
+
   @Test
   public void testMultiplexedSessionMaintenanceDuration() {
     assertEquals(
@@ -326,6 +372,10 @@ public void testToBuilder() {
         SessionPoolOptions.newBuilder()
             .setUseMultiplexedSession(ThreadLocalRandom.current().nextBoolean())
             .build());
+    assertToBuilderRoundtrip(
+        SessionPoolOptions.newBuilder()
+            .setUseMultiplexedSessionForRW(ThreadLocalRandom.current().nextBoolean())
+            .build());
     assertToBuilderRoundtrip(
         SessionPoolOptions.newBuilder()
             .setMinSessions(ThreadLocalRandom.current().nextInt(400))