[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.