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

Commit 72cafc8

Browse files
rahul2393rahul2393
committed
refactored classes
1 parent 7790a8a commit 72cafc8

7 files changed

Lines changed: 369 additions & 117 deletions

File tree

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java‎

Lines changed: 40 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,13 @@
1717
package com.google.cloud.spanner.spi.v1;
1818

1919
import com.google.api.core.InternalApi;
20+
import com.google.spanner.v1.BeginTransactionRequest;
2021
import com.google.spanner.v1.CacheUpdate;
2122
import com.google.spanner.v1.DirectedReadOptions;
2223
import com.google.spanner.v1.ExecuteSqlRequest;
2324
import com.google.spanner.v1.ReadRequest;
2425
import com.google.spanner.v1.RoutingHint;
26+
import com.google.spanner.v1.TransactionOptions;
2527
import com.google.spanner.v1.TransactionSelector;
2628
import java.util.Objects;
2729
import java.util.concurrent.atomic.AtomicLong;
@@ -37,11 +39,9 @@ public final class ChannelFinder {
3739
private final AtomicLong databaseId = new AtomicLong();
3840
private final KeyRecipeCache recipeCache = new KeyRecipeCache();
3941
private final KeyRangeCache rangeCache;
40-
private final String databaseUri;
4142

42-
public ChannelFinder(ChannelEndpointCache endpointCache, String databaseUri) {
43+
public ChannelFinder(ChannelEndpointCache endpointCache) {
4344
this.rangeCache = new KeyRangeCache(Objects.requireNonNull(endpointCache));
44-
this.databaseUri = Objects.requireNonNull(databaseUri);
4545
}
4646

4747
void useDeterministicRandom() {
@@ -83,18 +83,46 @@ public ChannelEndpoint findServer(ExecuteSqlRequest.Builder reqBuilder) {
8383
reqBuilder.getRoutingHintBuilder());
8484
}
8585

86+
public ChannelEndpoint findServer(BeginTransactionRequest.Builder reqBuilder) {
87+
if (!reqBuilder.hasMutationKey()) {
88+
return null;
89+
}
90+
TargetRange target = recipeCache.mutationToTargetRange(reqBuilder.getMutationKey());
91+
if (target == null) {
92+
return null;
93+
}
94+
RoutingHint.Builder hintBuilder = RoutingHint.newBuilder();
95+
hintBuilder.setKey(target.start);
96+
if (!target.limit.isEmpty()) {
97+
hintBuilder.setLimitKey(target.limit);
98+
}
99+
return fillRoutingHint(
100+
preferLeader(reqBuilder.getOptions()),
101+
KeyRangeCache.RangeMode.COVERING_SPLIT,
102+
DirectedReadOptions.getDefaultInstance(),
103+
hintBuilder);
104+
}
105+
86106
private ChannelEndpoint fillRoutingHint(
87107
TransactionSelector transactionSelector,
88108
DirectedReadOptions directedReadOptions,
89109
KeyRangeCache.RangeMode rangeMode,
90110
RoutingHint.Builder hintBuilder) {
111+
return fillRoutingHint(
112+
preferLeader(transactionSelector), rangeMode, directedReadOptions, hintBuilder);
113+
}
114+
115+
private ChannelEndpoint fillRoutingHint(
116+
boolean preferLeader,
117+
KeyRangeCache.RangeMode rangeMode,
118+
DirectedReadOptions directedReadOptions,
119+
RoutingHint.Builder hintBuilder) {
91120
long id = databaseId.get();
92121
if (id == 0) {
93122
return null;
94123
}
95124
hintBuilder.setDatabaseId(id);
96-
return rangeCache.fillRoutingHint(
97-
preferLeader(transactionSelector), rangeMode, directedReadOptions, hintBuilder);
125+
return rangeCache.fillRoutingHint(preferLeader, rangeMode, directedReadOptions, hintBuilder);
98126
}
99127

100128
private static boolean preferLeader(TransactionSelector selector) {
@@ -112,4 +140,11 @@ private static boolean preferLeader(TransactionSelector selector) {
112140
return true;
113141
}
114142
}
143+
144+
private static boolean preferLeader(TransactionOptions options) {
145+
if (options == null || !options.hasReadOnly()) {
146+
return true;
147+
}
148+
return options.getReadOnly().getStrong();
149+
}
115150
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java‎

Lines changed: 7 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,6 @@
5858
import com.google.api.gax.rpc.UnavailableException;
5959
import com.google.api.gax.rpc.WatchdogProvider;
6060
import com.google.api.pathtemplate.PathTemplate;
61-
import com.google.auth.Credentials;
6261
import com.google.cloud.RetryHelper;
6362
import com.google.cloud.RetryHelper.RetryHelperException;
6463
import com.google.cloud.grpc.GcpManagedChannel;
@@ -213,7 +212,6 @@
213212
import java.util.concurrent.ConcurrentLinkedDeque;
214213
import java.util.concurrent.ConcurrentMap;
215214
import java.util.concurrent.ExecutionException;
216-
import java.util.concurrent.Executor;
217215
import java.util.concurrent.ExecutorService;
218216
import java.util.concurrent.Executors;
219217
import java.util.concurrent.Future;
@@ -293,109 +291,6 @@ public class GapicSpannerRpc implements SpannerRpc {
293291

294292
private final GrpcCallContext baseGrpcCallContext;
295293

296-
private static final class KeyAwareTransportChannelProvider implements TransportChannelProvider {
297-
private final InstantiatingGrpcChannelProvider baseProvider;
298-
@Nullable private final ChannelEndpointCacheFactory endpointCacheFactory;
299-
300-
private KeyAwareTransportChannelProvider(
301-
InstantiatingGrpcChannelProvider.Builder builder,
302-
@Nullable ChannelEndpointCacheFactory endpointCacheFactory) {
303-
this.baseProvider = builder.build();
304-
this.endpointCacheFactory = endpointCacheFactory;
305-
}
306-
307-
private KeyAwareTransportChannelProvider(
308-
InstantiatingGrpcChannelProvider baseProvider,
309-
@Nullable ChannelEndpointCacheFactory endpointCacheFactory) {
310-
this.baseProvider = baseProvider;
311-
this.endpointCacheFactory = endpointCacheFactory;
312-
}
313-
314-
@Override
315-
public GrpcTransportChannel getTransportChannel() throws IOException {
316-
return GrpcTransportChannel.newBuilder()
317-
.setManagedChannel(KeyAwareChannel.create(baseProvider, endpointCacheFactory))
318-
.build();
319-
}
320-
321-
@Override
322-
public String getTransportName() {
323-
return baseProvider.getTransportName();
324-
}
325-
326-
@Override
327-
public boolean needsEndpoint() {
328-
return baseProvider.needsEndpoint();
329-
}
330-
331-
@Override
332-
public boolean needsCredentials() {
333-
return baseProvider.needsCredentials();
334-
}
335-
336-
@Override
337-
public boolean needsExecutor() {
338-
return baseProvider.needsExecutor();
339-
}
340-
341-
@Override
342-
public boolean needsHeaders() {
343-
return baseProvider.needsHeaders();
344-
}
345-
346-
@Override
347-
public boolean shouldAutoClose() {
348-
return baseProvider.shouldAutoClose();
349-
}
350-
351-
@Override
352-
public TransportChannelProvider withEndpoint(String endpoint) {
353-
return new KeyAwareTransportChannelProvider(
354-
(InstantiatingGrpcChannelProvider) baseProvider.withEndpoint(endpoint),
355-
endpointCacheFactory);
356-
}
357-
358-
@Override
359-
public TransportChannelProvider withCredentials(Credentials credentials) {
360-
return new KeyAwareTransportChannelProvider(
361-
(InstantiatingGrpcChannelProvider) baseProvider.withCredentials(credentials),
362-
endpointCacheFactory);
363-
}
364-
365-
@Override
366-
public TransportChannelProvider withHeaders(Map<String, String> headers) {
367-
return new KeyAwareTransportChannelProvider(
368-
(InstantiatingGrpcChannelProvider) baseProvider.withHeaders(headers),
369-
endpointCacheFactory);
370-
}
371-
372-
@Override
373-
public TransportChannelProvider withPoolSize(int poolSize) {
374-
return new KeyAwareTransportChannelProvider(
375-
(InstantiatingGrpcChannelProvider) baseProvider.withPoolSize(poolSize),
376-
endpointCacheFactory);
377-
}
378-
379-
@Override
380-
public TransportChannelProvider withExecutor(ScheduledExecutorService executor) {
381-
return new KeyAwareTransportChannelProvider(
382-
(InstantiatingGrpcChannelProvider) baseProvider.withExecutor(executor),
383-
endpointCacheFactory);
384-
}
385-
386-
@Override
387-
public TransportChannelProvider withExecutor(Executor executor) {
388-
return new KeyAwareTransportChannelProvider(
389-
(InstantiatingGrpcChannelProvider) baseProvider.withExecutor(executor),
390-
endpointCacheFactory);
391-
}
392-
393-
@Override
394-
public boolean acceptsPoolSize() {
395-
return baseProvider.acceptsPoolSize();
396-
}
397-
}
398-
399294
public static GapicSpannerRpc create(SpannerOptions options) {
400295
return new GapicSpannerRpc(options);
401296
}
@@ -506,12 +401,15 @@ public GapicSpannerRpc(final SpannerOptions options) {
506401

507402
boolean enableLocationApi =
508403
Boolean.parseBoolean(System.getenv(EXPERIMENTAL_LOCATION_API_ENV_VAR));
404+
TransportChannelProvider baseChannelProvider =
405+
MoreObjects.firstNonNull(
406+
options.getChannelProvider(), defaultChannelProviderBuilder.build());
509407
TransportChannelProvider channelProvider =
510-
enableLocationApi
408+
enableLocationApi && baseChannelProvider instanceof InstantiatingGrpcChannelProvider
511409
? new KeyAwareTransportChannelProvider(
512-
defaultChannelProviderBuilder, options.getChannelEndpointCacheFactory())
513-
: MoreObjects.firstNonNull(
514-
options.getChannelProvider(), defaultChannelProviderBuilder.build());
410+
(InstantiatingGrpcChannelProvider) baseChannelProvider,
411+
options.getChannelEndpointCacheFactory())
412+
: baseChannelProvider;
515413

516414
CredentialsProvider credentialsProvider =
517415
GrpcTransportOptions.setUpCredentialsProvider(options);

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ private ChannelFinder getOrCreateChannelFinder(String databaseId) {
111111
ref = channelFinders.get(databaseId);
112112
finder = (ref != null) ? ref.get() : null;
113113
if (finder == null) {
114-
finder = new ChannelFinder(endpointCache, databaseId);
114+
finder = new ChannelFinder(endpointCache);
115115
channelFinders.put(databaseId, new SoftReference<>(finder));
116116
}
117117
}
@@ -306,7 +306,18 @@ public void sendMessage(RequestT message) {
306306
}
307307
message = (RequestT) reqBuilder.build();
308308
} else if (message instanceof BeginTransactionRequest) {
309+
BeginTransactionRequest.Builder reqBuilder =
310+
((BeginTransactionRequest) message).toBuilder();
311+
String databaseId = parentChannel.extractDatabaseIdFromSession(reqBuilder.getSession());
312+
if (databaseId != null && reqBuilder.hasMutationKey()) {
313+
finder = parentChannel.getOrCreateChannelFinder(databaseId);
314+
ChannelEndpoint routed = finder.findServer(reqBuilder);
315+
if (endpoint == null) {
316+
endpoint = routed;
317+
}
318+
}
309319
allowDefaultAffinity = true;
320+
message = (RequestT) reqBuilder.build();
310321
} else if (message instanceof CommitRequest) {
311322
CommitRequest request = (CommitRequest) message;
312323
if (!request.getTransactionId().isEmpty()) {

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareTransportChannelProvider.java‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner.spi.v1;
18+
19+
import com.google.api.gax.grpc.GrpcTransportChannel;
20+
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
21+
import com.google.api.gax.rpc.TransportChannelProvider;
22+
import com.google.auth.Credentials;
23+
import java.io.IOException;
24+
import java.util.Map;
25+
import java.util.concurrent.Executor;
26+
import java.util.concurrent.ScheduledExecutorService;
27+
import javax.annotation.Nullable;
28+
29+
final class KeyAwareTransportChannelProvider implements TransportChannelProvider {
30+
private final InstantiatingGrpcChannelProvider baseProvider;
31+
@Nullable private final ChannelEndpointCacheFactory endpointCacheFactory;
32+
33+
KeyAwareTransportChannelProvider(
34+
InstantiatingGrpcChannelProvider.Builder builder,
35+
@Nullable ChannelEndpointCacheFactory endpointCacheFactory) {
36+
this.baseProvider = builder.build();
37+
this.endpointCacheFactory = endpointCacheFactory;
38+
}
39+
40+
KeyAwareTransportChannelProvider(
41+
InstantiatingGrpcChannelProvider baseProvider,
42+
@Nullable ChannelEndpointCacheFactory endpointCacheFactory) {
43+
this.baseProvider = baseProvider;
44+
this.endpointCacheFactory = endpointCacheFactory;
45+
}
46+
47+
@Override
48+
public GrpcTransportChannel getTransportChannel() throws IOException {
49+
return GrpcTransportChannel.newBuilder()
50+
.setManagedChannel(KeyAwareChannel.create(baseProvider, endpointCacheFactory))
51+
.build();
52+
}
53+
54+
@Override
55+
public String getTransportName() {
56+
return baseProvider.getTransportName();
57+
}
58+
59+
@Override
60+
public boolean needsEndpoint() {
61+
return baseProvider.needsEndpoint();
62+
}
63+
64+
@Override
65+
public boolean needsCredentials() {
66+
return baseProvider.needsCredentials();
67+
}
68+
69+
@Override
70+
public boolean needsExecutor() {
71+
return baseProvider.needsExecutor();
72+
}
73+
74+
@Override
75+
public boolean needsHeaders() {
76+
return baseProvider.needsHeaders();
77+
}
78+
79+
@Override
80+
public boolean shouldAutoClose() {
81+
return baseProvider.shouldAutoClose();
82+
}
83+
84+
@Override
85+
public TransportChannelProvider withEndpoint(String endpoint) {
86+
return new KeyAwareTransportChannelProvider(
87+
(InstantiatingGrpcChannelProvider) baseProvider.withEndpoint(endpoint),
88+
endpointCacheFactory);
89+
}
90+
91+
@Override
92+
public TransportChannelProvider withCredentials(Credentials credentials) {
93+
return new KeyAwareTransportChannelProvider(
94+
(InstantiatingGrpcChannelProvider) baseProvider.withCredentials(credentials),
95+
endpointCacheFactory);
96+
}
97+
98+
@Override
99+
public TransportChannelProvider withHeaders(Map<String, String> headers) {
100+
return new KeyAwareTransportChannelProvider(
101+
(InstantiatingGrpcChannelProvider) baseProvider.withHeaders(headers), endpointCacheFactory);
102+
}
103+
104+
@Override
105+
public TransportChannelProvider withPoolSize(int poolSize) {
106+
return new KeyAwareTransportChannelProvider(
107+
(InstantiatingGrpcChannelProvider) baseProvider.withPoolSize(poolSize),
108+
endpointCacheFactory);
109+
}
110+
111+
@Override
112+
public TransportChannelProvider withExecutor(ScheduledExecutorService executor) {
113+
return new KeyAwareTransportChannelProvider(
114+
(InstantiatingGrpcChannelProvider) baseProvider.withExecutor(executor),
115+
endpointCacheFactory);
116+
}
117+
118+
@Override
119+
public TransportChannelProvider withExecutor(Executor executor) {
120+
return new KeyAwareTransportChannelProvider(
121+
(InstantiatingGrpcChannelProvider) baseProvider.withExecutor(executor),
122+
endpointCacheFactory);
123+
}
124+
125+
@Override
126+
public boolean acceptsPoolSize() {
127+
return baseProvider.acceptsPoolSize();
128+
}
129+
}

0 commit comments

Comments
 (0)