fix(spanner): improve grpc-gcp affinity cleanup and location-aware re… · googleapis/google-cloud-java@aca0428 · GitHub
Skip to content

Commit aca0428

Browse files
rahul2393rahul2393
authored
fix(spanner): improve grpc-gcp affinity cleanup and location-aware retries (#12682)
2 parents ebad7e9 + 69fc961 commit aca0428

18 files changed

Lines changed: 771 additions & 51 deletions

‎java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java‎

Lines changed: 33 additions & 3 deletions

‎java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -398,6 +398,10 @@ private DatabaseClient createMultiplexedSessionTransaction(boolean singleUse) {
398398
private MultiplexedSessionTransaction createDirectMultiplexedSessionTransaction(
399399
boolean singleUse) {
400400
try {
401+
int singleUseChannelHint =
402+
singleUse && !sessionClient.getSpanner().getOptions().isGrpcGcpExtensionEnabled()
403+
? getSingleUseChannelHint()
404+
: NO_CHANNEL_HINT;
401405
return new MultiplexedSessionTransaction(
402406
this,
403407
tracer.getCurrentSpan(),
@@ -406,7 +410,7 @@ private MultiplexedSessionTransaction createDirectMultiplexedSessionTransaction(
406410
// session, such as for example a DatabaseNotFound exception. We therefore do not need
407411
// any special handling of such errors.
408412
multiplexedSessionReference.get().get(),
409-
singleUse ? getSingleUseChannelHint() : NO_CHANNEL_HINT,
413+
singleUseChannelHint,
410414
singleUse);
411415
} catch (ExecutionException executionException) {
412416
throw SpannerExceptionFactory.asSpannerException(executionException.getCause());

‎java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,10 @@ static SessionOption channelHint(long hint) {
8989
return new SessionOption(SpannerRpc.Option.CHANNEL_HINT, hint);
9090
}
9191

92+
static SessionOption unbindChannelHint() {
93+
return new SessionOption(SpannerRpc.Option.UNBIND_CHANNEL_HINT, Boolean.TRUE);
94+
}
95+
9296
SpannerRpc.Option rpcOption() {
9397
return rpcOption;
9498
}

‎java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
import java.util.List;
5252
import java.util.Map;
5353
import java.util.concurrent.ExecutionException;
54+
import java.util.concurrent.ThreadLocalRandom;
5455
import javax.annotation.Nullable;
5556

5657
/**
@@ -299,10 +300,11 @@ public CommitResponse writeAtLeastOnceWithOptions(
299300
}
300301
CommitRequest request = requestBuilder.build();
301302
ISpan span = tracer.spanBuilder(SpannerImpl.COMMIT);
303+
Map<SpannerRpc.Option, ?> singleUseWriteOptions = getSingleUseWriteChannelHintOptions();
302304

303305
try (IScope s = tracer.withSpan(span)) {
304306
return SpannerRetryHelper.runTxWithRetriesOnAborted(
305-
() -> new CommitResponse(spanner.getRpc().commit(request, getOptions())));
307+
() -> new CommitResponse(spanner.getRpc().commit(request, singleUseWriteOptions)));
306308
} catch (RuntimeException e) {
307309
span.setStatus(e);
308310
throw e;
@@ -326,6 +328,14 @@ private RequestOptions getRequestOptions(TransactionOption... transactionOptions
326328
return null;
327329
}
328330

331+
private Map<SpannerRpc.Option, ?> getSingleUseWriteChannelHintOptions() {
332+
if (!spanner.getOptions().isGrpcGcpExtensionEnabled()) {
333+
return getOptions();
334+
}
335+
return optionMap(
336+
SessionOption.channelHint(ThreadLocalRandom.current().nextLong(Long.MAX_VALUE)));
337+
}
338+
329339
@Override
330340
public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
331341
Iterable<MutationGroup> mutationGroups, TransactionOption... transactionOptions)

‎java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -232,7 +232,9 @@ private TransactionContextImpl(Builder builder) {
232232
this.clock = builder.clock;
233233
this.channelHint =
234234
getChannelHintOptions(
235-
session.getOptions(), ThreadLocalRandom.current().nextLong(Long.MAX_VALUE));
235+
session.getOptions(),
236+
ThreadLocalRandom.current().nextLong(Long.MAX_VALUE),
237+
session.getSpanner().getOptions().isGrpcGcpExtensionEnabled());
236238
this.previousTransactionId = builder.previousTransactionId;
237239
}
238240

‎java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,17 @@ long getNthRequest() {
8383
return this.nthRequest;
8484
}
8585

86+
/**
87+
* Returns a stable identifier for the logical request, excluding channel and attempt.
88+
*
89+
* <p>This can be used to associate retry-local state across attempts of the same RPC without
90+
* affecting unrelated requests.
91+
*/
92+
@InternalApi
93+
public String getLogicalRequestKey() {
94+
return this.nthClientId + "." + this.nthRequest;
95+
}
96+
8697
@VisibleForTesting
8798
static final Pattern REGEX =
8899
Pattern.compile("^(\\d)\\.([0-9a-z]{16})\\.(\\d+)\\.(\\d+)\\.(\\d+)\\.(\\d+)$");

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

Lines changed: 55 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import java.util.Objects;
3333
import java.util.concurrent.ThreadLocalRandom;
3434
import java.util.concurrent.atomic.AtomicLong;
35+
import java.util.function.Predicate;
3536

3637
/**
3738
* Finds a server for a request using location-aware routing metadata.
@@ -40,6 +41,8 @@
4041
*/
4142
@InternalApi
4243
public final class ChannelFinder {
44+
private static final Predicate<String> NO_EXCLUDED_ENDPOINTS = address -> false;
45+
4346
private final Object updateLock = new Object();
4447
private final AtomicLong databaseId = new AtomicLong();
4548
private final KeyRecipeCache recipeCache = new KeyRecipeCache();
@@ -71,47 +74,83 @@ public void update(CacheUpdate update) {
7174
}
7275

7376
public ChannelEndpoint findServer(ReadRequest.Builder reqBuilder) {
74-
return findServer(reqBuilder, preferLeader(reqBuilder.getTransaction()));
77+
return findServer(reqBuilder, preferLeader(reqBuilder.getTransaction()), NO_EXCLUDED_ENDPOINTS);
78+
}
79+
80+
public ChannelEndpoint findServer(
81+
ReadRequest.Builder reqBuilder, Predicate<String> excludedEndpoints) {
82+
return findServer(reqBuilder, preferLeader(reqBuilder.getTransaction()), excludedEndpoints);
7583
}
7684

7785
public ChannelEndpoint findServer(ReadRequest.Builder reqBuilder, boolean preferLeader) {
86+
return findServer(reqBuilder, preferLeader, NO_EXCLUDED_ENDPOINTS);
87+
}
88+
89+
public ChannelEndpoint findServer(
90+
ReadRequest.Builder reqBuilder, boolean preferLeader, Predicate<String> excludedEndpoints) {
7891
recipeCache.computeKeys(reqBuilder);
7992
return fillRoutingHint(
8093
preferLeader,
8194
KeyRangeCache.RangeMode.COVERING_SPLIT,
8295
reqBuilder.getDirectedReadOptions(),
83-
reqBuilder.getRoutingHintBuilder());
96+
reqBuilder.getRoutingHintBuilder(),
97+
excludedEndpoints);
8498
}
8599

86100
public ChannelEndpoint findServer(ExecuteSqlRequest.Builder reqBuilder) {
87-
return findServer(reqBuilder, preferLeader(reqBuilder.getTransaction()));
101+
return findServer(reqBuilder, preferLeader(reqBuilder.getTransaction()), NO_EXCLUDED_ENDPOINTS);
102+
}
103+
104+
public ChannelEndpoint findServer(
105+
ExecuteSqlRequest.Builder reqBuilder, Predicate<String> excludedEndpoints) {
106+
return findServer(reqBuilder, preferLeader(reqBuilder.getTransaction()), excludedEndpoints);
88107
}
89108

90109
public ChannelEndpoint findServer(ExecuteSqlRequest.Builder reqBuilder, boolean preferLeader) {
110+
return findServer(reqBuilder, preferLeader, NO_EXCLUDED_ENDPOINTS);
111+
}
112+
113+
public ChannelEndpoint findServer(
114+
ExecuteSqlRequest.Builder reqBuilder,
115+
boolean preferLeader,
116+
Predicate<String> excludedEndpoints) {
91117
recipeCache.computeKeys(reqBuilder);
92118
return fillRoutingHint(
93119
preferLeader,
94120
KeyRangeCache.RangeMode.PICK_RANDOM,
95121
reqBuilder.getDirectedReadOptions(),
96-
reqBuilder.getRoutingHintBuilder());
122+
reqBuilder.getRoutingHintBuilder(),
123+
excludedEndpoints);
97124
}
98125

99126
public ChannelEndpoint findServer(BeginTransactionRequest.Builder reqBuilder) {
127+
return findServer(reqBuilder, NO_EXCLUDED_ENDPOINTS);
128+
}
129+
130+
public ChannelEndpoint findServer(
131+
BeginTransactionRequest.Builder reqBuilder, Predicate<String> excludedEndpoints) {
100132
if (!reqBuilder.hasMutationKey()) {
101133
return null;
102134
}
103135
return routeMutation(
104136
reqBuilder.getMutationKey(),
105137
preferLeader(reqBuilder.getOptions()),
106-
reqBuilder.getRoutingHintBuilder());
138+
reqBuilder.getRoutingHintBuilder(),
139+
excludedEndpoints);
107140
}
108141

109142
public ChannelEndpoint fillRoutingHint(CommitRequest.Builder reqBuilder) {
143+
return fillRoutingHint(reqBuilder, NO_EXCLUDED_ENDPOINTS);
144+
}
145+
146+
public ChannelEndpoint fillRoutingHint(
147+
CommitRequest.Builder reqBuilder, Predicate<String> excludedEndpoints) {
110148
Mutation mutation = selectMutationForRouting(reqBuilder.getMutationsList());
111149
if (mutation == null) {
112150
return null;
113151
}
114-
return routeMutation(mutation, /* preferLeader= */ true, reqBuilder.getRoutingHintBuilder());
152+
return routeMutation(
153+
mutation, /* preferLeader= */ true, reqBuilder.getRoutingHintBuilder(), excludedEndpoints);
115154
}
116155

117156
private static Mutation selectMutationForRouting(List<Mutation> mutations) {
@@ -139,7 +178,10 @@ private static Mutation selectMutationForRouting(List<Mutation> mutations) {
139178
}
140179

141180
private ChannelEndpoint routeMutation(
142-
Mutation mutation, boolean preferLeader, RoutingHint.Builder hintBuilder) {
181+
Mutation mutation,
182+
boolean preferLeader,
183+
RoutingHint.Builder hintBuilder,
184+
Predicate<String> excludedEndpoints) {
143185
recipeCache.applySchemaGeneration(hintBuilder);
144186
TargetRange target = recipeCache.mutationToTargetRange(mutation);
145187
if (target == null) {
@@ -150,20 +192,23 @@ private ChannelEndpoint routeMutation(
150192
preferLeader,
151193
KeyRangeCache.RangeMode.COVERING_SPLIT,
152194
DirectedReadOptions.getDefaultInstance(),
153-
hintBuilder);
195+
hintBuilder,
196+
excludedEndpoints);
154197
}
155198

156199
private ChannelEndpoint fillRoutingHint(
157200
boolean preferLeader,
158201
KeyRangeCache.RangeMode rangeMode,
159202
DirectedReadOptions directedReadOptions,
160-
RoutingHint.Builder hintBuilder) {
203+
RoutingHint.Builder hintBuilder,
204+
Predicate<String> excludedEndpoints) {
161205
long id = databaseId.get();
162206
if (id == 0) {
163207
return null;
164208
}
165209
hintBuilder.setDatabaseId(id);
166-
return rangeCache.fillRoutingHint(preferLeader, rangeMode, directedReadOptions, hintBuilder);
210+
return rangeCache.fillRoutingHint(
211+
preferLeader, rangeMode, directedReadOptions, hintBuilder, excludedEndpoints);
167212
}
168213

169214
private static boolean preferLeader(TransactionSelector selector) {

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

Lines changed: 21 additions & 0 deletions

0 commit comments

Comments
 (0)