fix(spanner): release single-use channel hints for mutation-only transactions (P1-2, P1-3)

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

@@ -89,6 +89,8 @@ static class MultiplexedSessionTransaction extends SessionImpl {
 
     private boolean done;
 
+    private boolean channelHintReleased;
+
     MultiplexedSessionTransaction(
         MultiplexedSessionDatabaseClient client,
         ISpan span,
@@ -123,19 +125,18 @@ void onReadDone() {
       if (this.singleUse && getActiveTransaction() != null) {
         getActiveTransaction().close();
         setActive(null);
-        if (this.singleUseChannelHint != NO_CHANNEL_HINT) {
-          this.client.channelUsage.clear(this.singleUseChannelHint);
-        }
-        this.client.numCurrentSingleUseTransactions.decrementAndGet();
+        releaseSingleUseChannelHint();
       }
     }
 
     @Override
     public CommitResponse writeAtLeastOnceWithOptions(
         Iterable<Mutation> mutations, TransactionOption... options) throws SpannerException {
-      CommitResponse response = super.writeAtLeastOnceWithOptions(mutations, options);
-      onTransactionDone();
-      return response;
+      try {
+        return super.writeAtLeastOnceWithOptions(mutations, options);
+      } finally {
+        onTransactionDone();
+      }
     }
 
     @Override
@@ -150,6 +151,30 @@ void onTransactionDone() {
       if (markedDone) {
         client.numSessionsReleased.incrementAndGet();
       }
+      if (this.singleUse) {
+        // Mutation-only operations (writeAtLeastOnceWithOptions, batchWriteAtLeastOnce) never
+        // return a ResultSet and thus never call onReadDone(), so the channel hint must also be
+        // released here.
+        releaseSingleUseChannelHint();
+      }
+    }
+
+    /**
+     * Releases the channel hint that was reserved for this single-use transaction, so a subsequent
+     * single-use transaction can use the channel. This method is idempotent; the hint is released
+     * by whichever of {@link #onReadDone()} and {@link #onTransactionDone()} is called first.
+     */
+    private void releaseSingleUseChannelHint() {
+      boolean release = false;
+      synchronized (this) {
+        if (!this.channelHintReleased) {
+          this.channelHintReleased = true;
+          release = true;
+        }
+      }
+      if (release) {
+        this.client.releaseSingleUseChannelHint(this.singleUseChannelHint);
+      }
     }
 
     @Override
@@ -456,7 +481,11 @@ private DelayedMultiplexedSessionTransaction createDelayedMultiplexSessionTransa
   }
 
   private int getSingleUseChannelHint() {
+    // The counter tracks the number of single-use transactions that currently hold a channel hint,
+    // so it is only incremented when a hint is actually reserved, and any transient increment is
+    // rolled back before returning NO_CHANNEL_HINT.
     if (this.numCurrentSingleUseTransactions.incrementAndGet() > this.numChannels) {
+      this.numCurrentSingleUseTransactions.decrementAndGet();
       return NO_CHANNEL_HINT;
     }
     synchronized (this.channelUsage) {
@@ -466,13 +495,26 @@ private int getSingleUseChannelHint() {
       // This then means that all channels have already been assigned to single-use transactions,
       // and that we should not use a specific channel, but rather pick a random one.
       if (channel == this.numChannels) {
+        this.numCurrentSingleUseTransactions.decrementAndGet();
         return NO_CHANNEL_HINT;
       }
       this.channelUsage.set(channel);
       return channel;
     }
   }
 
+  private void releaseSingleUseChannelHint(int channelHint) {
+    if (channelHint == NO_CHANNEL_HINT) {
+      return;
+    }
+    // Take the same monitor as getSingleUseChannelHint(); the BitSet is not thread-safe and is
+    // shared by all clients of the same SpannerImpl.
+    synchronized (this.channelUsage) {
+      this.channelUsage.clear(channelHint);
+    }
+    this.numCurrentSingleUseTransactions.decrementAndGet();
+  }
+
   private final AbstractLazyInitializer<Dialect> dialectSupplier =
       new AbstractLazyInitializer<Dialect>() {
         @Override

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

@@ -535,6 +535,126 @@ public void testWriteAtLeastOnce() {
     assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
   }
 
+  @Test
+  public void testWriteAtLeastOnceReleasesSingleUseChannelHint() throws Exception {
+    // Channel hints are only used when the gRPC-GCP extension is disabled.
+    try (Spanner testSpanner = createSpannerWithoutGrpcGcp()) {
+      DatabaseClientImpl client =
+          (DatabaseClientImpl) testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+      int numChannels = testSpanner.getOptions().getNumChannels();
+
+      // Execute more mutation-only transactions than there are channels. Each transaction should
+      // release its channel hint when it is done, so the channel-spreading optimization keeps
+      // working for subsequent single-use transactions.
+      for (int i = 0; i < numChannels + 1; i++) {
+        assertNotNull(MockSpannerTestActions.writeAtLeastOnceInsertMutation(client));
+      }
+      assertEquals(0, getChannelUsage(client).cardinality());
+      assertEquals(0, getNumCurrentSingleUseTransactions(client).get());
+
+      // A subsequent single-use read-only transaction should still get a channel hint. The hint is
+      // held from the creation of the transaction until the result set has been consumed or closed.
+      try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) {
+        assertEquals(1, getChannelUsage(client).cardinality());
+        assertEquals(1, getNumCurrentSingleUseTransactions(client).get());
+        //noinspection StatementWithEmptyBody
+        while (resultSet.next()) {}
+      }
+      assertEquals(0, getChannelUsage(client).cardinality());
+      assertEquals(0, getNumCurrentSingleUseTransactions(client).get());
+    }
+  }
+
+  @Test
+  public void testFailedWriteAtLeastOnceReleasesSingleUseChannelHint() throws Exception {
+    try (Spanner testSpanner = createSpannerWithoutGrpcGcp()) {
+      DatabaseClientImpl client =
+          (DatabaseClientImpl) testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+      // Make sure that the multiplexed session has been created, so that the write transaction
+      // below does not use the delayed path (which never reserves a channel hint).
+      awaitMultiplexedSessionCreation(client);
+      mockSpanner.setCommitExecutionTime(
+          SimulatedExecutionTime.ofException(
+              Status.FAILED_PRECONDITION.withDescription("test").asRuntimeException()));
+
+      SpannerException exception =
+          assertThrows(
+              SpannerException.class,
+              () -> MockSpannerTestActions.writeAtLeastOnceInsertMutation(client));
+
+      assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode());
+      assertEquals(0, getChannelUsage(client).cardinality());
+      assertEquals(0, getNumCurrentSingleUseTransactions(client).get());
+    }
+  }
+
+  @Test
+  public void testBatchWriteAtLeastOnceReleasesSingleUseChannelHint() throws Exception {
+    try (Spanner testSpanner = createSpannerWithoutGrpcGcp()) {
+      DatabaseClientImpl client =
+          (DatabaseClientImpl) testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+      // Make sure that the multiplexed session has been created, so that the batch write below
+      // does not use the delayed path (which never reserves a channel hint).
+      awaitMultiplexedSessionCreation(client);
+
+      ServerStream<BatchWriteResponse> responseStream =
+          client.batchWriteAtLeastOnce(
+              ImmutableList.of(
+                  MutationGroup.of(
+                      Mutation.newInsertBuilder("FOO")
+                          .set("ID")
+                          .to(1L)
+                          .set("NAME")
+                          .to("Bar")
+                          .build())));
+      for (BatchWriteResponse response : responseStream) {
+        assertNotNull(response);
+      }
+
+      assertEquals(0, getChannelUsage(client).cardinality());
+      assertEquals(0, getNumCurrentSingleUseTransactions(client).get());
+    }
+  }
+
+  private void awaitMultiplexedSessionCreation(DatabaseClientImpl client) {
+    try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) {
+      //noinspection StatementWithEmptyBody
+      while (resultSet.next()) {}
+    }
+  }
+
+  private Spanner createSpannerWithoutGrpcGcp() {
+    return SpannerOptions.newBuilder()
+        .setProjectId("test-project")
+        .setChannelProvider(channelProvider)
+        .setCredentials(NoCredentials.getInstance())
+        .disableGrpcGcpExtension()
+        .setSessionPoolOption(
+            SessionPoolOptions.newBuilder()
+                .setUseMultiplexedSession(true)
+                .setUseMultiplexedSessionForRW(true)
+                .setUseMultiplexedSessionPartitionedOps(true)
+                .setFailOnSessionLeak()
+                .build())
+        .build()
+        .getService();
+  }
+
+  private static BitSet getChannelUsage(DatabaseClientImpl client) throws Exception {
+    java.lang.reflect.Field field =
+        MultiplexedSessionDatabaseClient.class.getDeclaredField("channelUsage");
+    field.setAccessible(true);
+    return (BitSet) field.get(client.multiplexedSessionDatabaseClient);
+  }
+
+  private static AtomicInteger getNumCurrentSingleUseTransactions(DatabaseClientImpl client)
+      throws Exception {
+    java.lang.reflect.Field field =
+        MultiplexedSessionDatabaseClient.class.getDeclaredField("numCurrentSingleUseTransactions");
+    field.setAccessible(true);
+    return (AtomicInteger) field.get(client.multiplexedSessionDatabaseClient);
+  }
+
   @Test
   public void testWriteAtLeastOnceWithCommitStats() {
     DatabaseClientImpl client =