chore: add a random hint for multi-use transactions when they are use… (#3058)

Author: arpan14

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

@@ -16,6 +16,7 @@
 
 package com.google.cloud.spanner;
 
+import static com.google.cloud.spanner.SessionClient.optionMap;
 import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
@@ -32,6 +33,7 @@
 import com.google.cloud.spanner.AsyncResultSet.ReadyCallback;
 import com.google.cloud.spanner.Options.QueryOption;
 import com.google.cloud.spanner.Options.ReadOption;
+import com.google.cloud.spanner.SessionClient.SessionOption;
 import com.google.cloud.spanner.SessionImpl.SessionTransaction;
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
 import com.google.common.annotations.VisibleForTesting;
@@ -52,6 +54,7 @@
 import com.google.spanner.v1.TransactionOptions;
 import com.google.spanner.v1.TransactionSelector;
 import java.util.Map;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicLong;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.GuardedBy;
@@ -180,9 +183,15 @@ static Builder newBuilder() {
     @GuardedBy("lock")
     private boolean used;
 
+    private final Map<SpannerRpc.Option, ?> channelHint;
+
     private SingleReadContext(Builder builder) {
       super(builder);
       this.bound = builder.bound;
+      // single use transaction have a single RPC and hence there is no need
+      // of a channel hint. GAX will automatically choose a hint when used
+      // with a multiplexed session.
+      this.channelHint = getChannelHintOptions(session.getOptions(), null);
     }
 
     @Override
@@ -209,6 +218,11 @@ TransactionSelector getTransactionSelector() {
           .setSingleUse(TransactionOptions.newBuilder().setReadOnly(bound.toProto()))
           .build();
     }
+
+    @Override
+    Map<SpannerRpc.Option, ?> getTransactionChannelHint() {
+      return channelHint;
+    }
   }
 
   private static void assertTimestampAvailable(boolean available) {
@@ -217,6 +231,7 @@ private static void assertTimestampAvailable(boolean available) {
 
   static class SingleUseReadOnlyTransaction extends SingleReadContext
       implements ReadOnlyTransaction {
+
     @GuardedBy("lock")
     private Timestamp timestamp;
 
@@ -300,6 +315,8 @@ static Builder newBuilder() {
     @GuardedBy("txnLock")
     private ByteString transactionId;
 
+    private final Map<SpannerRpc.Option, ?> channelHint;
+
     MultiUseReadOnlyTransaction(Builder builder) {
       super(builder);
       checkArgument(
@@ -318,6 +335,14 @@ static Builder newBuilder() {
         this.timestamp = builder.timestamp;
         this.transactionId = builder.transactionId;
       }
+      this.channelHint =
+          getChannelHintOptions(
+              session.getOptions(), ThreadLocalRandom.current().nextLong(Long.MAX_VALUE));
+    }
+
+    @Override
+    public Map<SpannerRpc.Option, ?> getTransactionChannelHint() {
+      return channelHint;
     }
 
     @Override
@@ -380,7 +405,7 @@ void initTransaction() {
                   .setOptions(options)
                   .build();
           Transaction transaction =
-              rpc.beginTransaction(request, session.getOptions(), isRouteToLeader());
+              rpc.beginTransaction(request, getTransactionChannelHint(), isRouteToLeader());
           if (!transaction.hasReadTimestamp()) {
             throw SpannerExceptionFactory.newSpannerException(
                 ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field");
@@ -727,7 +752,10 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
             }
             SpannerRpc.StreamingCall call =
                 rpc.executeQuery(
-                    request.build(), stream.consumer(), session.getOptions(), isRouteToLeader());
+                    request.build(),
+                    stream.consumer(),
+                    getTransactionChannelHint(),
+                    isRouteToLeader());
             session.markUsed(clock.instant());
             call.request(prefetchChunks);
             stream.setCall(call, request.getTransaction().hasBegin());
@@ -738,6 +766,16 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
         stream, this, options.hasDecodeMode() ? options.decodeMode() : defaultDecodeMode);
   }
 
+  Map<SpannerRpc.Option, ?> getChannelHintOptions(
+      Map<SpannerRpc.Option, ?> channelHintForSession, Long channelHintForTransaction) {
+    if (channelHintForSession != null) {
+      return channelHintForSession;
+    } else if (channelHintForTransaction != null) {
+      return optionMap(SessionOption.channelHint(channelHintForTransaction));
+    }
+    return null;
+  }
+
   /**
    * Called before any read or query is started to perform state checks and initializations.
    * Subclasses should call {@code super.beforeReadOrQuery()} if overriding.
@@ -782,6 +820,12 @@ public void close() {
   @Nullable
   abstract TransactionSelector getTransactionSelector();
 
+  /**
+   * Channel hint to be used for a transaction. This enables soft-stickiness per transaction by
+   * ensuring all RPCs within a transaction land up on the same channel.
+   */
+  abstract Map<SpannerRpc.Option, ?> getTransactionChannelHint();
+
   /**
    * Returns the transaction tag for this {@link AbstractReadContext} or <code>null</code> if this
    * {@link AbstractReadContext} does not have a transaction tag.
@@ -872,7 +916,10 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
             builder.setRequestOptions(buildRequestOptions(readOptions));
             SpannerRpc.StreamingCall call =
                 rpc.read(
-                    builder.build(), stream.consumer(), session.getOptions(), isRouteToLeader());
+                    builder.build(),
+                    stream.consumer(),
+                    getTransactionChannelHint(),
+                    isRouteToLeader());
             session.markUsed(clock.instant());
             call.request(prefetchChunks);
             stream.setCall(call, /* withBeginTransaction = */ builder.getTransaction().hasBegin());

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

@@ -30,6 +30,8 @@
 import com.google.cloud.spanner.Options.TransactionOption;
 import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.SessionImpl.SessionTransaction;
+import com.google.cloud.spanner.spi.v1.SpannerRpc;
+import com.google.cloud.spanner.spi.v1.SpannerRpc.Option;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
@@ -51,11 +53,13 @@
 import com.google.spanner.v1.TransactionSelector;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -198,13 +202,18 @@ public void removeListener(Runnable listener) {
     private CommitResponse commitResponse;
     private final Clock clock;
 
+    private final Map<SpannerRpc.Option, ?> channelHint;
+
     private TransactionContextImpl(Builder builder) {
       super(builder);
       this.transactionId = builder.transactionId;
       this.trackTransactionStarter = builder.trackTransactionStarter;
       this.options = builder.options;
       this.finishedAsyncOperations.set(null);
       this.clock = builder.clock;
+      this.channelHint =
+          getChannelHintOptions(
+              session.getOptions(), ThreadLocalRandom.current().nextLong(Long.MAX_VALUE));
     }
 
     @Override
@@ -559,6 +568,11 @@ TransactionSelector getTransactionSelector() {
       return TransactionSelector.newBuilder().setId(transactionId).build();
     }
 
+    @Override
+    Map<Option, ?> getTransactionChannelHint() {
+      return channelHint;
+    }
+
     @Override
     public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId) {
       Preconditions.checkNotNull(transaction);

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

@@ -25,6 +25,7 @@
 import com.google.api.gax.core.ExecutorProvider;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
+import com.google.cloud.spanner.spi.v1.SpannerRpc.Option;
 import com.google.spanner.v1.DirectedReadOptions;
 import com.google.spanner.v1.DirectedReadOptions.IncludeReplicas;
 import com.google.spanner.v1.DirectedReadOptions.ReplicaSelection;
@@ -39,6 +40,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -98,6 +100,11 @@ private final class TestReadContext extends AbstractReadContext {
     TransactionSelector getTransactionSelector() {
       return TransactionSelector.getDefaultInstance();
     }
+
+    @Override
+    Map<SpannerRpc.Option, ?> getTransactionChannelHint() {
+      return null;
+    }
   }
 
   private final class TestReadContextWithTag extends AbstractReadContext {
@@ -110,6 +117,11 @@ TransactionSelector getTransactionSelector() {
       return TransactionSelector.getDefaultInstance();
     }
 
+    @Override
+    Map<Option, ?> getTransactionChannelHint() {
+      return null;
+    }
+
     String getTransactionTag() {
       return "app=spanner,env=test";
     }