Merge branch 'googleapis:main' into main

arpan14 · Mar 25, 2024 · d6d8229 · d6d8229
2 parents fc49911 + e043de8
commit d6d8229
Show file tree

Hide file tree

Showing 14 changed files with 394 additions and 81 deletions.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java
@@ -257,7 +257,9 @@ private <T> T runWithSessionRetry(Function<Session, T> callable) {
       try {
         return callable.apply(session);
       } catch (SessionNotFoundException e) {
-        session = pool.replaceSession(e, session);
+        session =
+            (PooledSessionFuture)
+                pool.getPooledSessionReplacementHandler().replaceSession(e, session);
       }
     }
   }

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java
@@ -215,7 +215,8 @@ SessionImpl createSession() {
                   spanner.getOptions().getDatabaseRole(),
                   spanner.getOptions().getSessionLabels(),
                   options);
-      return new SessionImpl(spanner, session.getName(), options);
+      return new SessionImpl(
+          spanner, session.getName(), session.getCreateTime(), session.getMultiplexed(), options);
     } catch (RuntimeException e) {
       span.setStatus(e);
       throw e;
@@ -224,6 +225,39 @@ SessionImpl createSession() {
     }
   }
 
+  /**
+   * Create a multiplexed session and returns it to the given {@link SessionConsumer}. A multiplexed
+   * session is not affiliated with any GRPC channel. The given {@link SessionConsumer} is
+   * guaranteed to eventually get exactly 1 multiplexed session unless an error occurs. In case of
+   * an error on the gRPC calls, the consumer will receive one {@link
+   * SessionConsumer#onSessionCreateFailure(Throwable, int)} calls with the error.
+   *
+   * @param consumer The {@link SessionConsumer} to use for callbacks when sessions are available.
+   */
+  void createMultiplexedSession(SessionConsumer consumer) {
+    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, session.getName(), session.getCreateTime(), session.getMultiplexed(), null);
+      consumer.onSessionReady(sessionImpl);
+    } catch (Throwable t) {
+      span.setStatus(t);
+      consumer.onSessionCreateFailure(t, 1);
+    } finally {
+      span.end();
+    }
+  }
+
   /**
    * Asynchronously creates a batch of sessions and returns these to the given {@link
    * SessionConsumer}. This method may split the actual session creation over several gRPC calls in
@@ -311,7 +345,13 @@ private List<SessionImpl> internalBatchCreateSessions(
       span.end();
       List<SessionImpl> res = new ArrayList<>(sessionCount);
       for (com.google.spanner.v1.Session session : sessions) {
-        res.add(new SessionImpl(spanner, session.getName(), options));
+        res.add(
+            new SessionImpl(
+                spanner,
+                session.getName(),
+                session.getCreateTime(),
+                session.getMultiplexed(),
+                options));
       }
       return res;
     } catch (RuntimeException e) {

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java
@@ -98,6 +98,8 @@ interface SessionTransaction {
   ByteString readyTransactionId;
   private final Map<SpannerRpc.Option, ?> options;
   private volatile Instant lastUseTime;
+  @Nullable private final Instant createTime;
+  private final boolean isMultiplexed;
   private ISpan currentSpan;
 
   SessionImpl(SpannerImpl spanner, String name, Map<SpannerRpc.Option, ?> options) {
@@ -107,6 +109,24 @@ interface SessionTransaction {
     this.name = checkNotNull(name);
     this.databaseId = SessionId.of(name).getDatabaseId();
     this.lastUseTime = Instant.now();
+    this.createTime = null;
+    this.isMultiplexed = false;
+  }
+
+  SessionImpl(
+      SpannerImpl spanner,
+      String name,
+      com.google.protobuf.Timestamp createTime,
+      boolean isMultiplexed,
+      Map<SpannerRpc.Option, ?> options) {
+    this.spanner = spanner;
+    this.tracer = spanner.getTracer();
+    this.options = options;
+    this.name = checkNotNull(name);
+    this.databaseId = SessionId.of(name).getDatabaseId();
+    this.lastUseTime = Instant.now();
+    this.createTime = convert(createTime);
+    this.isMultiplexed = isMultiplexed;
   }
 
   @Override
@@ -130,6 +150,14 @@ Instant getLastUseTime() {
     return lastUseTime;
   }
 
+  Instant getCreateTime() {
+    return createTime;
+  }
+
+  boolean getIsMultiplexed() {
+    return isMultiplexed;
+  }
+
   void markUsed(Instant instant) {
     lastUseTime = instant;
   }
@@ -455,4 +483,11 @@ boolean hasReadyTransaction() {
   TraceWrapper getTracer() {
     return tracer;
   }
+
+  private Instant convert(com.google.protobuf.Timestamp timestamp) {
+    if (timestamp == null) {
+      return null;
+    }
+    return Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos());
+  }
 }