feat: Add x-goog-spanner-route-to-leader header to Spanner RPC contexts for RW/PDML transactions.

The header is added to support leader-aware-routing feature, which aims at reducing cross-regional latency for RW/PDML transactions in a multi-region instance.

Author: yifanzyifanz

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java

@@ -161,6 +161,7 @@ static Builder newBuilder() {
  private SingleReadContext(Builder builder) {
  super(builder);
  this.bound = builder.bound;
+ this.routeToLeader = false;
  }
 
  @GuardedBy("lock")
@@ -291,6 +292,7 @@ static Builder newBuilder() {
  this.timestamp = builder.timestamp;
  this.transactionId = builder.transactionId;
  }
+ this.routeToLeader = false;
  }
 
  @Override
@@ -347,7 +349,8 @@ void initTransaction() {
  .setSession(session.getName())
  .setOptions(options)
  .build();
- Transaction transaction = rpc.beginTransaction(request, session.getOptions());
+ Transaction transaction =
+ rpc.beginTransaction(request, session.getOptions(), routeToLeader);
  if (!transaction.hasReadTimestamp()) {
  throw SpannerExceptionFactory.newSpannerException(
  ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field");
@@ -380,6 +383,7 @@ void initTransaction() {
  Span span;
  private final int defaultPrefetchChunks;
  private final QueryOptions defaultQueryOptions;
+ protected boolean routeToLeader = false;
 
  @GuardedBy("lock")
  private boolean isValid = true;
@@ -664,7 +668,8 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
  request.setTransaction(selector);
  }
  SpannerRpc.StreamingCall call =
- rpc.executeQuery(request.build(), stream.consumer(), session.getOptions());
+ rpc.executeQuery(
+ request.build(), stream.consumer(), session.getOptions(), routeToLeader);
  call.request(prefetchChunks);
  stream.setCall(call, request.getTransaction().hasBegin());
  return stream;
@@ -792,7 +797,7 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
  }
  builder.setRequestOptions(buildRequestOptions(readOptions));
  SpannerRpc.StreamingCall call =
- rpc.read(builder.build(), stream.consumer(), session.getOptions());
+ rpc.read(builder.build(), stream.consumer(), session.getOptions(), routeToLeader);
  call.request(prefetchChunks);
  stream.setCall(call, /* withBeginTransaction = */ builder.getTransaction().hasBegin());
  return stream;

google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDmlTransaction.java

@@ -202,7 +202,7 @@ private ByteString initTransaction() {
  TransactionOptions.newBuilder()
  .setPartitionedDml(TransactionOptions.PartitionedDml.getDefaultInstance()))
  .build();
- Transaction tx = rpc.beginTransaction(request, session.getOptions());
+ Transaction tx = rpc.beginTransaction(request, session.getOptions(), true);
  if (tx.getId().isEmpty()) {
  throw SpannerExceptionFactory.newSpannerException(
  ErrorCode.INTERNAL,

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java

@@ -267,7 +267,7 @@ public AsyncTransactionManagerImpl transactionManagerAsync(TransactionOption...
  @Override
  public void prepareReadWriteTransaction() {
  setActive(null);
- readyTransactionId = beginTransaction();
+ readyTransactionId = beginTransaction(true);
  }
 
  @Override
@@ -288,17 +288,17 @@ public void close() {
  }
  }
 
- ByteString beginTransaction() {
+ ByteString beginTransaction(boolean routeToLeader) {
  try {
- return beginTransactionAsync().get();
+ return beginTransactionAsync(routeToLeader).get();
  } catch (ExecutionException e) {
  throw SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause());
  } catch (InterruptedException e) {
  throw SpannerExceptionFactory.propagateInterrupt(e);
  }
  }
 
- ApiFuture<ByteString> beginTransactionAsync() {
+ ApiFuture<ByteString> beginTransactionAsync(boolean routeToLeader) {
  final SettableApiFuture<ByteString> res = SettableApiFuture.create();
  final Span span = tracer.spanBuilder(SpannerImpl.BEGIN_TRANSACTION).startSpan();
  final BeginTransactionRequest request =
@@ -309,7 +309,7 @@ ApiFuture<ByteString> beginTransactionAsync() {
  .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance()))
  .build();
  final ApiFuture<Transaction> requestFuture =
- spanner.getRpc().beginTransactionAsync(request, options);
+ spanner.getRpc().beginTransactionAsync(request, options, routeToLeader);
  requestFuture.addListener(
  tracer.withSpan(
  span,

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java

@@ -196,6 +196,7 @@ private TransactionContextImpl(Builder builder) {
  this.trackTransactionStarter = builder.trackTransactionStarter;
  this.options = builder.options;
  this.finishedAsyncOperations.set(null);
+ this.routeToLeader = true;
  }
 
  private void increaseAsyncOperations() {
@@ -255,7 +256,7 @@ ApiFuture<Void> ensureTxnAsync() {
 
  private void createTxnAsync(final SettableApiFuture<Void> res) {
  span.addAnnotation("Creating Transaction");
- final ApiFuture<ByteString> fut = session.beginTransactionAsync();
+ final ApiFuture<ByteString> fut = session.beginTransactionAsync(routeToLeader);
  fut.addListener(
  () -> {
  try {
@@ -719,7 +720,7 @@ private ResultSet internalExecuteUpdate(
  /* withTransactionSelector = */ true);
  try {
  com.google.spanner.v1.ResultSet resultSet =
- rpc.executeQuery(builder.build(), session.getOptions());
+ rpc.executeQuery(builder.build(), session.getOptions(), routeToLeader);
  if (resultSet.getMetadata().hasTransaction()) {
  onTransactionMetadata(
  resultSet.getMetadata().getTransaction(), builder.getTransaction().hasBegin());
@@ -749,7 +750,7 @@ public ApiFuture<Long> executeUpdateAsync(Statement statement, UpdateOption... o
  // Register the update as an async operation that must finish before the transaction may
  // commit.
  increaseAsyncOperations();
- resultSet = rpc.executeQueryAsync(builder.build(), session.getOptions());
+ resultSet = rpc.executeQueryAsync(builder.build(), session.getOptions(), routeToLeader);
  } catch (Throwable t) {
  decreaseAsyncOperations();
  throw t;

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

@@ -1592,7 +1592,8 @@ public List<Session> batchCreateSessions(
  requestBuilder.setSessionTemplate(sessionBuilder);
  BatchCreateSessionsRequest request = requestBuilder.build();
  GrpcCallContext context =
- newCallContext(options, databaseName, request, SpannerGrpc.getBatchCreateSessionsMethod());
+ newCallContext(
+ options, databaseName, request, SpannerGrpc.getBatchCreateSessionsMethod(), true);
  return get(spannerStub.batchCreateSessionsCallable().futureCall(request, context))
  .getSessionList();
  }
@@ -1616,7 +1617,7 @@ public Session createSession(
  requestBuilder.setSession(sessionBuilder);
  CreateSessionRequest request = requestBuilder.build();
  GrpcCallContext context =
- newCallContext(options, databaseName, request, SpannerGrpc.getCreateSessionMethod());
+ newCallContext(options, databaseName, request, SpannerGrpc.getCreateSessionMethod(), true);
  return get(spannerStub.createSessionCallable().futureCall(request, context));
  }
 
@@ -1630,15 +1631,19 @@ public void deleteSession(String sessionName, @Nullable Map<Option, ?> options)
  public ApiFuture<Empty> asyncDeleteSession(String sessionName, @Nullable Map<Option, ?> options) {
  DeleteSessionRequest request = DeleteSessionRequest.newBuilder().setName(sessionName).build();
  GrpcCallContext context =
- newCallContext(options, sessionName, request, SpannerGrpc.getDeleteSessionMethod());
+ newCallContext(options, sessionName, request, SpannerGrpc.getDeleteSessionMethod(), false);
  return spannerStub.deleteSessionCallable().futureCall(request, context);
  }
 
  @Override
  public StreamingCall read(
- ReadRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options) {
+ ReadRequest request,
+ ResultStreamConsumer consumer,
+ @Nullable Map<Option, ?> options,
+ boolean routeToLeader) {
  GrpcCallContext context =
- newCallContext(options, request.getSession(), request, SpannerGrpc.getReadMethod());
+ newCallContext(
+ options, request.getSession(), request, SpannerGrpc.getReadMethod(), routeToLeader);
  SpannerResponseObserver responseObserver = new SpannerResponseObserver(consumer);
  spannerStub.streamingReadCallable().call(request, responseObserver, context);
  final StreamController controller = responseObserver.getController();
@@ -1658,13 +1663,14 @@ public void cancel(String message) {
  }
 
  @Override
- public ResultSet executeQuery(ExecuteSqlRequest request, @Nullable Map<Option, ?> options) {
- return get(executeQueryAsync(request, options));
+ public ResultSet executeQuery(
+ ExecuteSqlRequest request, @Nullable Map<Option, ?> options, boolean routeToLeader) {
+ return get(executeQueryAsync(request, options, routeToLeader));
  }
 
  @Override
  public ApiFuture<ResultSet> executeQueryAsync(
- ExecuteSqlRequest request, @Nullable Map<Option, ?> options) {
+ ExecuteSqlRequest request, @Nullable Map<Option, ?> options, boolean routeToLeader) {
  GrpcCallContext context =
  newCallContext(options, request.getSession(), request, SpannerGrpc.getExecuteSqlMethod());
  return spannerStub.executeSqlCallable().futureCall(request, context);
@@ -1674,7 +1680,8 @@ public ApiFuture<ResultSet> executeQueryAsync(
  public ResultSet executePartitionedDml(
  ExecuteSqlRequest request, @Nullable Map<Option, ?> options) {
  GrpcCallContext context =
- newCallContext(options, request.getSession(), request, SpannerGrpc.getExecuteSqlMethod());
+ newCallContext(
+ options, request.getSession(), request, SpannerGrpc.getExecuteSqlMethod(), true);
  return get(partitionedDmlStub.executeSqlCallable().futureCall(request, context));
  }
 
@@ -1688,18 +1695,29 @@ public ServerStream<PartialResultSet> executeStreamingPartitionedDml(
  ExecuteSqlRequest request, Map<Option, ?> options, Duration timeout) {
  GrpcCallContext context =
  newCallContext(
- options, request.getSession(), request, SpannerGrpc.getExecuteStreamingSqlMethod());
+ options,
+ request.getSession(),
+ request,
+ SpannerGrpc.getExecuteStreamingSqlMethod(),
+ true);
  // Override any timeout settings that might have been set on the call context.
  context = context.withTimeout(timeout).withStreamWaitTimeout(timeout);
  return partitionedDmlStub.executeStreamingSqlCallable().call(request, context);
  }
 
  @Override
  public StreamingCall executeQuery(
- ExecuteSqlRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options) {
+ ExecuteSqlRequest request,
+ ResultStreamConsumer consumer,
+ @Nullable Map<Option, ?> options,
+ boolean routeToLeader) {
  GrpcCallContext context =
  newCallContext(
- options, request.getSession(), request, SpannerGrpc.getExecuteStreamingSqlMethod());
+ options,
+ request.getSession(),
+ request,
+ SpannerGrpc.getExecuteStreamingSqlMethod(),
+ routeToLeader);
  SpannerResponseObserver responseObserver = new SpannerResponseObserver(consumer);
  spannerStub.executeStreamingSqlCallable().call(request, responseObserver, context);
  final StreamController controller = responseObserver.getController();
@@ -1729,30 +1747,35 @@ public ApiFuture<ExecuteBatchDmlResponse> executeBatchDmlAsync(
  ExecuteBatchDmlRequest request, @Nullable Map<Option, ?> options) {
  GrpcCallContext context =
  newCallContext(
- options, request.getSession(), request, SpannerGrpc.getExecuteBatchDmlMethod());
+ options, request.getSession(), request, SpannerGrpc.getExecuteBatchDmlMethod(), true);
  return spannerStub.executeBatchDmlCallable().futureCall(request, context);
  }
 
  @Override
  public ApiFuture<Transaction> beginTransactionAsync(
- BeginTransactionRequest request, @Nullable Map<Option, ?> options) {
+ BeginTransactionRequest request, @Nullable Map<Option, ?> options, boolean routeToLeader) {
  GrpcCallContext context =
  newCallContext(
- options, request.getSession(), request, SpannerGrpc.getBeginTransactionMethod());
+ options,
+ request.getSession(),
+ request,
+ SpannerGrpc.getBeginTransactionMethod(),
+ routeToLeader);
  return spannerStub.beginTransactionCallable().futureCall(request, context);
  }
 
  @Override
  public Transaction beginTransaction(
- BeginTransactionRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
- return get(beginTransactionAsync(request, options));
+ BeginTransactionRequest request, @Nullable Map<Option, ?> options, boolean routeToLeader)
+ throws SpannerException {
+ return get(beginTransactionAsync(request, options, routeToLeader));
  }
 
  @Override
  public ApiFuture<CommitResponse> commitAsync(
  CommitRequest request, @Nullable Map<Option, ?> options) {
  GrpcCallContext context =
- newCallContext(options, request.getSession(), request, SpannerGrpc.getCommitMethod());
+ newCallContext(options, request.getSession(), request, SpannerGrpc.getCommitMethod(), true);
  return spannerStub.commitCallable().futureCall(request, context);
  }
 
@@ -1765,7 +1788,8 @@ public CommitResponse commit(CommitRequest commitRequest, @Nullable Map<Option,
  @Override
  public ApiFuture<Empty> rollbackAsync(RollbackRequest request, @Nullable Map<Option, ?> options) {
  GrpcCallContext context =
- newCallContext(options, request.getSession(), request, SpannerGrpc.getRollbackMethod());
+ newCallContext(
+ options, request.getSession(), request, SpannerGrpc.getRollbackMethod(), true);
  return spannerStub.rollbackCallable().futureCall(request, context);
  }
 
@@ -1780,7 +1804,7 @@ public PartitionResponse partitionQuery(
  PartitionQueryRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
  GrpcCallContext context =
  newCallContext(
- options, request.getSession(), request, SpannerGrpc.getPartitionQueryMethod());
+ options, request.getSession(), request, SpannerGrpc.getPartitionQueryMethod(), true);
  return get(spannerStub.partitionQueryCallable().futureCall(request, context));
  }
 
@@ -1789,7 +1813,7 @@ public PartitionResponse partitionRead(
  PartitionReadRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
  GrpcCallContext context =
  newCallContext(
- options, request.getSession(), request, SpannerGrpc.getPartitionReadMethod());
+ options, request.getSession(), request, SpannerGrpc.getPartitionReadMethod(), true);
  return get(spannerStub.partitionReadCallable().futureCall(request, context));
  }
 
@@ -1898,6 +1922,16 @@ <ReqT, RespT> GrpcCallContext newCallContext(
  String resource,
  ReqT request,
  MethodDescriptor<ReqT, RespT> method) {
+ return newCallContext(options, resource, request, method, false);
+ }
+
+ @VisibleForTesting
+ <ReqT, RespT> GrpcCallContext newCallContext(
+ @Nullable Map<Option, ?> options,
+ String resource,
+ ReqT request,
+ MethodDescriptor<ReqT, RespT> method,
+ boolean routeToLeader) {
  GrpcCallContext context = GrpcCallContext.createDefault();
  if (options != null) {
  context = context.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue());
@@ -1907,6 +1941,9 @@ <ReqT, RespT> GrpcCallContext newCallContext(
  context = context.withCallOptions(context.getCallOptions().withCompression(compressorName));
  }
  context = context.withExtraHeaders(metadataProvider.newExtraHeaders(resource, projectName));
+ if (routeToLeader) {
+ context = context.withExtraHeaders(metadataProvider.newRouteToLeaderHeader());
+ }
  if (callCredentialsProvider != null) {
  CallCredentials callCredentials = callCredentialsProvider.getCallCredentials();
  if (callCredentials != null) {

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

@@ -28,7 +28,7 @@
 class SpannerMetadataProvider {
  private final Map<Metadata.Key<String>, String> headers;
  private final String resourceHeaderKey;
-
+ private static final String routeToLeaderHeaderKey = "x-goog-spanner-route-to-leader";
  private static final Pattern[] RESOURCE_TOKEN_PATTERNS = {
  Pattern.compile("^(?<headerValue>projects/[^/]*/instances/[^/]*/databases/[^/]*)(.*)?"),
  Pattern.compile("^(?<headerValue>projects/[^/]*/instances/[^/]*)(.*)?")
@@ -66,6 +66,12 @@ Map<String, List<String>> newExtraHeaders(
  .build();
  }
 
+ Map<String, List<String>> newRouteToLeaderHeader() {
+ return ImmutableMap.<String, List<String>>builder()
+ .put(routeToLeaderHeaderKey, Collections.singletonList("true"))
+ .build();
+ }
+
  private Map<Metadata.Key<String>, String> constructHeadersAsMetadata(
  Map<String, String> headers) {
  ImmutableMap.Builder<Metadata.Key<String>, String> headersAsMetadataBuilder =