all rpcs on explicit begin started transactions should use same server/channel conncetion

Author: rahul2393

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

@@ -404,6 +404,15 @@ ByteString getTransactionId() {
       }
     }
 
+    @Override
+    public void close() {
+      ByteString id = getTransactionId();
+      if (id != null && !id.isEmpty()) {
+        rpc.clearTransactionAffinity(id);
+      }
+      super.close();
+    }
+
     /**
      * Initializes the transaction with the timestamp specified within MultiUseReadOnlyTransaction.
      * This is used only for fallback of PartitionQueryRequest and PartitionReadRequest with

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

@@ -51,6 +51,7 @@
 import com.google.api.gax.rpc.StatusCode;
 import com.google.api.gax.rpc.StatusCode.Code;
 import com.google.api.gax.rpc.StreamController;
+import com.google.api.gax.rpc.TransportChannel;
 import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.api.gax.rpc.UnaryCallSettings;
 import com.google.api.gax.rpc.UnaryCallable;
@@ -105,6 +106,7 @@
 import com.google.longrunning.GetOperationRequest;
 import com.google.longrunning.Operation;
 import com.google.longrunning.OperationsGrpc;
+import com.google.protobuf.ByteString;
 import com.google.protobuf.Empty;
 import com.google.protobuf.FieldMask;
 import com.google.protobuf.InvalidProtocolBufferException;
@@ -187,6 +189,7 @@
 import com.google.spanner.v1.Transaction;
 import io.grpc.CallCredentials;
 import io.grpc.Context;
+import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.MethodDescriptor;
 import java.io.IOException;
@@ -286,6 +289,7 @@ public class GapicSpannerRpc implements SpannerRpc {
   private final int numChannels;
   private final boolean isGrpcGcpExtensionEnabled;
   private final boolean isDynamicChannelPoolEnabled;
+  @Nullable private final KeyAwareChannel keyAwareChannel;
 
   private final GrpcCallContext baseGrpcCallContext;
 
@@ -539,6 +543,7 @@ public GapicSpannerRpc(final SpannerOptions options) {
                         /* isAdminClient= */ false, isEmulatorEnabled(options, emulatorHost)))
                 .build();
         ClientContext clientContext = ClientContext.create(spannerStubSettings);
+        this.keyAwareChannel = extractKeyAwareChannel(clientContext.getTransportChannel());
         this.spannerStub =
             GrpcSpannerStubWithStubSettingsAndClientContext.create(
                 spannerStubSettings, clientContext);
@@ -658,6 +663,7 @@ public <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createUnaryCalla
         throw asSpannerException(e);
       }
     } else {
+      this.keyAwareChannel = null;
       this.databaseAdminStub = null;
       this.instanceAdminStub = null;
       this.spannerStub = null;
@@ -675,6 +681,23 @@ public <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createUnaryCalla
     }
   }
 
+  private static KeyAwareChannel extractKeyAwareChannel(TransportChannel transportChannel) {
+    if (transportChannel instanceof GrpcTransportChannel) {
+      ManagedChannel channel = ((GrpcTransportChannel) transportChannel).getChannel();
+      if (channel instanceof KeyAwareChannel) {
+        return (KeyAwareChannel) channel;
+      }
+    }
+    return null;
+  }
+
+  @Override
+  public void clearTransactionAffinity(ByteString transactionId) {
+    if (keyAwareChannel != null) {
+      keyAwareChannel.clearTransactionAffinity(transactionId);
+    }
+  }
+
   private static String parseGrpcGcpApiConfig() {
     try {
       return Resources.toString(

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

@@ -19,6 +19,7 @@
 import com.google.api.core.InternalApi;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.protobuf.ByteString;
+import com.google.spanner.v1.BeginTransactionRequest;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.PartialResultSet;
@@ -52,6 +53,8 @@ final class KeyAwareChannel extends ManagedChannel {
   private static final String STREAMING_SQL_METHOD =
       "google.spanner.v1.Spanner/ExecuteStreamingSql";
   private static final String UNARY_SQL_METHOD = "google.spanner.v1.Spanner/ExecuteSql";
+  private static final String BEGIN_TRANSACTION_METHOD =
+      "google.spanner.v1.Spanner/BeginTransaction";
   private static final String COMMIT_METHOD = "google.spanner.v1.Spanner/Commit";
   private static final String ROLLBACK_METHOD = "google.spanner.v1.Spanner/Rollback";
 
@@ -162,6 +165,7 @@ private static boolean isKeyAware(MethodDescriptor<?, ?> methodDescriptor) {
     return STREAMING_READ_METHOD.equals(method)
         || STREAMING_SQL_METHOD.equals(method)
         || UNARY_SQL_METHOD.equals(method)
+        || BEGIN_TRANSACTION_METHOD.equals(method)
         || COMMIT_METHOD.equals(method)
         || ROLLBACK_METHOD.equals(method);
   }
@@ -185,12 +189,17 @@ private void clearAffinity(ByteString transactionId) {
     transactionAffinities.remove(transactionId);
   }
 
-  private void recordAffinity(ByteString transactionId, @Nullable ChannelEndpoint endpoint) {
+  void clearTransactionAffinity(ByteString transactionId) {
+    clearAffinity(transactionId);
+  }
+
+  private void recordAffinity(
+      ByteString transactionId, @Nullable ChannelEndpoint endpoint, boolean allowDefault) {
     if (transactionId == null || transactionId.isEmpty() || endpoint == null) {
       return;
     }
     String address = endpoint.getAddress();
-    if (defaultEndpointAddress.equals(address)) {
+    if (!allowDefault && defaultEndpointAddress.equals(address)) {
       return;
     }
     transactionAffinities.put(transactionId, address);
@@ -238,6 +247,7 @@ static final class KeyAwareClientCall<RequestT, ResponseT>
     private ChannelFinder channelFinder;
     @Nullable private ChannelEndpoint selectedEndpoint;
     @Nullable private ByteString transactionIdToClear;
+    private boolean allowDefaultAffinity;
 
     KeyAwareClientCall(
         KeyAwareChannel parentChannel,
@@ -295,6 +305,8 @@ public void sendMessage(RequestT message) {
           }
         }
         message = (RequestT) reqBuilder.build();
+      } else if (message instanceof BeginTransactionRequest) {
+        allowDefaultAffinity = true;
       } else if (message instanceof CommitRequest) {
         CommitRequest request = (CommitRequest) message;
         if (!request.getTransactionId().isEmpty()) {
@@ -309,7 +321,8 @@ public void sendMessage(RequestT message) {
         }
       } else {
         throw new IllegalStateException(
-            "Only read, query, commit, and rollback requests are supported for key-aware calls.");
+            "Only read, query, begin transaction, commit, and rollback requests are supported for"
+                + " key-aware calls.");
       }
 
       if (endpoint == null) {
@@ -343,7 +356,7 @@ public void cancel(@Nullable String message, @Nullable Throwable cause) {
     }
 
     void maybeRecordAffinity(ByteString transactionId) {
-      parentChannel.recordAffinity(transactionId, selectedEndpoint);
+      parentChannel.recordAffinity(transactionId, selectedEndpoint, allowDefaultAffinity);
     }
 
     void maybeClearAffinity() {
@@ -378,6 +391,12 @@ public void onMessage(ResponseT message) {
         if (transactionId != null) {
           call.maybeRecordAffinity(transactionId);
         }
+      } else if (message instanceof Transaction) {
+        Transaction response = (Transaction) message;
+        ByteString transactionId = transactionIdFromTransaction(response);
+        if (transactionId != null) {
+          call.maybeRecordAffinity(transactionId);
+        }
       }
       super.onMessage(message);
     }

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

@@ -39,6 +39,7 @@
 import com.google.iam.v1.Policy;
 import com.google.iam.v1.TestIamPermissionsResponse;
 import com.google.longrunning.Operation;
+import com.google.protobuf.ByteString;
 import com.google.protobuf.Empty;
 import com.google.protobuf.FieldMask;
 import com.google.spanner.admin.database.v1.Backup;
@@ -194,6 +195,9 @@ default RequestIdCreator getRequestIdCreator() {
     throw new UnsupportedOperationException("Not implemented");
   }
 
+  /** Clears any client-side affinity associated with the given transaction id. */
+  default void clearTransactionAffinity(ByteString transactionId) {}
+
   // Instance admin APIs.
   Paginated<InstanceConfig> listInstanceConfigs(int pageSize, @Nullable String pageToken)
       throws SpannerException;