chore: route multi-use read-only txn reads via location-aware cache

[rahul2393]

• Added support for multi-use read-only transaction routing so reads/queries are sent to the cached routed server (based
  on routing hints/key ranges) instead of being pinned by transaction affinity/default endpoint.
• Introduced read-only transaction tracking in KeyAwareChannel to:
  • identify read-only txn IDs from BeginTransaction responses
  • preserve strong/stale leader preference for subsequent requests
  • bypass affinity lookup/recording for those txn IDs
• Extended ChannelFinder with overloads that allow explicit preferLeader routing for ReadRequest and ExecuteSqlRequest.
• Ensured cleanup removes both affinity and read-only tracking state when transaction is closed.
• Added unit tests in KeyAwareChannelTest for:
  • independent per-request routing within the same read-only txn
  • no affinity recording for read-only txns
  • cleanup path coverage.

[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 enhances the routing capabilities for read-only transactions by allowing them to leverage location-aware caching. Instead of being pinned to a single endpoint due to transaction affinity, read-only operations can now be routed independently based on key ranges and explicit leader preferences. This change aims to improve performance and resource utilization for multi-use read-only transactions by distributing their reads more efficiently across available servers.

Highlights

• Multi-use Read-Only Transaction Routing: Implemented support for routing reads and queries within multi-use read-only transactions to cached, location-aware servers, bypassing traditional transaction affinity.
• Read-Only Transaction Tracking: Introduced mechanisms in KeyAwareChannel to identify and track read-only transaction IDs, preserving strong/stale leader preferences and preventing affinity recording for these transactions.
• ChannelFinder Extensions: Extended ChannelFinder with new overloads for ReadRequest and ExecuteSqlRequest to allow explicit preferLeader routing, enabling more granular control over read-only transaction routing.
• State Cleanup: Ensured that both transaction affinity and the newly introduced read-only transaction tracking state are properly cleaned up when a transaction is closed.
• Comprehensive Unit Tests: Added new unit tests to validate independent per-request routing for read-only transactions, the bypass of affinity recording, and the correct cleanup of transaction state.

Changelog

• google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java
  • Added overloaded findServer methods for ReadRequest.Builder and ExecuteSqlRequest.Builder to include a preferLeader boolean parameter for explicit routing control.
• google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java
  • Introduced a readOnlyTransactions map to track read-only transaction IDs and their leader preference.
  • Modified clearAffinity to also remove read-only transaction tracking state.
  • Added helper methods isReadOnlyTransaction, readOnlyPreferLeader, and trackReadOnlyTransaction for managing read-only transaction state.
  • Added isReadOnlyBegin and readOnlyIsStrong fields to KeyAwareClientCall to capture read-only transaction properties during BeginTransaction.
  • Updated KeyAwareClientCall.sendMessage to identify and mark read-only BeginTransactionRequests.
  • Modified KeyAwareClientCall.routeFromRequest for ReadRequest and ExecuteSqlRequest to bypass affinity for read-only transactions and use the new ChannelFinder overloads for independent routing.
  • Updated KeyAwareClientCall.Listener.onMessage to track read-only transactions upon receiving BeginTransactionResponse.
• google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java
  • Imported necessary classes for testing new transaction and routing features.
  • Added readOnlyTransactionRoutesEachReadIndependently test to verify independent routing for read-only transactions.
  • Added readOnlyTransactionDoesNotRecordAffinity test to confirm that affinity is not recorded for read-only transactions.
  • Added readOnlyTransactionCleanupOnClose test to ensure proper cleanup of read-only transaction state.

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 implements location-aware routing for multi-use read-only transactions, which should improve performance by allowing individual reads and queries to be routed based on key ranges instead of being pinned by transaction affinity. The changes are well-structured, with clear separation of concerns between ChannelFinder and KeyAwareChannel. The addition of tracking for read-only transactions and the logic to bypass affinity are implemented correctly. The new unit tests cover the main scenarios well. I've only found one minor issue in a test case, which is missing assertions to fully verify the cleanup logic.

[olavloite]

I think that we should add similar logic for ExecuteSqlRequest and ReadRequest, as those can also include an inlined BeginTransactionOption. That might not be something that we use today for read-only transactions, but it is likely to change in the future. (Alternatively, we should have a test that fails if we see an ExecuteSqlRequest with an inlined begin. That way we get a signal if this part of the implementation changes in the future, and are forced to fix it here then)

[rahul2393]

added support for both explicit begin and inline begin with ExecuteSqlRequest and ReadRequest

[olavloite]

One limitation with this type of test is that they assume a given client behavior. That is: they assume that the client will call BeginTransaction for read-only transactions. If we change that in the future to use inlined-begin, then these tests will continue to pass, even though the actual feature will not work correctly. So I think that we also need something of an end-to-end test that really verifies that the feature works with client (and continues to work with the client, even if we make internal changes).

[rahul2393]

end-to-end test won't be possible here currently.