feat(spanner): Modified BatchClientImpl to store multiplexed session and create fresh session after expiration date.

Author: pratickchokhani

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

@@ -30,18 +30,53 @@
 import com.google.spanner.v1.PartitionReadRequest;
 import com.google.spanner.v1.PartitionResponse;
 import com.google.spanner.v1.TransactionSelector;
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.ReentrantLock;
 import javax.annotation.Nullable;
+import javax.annotation.concurrent.GuardedBy;
 
 /** Default implementation for Batch Client interface. */
 public class BatchClientImpl implements BatchClient {
   private final SessionClient sessionClient;
+
   private final boolean isMultiplexedSessionEnabled;
 
+  /** Lock to protect the multiplexed session. */
+  private final ReentrantLock multiplexedSessionLock = new ReentrantLock();
+
+  /** The duration before we try to replace the multiplexed session. The default is 7 days. */
+  private final Duration sessionExpirationDuration;
+
+  /** The expiration date/time of the current multiplexed session. */
+  @GuardedBy("multiplexedSessionLock")
+  private final AtomicReference<Instant> expirationDate;
+
+  @GuardedBy("multiplexedSessionLock")
+  private final AtomicReference<SessionImpl> multiplexedSessionReference;
+
+  private final Clock clock;
+
   BatchClientImpl(SessionClient sessionClient, boolean isMultiplexedSessionEnabled) {
     this.sessionClient = checkNotNull(sessionClient);
     this.isMultiplexedSessionEnabled = isMultiplexedSessionEnabled;
+    this.sessionExpirationDuration =
+        Duration.ofMillis(
+            sessionClient
+                .getSpanner()
+                .getOptions()
+                .getSessionPoolOptions()
+                .getMultiplexedSessionMaintenanceDuration()
+                .toMillis());
+    // Initialize the expiration date to the start of time to avoid unnecessary null checks.
+    // This also ensured that a new session is created on first request.
+    this.expirationDate = new AtomicReference<>(Instant.MIN);
+    this.multiplexedSessionReference = new AtomicReference<>();
+    clock = Clock.systemUTC();
   }
 
   @Override
@@ -54,7 +89,7 @@ public String getDatabaseRole() {
   public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) {
     SessionImpl session;
     if (isMultiplexedSessionEnabled) {
-      session = sessionClient.createMultiplexedSession();
+      session = getMultiplexedSession();
     } else {
       session = sessionClient.createSession();
     }
@@ -99,6 +134,20 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batc
         batchTransactionId);
   }
 
+  private SessionImpl getMultiplexedSession() {
+    this.multiplexedSessionLock.lock();
+    try {
+      if (this.clock.instant().isAfter(this.expirationDate.get())
+          || this.multiplexedSessionReference.get() == null) {
+        this.multiplexedSessionReference.set(this.sessionClient.createMultiplexedSession());
+        this.expirationDate.set(this.clock.instant().plus(this.sessionExpirationDuration));
+      }
+      return this.multiplexedSessionReference.get();
+    } finally {
+      this.multiplexedSessionLock.unlock();
+    }
+  }
+
   private static class BatchReadOnlyTransactionImpl extends MultiUseReadOnlyTransaction
       implements BatchReadOnlyTransaction {
     private final String sessionName;

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

@@ -51,6 +51,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -107,6 +108,11 @@ private static String nextDatabaseClientId(DatabaseId databaseId) {
   @GuardedBy("this")
   private final Map<DatabaseId, DatabaseClientImpl> dbClients = new HashMap<>();
 
+  @GuardedBy("dbBatchClientLock")
+  private final Map<DatabaseId, BatchClientImpl> dbBatchClients = new HashMap<>();
+
+  private final ReentrantLock dbBatchClientLock = new ReentrantLock();
+
   private final CloseableExecutorProvider asyncExecutorProvider;
 
   @GuardedBy("this")
@@ -336,9 +342,23 @@ DatabaseClientImpl createDatabaseClient(
 
   @Override
   public BatchClient getBatchClient(DatabaseId db) {
+    if (getOptions().getSessionPoolOptions().getUseMultiplexedSessionPartitionedOps()) {
+      this.dbBatchClientLock.lock();
+      try {
+        if (this.dbBatchClients.containsKey(db)) {
+          return this.dbBatchClients.get(db);
+        }
+        BatchClientImpl batchClient =
+            new BatchClientImpl(
+                getSessionClient(db), /*useMultiplexedSessionPartitionedOps=*/ true);
+        this.dbBatchClients.put(db, batchClient);
+        return batchClient;
+      } finally {
+        this.dbBatchClientLock.unlock();
+      }
+    }
     return new BatchClientImpl(
-        getSessionClient(db),
-        getOptions().getSessionPoolOptions().getUseMultiplexedSessionPartitionedOps());
+        getSessionClient(db), /*useMultiplexedSessionPartitionedOps=*/ false);
   }
 
   @Override

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

@@ -47,6 +47,7 @@
 import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
+import org.threeten.bp.Duration;
 
 /** Unit tests for {@link com.google.cloud.spanner.BatchClientImpl}. */
 @RunWith(JUnit4.class)
@@ -96,6 +97,7 @@ public void setUp() {
     when(sessionPoolOptions.getPoolMaintainerClock()).thenReturn(Clock.INSTANCE);
     when(sessionPoolOptions.getUseMultiplexedSessionPartitionedOps())
         .thenReturn(isMultiplexedSession);
+    when(sessionPoolOptions.getMultiplexedSessionMaintenanceDuration()).thenReturn(Duration.ZERO);
     when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions);
     @SuppressWarnings("resource")
     SpannerImpl spanner = new SpannerImpl(gapicRpc, spannerOptions);