Skip to content

Commit

Permalink
feat: Leader Aware Routing in Connection API
Browse files Browse the repository at this point in the history 
(cherry picked from commit 83ded36)
  • Loading branch information
@rajatbhatta
rajatbhatta committed Feb 27, 2023
1 parent 8372250 commit 13cb842
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 2 deletions.
20 changes: 20 additions & 0 deletions ...le-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:";
Expand All @@ -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. */
Expand Down Expand Up @@ -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)",
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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
Expand Down
11 changes: 9 additions & 2 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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
Expand All @@ -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
Expand All @@ -207,7 +210,8 @@ public int hashCode() {
this.numChannels,
this.usePlainText,
this.databaseRole,
this.userAgent);
this.userAgent,
this.routeToLeader);
}
}

Expand Down Expand Up @@ -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());
Expand Down
23 changes: 23 additions & 0 deletions ...loud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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;
Expand Down Expand Up @@ -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();
Expand Down
26 changes: 26 additions & 0 deletions google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SpannerPoolTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
}
}

0 comments on commit 13cb842

Please sign in to comment.