[sakthivelmanii]

Problem

• In MultiplexedSessionDatabaseClient, the initial future was exceptionally completed with setException(t) before multiplexedSessionReference was updated to the retry future, letting waiters observe the stale error.

Solution

• Reorder the logic in onSessionCreateFailure to instantiate and install the retry future into multiplexedSessionReference before exceptionally completing the current attempt.
• Add testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated to prevent regressions.

[sakthivelmanii]

/gemini review

[gemini-code-assist]

Code Review

This pull request modifies the multiplexed session creation failure handling in MultiplexedSessionDatabaseClient to install the retry future before propagating the failure to existing waiters, preventing subsequent transactions from failing prematurely. A unit test has been added to verify this behavior. The reviewer recommends wrapping the asyncCreateMultiplexedSession call in a try-catch block to exceptionally complete the newly installed future if a synchronous exception is thrown, preventing potential deadlocks.

[gemini-code-assist]

Code Review

This pull request ensures that when a multiplexed session creation fails, the future for the retry attempt is installed before the failure exception is propagated to existing waiters. This prevents new transactions initiated immediately after the failure from failing with the previous retryable error. A corresponding unit test has been added to verify this behavior. The reviewer suggested a refactoring to simplify the conditional logic and avoid duplicating the exception-setting call.

[sakthivelmanii]

/gemini review

[gemini-code-assist]

Code Review

This pull request modifies the session creation failure handling in MultiplexedSessionDatabaseClient to install the retry future before propagating the exception, ensuring subsequent transactions wait for the retry instead of failing immediately. A corresponding unit test has also been added. The review feedback recommends using compareAndSet instead of set when updating the multiplexed session reference to prevent potential race conditions, and wrapping the test client in a try-finally block to guarantee proper resource cleanup in case of test failures.

[sakthivelmanii]

/gemini review

[gemini-code-assist]

Code Review

This pull request updates MultiplexedSessionDatabaseClient to install a retry future for multiplexed session creation before propagating the failure of the current attempt, preventing subsequent transactions from failing with the previous error. A corresponding unit test has been added to verify this behavior. The review feedback points out a critical issue where a synchronous exception thrown during the retry initiation could leave the new future uncompleted indefinitely, and suggests wrapping the call in a try-catch block to safely propagate the exception.

[sakthivelmanii]

/gemini review

[gemini-code-assist]

Code Review

This pull request updates the multiplexed session creation retry logic in MultiplexedSessionDatabaseClient to install the retry future before propagating the failure of the current attempt, preventing subsequent transactions from failing with the retryable error. A new unit test was added to verify this behavior. The review feedback suggests wrapping the asynchronous session creation call in a try-catch block to handle potential synchronous exceptions, ensuring the newly installed future is completed exceptionally rather than hanging indefinitely.

[sakthivelmanii]

/gemini review

[gemini-code-assist]

Code Review

This pull request modifies MultiplexedSessionDatabaseClient to ensure that when a multiplexed session creation fails, the retry future is installed before the failure is propagated to waiters. This prevents subsequent transactions from failing with the transient error of the previous attempt, forcing them to wait for the retry instead. A corresponding unit test has been added to verify this behavior. There are no review comments to address, and I have no additional feedback to provide.

[olavloite]

Accepting that design premise makes sense from an operational perspective: fail-fast for in-flight waiters (giving them immediate feedback rather than stalling them across 10 consecutive attempts) while transparently establishing a retry for future requests.

If this is the intended contract, here is how to make the implementation 100% thread-safe, leak-free, and resilient against race conditions.

---

1. Requirements for a 100% Safe Implementation

To make this pattern safe, four critical concurrency guarantees must be met:

1. Strict Happens-Before Installation: The retry future (retryFuture) must be installed into multiplexedSessionReference before sessionReferenceFuture.setException(t) is invoked. This guarantees that any unblocking waiter that retries immediately will observe the new retry future, not the failed one.
2. Closing State Guard (No Leaked Uncompleted Futures & No Zombie RPCs):
   If client.close() is called concurrently right after CAS installs retryFuture:
   • An RPC must not be dispatched on a closed client.
   • retryFuture must not be left uncompleted (otherwise any transaction that grabbed retryFuture between CAS and close() will hang forever). It must be completed with newClosedException().
3. Proactive Failure on client.close():
   When close() is called, if the current session future in multiplexedSessionReference is still pending, close() should complete it with newClosedException() so existing waiters fail immediately rather than waiting for in-flight network timeouts.
4. No Reflection in Unit Tests: Expose a @VisibleForTesting package-private accessor instead of breaking encapsulation via Field.setAccessible(true).

---

2. The Recommended Implementation

Here is the robust, production-grade implementation for MultiplexedSessionDatabaseClient.java:

A. Updated onSessionCreateFailure in MultiplexedSessionDatabaseClient.java

          @Override
          public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) {
            SpannerException spannerException = SpannerExceptionFactory.asSpannerException(t);
            if (MultiplexedSessionDatabaseClient.this.resourceNotFoundException.get() == null
                && (spannerException instanceof DatabaseNotFoundException
                    || spannerException instanceof InstanceNotFoundException
                    || spannerException instanceof SessionNotFoundException)) {
              MultiplexedSessionDatabaseClient.this.resourceNotFoundException.set(
                  (ResourceNotFoundException) spannerException);
            }

            SettableApiFuture<SessionReference> retryFuture = null;
            if (remainingAttempts > 1
                && !isClientClosed()
                && RETRYABLE_ERROR_CODES.contains(spannerException.getErrorCode())) {
              final SettableApiFuture<SessionReference> nextFuture = SettableApiFuture.create();
              if (MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.compareAndSet(
                  sessionReferenceFuture, nextFuture)) {
                retryFuture = nextFuture;
              }
            }

            // 1. Unblock all waiters of the failed attempt first.
            // This is executed OUTSIDE any lock to avoid deadlocks in foreign listener callbacks.
            // Because the retryFuture was installed into multiplexedSessionReference above,
            // any waking waiter that retries immediately is guaranteed to observe the retry future.
            sessionReferenceFuture.setException(t);

            // 2. If a retry was scheduled, initiate it:
            if (retryFuture != null) {
              // Guard against concurrent client close that happened while waking waiters.
              if (isClientClosed()) {
                retryFuture.setException(newClosedException());
                return;
              }
              try {
                asyncCreateMultiplexedSession(retryFuture, remainingAttempts - 1);
              } catch (Throwable e) {
                // Failsafe: ensure waiting threads never hang if RPC scheduling throws synchronously.
                retryFuture.setException(e);
              }
            }
          }

B. Updated close() in MultiplexedSessionDatabaseClient.java

Ensure close() unblocks any currently waiting transaction:

  void close() {
    boolean releaseChannelUsage = false;
    ApiFuture<SessionReference> currentSessionFuture = null;
    synchronized (this) {
      if (!this.isClosed) {
        this.isClosed = true;
        this.maintainer.stop();
        releaseChannelUsage = true;
        currentSessionFuture = this.multiplexedSessionReference.get();
      }
    }
    if (releaseChannelUsage) {
      // If a session creation or retry is currently in flight, fail it immediately so
      // transactions waiting on the closed client do not hang.
      if (currentSessionFuture instanceof SettableApiFuture && !currentSessionFuture.isDone()) {
        ((SettableApiFuture<SessionReference>) currentSessionFuture)
            .setException(newClosedException());
      }
      spanner.getRpc().unregisterChannelPrimeSessionSource(this);
      synchronized (CHANNEL_USAGE) {
        SharedChannelUsage sharedChannelUsage = CHANNEL_USAGE.get(this.spanner);
        if (sharedChannelUsage != null) {
          sharedChannelUsage.referenceCount--;
          if (sharedChannelUsage.referenceCount == 0) {
            CHANNEL_USAGE.remove(this.spanner);
          }
        }
      }
    }
  }

C. Clean @VisibleForTesting Accessor (Removing Reflection)

Add this package-private method to MultiplexedSessionDatabaseClient.java:

  @VisibleForTesting
  ApiFuture<SessionReference> getMultiplexedSessionFuture() {
    return this.multiplexedSessionReference.get();
  }

---

3. Updated Unit Test (No Reflection)

In MultiplexedSessionDatabaseClientTest.java:

  @Test
  public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() throws Exception {
    Clock clock = mock(Clock.class);
    when(clock.instant()).thenReturn(Instant.now());
    SessionClient sessionClient = mock(SessionClient.class);
    SpannerImpl spanner = mock(SpannerImpl.class);
    SpannerOptions spannerOptions = mock(SpannerOptions.class);
    SessionPoolOptions sessionPoolOptions = mock(SessionPoolOptions.class);
    when(sessionClient.getSpanner()).thenReturn(spanner);
    when(sessionClient.getDatabaseId()).thenReturn(TEST_DATABASE_ID);
    when(spanner.getRpc()).thenReturn(mock(SpannerRpc.class));
    when(spanner.getOptions()).thenReturn(spannerOptions);
    when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions);
    when(sessionPoolOptions.getMultiplexedSessionMaintenanceDuration())
        .thenReturn(Duration.ofDays(7));
    when(sessionPoolOptions.getMultiplexedSessionMaintenanceLoopFrequency())
        .thenReturn(Duration.ofMinutes(10));
    SessionImpl session = mock(SessionImpl.class);
    SessionReference sessionReference = mock(SessionReference.class);
    when(session.getSessionReference()).thenReturn(sessionReference);

    AtomicReference<SessionConsumer> initialConsumer = new AtomicReference<>();
    AtomicReference<SessionConsumer> retryConsumer = new AtomicReference<>();
    doAnswer(
            (Answer<?>)
                invocationOnMock -> {
                  initialConsumer.set(invocationOnMock.getArgument(0));
                  return null;
                })
        .doAnswer(
            (Answer<?>)
                invocationOnMock -> {
                  retryConsumer.set(invocationOnMock.getArgument(0));
                  return null;
                })
        .when(sessionClient)
        .asyncCreateMultiplexedSession(any(SessionConsumer.class));

    MultiplexedSessionDatabaseClient client =
        new MultiplexedSessionDatabaseClient(sessionClient, clock);
    try {
      assertNotNull(initialConsumer.get());

      ApiFuture<SessionReference> initialFuture = client.getMultiplexedSessionFuture();
      AtomicReference<ApiFuture<SessionReference>> futureSeenByNewTransaction =
          new AtomicReference<>();
      initialFuture.addListener(
          () -> futureSeenByNewTransaction.set(client.getMultiplexedSessionFuture()),
          MoreExecutors.directExecutor());

      initialConsumer
          .get()
          .onSessionCreateFailure(
              SpannerExceptionFactory.newSpannerException(ErrorCode.DEADLINE_EXCEEDED, "test"), 1);

      // 1. Waiters for the initial attempt receive the error of that attempt.
      SpannerException exception =
          assertThrows(SpannerException.class, () -> SpannerApiFutures.get(initialFuture));
      assertEquals(ErrorCode.DEADLINE_EXCEEDED, exception.getErrorCode());

      // 2. New transactions arriving after the failure observe the retry future.
      assertNotNull(futureSeenByNewTransaction.get());
      assertFalse(futureSeenByNewTransaction.get().isDone());
      assertNotNull(retryConsumer.get());

      // 3. Completing the retry succeeds the pending future.
      retryConsumer.get().onSessionReady(session);
      assertEquals(sessionReference, futureSeenByNewTransaction.get().get());
      assertEquals(sessionReference, client.getCurrentSessionReference());
    } finally {
      client.close();
    }
  }

---

4. Summary of Fixes Over PR #14866

Hazard in PR #14866	Resolution in this Design
Race with close() after CAS	Before dispatching asyncCreateMultiplexedSession, checks if (isClientClosed()) and fails retryFuture with newClosedException() so waiters never hang and no RPC is initiated.
Hanging waiters on close()	In close(), proactively marks any in-flight uncompleted future as failed with newClosedException().
Unhandled CAS failure	Decouples CAS into if (compareAndSet(...)) { retryFuture = nextFuture; }, ensuring sessionReferenceFuture.setException(t) runs unconditionally while retryFuture is only scheduled if CAS succeeded.
Reflection in test code	Replaced getDeclaredField("multiplexedSessionReference") with a package-private @VisibleForTesting getter.

[sakthivelmanii]

/gemini review

[gemini-code-assist]

Code Review

This pull request improves the resilience of MultiplexedSessionDatabaseClient by ensuring that when a multiplexed session creation fails, any scheduled retry is registered before the failure is propagated to waiting threads. This prevents waking waiters from observing the failed session attempt. Additionally, the close() method is updated to immediately fail any in-flight session creation or retry to prevent waiting transactions from hanging. A unit test is added to verify this behavior. The review feedback suggests a minor improvement to guard against a race condition during concurrent client closure by checking if the retry future is already done before initiating the asynchronous session creation.

[olavloite]

1. assumeTrue(isJava8()) Skips the New Test in CI

In MultiplexedSessionDatabaseClientTest.java:253:

  @Test
  public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() throws Exception {
    assumeTrue(isJava8()); // <--- DELETE THIS LINE

• When this test originally used reflection on private fields, assumeTrue(isJava8()) was copied from environment-variable tests that fail on Java 9+.
• Now that we have the clean @VisibleForTesting ApiFuture<SessionReference> getMultiplexedSessionFuture() accessor, no reflection is used.
• Because CI and modern builds run on Java 11, 17, and 21, assumeTrue(isJava8()) causes Surefire to report:

  [WARNING] Tests run: 1, Failures: 0, Errors: 0, Skipped: 1
  

  The test is silently skipped and never executed in standard CI or local development!
• Fix: Remove assumeTrue(isJava8());. The test passes cleanly on Java 11/17/21 in under 1 second.

---

2. Race Condition on Concurrent close() (Zombie RPC Dispatch)

Regarding the discussion on whether retryFuture.isDone() needs to be checked alongside isClientClosed():

private boolean isClientClosed() {
  return isClosed;
}

1. isClientClosed() is not synchronized; it reads the volatile boolean isClosed.
2. Even if isClientClosed() were synchronized, calling a synchronized boolean getter releases the monitor immediately upon returning.
3. The Race Window:
   • Thread A (onSessionCreateFailure): passes if (isClientClosed()) (returns false).
   • Thread B (close()): acquires synchronized (this), sets this.isClosed = true, reads currentSessionFuture = retryFuture, marks retryFuture as failed with newClosedException(), and unregisters channel sources.
   • Thread A: resumes and calls asyncCreateMultiplexedSession(retryFuture, remainingAttempts - 1).
   • Result: An RPC is dispatched to Cloud Spanner on an already-closed client with an already-completed future. When the session arrives later, onSessionReady drops it, leaking the network call and server-side session.

---

3. Required Fixes in MultiplexedSessionDatabaseClient.java

A. Guard onSessionCreateFailure against already-completed futures

In onSessionCreateFailure:

            // 2. If a retry was scheduled, initiate it:
            if (retryFuture != null) {
              // Guard against concurrent client close that happened while waking waiters.
              if (isClientClosed() || retryFuture.isDone()) {
                retryFuture.setException(newClosedException());
                return;
              }
              try {
                asyncCreateMultiplexedSession(retryFuture, remainingAttempts - 1);
              } catch (Throwable e) {
                // Failsafe: ensure waiting threads never hang if RPC scheduling throws synchronously.
                retryFuture.setException(e);
              }
            }

B. Guard asyncCreateMultiplexedSession against dispatching on a closed client

In asyncCreateMultiplexedSession, add an upfront guard:

  private void asyncCreateMultiplexedSession(
      SettableApiFuture<SessionReference> sessionReferenceFuture, int remainingAttempts) {
    if (isClosed || sessionReferenceFuture.isDone()) {
      sessionReferenceFuture.setException(newClosedException());
      return;
    }
    this.sessionClient.asyncCreateMultiplexedSession(
        new SessionConsumer() { ... });

---

4. Real Tests to Add in MultiplexedSessionDatabaseClientMockServerTest.java

These are real integration tests exercising the client's public API (DatabaseClient.singleUse().executeQuery(...)) against the in-process MockSpanner gRPC server without arbitrary Thread.sleep calls.

Real Test 1: Verify close() unblocks transactions waiting for session creation

Verifies that when testSpanner.close() is called while a query is blocked waiting for session creation, the query does not hang on the frozen mock server, but unblocks immediately with FAILED_PRECONDITION: This client has been closed:

  @Test
  public void testCloseUnblocksWaitingTransaction() throws Exception {
    mockSpanner.freeze();

    Spanner testSpanner =
        SpannerOptions.newBuilder()
            .setProjectId("test-project")
            .setChannelProvider(channelProvider)
            .setCredentials(NoCredentials.getInstance())
            .build()
            .getService();
    DatabaseClient client = testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));

    // Wait until the initial CreateSession request is in flight and frozen on the server.
    mockSpanner.waitForRequestsToContain(CreateSessionRequest.class, 5000);

    // Start a query that blocks waiting for the session to be created.
    try (ReadContext readContext = client.singleUse();
        ResultSet resultSet = readContext.executeQuery(STATEMENT)) {
      // Close the client while the transaction is waiting for session creation.
      testSpanner.close();

      // The transaction must immediately fail with FAILED_PRECONDITION (client closed),
      // rather than hanging indefinitely waiting for mockSpanner to unfreeze.
      SpannerException exception = assertThrows(SpannerException.class, resultSet::next);
      assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode());
      assertTrue(exception.getMessage().contains("This client has been closed"));
    } finally {
      mockSpanner.unfreeze();
    }
  }

Real Test 2: Verify close() unblocks transactions waiting on a retry attempt

Verifies that if attempt 1 fails and a new transaction binds to the retry attempt (attempt 2), closing the client unblocks that transaction immediately without waiting for the server to unfreeze:

  @Test
  public void testCloseUnblocksTransactionWaitingOnRetry() throws Exception {
    // Attempt 1 fails with DEADLINE_EXCEEDED.
    mockSpanner.setCreateSessionExecutionTime(
        SimulatedExecutionTime.ofException(Status.DEADLINE_EXCEEDED.asRuntimeException()));
    mockSpanner.freeze();

    Spanner testSpanner =
        SpannerOptions.newBuilder()
            .setProjectId("test-project")
            .setChannelProvider(channelProvider)
            .setCredentials(NoCredentials.getInstance())
            .build()
            .getService();
    DatabaseClient client = testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));

    // Wait until the initial attempt is frozen.
    mockSpanner.waitForRequestsToContain(CreateSessionRequest.class, 5000);

    // Fail the initial attempt.
    try (ReadContext readContext = client.singleUse();
        ResultSet resultSet = readContext.executeQuery(STATEMENT)) {
      mockSpanner.unfreeze();
      SpannerException exception = assertThrows(SpannerException.class, resultSet::next);
      assertEquals(ErrorCode.DEADLINE_EXCEEDED, exception.getErrorCode());
    }

    // Freeze the mock server before the retry attempt completes.
    mockSpanner.freeze();

    // Wait until the second CreateSession request has reached mock server and is frozen.
    while (mockSpanner.countRequestsOfType(CreateSessionRequest.class) < 2) {
      Thread.yield();
    }

    // Start a new query that binds to the pending retry future (attempt 2).
    try (ReadContext readContext = client.singleUse();
        ResultSet resultSet = readContext.executeQuery(STATEMENT)) {
      // Close the client while query is waiting for attempt 2.
      testSpanner.close();

      // The transaction must immediately fail with FAILED_PRECONDITION (client closed),
      // instead of hanging while attempt 2 is frozen.
      SpannerException exception = assertThrows(SpannerException.class, resultSet::next);
      assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode());
      assertTrue(exception.getMessage().contains("This client has been closed"));
    } finally {
      mockSpanner.unfreeze();
    }
  }

[sakthivelmanii]

Addressed all feedback: added the retryFuture.isDone() / sessionReferenceFuture.isDone() / isClosed guards in onSessionCreateFailure and asyncCreateMultiplexedSession, and added testCloseUnblocksWaitingTransaction and testCloseUnblocksTransactionWaitingOnRetry in MultiplexedSessionDatabaseClientMockServerTest. Also switched assumeTrue(isJava8()) to assumeFalse(isNativeImage()) so testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated runs on Java 8/11/17/21+ JVMs while skipping Mockito dynamic mock creation in GraalVM native-image tests.

[sakthivelmanii]

/gemini review

[gemini-code-assist]

Code Review

This pull request improves the resilience and thread safety of multiplexed session creation and client closure in MultiplexedSessionDatabaseClient. Specifically, it ensures that closing the client immediately fails and unblocks any waiting or retrying transactions to prevent them from hanging indefinitely. It also refactors retry handling to install the retry future before propagating failures. Review feedback highlights a potential thread-safety issue where the raw isClosed field is accessed directly in asyncCreateMultiplexedSession instead of using the synchronized isClientClosed() helper method.