chore(spanner): support transaction channel hint for RW mux (googleapis#3317)

* chore(spanner): add transaction channel hint for RW to support multiplexed sessions * chore(spanner): lint fix

Author: harshachinta

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

@@ -426,7 +426,8 @@ public void close() {
  }
  }
 
- ApiFuture<ByteString> beginTransactionAsync(Options transactionOptions, boolean routeToLeader) {
+ ApiFuture<ByteString> beginTransactionAsync(
+ Options transactionOptions, boolean routeToLeader, Map<SpannerRpc.Option, ?> channelHint) {
  final SettableApiFuture<ByteString> res = SettableApiFuture.create();
  final ISpan span = tracer.spanBuilder(SpannerImpl.BEGIN_TRANSACTION);
  final BeginTransactionRequest request =
@@ -436,7 +437,7 @@ ApiFuture<ByteString> beginTransactionAsync(Options transactionOptions, boolean
  .build();
  final ApiFuture<Transaction> requestFuture;
  try (IScope ignore = tracer.withSpan(span)) {
- requestFuture = spanner.getRpc().beginTransactionAsync(request, getOptions(), routeToLeader);
+ requestFuture = spanner.getRpc().beginTransactionAsync(request, channelHint, routeToLeader);
  }
  requestFuture.addListener(
  () -> {

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

@@ -282,7 +282,8 @@ ApiFuture<Void> ensureTxnAsync() {
 
  private void createTxnAsync(final SettableApiFuture<Void> res) {
  span.addAnnotation("Creating Transaction");
- final ApiFuture<ByteString> fut = session.beginTransactionAsync(options, isRouteToLeader());
+ final ApiFuture<ByteString> fut =
+ session.beginTransactionAsync(options, isRouteToLeader(), getTransactionChannelHint());
  fut.addListener(
  () -> {
  try {
@@ -427,7 +428,7 @@ public void run() {
  final ApiFuture<com.google.spanner.v1.CommitResponse> commitFuture;
  final ISpan opSpan = tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span);
  try (IScope ignore = tracer.withSpan(opSpan)) {
- commitFuture = rpc.commitAsync(commitRequest, session.getOptions());
+ commitFuture = rpc.commitAsync(commitRequest, getTransactionChannelHint());
  }
  session.markUsed(clock.instant());
  commitFuture.addListener(
@@ -525,7 +526,7 @@ ApiFuture<Empty> rollbackAsync() {
  .setSession(session.getName())
  .setTransactionId(transactionId)
  .build(),
- session.getOptions());
+ getTransactionChannelHint());
  session.markUsed(clock.instant());
  return apiFuture;
  } else {
@@ -800,7 +801,7 @@ private ResultSet internalExecuteUpdate(
  statement, queryMode, options, /* withTransactionSelector = */ true);
  try {
  com.google.spanner.v1.ResultSet resultSet =
- rpc.executeQuery(builder.build(), session.getOptions(), isRouteToLeader());
+ rpc.executeQuery(builder.build(), getTransactionChannelHint(), isRouteToLeader());
  session.markUsed(clock.instant());
  if (resultSet.getMetadata().hasTransaction()) {
  onTransactionMetadata(
@@ -838,7 +839,8 @@ public ApiFuture<Long> executeUpdateAsync(Statement statement, UpdateOption... u
  // commit.
  increaseAsyncOperations();
  resultSet =
- rpc.executeQueryAsync(builder.build(), session.getOptions(), isRouteToLeader());
+ rpc.executeQueryAsync(
+ builder.build(), getTransactionChannelHint(), isRouteToLeader());
  session.markUsed(clock.instant());
  } catch (Throwable t) {
  decreaseAsyncOperations();
@@ -926,7 +928,7 @@ public long[] batchUpdate(Iterable<Statement> statements, UpdateOption... update
  getExecuteBatchDmlRequestBuilder(statements, options);
  try {
  com.google.spanner.v1.ExecuteBatchDmlResponse response =
- rpc.executeBatchDml(builder.build(), session.getOptions());
+ rpc.executeBatchDml(builder.build(), getTransactionChannelHint());
  session.markUsed(clock.instant());
  long[] results = new long[response.getResultSetsCount()];
  for (int i = 0; i < response.getResultSetsCount(); ++i) {
@@ -983,7 +985,7 @@ public ApiFuture<long[]> batchUpdateAsync(
  // Register the update as an async operation that must finish before the transaction may
  // commit.
  increaseAsyncOperations();
- response = rpc.executeBatchDmlAsync(builder.build(), session.getOptions());
+ response = rpc.executeBatchDmlAsync(builder.build(), getTransactionChannelHint());
  session.markUsed(clock.instant());
  } catch (Throwable t) {
  decreaseAsyncOperations();

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java

@@ -1497,7 +1497,7 @@ public void testSessionNotFoundReadWriteTransaction() {
  .thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
  when(closedSession.newTransaction(Options.fromTransactionOptions()))
  .thenReturn(closedTransactionContext);
- when(closedSession.beginTransactionAsync(any(), eq(true))).thenThrow(sessionNotFound);
+ when(closedSession.beginTransactionAsync(any(), eq(true), any())).thenThrow(sessionNotFound);
  when(closedSession.getTracer()).thenReturn(tracer);
  TransactionRunnerImpl closedTransactionRunner = new TransactionRunnerImpl(closedSession);
  closedTransactionRunner.setSpan(span);
@@ -1512,7 +1512,7 @@ public void testSessionNotFoundReadWriteTransaction() {
  final TransactionContextImpl openTransactionContext = mock(TransactionContextImpl.class);
  when(openSession.newTransaction(Options.fromTransactionOptions()))
  .thenReturn(openTransactionContext);
- when(openSession.beginTransactionAsync(any(), eq(true)))
+ when(openSession.beginTransactionAsync(any(), eq(true), any()))
  .thenReturn(ApiFutures.immediateFuture(ByteString.copyFromUtf8("open-txn")));
  when(openSession.getTracer()).thenReturn(tracer);
  TransactionRunnerImpl openTransactionRunner = new TransactionRunnerImpl(openSession);