chore(spanner): handle server side kill switch of R/W multiplexed sessions (googleapis#3441)

* chore(spanner): handle server side kill switch for multiplexed sessions with read-write * chore(spanner): lint fix * chore(spanner): add error message * chore(spanner): fix unit tests due to backgroung BeginTransaction RPC * chore(spanner): lint fix * chore(spanner): handle mock spanner logic * chore(spanner): testcase fix * chore(spanner): do not register backgroung begin txn in mock spanner * chore(spanner): revert unit test changes * chore(spanner): verify error message of the exeception

Author: harshachinta

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

@@ -92,7 +92,7 @@ DatabaseClient getMultiplexedSession() {
 
  @VisibleForTesting
  DatabaseClient getMultiplexedSessionForRW() {
- if (this.useMultiplexedSessionForRW) {
+ if (canUseMultiplexedSessionsForRW()) {
  return getMultiplexedSession();
  }
  return getSession();
@@ -107,6 +107,12 @@ private boolean canUseMultiplexedSessions() {
  && this.multiplexedSessionDatabaseClient.isMultiplexedSessionsSupported();
  }
 
+ private boolean canUseMultiplexedSessionsForRW() {
+ return this.useMultiplexedSessionForRW
+ && this.multiplexedSessionDatabaseClient != null
+ && this.multiplexedSessionDatabaseClient.isMultiplexedSessionsForRWSupported();
+ }
+
  @Override
  public Dialect getDialect() {
  return pool.getDialect();
@@ -129,7 +135,7 @@ public CommitResponse writeWithOptions(
  throws SpannerException {
  ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, options);
  try (IScope s = tracer.withSpan(span)) {
- if (this.useMultiplexedSessionForRW && getMultiplexedSessionDatabaseClient() != null) {
+ if (canUseMultiplexedSessionsForRW() && getMultiplexedSessionDatabaseClient() != null) {
  return getMultiplexedSessionDatabaseClient().writeWithOptions(mutations, options);
  }
  return runWithSessionRetry(session -> session.writeWithOptions(mutations, options));

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

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner;
 
 import static com.google.cloud.spanner.SessionImpl.NO_CHANNEL_HINT;
+import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException;
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
@@ -27,6 +28,10 @@
 import com.google.cloud.spanner.SpannerException.ResourceNotFoundException;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.util.concurrent.MoreExecutors;
+import com.google.spanner.v1.BeginTransactionRequest;
+import com.google.spanner.v1.RequestOptions;
+import com.google.spanner.v1.Transaction;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
@@ -92,6 +97,10 @@ void onError(SpannerException spannerException) {
  // synchronizing, as it does not really matter exactly which error is set.
  this.client.resourceNotFoundException.set((ResourceNotFoundException) spannerException);
  }
+ // Mark multiplexed sessions for RW as unimplemented and fall back to regular sessions if
+ // UNIMPLEMENTED with error message "Transaction type read_write not supported with
+ // multiplexed sessions" is returned.
+ this.client.maybeMarkUnimplementedForRW(spannerException);
  }
 
  @Override
@@ -164,6 +173,12 @@ public void close() {
  /** The current multiplexed session that is used by this client. */
  private final AtomicReference<ApiFuture<SessionReference>> multiplexedSessionReference;
 
+ /**
+ * The Transaction response returned by the BeginTransaction request with read-write when a
+ * multiplexed session is created during client initialization.
+ */
+ private final SettableApiFuture<Transaction> readWriteBeginTransactionReferenceFuture;
+
  /** The expiration date/time of the current multiplexed session. */
  private final AtomicReference<Instant> expirationDate;
 
@@ -190,6 +205,12 @@ public void close() {
  */
  private final AtomicBoolean unimplemented = new AtomicBoolean(false);
 
+ /**
+ * This flag is set to true if the server return UNIMPLEMENTED when a read-write transaction is
+ * executed on a multiplexed session. TODO: Remove once this is guaranteed to be available.
+ */
+ @VisibleForTesting final AtomicBoolean unimplementedForRW = new AtomicBoolean(false);
+
  MultiplexedSessionDatabaseClient(SessionClient sessionClient) {
  this(sessionClient, Clock.systemUTC());
  }
@@ -217,6 +238,7 @@ public void close() {
  this.tracer = sessionClient.getSpanner().getTracer();
  final SettableApiFuture<SessionReference> initialSessionReferenceFuture =
  SettableApiFuture.create();
+ this.readWriteBeginTransactionReferenceFuture = SettableApiFuture.create();
  this.multiplexedSessionReference = new AtomicReference<>(initialSessionReferenceFuture);
  this.sessionClient.asyncCreateMultiplexedSession(
  new SessionConsumer() {
@@ -226,6 +248,16 @@ public void onSessionReady(SessionImpl session) {
  // only start the maintainer if we actually managed to create a session in the first
  // place.
  maintainer.start();
+
+ // initiate a begin transaction request to verify if read-write transactions are
+ // supported using multiplexed sessions.
+ if (sessionClient
+ .getSpanner()
+ .getOptions()
+ .getSessionPoolOptions()
+ .getUseMultiplexedSessionForRW()) {
+ verifyBeginTransactionWithRWOnMultiplexedSessionAsync(session.getName());
+ }
  }
 
  @Override
@@ -267,6 +299,70 @@ private void maybeMarkUnimplemented(Throwable t) {
  }
  }
 
+ private void maybeMarkUnimplementedForRW(SpannerException spannerException) {
+ if (spannerException.getErrorCode() == ErrorCode.UNIMPLEMENTED
+ && verifyErrorMessage(
+ spannerException,
+ "Transaction type read_write not supported with multiplexed sessions")) {
+ unimplementedForRW.set(true);
+ }
+ }
+
+ private boolean verifyErrorMessage(SpannerException spannerException, String message) {
+ if (spannerException.getCause() == null) {
+ return false;
+ }
+ if (spannerException.getCause().getMessage() == null) {
+ return false;
+ }
+ return spannerException.getCause().getMessage().contains(message);
+ }
+
+ private void verifyBeginTransactionWithRWOnMultiplexedSessionAsync(String sessionName) {
+ // TODO: Remove once this is guaranteed to be available.
+ // annotate the explict BeginTransactionRequest with a transaction tag
+ // "multiplexed-rw-background-begin-txn" to avoid storing this request on mock spanner.
+ // this is to safeguard other mock spanner tests whose BeginTransaction request count will
+ // otherwise increase by 1. Modifying the unit tests do not seem valid since this code is
+ // temporary and will be removed once the read-write on multiplexed session looks stable at
+ // backend.
+ BeginTransactionRequest.Builder requestBuilder =
+ BeginTransactionRequest.newBuilder()
+ .setSession(sessionName)
+ .setOptions(
+ SessionImpl.createReadWriteTransactionOptions(
+ Options.fromTransactionOptions(), /* previousTransactionId = */ null))
+ .setRequestOptions(
+ RequestOptions.newBuilder()
+ .setTransactionTag("multiplexed-rw-background-begin-txn")
+ .build());
+ final BeginTransactionRequest request = requestBuilder.build();
+ final ApiFuture<Transaction> requestFuture;
+ requestFuture =
+ sessionClient
+ .getSpanner()
+ .getRpc()
+ .beginTransactionAsync(request, /* options = */ null, /* routeToLeader = */ true);
+ requestFuture.addListener(
+ () -> {
+ try {
+ Transaction txn = requestFuture.get();
+ if (txn.getId().isEmpty()) {
+ throw newSpannerException(
+ ErrorCode.INTERNAL, "Missing id in transaction\n" + sessionName);
+ }
+ readWriteBeginTransactionReferenceFuture.set(txn);
+ } catch (Exception e) {
+ SpannerException spannerException = SpannerExceptionFactory.newSpannerException(e);
+ // Mark multiplexed sessions for RW as unimplemented and fall back to regular sessions
+ // if UNIMPLEMENTED is returned.
+ maybeMarkUnimplementedForRW(spannerException);
+ readWriteBeginTransactionReferenceFuture.setException(e);
+ }
+ },
+ MoreExecutors.directExecutor());
+ }
+
  boolean isValid() {
  return resourceNotFoundException.get() == null;
  }
@@ -283,6 +379,10 @@ boolean isMultiplexedSessionsSupported() {
  return !this.unimplemented.get();
  }
 
+ boolean isMultiplexedSessionsForRWSupported() {
+ return !this.unimplementedForRW.get();
+ }
+
  void close() {
  synchronized (this) {
  if (!this.isClosed) {
@@ -308,6 +408,17 @@ SessionReference getCurrentSessionReference() {
  }
  }
 
+ @VisibleForTesting
+ Transaction getReadWriteBeginTransactionReference() {
+ try {
+ return this.readWriteBeginTransactionReferenceFuture.get();
+ } catch (ExecutionException executionException) {
+ throw SpannerExceptionFactory.asSpannerException(executionException.getCause());
+ } catch (InterruptedException interruptedException) {
+ throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
+ }
+ }
+
  /**
  * Returns true if the multiplexed session has been created. This client can be used before the
  * session has been created, and will in that case use a delayed transaction that contains a

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

@@ -1872,7 +1872,17 @@ private Transaction getTemporaryTransactionOrNull(TransactionSelector tx) {
  @Override
  public void beginTransaction(
  BeginTransactionRequest request, StreamObserver<Transaction> responseObserver) {
- requests.add(request);
+ // TODO: Remove once this is guaranteed to be available.
+ // Skip storing the explicit BeginTransactionRequest used to verify read-write transaction
+ // server availability on multiplexed sessions.
+ // This code will be removed once read-write multiplexed sessions are stable on the backend,
+ // hence the temporary trade-off.
+ if (!request
+ .getRequestOptions()
+ .getTransactionTag()
+ .equals("multiplexed-rw-background-begin-txn")) {
+ requests.add(request);
+ }
  Preconditions.checkNotNull(request.getSession());
  Session session = getSession(request.getSession());
  if (session == null) {