feat: add session timeout metric (#65)

* feat: add session timeout metric

* minor patch

* rename metric name and description

* fix potential flaky test

* fix code format

* update metric name

* minor nit

* fix review comment

* rename get_sessions_timeouts to get_session_timeouts

Author: mayurkale22

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

@@ -42,12 +42,16 @@ class MetricRegistryConstants {
   static final String COUNT = "1";
 
   // The Metric name and description
-  static final String MAX_IN_USE_SESSIONS = "cloud.google.com/java/spanner/max_in_use_session";
+  static final String MAX_IN_USE_SESSIONS = "cloud.google.com/java/spanner/max_in_use_sessions";
   static final String MAX_ALLOWED_SESSIONS = "cloud.google.com/java/spanner/max_allowed_sessions";
   static final String IN_USE_SESSIONS = "cloud.google.com/java/spanner/in_use_sessions";
+  static final String GET_SESSION_TIMEOUTS = "cloud.google.com/java/spanner/get_session_timeouts";
+
   static final String MAX_IN_USE_SESSIONS_DESCRIPTION =
       "The maximum number of sessions in use during the last 10 minute interval.";
   static final String MAX_ALLOWED_SESSIONS_DESCRIPTION =
       "The maximum number of sessions allowed. Configurable by the user.";
   static final String IN_USE_SESSIONS_DESCRIPTION = "The number of sessions currently in use.";
+  static final String SESSIONS_TIMEOUTS_DESCRIPTION =
+      "The number of get sessions timeouts due to pool exhaustion";
 }

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

@@ -17,12 +17,14 @@
 package com.google.cloud.spanner;
 
 import static com.google.cloud.spanner.MetricRegistryConstants.COUNT;
+import static com.google.cloud.spanner.MetricRegistryConstants.GET_SESSION_TIMEOUTS;
 import static com.google.cloud.spanner.MetricRegistryConstants.IN_USE_SESSIONS;
 import static com.google.cloud.spanner.MetricRegistryConstants.IN_USE_SESSIONS_DESCRIPTION;
 import static com.google.cloud.spanner.MetricRegistryConstants.MAX_ALLOWED_SESSIONS;
 import static com.google.cloud.spanner.MetricRegistryConstants.MAX_ALLOWED_SESSIONS_DESCRIPTION;
 import static com.google.cloud.spanner.MetricRegistryConstants.MAX_IN_USE_SESSIONS;
 import static com.google.cloud.spanner.MetricRegistryConstants.MAX_IN_USE_SESSIONS_DESCRIPTION;
+import static com.google.cloud.spanner.MetricRegistryConstants.SESSIONS_TIMEOUTS_DESCRIPTION;
 import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_DEFAULT_LABEL_VALUES;
 import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS;
 import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException;
@@ -50,6 +52,7 @@
 import com.google.protobuf.Empty;
 import io.opencensus.common.Scope;
 import io.opencensus.common.ToLongFunction;
+import io.opencensus.metrics.DerivedLongCumulative;
 import io.opencensus.metrics.DerivedLongGauge;
 import io.opencensus.metrics.LabelValue;
 import io.opencensus.metrics.MetricOptions;
@@ -1845,6 +1848,15 @@ private void initMetricsCollection(MetricRegistry metricRegistry, List<LabelValu
                 .setLabelKeys(SPANNER_LABEL_KEYS)
                 .build());
 
+    DerivedLongCumulative sessionsTimeouts =
+        metricRegistry.addDerivedLongCumulative(
+            GET_SESSION_TIMEOUTS,
+            MetricOptions.builder()
+                .setDescription(SESSIONS_TIMEOUTS_DESCRIPTION)
+                .setUnit(COUNT)
+                .setLabelKeys(SPANNER_LABEL_KEYS)
+                .build());
+
     // The value of a maxSessionsInUse is observed from a callback function. This function is
     // invoked whenever metrics are collected.
     maxInUseSessionsMetric.createTimeSeries(
@@ -1880,5 +1892,17 @@ public long applyAsLong(SessionPool sessionPool) {
             return sessionPool.numSessionsInUse;
           }
         });
+
+    // The value of a numWaiterTimeouts is observed from a callback function. This function is
+    // invoked whenever metrics are collected.
+    sessionsTimeouts.createTimeSeries(
+        labelValues,
+        this,
+        new ToLongFunction<SessionPool>() {
+          @Override
+          public long applyAsLong(SessionPool sessionPool) {
+            return sessionPool.getNumWaiterTimeouts();
+          }
+        });
   }
 }

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

@@ -96,6 +96,32 @@ public void removeTimeSeries(List<LabelValue> list) {}
     public void clear() {}
   }
 
+  public static final class FakeDerivedLongCumulative extends DerivedLongCumulative {
+    private final MetricsRecord record;
+    private final String name;
+    private final List<LabelKey> labelKeys;
+
+    private FakeDerivedLongCumulative(
+        FakeMetricRegistry metricRegistry, String name, List<LabelKey> labelKeys) {
+      this.record = metricRegistry.record;
+      this.labelKeys = labelKeys;
+      this.name = name;
+    }
+
+    @Override
+    public <T> void createTimeSeries(
+        List<LabelValue> labelValues, T t, ToLongFunction<T> toLongFunction) {
+      this.record.metrics.put(this.name, new PointWithFunction(t, toLongFunction));
+      this.record.labels.put(this.labelKeys, labelValues);
+    }
+
+    @Override
+    public void removeTimeSeries(List<LabelValue> list) {}
+
+    @Override
+    public void clear() {}
+  }
+
   /**
    * A {@link MetricRegistry} implementation that saves metrics records to be accessible from {@link
    * #pollRecord()}.
@@ -144,7 +170,7 @@ public DoubleCumulative addDoubleCumulative(String s, MetricOptions metricOption
 
     @Override
     public DerivedLongCumulative addDerivedLongCumulative(String s, MetricOptions metricOptions) {
-      throw new UnsupportedOperationException();
+      return new FakeDerivedLongCumulative(this, s, metricOptions.getLabelKeys());
     }
 
     @Override

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

@@ -1561,12 +1561,14 @@ public void run() {
   }
 
   @Test
-  public void testSessionMetrics() {
+  public void testSessionMetrics() throws Exception {
+    // Create a session pool with max 2 session and a low timeout for waiting for a session.
     options =
         SessionPoolOptions.newBuilder()
             .setMinSessions(1)
-            .setMaxSessions(3)
+            .setMaxSessions(2)
             .setMaxIdleSessions(0)
+            .setInitialWaitForSessionTimeoutMillis(20L)
             .build();
     FakeClock clock = new FakeClock();
     clock.currentTimeMillis = System.currentTimeMillis();
@@ -1583,16 +1585,46 @@ public void testSessionMetrics() {
     Session session2 = pool.getReadSession();
 
     MetricsRecord record = metricRegistry.pollRecord();
+    assertThat(record.getMetrics().size()).isEqualTo(4);
     assertThat(record.getMetrics()).containsEntry(MetricRegistryConstants.IN_USE_SESSIONS, 2L);
     assertThat(record.getMetrics()).containsEntry(MetricRegistryConstants.MAX_IN_USE_SESSIONS, 2L);
+    assertThat(record.getMetrics()).containsEntry(MetricRegistryConstants.GET_SESSION_TIMEOUTS, 0L);
     assertThat(record.getMetrics())
         .containsEntry(
             MetricRegistryConstants.MAX_ALLOWED_SESSIONS, (long) options.getMaxSessions());
     assertThat(record.getLabels()).containsEntry(SPANNER_LABEL_KEYS, labelValues);
 
+    final CountDownLatch latch = new CountDownLatch(1);
+    // Try asynchronously to take another session. This attempt should time out.
+    Future<Void> fut =
+        executor.submit(
+            new Callable<Void>() {
+              @Override
+              public Void call() {
+                latch.countDown();
+                Session session = pool.getReadSession();
+                session.close();
+                return null;
+              }
+            });
+    // Wait until the background thread is actually waiting for a session.
+    latch.await();
+    // Wait until the request has timed out.
+    int waitCount = 0;
+    while (pool.getNumWaiterTimeouts() == 0L && waitCount < 1000) {
+      Thread.sleep(5L);
+      waitCount++;
+    }
+    // Return the checked out session to the pool so the async request will get a session and
+    // finish.
     session2.close();
-    session1.close();
+    // Verify that the async request also succeeds.
+    fut.get(10L, TimeUnit.SECONDS);
+    executor.shutdown();
 
+    session1.close();
+    assertThat(record.getMetrics().get(MetricRegistryConstants.GET_SESSION_TIMEOUTS).longValue())
+        .isAtLeast(1L);
     assertThat(record.getMetrics()).containsEntry(MetricRegistryConstants.IN_USE_SESSIONS, 0L);
     assertThat(record.getMetrics()).containsEntry(MetricRegistryConstants.MAX_IN_USE_SESSIONS, 2L);
   }