feat: Leader Aware Routing in Connection API

Author: rajatbhatta

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java

@@ -172,6 +172,7 @@ public String[] getValidValues() {
   private static final RpcPriority DEFAULT_RPC_PRIORITY = null;
   private static final boolean DEFAULT_RETURN_COMMIT_STATS = false;
   private static final boolean DEFAULT_LENIENT = false;
+  private static final boolean DEFAULT_ROUTE_TO_LEADER = true;
 
   private static final String PLAIN_TEXT_PROTOCOL = "http:";
   private static final String HOST_PROTOCOL = "https:";
@@ -183,6 +184,8 @@ public String[] getValidValues() {
   public static final String AUTOCOMMIT_PROPERTY_NAME = "autocommit";
   /** Name of the 'readonly' connection property. */
   public static final String READONLY_PROPERTY_NAME = "readonly";
+  /** Name of the 'routeToLeader' connection property. */
+  public static final String ROUTE_TO_LEADER_PROPERTY_NAME = "routeToLeader";
   /** Name of the 'retry aborts internally' connection property. */
   public static final String RETRY_ABORTS_INTERNALLY_PROPERTY_NAME = "retryAbortsInternally";
   /** Name of the 'credentials' connection property. */
@@ -231,6 +234,10 @@ public String[] getValidValues() {
                       READONLY_PROPERTY_NAME,
                       "Should the connection start in read-only mode (true/false)",
                       DEFAULT_READONLY),
+                  ConnectionProperty.createBooleanProperty(
+                      ROUTE_TO_LEADER_PROPERTY_NAME,
+                      "Should the RW/PDML requests be routed to leader region (true/false)",
+                      DEFAULT_ROUTE_TO_LEADER),
                   ConnectionProperty.createBooleanProperty(
                       RETRY_ABORTS_INTERNALLY_PROPERTY_NAME,
                       "Should the connection automatically retry Aborted errors (true/false)",
@@ -426,6 +433,8 @@ private boolean isValidUri(String uri) {
      *       created on the emulator if any of them do not yet exist. Any existing instance or
      *       database on the emulator will remain untouched. No other configuration is needed in
      *       order to connect to the emulator than setting this property.
+     *   <li>routeToLeader (boolean): Sets the routeToLeader flag to route requests to leader (true)
+     *       or any region (false) in RW/PDML transactions. Default is true.
      * </ul>
      *
      * @param uri The URI of the Spanner database to connect to.
@@ -547,6 +556,7 @@ public static Builder newBuilder() {
 
   private final boolean autocommit;
   private final boolean readOnly;
+  private final boolean routeToLeader;
   private final boolean retryAbortsInternally;
   private final List<StatementExecutionInterceptor> statementExecutionInterceptors;
   private final SpannerOptionsConfigurator configurator;
@@ -636,6 +646,7 @@ private ConnectionOptions(Builder builder) {
 
     this.autocommit = parseAutocommit(this.uri);
     this.readOnly = parseReadOnly(this.uri);
+    this.routeToLeader = parseRouteToLeader(this.uri);
     this.retryAbortsInternally = parseRetryAbortsInternally(this.uri);
     this.statementExecutionInterceptors =
         Collections.unmodifiableList(builder.statementExecutionInterceptors);
@@ -719,6 +730,11 @@ static boolean parseReadOnly(String uri) {
     return value != null ? Boolean.parseBoolean(value) : DEFAULT_READONLY;
   }
 
+  static boolean parseRouteToLeader(String uri) {
+    String value = parseUriProperty(uri, ROUTE_TO_LEADER_PROPERTY_NAME);
+    return value != null ? Boolean.parseBoolean(value) : DEFAULT_ROUTE_TO_LEADER;
+  }
+
   @VisibleForTesting
   static boolean parseRetryAbortsInternally(String uri) {
     String value = parseUriProperty(uri, RETRY_ABORTS_INTERNALLY_PROPERTY_NAME);
@@ -1025,6 +1041,10 @@ public boolean isAutocommit() {
   public boolean isReadOnly() {
     return readOnly;
   }
+  /** Whether RW/PDML requests are preferred to be routed to the leader region. */
+  public boolean isRouteToLeader() {
+    return routeToLeader;
+  }
 
   /**
    * The initial retryAbortsInternally value for connections created by this {@link

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java

@@ -154,6 +154,7 @@ static class SpannerPoolKey {
     private final boolean usePlainText;
     private final String userAgent;
     private final String databaseRole;
+    private final boolean routeToLeader;
 
     @VisibleForTesting
     static SpannerPoolKey of(ConnectionOptions options) {
@@ -179,6 +180,7 @@ private SpannerPoolKey(ConnectionOptions options) throws IOException {
       this.numChannels = options.getNumChannels();
       this.usePlainText = options.isUsePlainText();
       this.userAgent = options.getUserAgent();
+      this.routeToLeader = options.isRouteToLeader();
     }
 
     @Override
@@ -194,7 +196,8 @@ public boolean equals(Object o) {
           && Objects.equals(this.numChannels, other.numChannels)
           && Objects.equals(this.databaseRole, other.databaseRole)
           && Objects.equals(this.usePlainText, other.usePlainText)
-          && Objects.equals(this.userAgent, other.userAgent);
+          && Objects.equals(this.userAgent, other.userAgent)
+          && Objects.equals(this.routeToLeader, other.routeToLeader);
     }
 
     @Override
@@ -207,7 +210,8 @@ public int hashCode() {
           this.numChannels,
           this.usePlainText,
           this.databaseRole,
-          this.userAgent);
+          this.userAgent,
+          this.routeToLeader);
     }
   }
 
@@ -342,6 +346,9 @@ Spanner createSpanner(SpannerPoolKey key, ConnectionOptions options) {
     if (options.getChannelProvider() != null) {
       builder.setChannelProvider(options.getChannelProvider());
     }
+    if (!options.isRouteToLeader()) {
+      builder.disableLeaderAwareRouting();
+    }
     if (key.usePlainText) {
       // Credentials may not be sent over a plain text channel.
       builder.setCredentials(NoCredentials.getInstance());

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java

@@ -18,6 +18,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
@@ -30,6 +31,7 @@
 import com.google.auth.oauth2.ServiceAccountCredentials;
 import com.google.cloud.NoCredentials;
 import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.Spanner;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerOptions;
 import com.google.common.io.BaseEncoding;
@@ -148,6 +150,27 @@ public void testBuildWithAutoConfigEmulator() {
     assertTrue(options.isUsePlainText());
   }
 
+  @Test
+  public void testBuildWithRouteToLeader() {
+    ConnectionOptions.Builder builder = ConnectionOptions.newBuilder();
+    builder.setUri(
+        "cloudspanner:/projects/test-project-123/instances/test-instance-123/databases/test-database-123?routeToLeader=false");
+    ConnectionOptions options = builder.build();
+    assertEquals(options.getHost(), DEFAULT_HOST);
+    assertEquals(options.getProjectId(), "test-project-123");
+    assertEquals(options.getInstanceId(), "test-instance-123");
+    assertEquals(options.getDatabaseName(), "test-database-123");
+    assertFalse(options.isRouteToLeader());
+
+    // Test for default behavior for routeToLeader property.
+    builder = ConnectionOptions
+        .newBuilder()
+        .setUri(
+            "cloudspanner:/projects/test-project-123/instances/test-instance-123/databases/test-database-123");
+    options = builder.build();
+    assertTrue(options.isRouteToLeader());
+  }
+
   @Test
   public void testBuildWithAutoConfigEmulatorAndHost() {
     ConnectionOptions.Builder builder = ConnectionOptions.newBuilder();

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SpannerPoolTest.java

@@ -18,7 +18,9 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -65,6 +67,8 @@ public class SpannerPoolTest {
 
   private ConnectionOptions options5 = mock(ConnectionOptions.class);
   private ConnectionOptions options6 = mock(ConnectionOptions.class);
+  private ConnectionOptions options7 = mock(ConnectionOptions.class);
+  private ConnectionOptions options8 = mock(ConnectionOptions.class);
 
   private SpannerPool createSubjectAndMocks() {
     return createSubjectAndMocks(0L, Ticker.systemTicker());
@@ -93,6 +97,10 @@ Spanner createSpanner(SpannerPoolKey key, ConnectionOptions options) {
     // ConnectionOptions with no specific credentials.
     when(options5.getProjectId()).thenReturn("test-project-3");
     when(options6.getProjectId()).thenReturn("test-project-3");
+    when(options7.getProjectId()).thenReturn("test-project-3");
+    when(options7.isRouteToLeader()).thenReturn(true);
+    when(options8.getProjectId()).thenReturn("test-project-3");
+    when(options8.isRouteToLeader()).thenReturn(false);
 
     return pool;
   }
@@ -145,6 +153,9 @@ public void testGetSpanner() {
     spanner1 = pool.getSpanner(options3, connection1);
     spanner2 = pool.getSpanner(options4, connection2);
     assertThat(spanner1).isNotEqualTo(spanner2);
+    spanner1 = pool.getSpanner(options7, connection1);
+    spanner2 = pool.getSpanner(options8, connection2);
+    assertNotEquals(spanner1, spanner2);
   }
 
   @Test
@@ -460,14 +471,29 @@ public void testSpannerPoolKeyEquality() {
             .setUri("cloudspanner:/projects/p/instances/i/databases/d")
             .setCredentials(NoCredentials.getInstance())
             .build();
+    ConnectionOptions options4 =
+        ConnectionOptions.newBuilder()
+            .setUri("cloudspanner:/projects/p/instances/i/databases/d?routeToLeader=true")
+            .setCredentials(NoCredentials.getInstance())
+            .build();
+    ConnectionOptions options5 =
+        ConnectionOptions.newBuilder()
+            .setUri("cloudspanner:/projects/p/instances/i/databases/d?routeToLeader=false")
+            .setCredentials(NoCredentials.getInstance())
+            .build();
 
     SpannerPoolKey key1 = SpannerPoolKey.of(options1);
     SpannerPoolKey key2 = SpannerPoolKey.of(options2);
     SpannerPoolKey key3 = SpannerPoolKey.of(options3);
+    SpannerPoolKey key4 = SpannerPoolKey.of(options4);
+    SpannerPoolKey key5 = SpannerPoolKey.of(options5);
 
     assertNotEquals(key1, key2);
     assertEquals(key2, key3);
     assertNotEquals(key1, key3);
     assertNotEquals(key1, new Object());
+    assertEquals(key3, key4);
+    assertNotEquals(key4, key5);
+    assertEquals(key2, key4);
   }
 }