chore: handle unary gRPC call ordering in KeyAwareChannel

[rahul2393]

• Fixes KeyAwareChannel unary-call initialization ordering and hardens non-happy-path behavior (cancel/timeout/error).
• Expands test coverage for KeyAwareChannel and location-aware integration flows.

Motivation

• Unary RPCs can invoke request() before sendMessage(), which could throw:
  Delegate call not initialized before use. sendMessage was likely not called.
• Cancel/error/timeout paths needed stronger guarantees and tests.

What Changed

• KeyAwareChannel:
  • Buffers request(n), halfClose, and setMessageCompression until delegate creation, then replays.
  • Simplifies cancel state to a single cancelledStatus (+ preserved trailers), reused across paths.
  • Preserves trailers from cancel causes (Status.trailersFromThrowable(cause)).
  • Clears pending message compression after replay.
  • Removes redundant endpoint null-check in mutation-key routing path.

[gemini-code-assist]

Summary of Changes

Hello @rahul2393, 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 resolves a critical issue in the KeyAwareChannel where unary gRPC calls could crash due to an IllegalStateException caused by an incorrect ordering of request() and sendMessage() calls. The solution involves implementing a buffering mechanism within KeyAwareClientCall to temporarily hold request, halfClose, and setMessageCompression operations until the underlying gRPC delegate is fully initialized. This ensures that these operations are executed in the correct sequence, preventing crashes and improving the robustness of unary gRPC interactions without affecting streaming calls.

Highlights

• Buffering gRPC Call Operations: The KeyAwareClientCall now buffers request(), halfClose(), and setMessageCompression() calls if the underlying gRPC delegate is not yet initialized. These buffered operations are replayed once the delegate is ready.
• Fix for Unary Call Crash: Addresses an IllegalStateException that occurred in unary gRPC calls when request() was invoked before sendMessage(), specifically in scenarios like BeginTransaction or ExecuteSql.
• Improved Cancellation Handling: The cancel() method now correctly handles early cancellations by storing the cancellation state and immediately closing the response listener if start() has already been called, or buffering the cancellation if not.
• State Management for Delegate: New internal fields (pendingRequests, pendingHalfClose, pendingMessageCompression, cancelled, cancelMessage, cancelCause) have been introduced to manage the state of these buffered operations and cancellation.

Changelog

• google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java
  • Added new fields (pendingRequests, pendingHalfClose, pendingMessageCompression, cancelled, cancelMessage, cancelCause) to KeyAwareClientCall for state management.
  • Modified start() to handle early cancellations by immediately closing the response listener if cancelled.
  • Updated sendMessage() to include checks for cancellation and proper initialization, and to apply buffered message compression, drain pending requests, and apply buffered half-close after delegate creation.
  • Changed halfClose() to buffer the call if the delegate is not yet initialized, instead of throwing an exception.
  • Overrode request() to buffer requests if the delegate is not ready.
  • Overrode setMessageCompression() to buffer the compression setting if the delegate is not ready.
  • Added drainPendingRequests() method to replay buffered request() calls.
  • Modified cancel() to buffer cancellation details if the call has not yet started.

Activity

• No human activity has been recorded for this pull request yet.

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 effectively addresses a crash in KeyAwareChannel caused by out-of-order gRPC calls for unary operations. The approach of buffering requests, half-close, and message compression until the delegate is initialized is sound and well-implemented. The changes are clear and align with the problem description. I have one minor suggestion to refactor a small piece of duplicated code to improve maintainability. Overall, this is a solid fix.

[olavloite]

I think that this class could use some more tests, especially for the non-happy path. Are we sure that everything works as we expect in case of errors? And cancelled calls? And timeouts? etc.

[olavloite]

this.pendingMessageCompression is never set back to null, meaning that once it has been set, this if statement will be true for every invocation.

[rahul2393]

Updated to set it to null

[olavloite]

(Not changed in this pull request, but still relevant): This if statement is redundant. endpoint is always null at this point, so the entire if statement can be removed. And if that is not intentional, then that might be an indication of a bug.

[rahul2393]

this was for my debugging and logging the endpoint whn validating the location aware flow, removed it.