chore(spanner): add multiplexed session flag in executor code (#3030)

* chore(spanner): add multiplexed session flag for executor code

* chore(spanner): clirr fix

* chore(spanner): refactor doublt to boolean flag

* chore(spanner): clirr fix

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* chore(spanner): code refactoring

* chore(spanner): add validation for ratio flag

* chore(spanner): add logs

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: harshachinta

google-cloud-spanner-executor/clirr-ignored-differences.xml

@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- see http://www.mojohaus.org/clirr-maven-plugin/examples/ignored-differences.html -->
+<differences>
+  <difference>
+    <differenceType>7004</differenceType>
+    <className>com/google/cloud/executor/spanner/CloudExecutorImpl</className>
+    <method>CloudExecutorImpl(boolean)</method>
+    <to>CloudExecutorImpl(boolean, double)</to>
+  </difference>
+</differences>

google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudExecutorImpl.java

@@ -18,9 +18,12 @@
 
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.SpannerExceptionFactory;
+import com.google.spanner.executor.v1.SessionPoolOptions;
+import com.google.spanner.executor.v1.SpannerAction;
 import com.google.spanner.executor.v1.SpannerAsyncActionRequest;
 import com.google.spanner.executor.v1.SpannerAsyncActionResponse;
 import com.google.spanner.executor.v1.SpannerExecutorProxyGrpc;
+import com.google.spanner.executor.v1.SpannerOptions;
 import io.grpc.Status;
 import io.grpc.stub.StreamObserver;
 import java.util.logging.Level;
@@ -34,8 +37,13 @@ public class CloudExecutorImpl extends SpannerExecutorProxyGrpc.SpannerExecutorP
   // Executors to proxy.
   private final CloudClientExecutor clientExecutor;
 
-  public CloudExecutorImpl(boolean enableGrpcFaultInjector) {
+  // Ratio of operations to use multiplexed sessions.
+  private final double multiplexedSessionOperationsRatio;
+
+  public CloudExecutorImpl(
+      boolean enableGrpcFaultInjector, double multiplexedSessionOperationsRatio) {
     clientExecutor = new CloudClientExecutor(enableGrpcFaultInjector);
+    this.multiplexedSessionOperationsRatio = multiplexedSessionOperationsRatio;
   }
 
   /** Execute SpannerAsync action requests. */
@@ -48,6 +56,36 @@ public StreamObserver<SpannerAsyncActionRequest> executeActionAsync(
       @Override
       public void onNext(SpannerAsyncActionRequest request) {
         LOGGER.log(Level.INFO, String.format("Receiving request: \n%s", request));
+
+        // Use Multiplexed sessions for all supported operations if the
+        // multiplexedSessionOperationsRatio from command line is > 0.0
+        if (multiplexedSessionOperationsRatio > 0.0) {
+          SessionPoolOptions.Builder sessionPoolOptionsBuilder;
+          if (request.getAction().getSpannerOptions().hasSessionPoolOptions()) {
+            sessionPoolOptionsBuilder =
+                request
+                    .getAction()
+                    .getSpannerOptions()
+                    .getSessionPoolOptions()
+                    .toBuilder()
+                    .setUseMultiplexed(true);
+          } else {
+            sessionPoolOptionsBuilder = SessionPoolOptions.newBuilder().setUseMultiplexed(true);
+          }
+
+          SpannerOptions.Builder optionsBuilder =
+              request
+                  .getAction()
+                  .getSpannerOptions()
+                  .toBuilder()
+                  .setSessionPoolOptions(sessionPoolOptionsBuilder);
+          SpannerAction.Builder actionBuilder =
+              request.getAction().toBuilder().setSpannerOptions(optionsBuilder);
+          request = request.toBuilder().setAction(actionBuilder).build();
+          LOGGER.log(
+              Level.INFO,
+              String.format("Updated request to set multiplexed session flag: \n%s", request));
+        }
         Status status = clientExecutor.startHandlingRequest(request, executionContext);
         if (!status.isOk()) {
           LOGGER.log(

google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/WorkerProxy.java

@@ -45,17 +45,21 @@ public class WorkerProxy {
   private static final String OPTION_SERVICE_KEY_FILE = "service_key_file";
   private static final String OPTION_USE_PLAIN_TEXT_CHANNEL = "use_plain_text_channel";
   private static final String OPTION_ENABLE_GRPC_FAULT_INJECTOR = "enable_grpc_fault_injector";
+  private static final String OPTION_MULTIPLEXED_SESSION_OPERATIONS_RATIO =
+      "multiplexed_session_operations_ratio";
 
   public static int spannerPort = 0;
   public static int proxyPort = 0;
   public static String cert = "";
   public static String serviceKeyFile = "";
+  public static double multiplexedSessionOperationsRatio = 0.0;
   public static boolean usePlainTextChannel = false;
   public static boolean enableGrpcFaultInjector = false;
 
   public static CommandLine commandLine;
 
   private static final int MIN_PORT = 0, MAX_PORT = 65535;
+  private static final double MIN_RATIO = 0.0, MAX_RATIO = 1.0;
 
   public static void main(String[] args) throws Exception {
     commandLine = buildOptions(args);
@@ -95,10 +99,30 @@ public static void main(String[] args) throws Exception {
     usePlainTextChannel = commandLine.hasOption(OPTION_USE_PLAIN_TEXT_CHANNEL);
     enableGrpcFaultInjector = commandLine.hasOption(OPTION_ENABLE_GRPC_FAULT_INJECTOR);
 
+    if (commandLine.hasOption(OPTION_MULTIPLEXED_SESSION_OPERATIONS_RATIO)) {
+      multiplexedSessionOperationsRatio =
+          Double.parseDouble(
+              commandLine.getOptionValue(OPTION_MULTIPLEXED_SESSION_OPERATIONS_RATIO));
+      LOGGER.log(
+          Level.INFO,
+          String.format(
+              "Multiplexed session ratio from commandline arg: \n%s",
+              multiplexedSessionOperationsRatio));
+      if (multiplexedSessionOperationsRatio < MIN_RATIO
+          || multiplexedSessionOperationsRatio > MAX_RATIO) {
+        throw new IllegalArgumentException(
+            "Spanner multiplexedSessionOperationsRatio must be between "
+                + MIN_RATIO
+                + " and "
+                + MAX_RATIO);
+      }
+    }
+
     Server server;
     while (true) {
       try {
-        CloudExecutorImpl cloudExecutorImpl = new CloudExecutorImpl(enableGrpcFaultInjector);
+        CloudExecutorImpl cloudExecutorImpl =
+            new CloudExecutorImpl(enableGrpcFaultInjector, multiplexedSessionOperationsRatio);
         HealthStatusManager healthStatusManager = new HealthStatusManager();
         // Set up Cloud server.
         server =
@@ -139,6 +163,11 @@ private static CommandLine buildOptions(String[] args) {
         OPTION_ENABLE_GRPC_FAULT_INJECTOR,
         false,
         "Enable grpc fault injector in cloud client executor.");
+    options.addOption(
+        null,
+        OPTION_MULTIPLEXED_SESSION_OPERATIONS_RATIO,
+        true,
+        "Ratio of operations to use multiplexed sessions.");
 
     CommandLineParser parser = new DefaultParser();
     try {