fix: Retry multiplexed session failures

[sakthivelmanii]

Description:

Currently when multiplexed session fails with any error, we are storing the exception in the session reference and re-throwing that error to all the subsequent requests. This will cause the library to stall since no further requests will be processed successfully. It's a general expectation that all RPC requests are expected due to CPU, Network, GFE and other factors.

Proposed solution:
If session creation failed with an any error, it should retry the error. So we will be checking the following before initiating the retry.

• Retry already in progress - This is make sure only one RPC call is sent at a time
• If it fails with UNIMPLEMENTED - We shouldn't retry UNIMPLEMENTED error. SpanFE will throw UNIMPLEMENTED error only if the multiplexed session is not enabled.

Old Behaviour:

1. Multiplexed session failed during initial creation
2. RPC1 - Replay the same error
3. RPC2 - Replay the same error
4. RPC3 - Replay the same error

Updated behaviour:

State - 1

1. Multiplexed session failed during initial creation
2. RPC1 - Will retry the RPC.
3. RPC2 - Wait for RPC
4. RPC3 - Wait for RPC

State - 2

1. Multiplexed session failed during initial creation
2. RPC1 - Will retry the RPC -> RPC Failed
3. RPC2 - Wait for RPC -> RPC Failed
4. RPC3 - Wait for RPC -> RPC Failed

State - 3

1. Multiplexed session failed during initial creation
2. RPC1 - Will retry the RPC -> RPC Failed
3. RPC2 - Wait for RPC -> RPC Failed
4. RPC3 - Wait for RPC -> RPC Failed
5. RPC4 - Will retry the RPC
6. RPC5 - Wait for RPC
7. RPC6 - Wait for RPC

[gemini-code-assist]

Summary of Changes

Hello @sakthivelmanii, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a crucial retry mechanism for multiplexed session creation failures in the Spanner client library. Previously, a single failure during session initialization would cause the library to stall, re-throwing the same error for all subsequent requests. The new logic ensures that session creation is retried upon encountering transient errors (excluding UNIMPLEMENTED errors), significantly improving the client's resilience and preventing service disruption.

Highlights

• Enhanced Multiplexed Session Resilience: Implemented a retry mechanism for multiplexed session creation failures, preventing the client library from stalling when initial session setup encounters transient errors.
• Controlled Retry Logic: Introduced safeguards to ensure only one RPC call attempts a session creation retry at a time and explicitly avoids retrying UNIMPLEMENTED errors, which signify that multiplexed sessions are not enabled.
• Concurrency Management: Utilized AtomicBoolean and ReentrantLock to manage the state of session creation retries and ensure thread-safe operations.
• Updated Session Retrieval: Modified the internal session retrieval logic to incorporate the new retry mechanism, ensuring that subsequent requests can trigger a session re-creation if the current one is broken.
• Comprehensive Testing: Added new test cases to validate the retry behavior for DEADLINE_EXCEEDED errors, covering both single and parallel request scenarios, ensuring the robustness of the fix.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a retry mechanism for multiplexed session creation failures, which is a valuable improvement to prevent the client library from stalling. The use of a lock to manage retries and prevent concurrent session creation attempts is a good approach.

I've identified a critical issue where a change intended to be non-blocking has become blocking, which could impact performance and behavior. I've also included a suggestion to improve the robustness of the locking mechanism by using a try-finally block, which is a standard best practice.

The added tests effectively cover the new retry logic for both single and parallel requests. Addressing the identified issues will make this a solid contribution.

[gemini-code-assist]

The current implementation of getCurrentSessionReferenceFuture() makes it a blocking method because it directly calls the blocking method getCurrentSessionReference(). This negates the purpose of DelayedMultiplexedSessionTransaction, which relies on a non-blocking future to defer session acquisition. The current implementation will cause createDelayedMultiplexedSessionTransaction() to block until a session is available, which is a significant change in behavior.

The method should be implemented asynchronously to chain the retry logic without blocking. Here's a suggested implementation that preserves the non-blocking nature:

  ApiFuture<SessionReference> getCurrentSessionReferenceFuture() {
    return ApiFutures.catchingAsync(
        this.multiplexedSessionReference.get(),
        Throwable.class,
        (throwable) -> {
          sessionCreationLock.lock();
          try {
            if (isMultiplexedSessionsSupported()
                && retryingSessionCreation.compareAndSet(false, true)) {
              SettableApiFuture<SessionReference> settableApiFuture = SettableApiFuture.create();
              asyncCreateMultiplexedSession(settableApiFuture);
              multiplexedSessionReference.set(settableApiFuture);
            }
          } finally {
            sessionCreationLock.unlock();
          }
          return this.multiplexedSessionReference.get();
        },
        MoreExecutors.directExecutor());
  }

This change ensures that DelayedMultiplexedSessionTransaction remains non-blocking and that the retry logic is correctly applied to it.

[sakthivelmanii]

Even current logic will ensure that only DelayedMultiplexedSessionTransaction remains non blocking during getDatabaseClient creation.

[sakthivelmanii]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a retry mechanism for multiplexed session creation failures, which is a valuable improvement for the library's resilience. The changes correctly identify UNIMPLEMENTED errors as non-retryable and use a ReentrantLock to prevent multiple concurrent session creation retries. New test cases have been added to validate this behavior, including scenarios with parallel requests and DEADLINE_EXCEEDED errors. However, there are a couple of areas that could be improved for clarity and efficiency.

[gemini-code-assist]

The retryingSessionCreation flag is initialized to true here, but it's immediately set to false in both onSessionReady and onSessionCreateFailure callbacks. Additionally, the compareAndSet(false, true) logic in maybeRetrySessionCreation expects the flag to be false to initiate a new retry.

Initializing it to false would align better with its intended state (not retrying initially) and the subsequent logic.

Suggested change

	private final AtomicBoolean retryingSessionCreation = new AtomicBoolean(true);
	private final AtomicBoolean retryingSessionCreation = new AtomicBoolean(false);