chore: handle unary gRPC call ordering in KeyAwareChannel

[rahul2393]

• Fixes a crash in KeyAwareChannel where unary gRPC calls can invoke request() before sendMessage().
• Buffers request/halfClose/setMessageCompression until the delegate is created, then replays them.

Motivation

• BeginTransaction/ExecuteSql/Commit/Rollback use unary call paths.
• gRPC calls request() in onStart for unary calls, which previously hit delegate() before sendMessage(), causing:
  "Delegate call not initialized before use. sendMessage was likely not called."

What changed

• KeyAwareClientCall now:
• buffers request(n) counts
• buffers halfClose and message compression
• handles early cancel()
• drains buffered requests after delegate.start()

Impact

• Removes IllegalStateException for unary calls.
• No behavior change for streaming calls.
• Minimal overhead (small in‑memory counters/flags).

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

The || Boolean.parseBoolean(System.getenv(EXPERIMENTAL_LOCATION_API_ENV_VAR) should not be necessary (and is also slightly counter-productive). The default for options.isEnableLocationApi() is to read the value from the environment variable. Only if a test has specified that SpannerOptions should use a fake environment, will it not read the environment variable. And in a case like that, the fake environment should also be able to override any environment variable that has been set, in order to get deterministic behavior (that is: The fake environment should be able to set the value to false, and it should not be overridden by what might have been set in a system environment variable).

[olavloite]

This can be simplified into one field: cancelledStatus. The if (cancelled) checks in the other places in this file can then be replaced with if (this.cancelledStatus != null)

[olavloite]

I think that we should try to preserve any trailers that might be present in the cause. And just create the Cancelled status once here, and re-use everywhere it is needed.

Suggested change

	cancelled = true;
	cancelMessage = message;
	cancelCause = cause;
	if (responseListener != null) {
	responseListener.onClose(
	io.grpc.Status.CANCELLED.withDescription(message).withCause(cause), new Metadata());
	}
	this.cancelledStatus = CANCELLED.withDescription(message).withCause(cause);
	Metadata trailers = cause == null ? new Metadata() : Status.trailersFromThrowable(cause);
	this.cancelledTrailers = trailers == null ? new Metadata() : trailers;
	if (responseListener != null) {
	responseListener.onClose(this.cancelledStatus, this.cancelledTrailers);
	}

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

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