Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Leader Aware Routing in Connection API #2308

Merged
merged 9 commits into from
Aug 4, 2023
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 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) {
rajatbhatta marked this conversation as resolved.
Show resolved Hide resolved
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
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
Original file line number Diff line number Diff line change
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 =
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Can be defined as private static.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are several places in this file where this can be done and replace the hard-codings with constants. We can take it as a follow-up task in a separate PR.

"cloudspanner:/projects/test-project-123/instances/test-instance-123/databases/test-database-123";
ConnectionOptions.Builder builder = ConnectionOptions.newBuilder();
builder.setUri(BASE_URI + "?routeToLeader=false");
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Prefer StringBuilder() . + is a more expensive way of concatenating strings.

Copy link
Contributor Author

@rajatbhatta rajatbhatta Jul 21, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think for single statement, the compiler would use StringBuilder automatically: https://stackoverflow.com/a/4645155.

builder.setCredentialsUrl(FILE_TEST_PATH);
ConnectionOptions options = builder.build();
assertEquals(options.getHost(), DEFAULT_HOST);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While the asserts here are consistent with what is used in the test class, we should try having a single assert statement. For ex - assertEquals(expectedConnectionOptions, actualConnectionOptions)

This provides benefits where you asserting the entire object and not picking few of the members and can help to pick code side-effects where any member got modified by mistake.

Copy link
Contributor Author

@rajatbhatta rajatbhatta Mar 6, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I understand your concern here. I'll want to take this up as a separate effort, as it's changing the way we have been writing unit tests so far. Let's discuss more in details offline, and take it forward from there. Keeping this thread open until then.

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
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,8 @@ public class SpannerPoolTest {

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Query : Do we have a test which fails when the new property is not added to the equals() method?

Copy link
Contributor Author

@rajatbhatta rajatbhatta Mar 1, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As discussed offline, we can take it as a separate exercise, as we'll need this across classes. Keeping the thread open for the time being.

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);
surbhigarg92 marked this conversation as resolved.
Show resolved Hide resolved
spanner2 = pool.getSpanner(options8, connection2);
assertNotEquals(spanner1, spanner2);
rajatbhatta marked this conversation as resolved.
Show resolved Hide resolved
}

@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);
}
}