Skip to content
This repository was archived by the owner on Apr 7, 2026. It is now read-only.

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

Merged
rahul2393 merged 2 commits into from
Feb 11, 2026
Merged

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
540f387
chore: route multi-use read-only txn reads via location-aware cache
rahul2393 Feb 11, 2026
519f0ef
incorporate changes
rahul2393 Feb 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,15 @@ public ChannelEndpoint findServer(ReadRequest.Builder reqBuilder) {
reqBuilder.getRoutingHintBuilder());
}

public ChannelEndpoint findServer(ReadRequest.Builder reqBuilder, boolean preferLeader) {
recipeCache.computeKeys(reqBuilder);
return fillRoutingHint(
preferLeader,
KeyRangeCache.RangeMode.COVERING_SPLIT,
reqBuilder.getDirectedReadOptions(),
reqBuilder.getRoutingHintBuilder());
}

public ChannelEndpoint findServer(ExecuteSqlRequest.Builder reqBuilder) {
recipeCache.computeKeys(reqBuilder);
return fillRoutingHint(
Expand All @@ -83,6 +92,15 @@ public ChannelEndpoint findServer(ExecuteSqlRequest.Builder reqBuilder) {
reqBuilder.getRoutingHintBuilder());
}

public ChannelEndpoint findServer(ExecuteSqlRequest.Builder reqBuilder, boolean preferLeader) {
recipeCache.computeKeys(reqBuilder);
return fillRoutingHint(
preferLeader,
KeyRangeCache.RangeMode.PICK_RANDOM,
reqBuilder.getDirectedReadOptions(),
reqBuilder.getRoutingHintBuilder());
}

public ChannelEndpoint findServer(BeginTransactionRequest.Builder reqBuilder) {
if (!reqBuilder.hasMutationKey()) {
return null;
Expand Down
62 changes: 56 additions & 6 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,9 @@ final class KeyAwareChannel extends ManagedChannel {
private final Map<String, SoftReference<ChannelFinder>> channelFinders =
new ConcurrentHashMap<>();
private final Map<ByteString, String> transactionAffinities = new ConcurrentHashMap<>();
// Maps read-only transaction IDs to their preferLeader value.
// Strong reads → true (prefer leader), Stale reads → false (any replica).
private final Map<ByteString, Boolean> readOnlyTransactions = new ConcurrentHashMap<>();
Comment thread
rahul2393 marked this conversation as resolved.
Outdated

private KeyAwareChannel(
InstantiatingGrpcChannelProvider channelProvider,
Expand Down Expand Up @@ -184,12 +187,34 @@ private void clearAffinity(ByteString transactionId) {
return;
}
transactionAffinities.remove(transactionId);
readOnlyTransactions.remove(transactionId);
Comment thread
rahul2393 marked this conversation as resolved.
Outdated
}

void clearTransactionAffinity(ByteString transactionId) {
clearAffinity(transactionId);
}

private boolean isReadOnlyTransaction(ByteString transactionId) {
return transactionId != null
&& !transactionId.isEmpty()
&& readOnlyTransactions.containsKey(transactionId);
}

@Nullable
private Boolean readOnlyPreferLeader(ByteString transactionId) {
if (transactionId == null || transactionId.isEmpty()) {
return null;
}
return readOnlyTransactions.get(transactionId);
}

private void trackReadOnlyTransaction(ByteString transactionId, boolean preferLeader) {
if (transactionId == null || transactionId.isEmpty()) {
return;
}
readOnlyTransactions.put(transactionId, preferLeader);
}

private void recordAffinity(
ByteString transactionId, @Nullable ChannelEndpoint endpoint, boolean allowDefault) {
if (transactionId == null || transactionId.isEmpty() || endpoint == null) {
Expand Down Expand Up @@ -250,6 +275,8 @@ static final class KeyAwareClientCall<RequestT, ResponseT>
@Nullable private Boolean pendingMessageCompression;
@Nullable private io.grpc.Status cancelledStatus;
@Nullable private Metadata cancelledTrailers;
private boolean isReadOnlyBegin;
private boolean readOnlyIsStrong;
private final Object lock = new Object();

KeyAwareClientCall(
Expand Down Expand Up @@ -325,7 +352,12 @@ public void sendMessage(RequestT message) {
finder = parentChannel.getOrCreateChannelFinder(databaseId);
endpoint = finder.findServer(reqBuilder);
}
allowDefaultAffinity = true;
if (reqBuilder.hasOptions() && reqBuilder.getOptions().hasReadOnly()) {

@olavloite olavloite Feb 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 rahul2393 Feb 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

isReadOnlyBegin = true;
readOnlyIsStrong = reqBuilder.getOptions().getReadOnly().getStrong();
} else {
allowDefaultAffinity = true;
}
message = (RequestT) reqBuilder.build();
} else if (message instanceof CommitRequest) {
CommitRequest request = (CommitRequest) message;
Expand Down Expand Up @@ -486,11 +518,17 @@ void maybeClearAffinity() {
private RoutingDecision routeFromRequest(ReadRequest.Builder reqBuilder) {
String databaseId = parentChannel.extractDatabaseIdFromSession(reqBuilder.getSession());
ByteString transactionId = transactionIdFromSelector(reqBuilder.getTransaction());
ChannelEndpoint endpoint = parentChannel.affinityEndpoint(transactionId);
// Skip affinity for read-only transactions so each read routes independently.
boolean isReadOnly = parentChannel.isReadOnlyTransaction(transactionId);
ChannelEndpoint endpoint = isReadOnly ? null : parentChannel.affinityEndpoint(transactionId);
ChannelFinder finder = null;
if (databaseId != null) {
finder = parentChannel.getOrCreateChannelFinder(databaseId);
ChannelEndpoint routed = finder.findServer(reqBuilder);
Boolean preferLeaderOverride = parentChannel.readOnlyPreferLeader(transactionId);
ChannelEndpoint routed =
preferLeaderOverride != null
? finder.findServer(reqBuilder, preferLeaderOverride)
: finder.findServer(reqBuilder);
Comment thread
rahul2393 marked this conversation as resolved.
if (endpoint == null) {
endpoint = routed;
}
Expand All @@ -501,11 +539,17 @@ private RoutingDecision routeFromRequest(ReadRequest.Builder reqBuilder) {
private RoutingDecision routeFromRequest(ExecuteSqlRequest.Builder reqBuilder) {
String databaseId = parentChannel.extractDatabaseIdFromSession(reqBuilder.getSession());
ByteString transactionId = transactionIdFromSelector(reqBuilder.getTransaction());
ChannelEndpoint endpoint = parentChannel.affinityEndpoint(transactionId);
// Skip affinity for read-only transactions so each query routes independently.
boolean isReadOnly = parentChannel.isReadOnlyTransaction(transactionId);
ChannelEndpoint endpoint = isReadOnly ? null : parentChannel.affinityEndpoint(transactionId);
ChannelFinder finder = null;
if (databaseId != null) {
finder = parentChannel.getOrCreateChannelFinder(databaseId);
ChannelEndpoint routed = finder.findServer(reqBuilder);
Boolean preferLeaderOverride = parentChannel.readOnlyPreferLeader(transactionId);
ChannelEndpoint routed =
preferLeaderOverride != null
? finder.findServer(reqBuilder, preferLeaderOverride)
: finder.findServer(reqBuilder);
if (endpoint == null) {
endpoint = routed;
}
Expand Down Expand Up @@ -554,7 +598,13 @@ public void onMessage(ResponseT message) {
transactionId = transactionIdFromTransaction(response);
}
if (transactionId != null) {
call.maybeRecordAffinity(transactionId);
if (call.isReadOnlyBegin) {
Comment thread
rahul2393 marked this conversation as resolved.
// Track the read-only transaction so subsequent reads skip affinity
// and route independently based on key-based routing.
call.parentChannel.trackReadOnlyTransaction(transactionId, call.readOnlyIsStrong);
} else if (!call.parentChannel.isReadOnlyTransaction(transactionId)) {
call.maybeRecordAffinity(transactionId);
}
}
super.onMessage(message);
}
Expand Down
215 changes: 215 additions & 0 deletions google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,18 @@
import com.google.spanner.v1.CommitResponse;
import com.google.spanner.v1.ExecuteSqlRequest;
import com.google.spanner.v1.Group;
import com.google.spanner.v1.PartialResultSet;
import com.google.spanner.v1.Range;
import com.google.spanner.v1.ReadRequest;
import com.google.spanner.v1.ResultSet;
import com.google.spanner.v1.ResultSetMetadata;
import com.google.spanner.v1.RollbackRequest;
import com.google.spanner.v1.RoutingHint;
import com.google.spanner.v1.SpannerGrpc;
import com.google.spanner.v1.Tablet;
import com.google.spanner.v1.Transaction;
import com.google.spanner.v1.TransactionOptions;
import com.google.spanner.v1.TransactionSelector;
import io.grpc.CallOptions;
import io.grpc.ClientCall;
import io.grpc.ManagedChannel;
Expand Down Expand Up @@ -269,6 +274,216 @@ public void resultSetCacheUpdateRoutesSubsequentRequest() throws Exception {
assertThat(harness.endpointCache.callCountForAddress("routed:1234")).isEqualTo(1);
}

@Test
public void readOnlyTransactionRoutesEachReadIndependently() throws Exception {

@olavloite olavloite Feb 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 rahul2393 Feb 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

TestHarness harness = createHarness();
ByteString transactionId = ByteString.copyFromUtf8("ro-tx-1");

// 1. Begin a read-only transaction (stale read).
ClientCall<BeginTransactionRequest, Transaction> beginCall =
harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
CapturingListener<Transaction> beginListener = new CapturingListener<>();
beginCall.start(beginListener, new Metadata());
beginCall.sendMessage(
BeginTransactionRequest.newBuilder()
.setSession(SESSION)
.setOptions(
TransactionOptions.newBuilder()
.setReadOnly(
TransactionOptions.ReadOnly.newBuilder()
.setReturnReadTimestamp(true)
.build()))
.build());

// BeginTransaction goes to default channel.
assertThat(harness.defaultManagedChannel.callCount()).isEqualTo(1);

@SuppressWarnings("unchecked")
RecordingClientCall<BeginTransactionRequest, Transaction> beginDelegate =
(RecordingClientCall<BeginTransactionRequest, Transaction>)
harness.defaultManagedChannel.latestCall();
beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build());
beginDelegate.emitOnClose(Status.OK, new Metadata());

// 2. Populate cache with routing data for two different key ranges.
CacheUpdate cacheUpdate =
CacheUpdate.newBuilder()
.setDatabaseId(7L)
.addRange(
Range.newBuilder()
.setStartKey(bytes("a"))
.setLimitKey(bytes("m"))
.setGroupUid(1L)
.setSplitId(1L)
.setGeneration(bytes("1")))
.addRange(
Range.newBuilder()
.setStartKey(bytes("m"))
.setLimitKey(bytes("z"))
.setGroupUid(2L)
.setSplitId(2L)
.setGeneration(bytes("1")))
.addGroup(
Group.newBuilder()
.setGroupUid(1L)
.setGeneration(bytes("1"))
.addTablets(
Tablet.newBuilder()
.setTabletUid(1L)
.setServerAddress("server-a:1234")
.setIncarnation(bytes("1"))
.setDistance(0)))
.addGroup(
Group.newBuilder()
.setGroupUid(2L)
.setGeneration(bytes("1"))
.addTablets(
Tablet.newBuilder()
.setTabletUid(2L)
.setServerAddress("server-b:1234")
.setIncarnation(bytes("1"))
.setDistance(0)))
.build();

// Seed the cache via a dummy query response with cache update.
ClientCall<ExecuteSqlRequest, ResultSet> seedCall =
harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT);
seedCall.start(new CapturingListener<ResultSet>(), new Metadata());
seedCall.sendMessage(
ExecuteSqlRequest.newBuilder()
.setSession(SESSION)
.setRoutingHint(RoutingHint.newBuilder().setKey(bytes("a")).build())
.build());
@SuppressWarnings("unchecked")
RecordingClientCall<ExecuteSqlRequest, ResultSet> seedDelegate =
(RecordingClientCall<ExecuteSqlRequest, ResultSet>)
harness.defaultManagedChannel.latestCall();
seedDelegate.emitOnMessage(ResultSet.newBuilder().setCacheUpdate(cacheUpdate).build());

// 3. Send a streaming read with key in range [a, m) → should go to server-a.
ClientCall<ReadRequest, PartialResultSet> readCallA =
harness.channel.newCall(SpannerGrpc.getStreamingReadMethod(), CallOptions.DEFAULT);
readCallA.start(new CapturingListener<PartialResultSet>(), new Metadata());
readCallA.sendMessage(
ReadRequest.newBuilder()
.setSession(SESSION)
.setTransaction(TransactionSelector.newBuilder().setId(transactionId))
.setRoutingHint(RoutingHint.newBuilder().setKey(bytes("b")).build())
.build());

assertThat(harness.endpointCache.callCountForAddress("server-a:1234")).isEqualTo(1);

// 4. Send an ExecuteStreamingSql with key in range [m, z) → should go to server-b.
ClientCall<ExecuteSqlRequest, PartialResultSet> queryCallB =
harness.channel.newCall(SpannerGrpc.getExecuteStreamingSqlMethod(), CallOptions.DEFAULT);
queryCallB.start(new CapturingListener<PartialResultSet>(), new Metadata());
queryCallB.sendMessage(
ExecuteSqlRequest.newBuilder()
.setSession(SESSION)
.setTransaction(TransactionSelector.newBuilder().setId(transactionId))
.setRoutingHint(RoutingHint.newBuilder().setKey(bytes("n")).build())
.build());

assertThat(harness.endpointCache.callCountForAddress("server-b:1234")).isEqualTo(1);

// Neither read was pinned to the default host (besides the initial begin + seed).
// default had: 1 begin + 1 seed = 2 calls
assertThat(harness.defaultManagedChannel.callCount()).isEqualTo(2);
}

@Test
public void readOnlyTransactionDoesNotRecordAffinity() throws Exception {
TestHarness harness = createHarness();
ByteString transactionId = ByteString.copyFromUtf8("ro-tx-2");

// Begin a read-only transaction.
ClientCall<BeginTransactionRequest, Transaction> beginCall =
harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
beginCall.start(new CapturingListener<Transaction>(), new Metadata());
beginCall.sendMessage(
BeginTransactionRequest.newBuilder()
.setSession(SESSION)
.setOptions(
TransactionOptions.newBuilder()
.setReadOnly(
TransactionOptions.ReadOnly.newBuilder()
.setReturnReadTimestamp(true)
.build()))
.build());

@SuppressWarnings("unchecked")
RecordingClientCall<BeginTransactionRequest, Transaction> beginDelegate =
(RecordingClientCall<BeginTransactionRequest, Transaction>)
harness.defaultManagedChannel.latestCall();
beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build());
beginDelegate.emitOnClose(Status.OK, new Metadata());

// No affinity should be recorded for the default endpoint.
// Verify by checking that the endpoint cache was never queried for affinity lookup.
// The default endpoint getCount tracks affinity lookups.
assertThat(harness.endpointCache.getCount(DEFAULT_ADDRESS)).isEqualTo(0);

// Send a read using the transaction ID (no cache populated, so falls back to default).
ClientCall<ExecuteSqlRequest, ResultSet> readCall =
harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT);
readCall.start(new CapturingListener<ResultSet>(), new Metadata());
readCall.sendMessage(
ExecuteSqlRequest.newBuilder()
.setSession(SESSION)
.setTransaction(TransactionSelector.newBuilder().setId(transactionId))
.build());

// The read goes to default (no cache data), but NOT because of affinity.
// No affinity lookup should have been performed for the read-only txn.
assertThat(harness.endpointCache.getCount(DEFAULT_ADDRESS)).isEqualTo(0);

// Now receive a response with the transaction ID — should NOT record affinity.
@SuppressWarnings("unchecked")
RecordingClientCall<ExecuteSqlRequest, ResultSet> readDelegate =
(RecordingClientCall<ExecuteSqlRequest, ResultSet>)
harness.defaultManagedChannel.latestCall();
readDelegate.emitOnMessage(
ResultSet.newBuilder()
.setMetadata(
ResultSetMetadata.newBuilder()
.setTransaction(Transaction.newBuilder().setId(transactionId)))
.build());

// Still no affinity recorded.
assertThat(harness.endpointCache.getCount(DEFAULT_ADDRESS)).isEqualTo(0);
}

@Test
public void readOnlyTransactionCleanupOnClose() throws Exception {
TestHarness harness = createHarness();
ByteString transactionId = ByteString.copyFromUtf8("ro-tx-3");

// Begin a read-only transaction.
ClientCall<BeginTransactionRequest, Transaction> beginCall =
harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
beginCall.start(new CapturingListener<Transaction>(), new Metadata());
beginCall.sendMessage(
BeginTransactionRequest.newBuilder()
.setSession(SESSION)
.setOptions(
TransactionOptions.newBuilder()
.setReadOnly(
TransactionOptions.ReadOnly.newBuilder()
.setReturnReadTimestamp(true)
.build()))
.build());

@SuppressWarnings("unchecked")
RecordingClientCall<BeginTransactionRequest, Transaction> beginDelegate =
(RecordingClientCall<BeginTransactionRequest, Transaction>)
harness.defaultManagedChannel.latestCall();
beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build());
beginDelegate.emitOnClose(Status.OK, new Metadata());

// Clear transaction affinity (simulates MultiUseReadOnlyTransaction.close()).
harness.channel.clearTransactionAffinity(transactionId);
}

private static TestHarness createHarness() throws IOException {
FakeEndpointCache endpointCache = new FakeEndpointCache(DEFAULT_ADDRESS);
InstantiatingGrpcChannelProvider provider =
Expand Down
Loading