feat(spanner): support multiplexed session for Partitioned operations

.github/workflows/ci.yaml

@@ -53,6 +53,7 @@ jobs:
       env:
         JOB_TYPE: test
         GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true
+        GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true
   units-java8:
     # Building using Java 17 and run the tests with Java 8 runtime
     name: "units (8)"
@@ -92,6 +93,7 @@ jobs:
         env:
           JOB_TYPE: test
           GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true
+          GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true
   windows:
     runs-on: windows-latest
     steps:

.kokoro/presubmit/integration-multiplexed-sessions-enabled.cfg

@@ -36,3 +36,8 @@ env_vars: {
     key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS"
     value: "true"
 }
+
+env_vars: {
+    key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS"
+    value: "true"
+}

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

@@ -19,7 +19,6 @@
 import com.google.api.gax.rpc.ServerStream;
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.Options.TransactionOption;
-import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.spanner.v1.BatchWriteResponse;
 
 /**
@@ -51,9 +50,4 @@ public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
       throws SpannerException {
     throw new UnsupportedOperationException();
   }
-
-  @Override
-  public long executePartitionedUpdate(Statement stmt, UpdateOption... options) {
-    throw new UnsupportedOperationException();
-  }
 }

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

@@ -37,9 +37,11 @@
 /** Default implementation for Batch Client interface. */
 public class BatchClientImpl implements BatchClient {
   private final SessionClient sessionClient;
+  private final boolean isMultiplexedSessionEnabled;
 
-  BatchClientImpl(SessionClient sessionClient) {
+  BatchClientImpl(SessionClient sessionClient, boolean isMultiplexedSessionEnabled) {
     this.sessionClient = checkNotNull(sessionClient);
+    this.isMultiplexedSessionEnabled = isMultiplexedSessionEnabled;
   }
 
   @Override
@@ -50,7 +52,12 @@ public String getDatabaseRole() {
 
   @Override
   public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) {
-    SessionImpl session = sessionClient.createSession();
+    SessionImpl session;
+    if (isMultiplexedSessionEnabled) {
+      session = sessionClient.createMultiplexedSession();
+    } else {
+      session = sessionClient.createSession();
+    }
     return new BatchReadOnlyTransactionImpl(
         MultiUseReadOnlyTransaction.newBuilder()
             .setSession(session)
@@ -72,7 +79,9 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) {
   @Override
   public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batchTransactionId) {
     SessionImpl session =
-        sessionClient.sessionWithId(checkNotNull(batchTransactionId).getSessionId());
+        sessionClient.sessionWithId(
+            checkNotNull(batchTransactionId).getSessionId(),
+            batchTransactionId.isMultiplexedSession());
     return new BatchReadOnlyTransactionImpl(
         MultiUseReadOnlyTransaction.newBuilder()
             .setSession(session)
@@ -95,12 +104,14 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batc
   private static class BatchReadOnlyTransactionImpl extends MultiUseReadOnlyTransaction
       implements BatchReadOnlyTransaction {
     private final String sessionName;
+    private final boolean isMultiplexedSession;
     private final Map<SpannerRpc.Option, ?> options;
 
     BatchReadOnlyTransactionImpl(
         MultiUseReadOnlyTransaction.Builder builder, TimestampBound bound) {
       super(builder.setTimestampBound(bound));
       this.sessionName = session.getName();
+      this.isMultiplexedSession = session.getIsMultiplexed();
       this.options = session.getOptions();
       initTransaction();
     }
@@ -110,11 +121,13 @@ private static class BatchReadOnlyTransactionImpl extends MultiUseReadOnlyTransa
       super(builder.setTransactionId(batchTransactionId.getTransactionId()));
       this.sessionName = session.getName();
       this.options = session.getOptions();
+      this.isMultiplexedSession = session.getIsMultiplexed();
     }
 
     @Override
     public BatchTransactionId getBatchTransactionId() {
-      return new BatchTransactionId(sessionName, getTransactionId(), getReadTimestamp());
+      return new BatchTransactionId(
+          sessionName, getTransactionId(), getReadTimestamp(), session.getIsMultiplexed());
     }
 
     @Override

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

@@ -32,12 +32,18 @@ public class BatchTransactionId implements Serializable {
   private final ByteString transactionId;
   private final String sessionId;
   private final Timestamp timestamp;
+  private final boolean isMultiplexedSession;
   private static final long serialVersionUID = 8067099123096783939L;
 
-  BatchTransactionId(String sessionId, ByteString transactionId, Timestamp timestamp) {
+  BatchTransactionId(
+      String sessionId,
+      ByteString transactionId,
+      Timestamp timestamp,
+      boolean isMultiplexedSession) {
     this.transactionId = Preconditions.checkNotNull(transactionId);
     this.sessionId = Preconditions.checkNotNull(sessionId);
     this.timestamp = Preconditions.checkNotNull(timestamp);
+    this.isMultiplexedSession = isMultiplexedSession;
   }
 
   ByteString getTransactionId() {
@@ -52,11 +58,15 @@ Timestamp getTimestamp() {
     return timestamp;
   }
 
+  public boolean isMultiplexedSession() {
+    return isMultiplexedSession;
+  }
+
   @Override
   public String toString() {
     return String.format(
-        "transactionId: %s, sessionId: %s, timestamp: %s",
-        transactionId.toStringUtf8(), sessionId, timestamp);
+        "transactionId: %s, sessionId: %s, timestamp: %s, isMultiplexedSession: %s",
+        transactionId.toStringUtf8(), sessionId, timestamp, isMultiplexedSession);
   }
 
   @Override

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

@@ -36,6 +36,7 @@ class DatabaseClientImpl implements DatabaseClient {
   @VisibleForTesting final String clientId;
   @VisibleForTesting final SessionPool pool;
   @VisibleForTesting final MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient;
+  @VisibleForTesting final boolean useMultiplexedSessionPartitionedOps;
   @VisibleForTesting final boolean useMultiplexedSessionForRW;
 
   final boolean useMultiplexedSessionBlindWrite;
@@ -47,6 +48,7 @@ class DatabaseClientImpl implements DatabaseClient {
         pool,
         /* useMultiplexedSessionBlindWrite = */ false,
         /* multiplexedSessionDatabaseClient = */ null,
+        /* useMultiplexedSessionPartitionedOps= */ false,
         tracer,
         /* useMultiplexedSessionForRW = */ false);
   }
@@ -58,6 +60,7 @@ class DatabaseClientImpl implements DatabaseClient {
         pool,
         /* useMultiplexedSessionBlindWrite = */ false,
         /* multiplexedSessionDatabaseClient = */ null,
+        /* useMultiplexedSessionPartitionedOps= */ false,
         tracer,
         /* useMultiplexedSessionForRW = */ false);
   }
@@ -67,12 +70,14 @@ class DatabaseClientImpl implements DatabaseClient {
       SessionPool pool,
       boolean useMultiplexedSessionBlindWrite,
       @Nullable MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient,
+      boolean useMultiplexedSessionPartitionedOps,
       TraceWrapper tracer,
       boolean useMultiplexedSessionForRW) {
     this.clientId = clientId;
     this.pool = pool;
     this.useMultiplexedSessionBlindWrite = useMultiplexedSessionBlindWrite;
     this.multiplexedSessionDatabaseClient = multiplexedSessionDatabaseClient;
+    this.useMultiplexedSessionPartitionedOps = useMultiplexedSessionPartitionedOps;
     this.tracer = tracer;
     this.useMultiplexedSessionForRW = useMultiplexedSessionForRW;
   }
@@ -309,6 +314,14 @@ public AsyncTransactionManager transactionManagerAsync(TransactionOption... opti
 
   @Override
   public long executePartitionedUpdate(final Statement stmt, final UpdateOption... options) {
+    if (useMultiplexedSessionPartitionedOps) {
+      return getMultiplexedSession().executePartitionedUpdate(stmt, options);
+    }
+    return executePartitionedUpdateSession(stmt, options);
+  }
+
+  private long executePartitionedUpdateSession(
+      final Statement stmt, final UpdateOption... options) {
     ISpan span = tracer.spanBuilder(PARTITION_DML_TRANSACTION);
     try (IScope s = tracer.withSpan(span)) {
       return runWithSessionRetry(session -> session.executePartitionedUpdate(stmt, options));

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

@@ -24,6 +24,7 @@
 import com.google.cloud.spanner.DelayedReadContext.DelayedReadOnlyTransaction;
 import com.google.cloud.spanner.MultiplexedSessionDatabaseClient.MultiplexedSessionTransaction;
 import com.google.cloud.spanner.Options.TransactionOption;
+import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.common.util.concurrent.MoreExecutors;
 import java.util.concurrent.ExecutionException;
 
@@ -224,4 +225,22 @@ private SessionReference getSessionReference() {
       throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
     }
   }
+
+  /**
+   * Execute `stmt` within PARTITIONED_DML transaction using multiplexed session. This method is a
+   * blocking call as the interface expects to return the output of the `stmt`.
+   */
+  @Override
+  public long executePartitionedUpdate(Statement stmt, UpdateOption... options) {
+    try {
+      SessionReference sessionReference = this.sessionFuture.get();
+      return new MultiplexedSessionTransaction(
+              client, span, sessionReference, NO_CHANNEL_HINT, true)
+          .executePartitionedUpdate(stmt, options);
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e);
+    } catch (ExecutionException e) {
+      throw new RuntimeException(e);
+    }
+  }
 }

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

@@ -24,6 +24,7 @@
 import com.google.api.core.SettableApiFuture;
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.Options.TransactionOption;
+import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.SessionClient.SessionConsumer;
 import com.google.cloud.spanner.SpannerException.ResourceNotFoundException;
 import com.google.common.annotations.VisibleForTesting;
@@ -553,6 +554,11 @@ public AsyncTransactionManager transactionManagerAsync(TransactionOption... opti
         .transactionManagerAsync(options);
   }
 
+  @Override
+  public long executePartitionedUpdate(Statement stmt, UpdateOption... options) {
+    return createMultiplexedSessionTransaction(true).executePartitionedUpdate(stmt, options);
+  }
+
   /**
    * It is enough with one executor to maintain the multiplexed sessions in all the clients, as they
    * do not need to be updated often, and the maintenance task is light. The core pool size is set

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

@@ -237,6 +237,19 @@ SessionImpl createSession() {
    * @param consumer The {@link SessionConsumer} to use for callbacks when sessions are available.
    */
   void createMultiplexedSession(SessionConsumer consumer) {
+    try {
+      SessionImpl sessionImpl = createMultiplexedSession();
+      consumer.onSessionReady(sessionImpl);
+    } catch (Throwable t) {
+      consumer.onSessionCreateFailure(t, 1);
+    }
+  }
+
+  /**
+   * Create a multiplexed session and returns it. A multiplexed session is not affiliated with any
+   * GRPC channel. In case of an error on the gRPC calls, the error will be returned.
+   */
+  SessionImpl createMultiplexedSession() {
     ISpan span = spanner.getTracer().spanBuilder(SpannerImpl.CREATE_MULTIPLEXED_SESSION);
     try (IScope s = spanner.getTracer().withSpan(span)) {
       com.google.spanner.v1.Session session =
@@ -253,10 +266,12 @@ void createMultiplexedSession(SessionConsumer consumer) {
               spanner,
               new SessionReference(
                   session.getName(), session.getCreateTime(), session.getMultiplexed(), null));
-      consumer.onSessionReady(sessionImpl);
+      span.addAnnotation(
+          String.format("Request for %d multiplexed session returned %d session", 1, 1));
+      return sessionImpl;
     } catch (Throwable t) {
       span.setStatus(t);
-      consumer.onSessionCreateFailure(t, 1);
+      throw t;
     } finally {
       span.end();
     }
@@ -289,31 +304,7 @@ private CreateMultiplexedSessionsRunnable(SessionConsumer consumer) {
 
     @Override
     public void run() {
-      ISpan span = spanner.getTracer().spanBuilder(SpannerImpl.CREATE_MULTIPLEXED_SESSION);
-      try (IScope s = spanner.getTracer().withSpan(span)) {
-        com.google.spanner.v1.Session session =
-            spanner
-                .getRpc()
-                .createSession(
-                    db.getName(),
-                    spanner.getOptions().getDatabaseRole(),
-                    spanner.getOptions().getSessionLabels(),
-                    null,
-                    true);
-        SessionImpl sessionImpl =
-            new SessionImpl(
-                spanner,
-                new SessionReference(
-                    session.getName(), session.getCreateTime(), session.getMultiplexed(), null));
-        span.addAnnotation(
-            String.format("Request for %d multiplexed session returned %d session", 1, 1));
-        consumer.onSessionReady(sessionImpl);
-      } catch (Throwable t) {
-        span.setStatus(t);
-        consumer.onSessionCreateFailure(t, 1);
-      } finally {
-        span.end();
-      }
+      createMultiplexedSession(consumer);
     }
   }
 
@@ -423,10 +414,16 @@ private List<SessionImpl> internalBatchCreateSessions(
 
   /** Returns a {@link SessionImpl} that references the existing session with the given name. */
   SessionImpl sessionWithId(String name) {
+    return sessionWithId(name, /*isMultiplexedSession= */ false);
+  }
+
+  /** Returns a {@link SessionImpl} that references the existing session with the given name. */
+  SessionImpl sessionWithId(String name, boolean isMultiplexedSession) {
     final Map<SpannerRpc.Option, ?> options;
     synchronized (this) {
       options = optionMap(SessionOption.channelHint(sessionChannelCounter++));
     }
-    return new SessionImpl(spanner, new SessionReference(name, options));
+    return new SessionImpl(
+        spanner, new SessionReference(name, /*createTime= */ null, isMultiplexedSession, options));
   }
 }