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

* chore(spanner): add transaction channel hint for RW to support multiplexed sessions * chore(spanner): lint fix
googleapis · Sep 4, 2024 · 6004c9f · 6004c9f
1 parent cce008d
commit 6004c9f
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 11 deletions.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/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(
         () -> {

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/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();

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/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);