Skip to content
This repository was archived by the owner on Apr 7, 2026. It is now read-only.

Commit 82c3244

Browse files
rahul2393rahul2393
committed
chore(spanner): add routing hints and cache updates for begin/commit
1 parent bb63e92 commit 82c3244

4 files changed

Lines changed: 249 additions & 20 deletions

File tree

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java‎

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,10 @@
1919
import com.google.api.core.InternalApi;
2020
import com.google.spanner.v1.BeginTransactionRequest;
2121
import com.google.spanner.v1.CacheUpdate;
22+
import com.google.spanner.v1.CommitRequest;
2223
import com.google.spanner.v1.DirectedReadOptions;
2324
import com.google.spanner.v1.ExecuteSqlRequest;
25+
import com.google.spanner.v1.Mutation;
2426
import com.google.spanner.v1.ReadRequest;
2527
import com.google.spanner.v1.RoutingHint;
2628
import com.google.spanner.v1.TransactionOptions;
@@ -92,23 +94,31 @@ public ChannelEndpoint findServer(ExecuteSqlRequest.Builder reqBuilder, boolean
9294
}
9395

9496
public ChannelEndpoint findServer(BeginTransactionRequest.Builder reqBuilder) {
95-
if (!reqBuilder.hasMutationKey()) {
97+
if (!reqBuilder.hasMutationKey()
98+
|| !recipeCache.computeKeys(
99+
reqBuilder.getMutationKey(), reqBuilder.getRoutingHintBuilder())) {
96100
return null;
97101
}
98-
TargetRange target = recipeCache.mutationToTargetRange(reqBuilder.getMutationKey());
99-
if (target == null) {
100-
return null;
101-
}
102-
RoutingHint.Builder hintBuilder = RoutingHint.newBuilder();
103-
hintBuilder.setKey(target.start);
104-
if (!target.limit.isEmpty()) {
105-
hintBuilder.setLimitKey(target.limit);
106-
}
107102
return fillRoutingHint(
108103
preferLeader(reqBuilder.getOptions()),
109104
KeyRangeCache.RangeMode.COVERING_SPLIT,
110105
DirectedReadOptions.getDefaultInstance(),
111-
hintBuilder);
106+
reqBuilder.getRoutingHintBuilder());
107+
}
108+
109+
public void fillRoutingHint(CommitRequest.Builder reqBuilder) {
110+
if (reqBuilder.getMutationsCount() == 0) {
111+
return;
112+
}
113+
Mutation mutation = reqBuilder.getMutations(0);
114+
if (!recipeCache.computeKeys(mutation, reqBuilder.getRoutingHintBuilder())) {
115+
return;
116+
}
117+
fillRoutingHint(
118+
/* preferLeader= */ true,
119+
KeyRangeCache.RangeMode.COVERING_SPLIT,
120+
DirectedReadOptions.getDefaultInstance(),
121+
reqBuilder.getRoutingHintBuilder());
112122
}
113123

114124
private ChannelEndpoint fillRoutingHint(

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java‎

Lines changed: 28 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.google.protobuf.ByteString;
2424
import com.google.spanner.v1.BeginTransactionRequest;
2525
import com.google.spanner.v1.CommitRequest;
26+
import com.google.spanner.v1.CommitResponse;
2627
import com.google.spanner.v1.ExecuteSqlRequest;
2728
import com.google.spanner.v1.PartialResultSet;
2829
import com.google.spanner.v1.ReadRequest;
@@ -47,9 +48,10 @@
4748
/**
4849
* ManagedChannel that routes eligible requests using location-aware routing hints.
4950
*
50-
* <p>Routing hints are applied to streaming read/query and unary ExecuteSql. Commit/Rollback use
51-
* transaction affinity when available. BeginTransaction is routed only when a mutation key is
52-
* provided.
51+
* <p>Routing hints are applied to streaming read/query and unary ExecuteSql. Mutation-based
52+
* BeginTransaction and Commit requests also carry routing hints when recipes are available.
53+
* Commit/Rollback use transaction affinity when available. BeginTransaction is routed only when a
54+
* mutation key is provided.
5355
*/
5456
@InternalApi
5557
final class KeyAwareChannel extends ManagedChannel {
@@ -355,8 +357,10 @@ public void sendMessage(RequestT message) {
355357
BeginTransactionRequest.Builder reqBuilder =
356358
((BeginTransactionRequest) message).toBuilder();
357359
String databaseId = parentChannel.extractDatabaseIdFromSession(reqBuilder.getSession());
358-
if (databaseId != null && reqBuilder.hasMutationKey()) {
360+
if (databaseId != null) {
359361
finder = parentChannel.getOrCreateChannelFinder(databaseId);
362+
}
363+
if (finder != null && reqBuilder.hasMutationKey()) {
360364
endpoint = finder.findServer(reqBuilder);
361365
}
362366
if (reqBuilder.hasOptions() && reqBuilder.getOptions().hasReadOnly()) {
@@ -367,11 +371,19 @@ public void sendMessage(RequestT message) {
367371
}
368372
message = (RequestT) reqBuilder.build();
369373
} else if (message instanceof CommitRequest) {
370-
CommitRequest request = (CommitRequest) message;
371-
if (!request.getTransactionId().isEmpty()) {
372-
endpoint = parentChannel.affinityEndpoint(request.getTransactionId());
373-
transactionIdToClear = request.getTransactionId();
374+
CommitRequest.Builder reqBuilder = ((CommitRequest) message).toBuilder();
375+
String databaseId = parentChannel.extractDatabaseIdFromSession(reqBuilder.getSession());
376+
if (databaseId != null) {
377+
finder = parentChannel.getOrCreateChannelFinder(databaseId);
378+
}
379+
if (finder != null) {
380+
finder.fillRoutingHint(reqBuilder);
381+
}
382+
if (!reqBuilder.getTransactionId().isEmpty()) {
383+
endpoint = parentChannel.affinityEndpoint(reqBuilder.getTransactionId());
384+
transactionIdToClear = reqBuilder.getTransactionId();
374385
}
386+
message = (RequestT) reqBuilder.build();
375387
} else if (message instanceof RollbackRequest) {
376388
RollbackRequest request = (RollbackRequest) message;
377389
if (!request.getTransactionId().isEmpty()) {
@@ -610,7 +622,15 @@ public void onMessage(ResponseT message) {
610622
transactionId = transactionIdFromMetadata(response);
611623
} else if (message instanceof Transaction) {
612624
Transaction response = (Transaction) message;
625+
if (response.hasCacheUpdate() && call.channelFinder != null) {
626+
call.channelFinder.update(response.getCacheUpdate());
627+
}
613628
transactionId = transactionIdFromTransaction(response);
629+
} else if (message instanceof CommitResponse) {
630+
CommitResponse response = (CommitResponse) message;
631+
if (response.hasCacheUpdate() && call.channelFinder != null) {
632+
call.channelFinder.update(response.getCacheUpdate());
633+
}
614634
}
615635
if (transactionId != null) {
616636
if (call.isReadOnlyBegin) {

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRecipeCache.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,23 @@ public void computeKeys(ExecuteSqlRequest.Builder reqBuilder) {
230230
}
231231
}
232232

233+
boolean computeKeys(Mutation mutation, RoutingHint.Builder hintBuilder) {
234+
if (!schemaGeneration.isEmpty()) {
235+
hintBuilder.setSchemaGeneration(schemaGeneration);
236+
}
237+
238+
TargetRange target = mutationToTargetRange(mutation);
239+
if (target == null) {
240+
return false;
241+
}
242+
243+
hintBuilder.setKey(target.start);
244+
if (!target.limit.isEmpty()) {
245+
hintBuilder.setLimitKey(target.limit);
246+
}
247+
return true;
248+
}
249+
233250
public TargetRange mutationToTargetRange(Mutation mutation) {
234251
if (mutation == null) {
235252
return null;

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java‎

Lines changed: 183 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,15 +22,20 @@
2222
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
2323
import com.google.protobuf.ByteString;
2424
import com.google.protobuf.Empty;
25+
import com.google.protobuf.ListValue;
26+
import com.google.protobuf.TextFormat;
27+
import com.google.protobuf.Value;
2528
import com.google.spanner.v1.BeginTransactionRequest;
2629
import com.google.spanner.v1.CacheUpdate;
2730
import com.google.spanner.v1.CommitRequest;
2831
import com.google.spanner.v1.CommitResponse;
2932
import com.google.spanner.v1.ExecuteSqlRequest;
3033
import com.google.spanner.v1.Group;
34+
import com.google.spanner.v1.Mutation;
3135
import com.google.spanner.v1.PartialResultSet;
3236
import com.google.spanner.v1.Range;
3337
import com.google.spanner.v1.ReadRequest;
38+
import com.google.spanner.v1.RecipeList;
3439
import com.google.spanner.v1.ResultSet;
3540
import com.google.spanner.v1.ResultSetMetadata;
3641
import com.google.spanner.v1.RollbackRequest;
@@ -274,6 +279,124 @@ public void resultSetCacheUpdateRoutesSubsequentRequest() throws Exception {
274279
assertThat(harness.endpointCache.callCountForAddress("routed:1234")).isEqualTo(1);
275280
}
276281

282+
@Test
283+
public void beginTransactionWithMutationKeyAddsRoutingHint() throws Exception {
284+
TestHarness harness = createHarness();
285+
seedCache(harness, createMutationRoutingCacheUpdate());
286+
287+
Mutation mutation = createInsertMutation("b");
288+
ClientCall<BeginTransactionRequest, Transaction> beginCall =
289+
harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
290+
beginCall.start(new CapturingListener<Transaction>(), new Metadata());
291+
beginCall.sendMessage(
292+
BeginTransactionRequest.newBuilder().setSession(SESSION).setMutationKey(mutation).build());
293+
294+
@SuppressWarnings("unchecked")
295+
RecordingClientCall<BeginTransactionRequest, Transaction> beginDelegate =
296+
(RecordingClientCall<BeginTransactionRequest, Transaction>)
297+
harness.defaultManagedChannel.latestCall();
298+
299+
assertThat(beginDelegate.lastMessage).isNotNull();
300+
assertThat(beginDelegate.lastMessage.getRoutingHint().getDatabaseId()).isEqualTo(7L);
301+
assertThat(beginDelegate.lastMessage.getRoutingHint().getSchemaGeneration().toStringUtf8())
302+
.isEqualTo("1");
303+
assertThat(beginDelegate.lastMessage.getRoutingHint().getKey().isEmpty()).isFalse();
304+
}
305+
306+
@Test
307+
public void transactionCacheUpdateEnablesCommitRoutingHint() throws Exception {
308+
TestHarness harness = createHarness();
309+
ByteString transactionId = ByteString.copyFromUtf8("tx-with-cache-update");
310+
311+
ClientCall<BeginTransactionRequest, Transaction> beginCall =
312+
harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
313+
beginCall.start(new CapturingListener<Transaction>(), new Metadata());
314+
beginCall.sendMessage(BeginTransactionRequest.newBuilder().setSession(SESSION).build());
315+
316+
@SuppressWarnings("unchecked")
317+
RecordingClientCall<BeginTransactionRequest, Transaction> beginDelegate =
318+
(RecordingClientCall<BeginTransactionRequest, Transaction>)
319+
harness.defaultManagedChannel.latestCall();
320+
beginDelegate.emitOnMessage(
321+
Transaction.newBuilder()
322+
.setId(transactionId)
323+
.setCacheUpdate(createMutationRoutingCacheUpdate())
324+
.build());
325+
beginDelegate.emitOnClose(Status.OK, new Metadata());
326+
327+
ClientCall<CommitRequest, CommitResponse> commitCall =
328+
harness.channel.newCall(SpannerGrpc.getCommitMethod(), CallOptions.DEFAULT);
329+
commitCall.start(new CapturingListener<CommitResponse>(), new Metadata());
330+
commitCall.sendMessage(
331+
CommitRequest.newBuilder()
332+
.setSession(SESSION)
333+
.setTransactionId(transactionId)
334+
.addMutations(createInsertMutation("b"))
335+
.build());
336+
337+
@SuppressWarnings("unchecked")
338+
RecordingClientCall<CommitRequest, CommitResponse> commitDelegate =
339+
(RecordingClientCall<CommitRequest, CommitResponse>)
340+
harness.defaultManagedChannel.latestCall();
341+
342+
assertThat(commitDelegate.lastMessage).isNotNull();
343+
assertThat(commitDelegate.lastMessage.getRoutingHint().getDatabaseId()).isEqualTo(7L);
344+
assertThat(commitDelegate.lastMessage.getRoutingHint().getSchemaGeneration().toStringUtf8())
345+
.isEqualTo("1");
346+
assertThat(commitDelegate.lastMessage.getRoutingHint().getKey().isEmpty()).isFalse();
347+
}
348+
349+
@Test
350+
public void commitResponseCacheUpdateEnablesSubsequentBeginRoutingHint() throws Exception {
351+
TestHarness harness = createHarness();
352+
ByteString transactionId = ByteString.copyFromUtf8("tx-before-commit-cache-update");
353+
354+
ClientCall<BeginTransactionRequest, Transaction> beginCall =
355+
harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
356+
beginCall.start(new CapturingListener<Transaction>(), new Metadata());
357+
beginCall.sendMessage(BeginTransactionRequest.newBuilder().setSession(SESSION).build());
358+
359+
@SuppressWarnings("unchecked")
360+
RecordingClientCall<BeginTransactionRequest, Transaction> beginDelegate =
361+
(RecordingClientCall<BeginTransactionRequest, Transaction>)
362+
harness.defaultManagedChannel.latestCall();
363+
beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build());
364+
beginDelegate.emitOnClose(Status.OK, new Metadata());
365+
366+
ClientCall<CommitRequest, CommitResponse> commitCall =
367+
harness.channel.newCall(SpannerGrpc.getCommitMethod(), CallOptions.DEFAULT);
368+
commitCall.start(new CapturingListener<CommitResponse>(), new Metadata());
369+
commitCall.sendMessage(
370+
CommitRequest.newBuilder().setSession(SESSION).setTransactionId(transactionId).build());
371+
372+
@SuppressWarnings("unchecked")
373+
RecordingClientCall<CommitRequest, CommitResponse> commitDelegate =
374+
(RecordingClientCall<CommitRequest, CommitResponse>)
375+
harness.defaultManagedChannel.latestCall();
376+
commitDelegate.emitOnMessage(
377+
CommitResponse.newBuilder().setCacheUpdate(createMutationRoutingCacheUpdate()).build());
378+
commitDelegate.emitOnClose(Status.OK, new Metadata());
379+
380+
Mutation mutation = createInsertMutation("b");
381+
ClientCall<BeginTransactionRequest, Transaction> secondBeginCall =
382+
harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
383+
secondBeginCall.start(new CapturingListener<Transaction>(), new Metadata());
384+
secondBeginCall.sendMessage(
385+
BeginTransactionRequest.newBuilder().setSession(SESSION).setMutationKey(mutation).build());
386+
387+
@SuppressWarnings("unchecked")
388+
RecordingClientCall<BeginTransactionRequest, Transaction> routedBeginDelegate =
389+
(RecordingClientCall<BeginTransactionRequest, Transaction>)
390+
harness.defaultManagedChannel.latestCall();
391+
392+
assertThat(routedBeginDelegate.lastMessage).isNotNull();
393+
assertThat(routedBeginDelegate.lastMessage.getRoutingHint().getDatabaseId()).isEqualTo(7L);
394+
assertThat(
395+
routedBeginDelegate.lastMessage.getRoutingHint().getSchemaGeneration().toStringUtf8())
396+
.isEqualTo("1");
397+
assertThat(routedBeginDelegate.lastMessage.getRoutingHint().getKey().isEmpty()).isFalse();
398+
}
399+
277400
@Test
278401
public void readOnlyTransactionRoutesEachReadIndependently() throws Exception {
279402
TestHarness harness = createHarness();
@@ -635,6 +758,43 @@ private static CacheUpdate createTwoRangeCacheUpdate() {
635758
.build();
636759
}
637760

761+
private static CacheUpdate createMutationRoutingCacheUpdate() throws TextFormat.ParseException {
762+
RecipeList keyRecipes =
763+
parseRecipeList(
764+
"schema_generation: \"1\"\n"
765+
+ "recipe {\n"
766+
+ " table_name: \"T\"\n"
767+
+ " part { tag: 1 }\n"
768+
+ " part {\n"
769+
+ " order: ASCENDING\n"
770+
+ " null_order: NULLS_FIRST\n"
771+
+ " type { code: STRING }\n"
772+
+ " identifier: \"k\"\n"
773+
+ " }\n"
774+
+ "}\n");
775+
return CacheUpdate.newBuilder()
776+
.setDatabaseId(7L)
777+
.setKeyRecipes(keyRecipes)
778+
.addRange(
779+
Range.newBuilder()
780+
.setStartKey(bytes("a"))
781+
.setLimitKey(bytes("m"))
782+
.setGroupUid(1L)
783+
.setSplitId(1L)
784+
.setGeneration(bytes("1")))
785+
.addGroup(
786+
Group.newBuilder()
787+
.setGroupUid(1L)
788+
.setGeneration(bytes("1"))
789+
.addTablets(
790+
Tablet.newBuilder()
791+
.setTabletUid(1L)
792+
.setServerAddress("server-a:1234")
793+
.setIncarnation(bytes("1"))
794+
.setDistance(0)))
795+
.build();
796+
}
797+
638798
private static void seedCache(TestHarness harness, CacheUpdate cacheUpdate) {
639799
ClientCall<ExecuteSqlRequest, ResultSet> seedCall =
640800
harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT);
@@ -652,6 +812,25 @@ private static void seedCache(TestHarness harness, CacheUpdate cacheUpdate) {
652812
seedDelegate.emitOnMessage(ResultSet.newBuilder().setCacheUpdate(cacheUpdate).build());
653813
}
654814

815+
private static Mutation createInsertMutation(String keyValue) {
816+
return Mutation.newBuilder()
817+
.setInsert(
818+
Mutation.Write.newBuilder()
819+
.setTable("T")
820+
.addColumns("k")
821+
.addValues(
822+
ListValue.newBuilder()
823+
.addValues(Value.newBuilder().setStringValue(keyValue).build())
824+
.build()))
825+
.build();
826+
}
827+
828+
private static RecipeList parseRecipeList(String text) throws TextFormat.ParseException {
829+
RecipeList.Builder builder = RecipeList.newBuilder();
830+
TextFormat.merge(text, builder);
831+
return builder.build();
832+
}
833+
655834
private static TestHarness createHarness() throws IOException {
656835
FakeEndpointCache endpointCache = new FakeEndpointCache(DEFAULT_ADDRESS);
657836
InstantiatingGrpcChannelProvider provider =
@@ -841,6 +1020,7 @@ int callCount() {
8411020
private static final class RecordingClientCall<RequestT, ResponseT>
8421021
extends ClientCall<RequestT, ResponseT> {
8431022
@Nullable private ClientCall.Listener<ResponseT> listener;
1023+
@Nullable private RequestT lastMessage;
8441024
private boolean cancelCalled;
8451025
@Nullable private String cancelMessage;
8461026
@Nullable private Throwable cancelCause;
@@ -864,7 +1044,9 @@ public void cancel(@Nullable String message, @Nullable Throwable cause) {
8641044
public void halfClose() {}
8651045

8661046
@Override
867-
public void sendMessage(RequestT message) {}
1047+
public void sendMessage(RequestT message) {
1048+
this.lastMessage = message;
1049+
}
8681050

8691051
void emitOnMessage(ResponseT response) {
8701052
if (listener != null) {

0 commit comments

Comments
 (0)