chore: add a random hint for multi-use transactions when they are use… (googleapis#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";
  }