fix: multiplexed session metrics were not included in refactor move (#3088)

The metrics for multiplexed sessions were not included in the
refactoring that moved multiplexed sessions out of the session pool.
This change re-adds those metrics based on the new client for
multiplexed sessions.

Author: olavloite

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

@@ -809,6 +809,7 @@ public final void invalidate() {
 
   @Override
   public void close() {
+    session.onTransactionDone();
     span.end();
     synchronized (lock) {
       isClosed = true;

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

@@ -130,7 +130,14 @@ public ResultSet analyzeQuery(Statement statement, QueryAnalyzeMode queryMode) {
   }
 
   @Override
-  public void close() {}
+  public void close() {
+    try {
+      this.readContextFuture.get().close();
+    } catch (Throwable ignore) {
+      // Ignore any errors during close, as this error has already propagated to the user through
+      // other means.
+    }
+  }
 
   /**
    * Represents a {@link ReadContext} using a multiplexed session that is not yet ready. The

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

@@ -39,6 +39,7 @@
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 
 /**
@@ -63,6 +64,8 @@ static class MultiplexedSessionTransaction extends SessionImpl {
 
     private final int singleUseChannelHint;
 
+    private boolean done;
+
     MultiplexedSessionTransaction(
         MultiplexedSessionDatabaseClient client,
         ISpan span,
@@ -73,6 +76,7 @@ static class MultiplexedSessionTransaction extends SessionImpl {
       this.client = client;
       this.singleUse = singleUse;
       this.singleUseChannelHint = singleUseChannelHint;
+      this.client.numSessionsAcquired.incrementAndGet();
       setCurrentSpan(span);
     }
 
@@ -103,6 +107,20 @@ void onReadDone() {
       }
     }
 
+    @Override
+    void onTransactionDone() {
+      boolean markedDone = false;
+      synchronized (this) {
+        if (!this.done) {
+          this.done = true;
+          markedDone = true;
+        }
+      }
+      if (markedDone) {
+        client.numSessionsReleased.incrementAndGet();
+      }
+    }
+
     @Override
     public void close() {
       // no-op, we don't want to delete the multiplexed session.
@@ -152,6 +170,10 @@ public void close() {
   private final AtomicReference<ResourceNotFoundException> resourceNotFoundException =
       new AtomicReference<>();
 
+  private final AtomicLong numSessionsAcquired = new AtomicLong();
+
+  private final AtomicLong numSessionsReleased = new AtomicLong();
+
   /**
    * This flag is set to true if the server return UNIMPLEMENTED when we try to create a multiplexed
    * session. TODO: Remove once this is guaranteed to be available.
@@ -239,6 +261,14 @@ boolean isValid() {
     return resourceNotFoundException.get() == null;
   }
 
+  AtomicLong getNumSessionsAcquired() {
+    return this.numSessionsAcquired;
+  }
+
+  AtomicLong getNumSessionsReleased() {
+    return this.numSessionsReleased;
+  }
+
   boolean isMultiplexedSessionsSupported() {
     return !this.unimplemented.get();
   }

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

@@ -484,6 +484,8 @@ void onError(SpannerException spannerException) {}
 
   void onReadDone() {}
 
+  void onTransactionDone() {}
+
   TraceWrapper getTracer() {
     return tracer;
   }

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

@@ -68,7 +68,6 @@
 import com.google.common.base.Function;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
-import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ForwardingListenableFuture;
 import com.google.common.util.concurrent.ForwardingListenableFuture.SimpleForwardingListenableFuture;
@@ -156,7 +155,8 @@ void maybeWaitOnMinSessions() {
     }
   }
 
-  private abstract static class CachedResultSetSupplier implements Supplier<ResultSet> {
+  private abstract static class CachedResultSetSupplier
+      implements com.google.common.base.Supplier<ResultSet> {
 
     private ResultSet cached;
 
@@ -2265,7 +2265,6 @@ public String getName() {
     @Override
     public void close() {
       synchronized (lock) {
-        numMultiplexedSessionsReleased++;
         if (lastException != null && isDatabaseOrInstanceNotFound(lastException)) {
           SessionPool.this.resourceNotFoundException =
               MoreObjects.firstNonNull(
@@ -2771,15 +2770,9 @@ enum Position {
   @GuardedBy("lock")
   private long numSessionsAcquired = 0;
 
-  @GuardedBy("lock")
-  private long numMultiplexedSessionsAcquired = 0;
-
   @GuardedBy("lock")
   private long numSessionsReleased = 0;
 
-  @GuardedBy("lock")
-  private long numMultiplexedSessionsReleased = 0;
-
   @GuardedBy("lock")
   private long numIdleSessionsRemoved = 0;
 
@@ -2830,7 +2823,9 @@ static SessionPool createPool(
       SessionClient sessionClient,
       TraceWrapper tracer,
       List<LabelValue> labelValues,
-      Attributes attributes) {
+      Attributes attributes,
+      AtomicLong numMultiplexedSessionsAcquired,
+      AtomicLong numMultiplexedSessionsReleased) {
     final SessionPoolOptions sessionPoolOptions = spannerOptions.getSessionPoolOptions();
 
     // A clock instance is passed in {@code SessionPoolOptions} in order to allow mocking via tests.
@@ -2846,7 +2841,9 @@ static SessionPool createPool(
         tracer,
         labelValues,
         spannerOptions.getOpenTelemetry(),
-        attributes);
+        attributes,
+        numMultiplexedSessionsAcquired,
+        numMultiplexedSessionsReleased);
   }
 
   static SessionPool createPool(
@@ -2884,7 +2881,9 @@ static SessionPool createPool(
         tracer,
         SPANNER_DEFAULT_LABEL_VALUES,
         openTelemetry,
-        null);
+        null,
+        new AtomicLong(),
+        new AtomicLong());
   }
 
   static SessionPool createPool(
@@ -2898,7 +2897,9 @@ static SessionPool createPool(
       TraceWrapper tracer,
       List<LabelValue> labelValues,
       OpenTelemetry openTelemetry,
-      Attributes attributes) {
+      Attributes attributes,
+      AtomicLong numMultiplexedSessionsAcquired,
+      AtomicLong numMultiplexedSessionsReleased) {
     SessionPool pool =
         new SessionPool(
             poolOptions,
@@ -2912,7 +2913,9 @@ static SessionPool createPool(
             tracer,
             labelValues,
             openTelemetry,
-            attributes);
+            attributes,
+            numMultiplexedSessionsAcquired,
+            numMultiplexedSessionsReleased);
     pool.initPool();
     return pool;
   }
@@ -2929,7 +2932,9 @@ private SessionPool(
       TraceWrapper tracer,
       List<LabelValue> labelValues,
       OpenTelemetry openTelemetry,
-      Attributes attributes) {
+      Attributes attributes,
+      AtomicLong numMultiplexedSessionsAcquired,
+      AtomicLong numMultiplexedSessionsReleased) {
     this.options = options;
     this.databaseRole = databaseRole;
     this.executorFactory = executorFactory;
@@ -2940,8 +2945,13 @@ private SessionPool(
     this.initialReleasePosition = initialReleasePosition;
     this.poolMaintainer = new PoolMaintainer();
     this.tracer = tracer;
-    this.initOpenCensusMetricsCollection(metricRegistry, labelValues);
-    this.initOpenTelemetryMetricsCollection(openTelemetry, attributes);
+    this.initOpenCensusMetricsCollection(
+        metricRegistry,
+        labelValues,
+        numMultiplexedSessionsAcquired,
+        numMultiplexedSessionsReleased);
+    this.initOpenTelemetryMetricsCollection(
+        openTelemetry, attributes, numMultiplexedSessionsAcquired, numMultiplexedSessionsReleased);
     this.waitOnMinSessionsLatch =
         options.getMinSessions() > 0 ? new CountDownLatch(1) : new CountDownLatch(0);
     this.waitOnMultiplexedSessionsLatch = new CountDownLatch(1);
@@ -3143,7 +3153,7 @@ boolean isValid() {
 
   /**
    * Returns a multiplexed session. The method fallbacks to a regular session if {@link
-   * SessionPoolOptions#useMultiplexedSession} is not set.
+   * SessionPoolOptions#getUseMultiplexedSession} is not set.
    */
   SessionFutureWrapper getMultiplexedSessionWithFallback() throws SpannerException {
     if (useMultiplexedSessions()) {
@@ -3250,8 +3260,6 @@ private void incrementNumSessionsInUse(boolean isMultiplexed) {
           maxSessionsInUse = numSessionsInUse;
         }
         numSessionsAcquired++;
-      } else {
-        numMultiplexedSessionsAcquired++;
       }
     }
   }
@@ -3775,7 +3783,10 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount
    * exporter, it allows users to monitor client behavior.
    */
   private void initOpenCensusMetricsCollection(
-      MetricRegistry metricRegistry, List<LabelValue> labelValues) {
+      MetricRegistry metricRegistry,
+      List<LabelValue> labelValues,
+      AtomicLong numMultiplexedSessionsAcquired,
+      AtomicLong numMultiplexedSessionsReleased) {
     if (!SpannerOptions.isEnabledOpenCensusMetrics()) {
       return;
     }
@@ -3860,18 +3871,14 @@ private void initOpenCensusMetricsCollection(
         labelValuesWithRegularSessions, this, sessionPool -> sessionPool.numSessionsAcquired);
     numAcquiredSessionsMetric.removeTimeSeries(labelValuesWithMultiplexedSessions);
     numAcquiredSessionsMetric.createTimeSeries(
-        labelValuesWithMultiplexedSessions,
-        this,
-        sessionPool -> sessionPool.numMultiplexedSessionsAcquired);
+        labelValuesWithMultiplexedSessions, this, unused -> numMultiplexedSessionsAcquired.get());
 
     numReleasedSessionsMetric.removeTimeSeries(labelValuesWithRegularSessions);
     numReleasedSessionsMetric.createTimeSeries(
         labelValuesWithRegularSessions, this, sessionPool -> sessionPool.numSessionsReleased);
     numReleasedSessionsMetric.removeTimeSeries(labelValuesWithMultiplexedSessions);
     numReleasedSessionsMetric.createTimeSeries(
-        labelValuesWithMultiplexedSessions,
-        this,
-        sessionPool -> sessionPool.numMultiplexedSessionsReleased);
+        labelValuesWithMultiplexedSessions, this, unused -> numMultiplexedSessionsReleased.get());
 
     List<LabelValue> labelValuesWithBeingPreparedType = new ArrayList<>(labelValues);
     labelValuesWithBeingPreparedType.add(NUM_SESSIONS_BEING_PREPARED);
@@ -3909,7 +3916,10 @@ private void initOpenCensusMetricsCollection(
    * an exporter, it allows users to monitor client behavior.
    */
   private void initOpenTelemetryMetricsCollection(
-      OpenTelemetry openTelemetry, Attributes attributes) {
+      OpenTelemetry openTelemetry,
+      Attributes attributes,
+      AtomicLong numMultiplexedSessionsAcquired,
+      AtomicLong numMultiplexedSessionsReleased) {
     if (openTelemetry == null || !SpannerOptions.isEnabledOpenTelemetryMetrics()) {
       return;
     }
@@ -3981,7 +3991,8 @@ private void initOpenTelemetryMetricsCollection(
         .buildWithCallback(
             measurement -> {
               measurement.record(this.numSessionsAcquired, attributesRegularSession);
-              measurement.record(this.numMultiplexedSessionsAcquired, attributesMultiplexedSession);
+              measurement.record(
+                  numMultiplexedSessionsAcquired.get(), attributesMultiplexedSession);
             });
 
     meter
@@ -3991,7 +4002,8 @@ private void initOpenTelemetryMetricsCollection(
         .buildWithCallback(
             measurement -> {
               measurement.record(this.numSessionsReleased, attributesRegularSession);
-              measurement.record(this.numMultiplexedSessionsReleased, attributesMultiplexedSession);
+              measurement.record(
+                  numMultiplexedSessionsReleased.get(), attributesMultiplexedSession);
             });
   }
 }

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

@@ -50,6 +50,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -271,17 +272,29 @@ public DatabaseClient getDatabaseClient(DatabaseId db) {
         attributesBuilder.put("database", db.getDatabase());
         attributesBuilder.put("instance_id", db.getInstanceId().getName());
 
+        boolean useMultiplexedSession =
+            getOptions().getSessionPoolOptions().getUseMultiplexedSession();
+        MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient =
+            useMultiplexedSession
+                ? new MultiplexedSessionDatabaseClient(SpannerImpl.this.getSessionClient(db))
+                : null;
+        AtomicLong numMultiplexedSessionsAcquired =
+            useMultiplexedSession
+                ? multiplexedSessionDatabaseClient.getNumSessionsAcquired()
+                : new AtomicLong();
+        AtomicLong numMultiplexedSessionsReleased =
+            useMultiplexedSession
+                ? multiplexedSessionDatabaseClient.getNumSessionsReleased()
+                : new AtomicLong();
         SessionPool pool =
             SessionPool.createPool(
                 getOptions(),
                 SpannerImpl.this.getSessionClient(db),
                 this.tracer,
                 labelValues,
-                attributesBuilder.build());
-        MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient =
-            getOptions().getSessionPoolOptions().getUseMultiplexedSession()
-                ? new MultiplexedSessionDatabaseClient(SpannerImpl.this.getSessionClient(db))
-                : null;
+                attributesBuilder.build(),
+                numMultiplexedSessionsAcquired,
+                numMultiplexedSessionsReleased);
         pool.maybeWaitOnMinSessions();
         DatabaseClientImpl dbClient =
             createDatabaseClient(clientId, pool, multiplexedSessionDatabaseClient);

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

@@ -3859,12 +3859,12 @@ public void testCreateSessionsFailure_shouldNotPropagateToCloseMethod() {
       // Simulate session creation failures on the backend.
       mockSpanner.setCreateSessionExecutionTime(
           SimulatedExecutionTime.ofStickyException(Status.RESOURCE_EXHAUSTED.asRuntimeException()));
-      DatabaseClient client =
-          spannerWithEmptySessionPool.getDatabaseClient(
-              DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
       // This will not cause any failure as getting a session from the pool is guaranteed to be
       // non-blocking, and any exceptions will be delayed until actual query execution.
       mockSpanner.freeze();
+      DatabaseClient client =
+          spannerWithEmptySessionPool.getDatabaseClient(
+              DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
       try (ResultSet rs = client.singleUse().executeQuery(SELECT1)) {
         mockSpanner.unfreeze();
         SpannerException e = assertThrows(SpannerException.class, rs::next);