chore: cleanup

Author: olavloite

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

@@ -22,6 +22,12 @@
 import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.spanner.v1.BatchWriteResponse;
 
+/**
+ * Base class for the Multiplexed Session {@link DatabaseClient} implementation. Throws {@link
+ * UnsupportedOperationException} for all methods that are currently not supported for multiplexed
+ * sessions. The concrete implementation implements the methods that are supported with multiplexed
+ * sessions.
+ */
 abstract class AbstractMultiplexedSessionDatabaseClient implements DatabaseClient {
 
   @Override

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

@@ -157,20 +157,13 @@ static class SingleReadContext extends AbstractReadContext {
     static class Builder extends AbstractReadContext.Builder<Builder, SingleReadContext> {
       private TimestampBound bound;
 
-      private boolean useRandomChannelHint;
-
       private Builder() {}
 
       Builder setTimestampBound(TimestampBound bound) {
         this.bound = bound;
         return self();
       }
 
-      Builder setUseRandomChannelHint(boolean useRandomChannelHint) {
-        this.useRandomChannelHint = useRandomChannelHint;
-        return self();
-      }
-
       @Override
       SingleReadContext build() {
         return new SingleReadContext(this);
@@ -195,14 +188,13 @@ static Builder newBuilder() {
     private SingleReadContext(Builder builder) {
       super(builder);
       this.bound = builder.bound;
-      // Use the channel hint that is stored with the regular session if that is being used.
-      // Otherwise, either use a random channel hint or no hint at all.
+      // single use transaction have a single RPC and hence there is no need
+      // of a channel hint. GAX will automatically choose a hint when used
+      // with a multiplexed session to perform a round-robin channel selection. We are
+      // passing a hint here to prefer random channel selection instead of doing GAX round-robin.
       this.channelHint =
           getChannelHintOptions(
-              session.getOptions(),
-              builder.useRandomChannelHint
-                  ? ThreadLocalRandom.current().nextLong(Long.MAX_VALUE)
-                  : null);
+              session.getOptions(), ThreadLocalRandom.current().nextLong(Long.MAX_VALUE));
     }
 
     @Override

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

@@ -39,12 +39,12 @@ class DatabaseClientImpl implements DatabaseClient {
 
   @VisibleForTesting
   DatabaseClientImpl(SessionPool pool, TraceWrapper tracer) {
-    this("", pool, null, tracer);
+    this("", pool, /* multiplexedSessionDatabaseClient = */ null, tracer);
   }
 
   @VisibleForTesting
   DatabaseClientImpl(String clientId, SessionPool pool, TraceWrapper tracer) {
-    this(clientId, pool, null, tracer);
+    this(clientId, pool, /* multiplexedSessionDatabaseClient = */ null, tracer);
   }
 
   DatabaseClientImpl(
@@ -286,7 +286,9 @@ private <T> T runWithSessionRetry(Function<Session, T> callable) {
 
   boolean isValid() {
     return pool.isValid()
-        && (multiplexedSessionDatabaseClient == null || multiplexedSessionDatabaseClient.isValid());
+        && (multiplexedSessionDatabaseClient == null
+            || multiplexedSessionDatabaseClient.isValid()
+            || !multiplexedSessionDatabaseClient.isMultiplexedSessionsSupported());
   }
 
   ListenableFuture<Void> closeAsync(ClosedException closedException) {

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

@@ -90,6 +90,7 @@ public boolean next() throws SpannerException {
       boolean hasNext = currRow.consumeRow(iterator);
       if (!hasNext) {
         statistics = iterator.getStats();
+        // Close the ResultSet when there is no more data.
         close();
       }
       return hasNext;
@@ -114,11 +115,14 @@ public ResultSetMetadata getMetadata() {
 
   @Override
   public void close() {
-    if (!closed) {
-      listener.onDone(iterator.isWithBeginTransaction());
-      iterator.close("ResultSet closed");
+    synchronized (this) {
+      if (closed) {
+        return;
+      }
       closed = true;
     }
+    listener.onDone(iterator.isWithBeginTransaction());
+    iterator.close("ResultSet closed");
   }
 
   @Override

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

@@ -109,31 +109,53 @@ public void close() {
     }
   }
 
+  /**
+   * Keeps track of which channels have been 'given' to single-use transactions for a given Spanner
+   * instance.
+   */
   private static final Map<SpannerImpl, BitSet> CHANNEL_USAGE = new HashMap<>();
 
   private final BitSet channelUsage;
 
   private final int numChannels;
 
+  /**
+   * The number of single-use read-only transactions currently running on this multiplexed session.
+   */
   private final AtomicInteger numCurrentSingleUseTransactions = new AtomicInteger();
 
   private boolean isClosed;
 
+  /** The duration before we try to replace the multiplexed session. The default is 7 days. */
   private final Duration sessionExpirationDuration;
 
   private final SessionClient sessionClient;
 
   private final TraceWrapper tracer;
 
+  /** The current multiplexed session that is used by this client. */
   private final AtomicReference<ApiFuture<SessionReference>> multiplexedSessionReference;
 
+  /** The expiration date/time of the current multiplexed session. */
   private final AtomicReference<Instant> expirationDate;
 
+  /**
+   * The maintainer runs every 10 minutes to check whether the multiplexed session should be
+   * refreshed.
+   */
   private final MultiplexedSessionMaintainer maintainer;
 
+  /**
+   * If a {@link DatabaseNotFoundException} or {@link InstanceNotFoundException} is returned by the
+   * server, then we set this field to mark the client as invalid.
+   */
   private final AtomicReference<ResourceNotFoundException> resourceNotFoundException =
       new AtomicReference<>();
 
+  /**
+   * This flag is set to true if the server return UNIMPLEMENTED when we try to create a multiplexed
+   * session.
+   */
   private final AtomicBoolean unimplemented = new AtomicBoolean(false);
 
   MultiplexedSessionDatabaseClient(SessionClient sessionClient) {
@@ -293,7 +315,11 @@ private int getSingleUseChannelHint() {
       return NO_CHANNEL_HINT;
     }
     synchronized (this.channelUsage) {
-      int channel = this.channelUsage.nextClearBit(0);
+      // Get the first unused channel.
+      int channel = this.channelUsage.nextClearBit(/* fromIndex = */ 0);
+      // BitSet returns an index larger than its original size if all the bits are set.
+      // 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) {
         return NO_CHANNEL_HINT;
       }

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

@@ -307,8 +307,6 @@ public ReadContext singleUse(TimestampBound bound) {
         SingleReadContext.newBuilder()
             .setSession(this)
             .setTimestampBound(bound)
-            .setUseRandomChannelHint(
-                spanner.getOptions().getSessionPoolOptions().isUseRandomChannelHint())
             .setRpc(spanner.getRpc())
             .setDefaultQueryOptions(spanner.getDefaultQueryOptions(getDatabaseId()))
             .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks())
@@ -332,8 +330,6 @@ public ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) {
         SingleUseReadOnlyTransaction.newBuilder()
             .setSession(this)
             .setTimestampBound(bound)
-            .setUseRandomChannelHint(
-                spanner.getOptions().getSessionPoolOptions().isUseRandomChannelHint())
             .setRpc(spanner.getRpc())
             .setDefaultQueryOptions(spanner.getDefaultQueryOptions(getDatabaseId()))
             .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks())

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

@@ -2947,7 +2947,7 @@ private SessionPool(
     this.waitOnMultiplexedSessionsLatch = new CountDownLatch(1);
   }
 
-  // TODO: Remove once all code for multiplexed sessions have been removed from the pool.
+  // TODO: Remove once all code for multiplexed sessions has been removed from the pool.
   private boolean useMultiplexedSessions() {
     // Multiplexed sessions have moved to MultiplexedSessionDatabaseClient
     return false;

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

@@ -3045,8 +3045,7 @@ public void testDatabaseOrInstanceDoesNotExistOnCreate() {
             DatabaseClientImpl dbClient =
                 (DatabaseClientImpl) spanner.getDatabaseClient(databaseId);
             // The CreateSession / BatchCreateSessions failure should propagate to the client and
-            // not
-            // retry.
+            // not retry.
             try (ResultSet rs = dbClient.singleUse().executeQuery(SELECT1)) {
               mockSpanner.unfreeze();
               assertThrows(ResourceNotFoundException.class, rs::next);