Skip to content

Commit

Permalink
feat: Leader Aware Routing in Connection API (#2308)
Browse files Browse the repository at this point in the history 
* feat: Leader Aware Routing in Connection API

(cherry picked from commit 83ded36)

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

Co-authored-by: Knut Olav Løite <koloite@gmail.com>

* add verbose documentation

* address review comments

* change assertThat to assertEquals/assertNotEquals within testGetSpanner method

* add credentials url

* add credentials url

---------

Co-authored-by: Knut Olav Løite <koloite@gmail.com>
  • Loading branch information
@rajatbhatta @olavloite
rajatbhatta and olavloite committed Aug 4, 2023
1 parent 441c1b0 commit 2a85446
Show file tree
Hide file tree
Showing 4 changed files with 93 additions and 13 deletions.
24 changes: 24 additions & 0 deletions ...le-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java
View file
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 boolean DEFAULT_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE = false;
private static final boolean DEFAULT_TRACK_SESSION_LEAKS = true;
private static final boolean DEFAULT_TRACK_CONNECTION_LEAKS = true;
Expand All @@ -186,6 +187,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 @@ -241,6 +244,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 read/write transactions and partitioned DML 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 @@ -462,6 +469,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 read/write transactions and Partitioned DML. Default is true.
* </ul>
*
* @param uri The URI of the Spanner database to connect to.
Expand Down Expand Up @@ -586,6 +595,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 @@ -678,6 +688,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 @@ -762,6 +773,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 @@ -1089,6 +1105,14 @@ public boolean isReadOnly() {
return readOnly;
}

/**
* Whether read/write transactions and partitioned DML are preferred to be routed to the leader
* region.
*/
public boolean isRouteToLeader() {
return routeToLeader;
}

/**
* The initial retryAbortsInternally value for connections created by this {@link
* ConnectionOptions}
Expand Down
11 changes: 9 additions & 2 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java
View file
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
24 changes: 24 additions & 0 deletions ...loud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java
View file
Expand Up @@ -49,6 +49,9 @@ public class ConnectionOptionsTest {
private static final String FILE_TEST_PATH =
ConnectionOptionsTest.class.getResource("test-key.json").getFile();
private static final String DEFAULT_HOST = "https://spanner.googleapis.com";
private static final String TEST_PROJECT = "test-project-123";
private static final String TEST_INSTANCE = "test-instance-123";
private static final String TEST_DATABASE = "test-database-123";

@Test
public void testBuildWithURIWithDots() {
Expand Down Expand Up @@ -149,6 +152,27 @@ public void testBuildWithAutoConfigEmulator() {
assertTrue(options.isUsePlainText());
}

@Test
public void testBuildWithRouteToLeader() {
final String BASE_URI =
"cloudspanner:/projects/test-project-123/instances/test-instance-123/databases/test-database-123";
ConnectionOptions.Builder builder = ConnectionOptions.newBuilder();
builder.setUri(BASE_URI + "?routeToLeader=false");
builder.setCredentialsUrl(FILE_TEST_PATH);
ConnectionOptions options = builder.build();
assertEquals(options.getHost(), DEFAULT_HOST);
assertEquals(options.getProjectId(), TEST_PROJECT);
assertEquals(options.getInstanceId(), TEST_INSTANCE);
assertEquals(options.getDatabaseName(), TEST_DATABASE);
assertFalse(options.isRouteToLeader());

// Test for default behavior for routeToLeader property.
builder = ConnectionOptions.newBuilder().setUri(BASE_URI);
builder.setCredentialsUrl(FILE_TEST_PATH);
options = builder.build();
assertTrue(options.isRouteToLeader());
}

@Test
public void testBuildWithAutoConfigEmulatorAndHost() {
ConnectionOptions.Builder builder = ConnectionOptions.newBuilder();
Expand Down
47 changes: 36 additions & 11 deletions google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SpannerPoolTest.java
View file
Expand Up @@ -65,6 +65,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 +95,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 All @@ -111,40 +117,43 @@ public void testGetSpanner() {
// assert equal
spanner1 = pool.getSpanner(options1, connection1);
spanner2 = pool.getSpanner(options1, connection2);
assertThat(spanner1).isEqualTo(spanner2);
assertEquals(spanner1, spanner2);
spanner1 = pool.getSpanner(options2, connection1);
spanner2 = pool.getSpanner(options2, connection2);
assertThat(spanner1).isEqualTo(spanner2);
assertEquals(spanner1, spanner2);
spanner1 = pool.getSpanner(options3, connection1);
spanner2 = pool.getSpanner(options3, connection2);
assertThat(spanner1).isEqualTo(spanner2);
assertEquals(spanner1, spanner2);
spanner1 = pool.getSpanner(options4, connection1);
spanner2 = pool.getSpanner(options4, connection2);
assertThat(spanner1).isEqualTo(spanner2);
assertEquals(spanner1, spanner2);
// Options 5 and 6 both use default credentials.
spanner1 = pool.getSpanner(options5, connection1);
spanner2 = pool.getSpanner(options6, connection2);
assertThat(spanner1).isEqualTo(spanner2);
assertEquals(spanner1, spanner2);

// assert not equal
spanner1 = pool.getSpanner(options1, connection1);
spanner2 = pool.getSpanner(options2, connection2);
assertThat(spanner1).isNotEqualTo(spanner2);
assertNotEquals(spanner1, spanner2);
spanner1 = pool.getSpanner(options1, connection1);
spanner2 = pool.getSpanner(options3, connection2);
assertThat(spanner1).isNotEqualTo(spanner2);
assertNotEquals(spanner1, spanner2);
spanner1 = pool.getSpanner(options1, connection1);
spanner2 = pool.getSpanner(options4, connection2);
assertThat(spanner1).isNotEqualTo(spanner2);
assertNotEquals(spanner1, spanner2);
spanner1 = pool.getSpanner(options2, connection1);
spanner2 = pool.getSpanner(options3, connection2);
assertThat(spanner1).isNotEqualTo(spanner2);
assertNotEquals(spanner1, spanner2);
spanner1 = pool.getSpanner(options2, connection1);
spanner2 = pool.getSpanner(options4, connection2);
assertThat(spanner1).isNotEqualTo(spanner2);
assertNotEquals(spanner1, spanner2);
spanner1 = pool.getSpanner(options3, connection1);
spanner2 = pool.getSpanner(options4, connection2);
assertThat(spanner1).isNotEqualTo(spanner2);
assertNotEquals(spanner1, spanner2);
spanner1 = pool.getSpanner(options7, connection1);
spanner2 = pool.getSpanner(options8, connection2);
assertNotEquals(spanner1, spanner2);
}

@Test
Expand Down Expand Up @@ -460,14 +469,30 @@ public void testSpannerPoolKeyEquality() {
.setUri("cloudspanner:/projects/p/instances/i/databases/d")
.setCredentials(NoCredentials.getInstance())
.build();
// Not passing in routeToLeader in Connection URI is equivalent to passing it as true,
// as routeToLeader is true by default.
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);
}
}

0 comments on commit 2a85446

Please sign in to comment.