Skip to content

Commit 95da1ed

Browse files
olavloiteolavloite
authored
feat: add option for cancelling queries when closing client (#3276)
* feat: add option for cancelling queries when closing client Adds an option to keep track of all running queries and cancel these when the client is closed. * test: unflake test cases * fix: cancel query by default for batch transactions * chore: cleanup
1 parent c449a91 commit 95da1ed

File tree

11 files changed

+318
-12
lines changed

11 files changed

+318
-12
lines changed

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ abstract class AbstractReadContext
6969

7070
abstract static class Builder<B extends Builder<?, T>, T extends AbstractReadContext> {
7171
private SessionImpl session;
72+
private boolean cancelQueryWhenClientIsClosed;
7273
private SpannerRpc rpc;
7374
private ISpan span;
7475
private TraceWrapper tracer;
@@ -91,6 +92,11 @@ B setSession(SessionImpl session) {
9192
return self();
9293
}
9394

95+
B setCancelQueryWhenClientIsClosed(boolean cancelQueryWhenClientIsClosed) {
96+
this.cancelQueryWhenClientIsClosed = cancelQueryWhenClientIsClosed;
97+
return self();
98+
}
99+
94100
B setRpc(SpannerRpc rpc) {
95101
this.rpc = rpc;
96102
return self();
@@ -450,6 +456,7 @@ void initTransaction() {
450456

451457
final Object lock = new Object();
452458
final SessionImpl session;
459+
final boolean cancelQueryWhenClientIsClosed;
453460
final SpannerRpc rpc;
454461
final ExecutorProvider executorProvider;
455462
ISpan span;
@@ -479,6 +486,7 @@ void initTransaction() {
479486

480487
AbstractReadContext(Builder<?, ?> builder) {
481488
this.session = builder.session;
489+
this.cancelQueryWhenClientIsClosed = builder.cancelQueryWhenClientIsClosed;
482490
this.rpc = builder.rpc;
483491
this.defaultPrefetchChunks = builder.defaultPrefetchChunks;
484492
this.defaultQueryOptions = builder.defaultQueryOptions;
@@ -760,7 +768,8 @@ ResultSet executeQueryInternalWithOptions(
760768
rpc.getExecuteQueryRetryableCodes()) {
761769
@Override
762770
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
763-
GrpcStreamIterator stream = new GrpcStreamIterator(statement, prefetchChunks);
771+
GrpcStreamIterator stream =
772+
new GrpcStreamIterator(statement, prefetchChunks, cancelQueryWhenClientIsClosed);
764773
if (partitionToken != null) {
765774
request.setPartitionToken(partitionToken);
766775
}
@@ -943,7 +952,8 @@ ResultSet readInternalWithOptions(
943952
rpc.getReadRetryableCodes()) {
944953
@Override
945954
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
946-
GrpcStreamIterator stream = new GrpcStreamIterator(prefetchChunks);
955+
GrpcStreamIterator stream =
956+
new GrpcStreamIterator(prefetchChunks, cancelQueryWhenClientIsClosed);
947957
TransactionSelector selector = null;
948958
if (resumeToken != null) {
949959
builder.setResumeToken(resumeToken);

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) {
5454
return new BatchReadOnlyTransactionImpl(
5555
MultiUseReadOnlyTransaction.newBuilder()
5656
.setSession(session)
57+
.setCancelQueryWhenClientIsClosed(true)
5758
.setRpc(sessionClient.getSpanner().getRpc())
5859
.setTimestampBound(bound)
5960
.setDefaultQueryOptions(
@@ -75,6 +76,7 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batc
7576
return new BatchReadOnlyTransactionImpl(
7677
MultiUseReadOnlyTransaction.newBuilder()
7778
.setSession(session)
79+
.setCancelQueryWhenClientIsClosed(true)
7880
.setRpc(sessionClient.getSpanner().getRpc())
7981
.setTransactionId(batchTransactionId.getTransactionId())
8082
.setTimestamp(batchTransactionId.getTimestamp())

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

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
3838
private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName());
3939
private static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build();
4040

41-
private final ConsumerImpl consumer = new ConsumerImpl();
41+
private final ConsumerImpl consumer;
4242
private final BlockingQueue<PartialResultSet> stream;
4343
private final Statement statement;
4444

@@ -49,13 +49,15 @@ class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
4949
private SpannerException error;
5050

5151
@VisibleForTesting
52-
GrpcStreamIterator(int prefetchChunks) {
53-
this(null, prefetchChunks);
52+
GrpcStreamIterator(int prefetchChunks, boolean cancelQueryWhenClientIsClosed) {
53+
this(null, prefetchChunks, cancelQueryWhenClientIsClosed);
5454
}
5555

5656
@VisibleForTesting
57-
GrpcStreamIterator(Statement statement, int prefetchChunks) {
57+
GrpcStreamIterator(
58+
Statement statement, int prefetchChunks, boolean cancelQueryWhenClientIsClosed) {
5859
this.statement = statement;
60+
this.consumer = new ConsumerImpl(cancelQueryWhenClientIsClosed);
5961
// One extra to allow for END_OF_STREAM message.
6062
this.stream = new LinkedBlockingQueue<>(prefetchChunks + 1);
6163
}
@@ -136,6 +138,12 @@ private void addToStream(PartialResultSet results) {
136138
}
137139

138140
private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer {
141+
private final boolean cancelQueryWhenClientIsClosed;
142+
143+
ConsumerImpl(boolean cancelQueryWhenClientIsClosed) {
144+
this.cancelQueryWhenClientIsClosed = cancelQueryWhenClientIsClosed;
145+
}
146+
139147
@Override
140148
public void onPartialResultSet(PartialResultSet results) {
141149
addToStream(results);
@@ -168,5 +176,10 @@ public void onError(SpannerException e) {
168176
error = e;
169177
addToStream(END_OF_STREAM);
170178
}
179+
180+
@Override
181+
public boolean cancelQueryWhenClientIsClosed() {
182+
return this.cancelQueryWhenClientIsClosed;
183+
}
171184
}
172185
}

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
/*
2+
* Copyright 2024 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;
18+
19+
import com.google.api.gax.rpc.UnavailableException;
20+
import com.google.common.base.Predicate;
21+
import io.grpc.Status.Code;
22+
import io.grpc.StatusRuntimeException;
23+
24+
/**
25+
* Predicate that checks whether an exception is a ChannelShutdownException. This exception is
26+
* thrown by gRPC if the underlying gRPC stub has been shut down and uses the UNAVAILABLE error
27+
* code. This means that it would normally be retried by the Spanner client, but this specific
28+
* UNAVAILABLE error should not be retried, as it would otherwise directly return the same error.
29+
*/
30+
class IsChannelShutdownException implements Predicate<Throwable> {
31+
32+
@Override
33+
public boolean apply(Throwable input) {
34+
Throwable cause = input;
35+
do {
36+
if (isUnavailableError(cause)
37+
&& (cause.getMessage().contains("Channel shutdown invoked")
38+
|| cause.getMessage().contains("Channel shutdownNow invoked"))) {
39+
return true;
40+
}
41+
} while ((cause = cause.getCause()) != null);
42+
return false;
43+
}
44+
45+
private boolean isUnavailableError(Throwable cause) {
46+
return (cause instanceof UnavailableException)
47+
|| (cause instanceof StatusRuntimeException
48+
&& ((StatusRuntimeException) cause).getStatus().getCode() == Code.UNAVAILABLE);
49+
}
50+
}

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -333,7 +333,9 @@ private static boolean isRetryable(ErrorCode code, @Nullable Throwable cause) {
333333
case UNAVAILABLE:
334334
// SSLHandshakeException is (probably) not retryable, as it is an indication that the server
335335
// certificate was not accepted by the client.
336-
return !hasCauseMatching(cause, Matchers.isSSLHandshakeException);
336+
// Channel shutdown is also not a retryable exception.
337+
return !(hasCauseMatching(cause, Matchers.isSSLHandshakeException)
338+
|| hasCauseMatching(cause, Matchers.IS_CHANNEL_SHUTDOWN_EXCEPTION));
337339
case RESOURCE_EXHAUSTED:
338340
return SpannerException.extractRetryDelay(cause) > 0;
339341
default:
@@ -356,5 +358,8 @@ private static class Matchers {
356358

357359
static final Predicate<Throwable> isRetryableInternalError = new IsRetryableInternalError();
358360
static final Predicate<Throwable> isSSLHandshakeException = new IsSslHandshakeException();
361+
362+
static final Predicate<Throwable> IS_CHANNEL_SHUTDOWN_EXCEPTION =
363+
new IsChannelShutdownException();
359364
}
360365
}

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

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,7 @@
201201
import java.util.concurrent.Callable;
202202
import java.util.concurrent.CancellationException;
203203
import java.util.concurrent.ConcurrentHashMap;
204+
import java.util.concurrent.ConcurrentLinkedDeque;
204205
import java.util.concurrent.ConcurrentMap;
205206
import java.util.concurrent.ExecutionException;
206207
import java.util.concurrent.ExecutorService;
@@ -262,6 +263,9 @@ public class GapicSpannerRpc implements SpannerRpc {
262263

263264
private final ScheduledExecutorService spannerWatchdog;
264265

266+
private final ConcurrentLinkedDeque<SpannerResponseObserver> responseObservers =
267+
new ConcurrentLinkedDeque<>();
268+
265269
private final boolean throttleAdministrativeRequests;
266270
private final RetrySettings retryAdministrativeRequestsSettings;
267271
private static final double ADMINISTRATIVE_REQUESTS_RATE_LIMIT = 1.0D;
@@ -2004,9 +2008,29 @@ <ReqT, RespT> GrpcCallContext newCallContext(
20042008
return (GrpcCallContext) context.merge(apiCallContextFromContext);
20052009
}
20062010

2011+
void registerResponseObserver(SpannerResponseObserver responseObserver) {
2012+
responseObservers.add(responseObserver);
2013+
}
2014+
2015+
void unregisterResponseObserver(SpannerResponseObserver responseObserver) {
2016+
responseObservers.remove(responseObserver);
2017+
}
2018+
2019+
void closeResponseObservers() {
2020+
responseObservers.forEach(SpannerResponseObserver::close);
2021+
responseObservers.clear();
2022+
}
2023+
2024+
@InternalApi
2025+
@VisibleForTesting
2026+
public int getNumActiveResponseObservers() {
2027+
return responseObservers.size();
2028+
}
2029+
20072030
@Override
20082031
public void shutdown() {
20092032
this.rpcIsClosed = true;
2033+
closeResponseObservers();
20102034
if (this.spannerStub != null) {
20112035
this.spannerStub.close();
20122036
this.partitionedDmlStub.close();
@@ -2028,6 +2052,7 @@ public void shutdown() {
20282052

20292053
public void shutdownNow() {
20302054
this.rpcIsClosed = true;
2055+
closeResponseObservers();
20312056
this.spannerStub.close();
20322057
this.partitionedDmlStub.close();
20332058
this.instanceAdminStub.close();
@@ -2085,7 +2110,7 @@ public void cancel(@Nullable String message) {
20852110
* A {@code ResponseObserver} that exposes the {@code StreamController} and delegates callbacks to
20862111
* the {@link ResultStreamConsumer}.
20872112
*/
2088-
private static class SpannerResponseObserver implements ResponseObserver<PartialResultSet> {
2113+
private class SpannerResponseObserver implements ResponseObserver<PartialResultSet> {
20892114

20902115
private StreamController controller;
20912116
private final ResultStreamConsumer consumer;
@@ -2094,13 +2119,21 @@ public SpannerResponseObserver(ResultStreamConsumer consumer) {
20942119
this.consumer = consumer;
20952120
}
20962121

2122+
void close() {
2123+
if (this.controller != null) {
2124+
this.controller.cancel();
2125+
}
2126+
}
2127+
20972128
@Override
20982129
public void onStart(StreamController controller) {
2099-
21002130
// Disable the auto flow control to allow client library
21012131
// set the number of messages it prefers to request
21022132
controller.disableAutoInboundFlowControl();
21032133
this.controller = controller;
2134+
if (this.consumer.cancelQueryWhenClientIsClosed()) {
2135+
registerResponseObserver(this);
2136+
}
21042137
}
21052138

21062139
@Override
@@ -2110,11 +2143,19 @@ public void onResponse(PartialResultSet response) {
21102143

21112144
@Override
21122145
public void onError(Throwable t) {
2146+
// Unregister the response observer when the query has completed with an error.
2147+
if (this.consumer.cancelQueryWhenClientIsClosed()) {
2148+
unregisterResponseObserver(this);
2149+
}
21132150
consumer.onError(newSpannerException(t));
21142151
}
21152152

21162153
@Override
21172154
public void onComplete() {
2155+
// Unregister the response observer when the query has completed normally.
2156+
if (this.consumer.cancelQueryWhenClientIsClosed()) {
2157+
unregisterResponseObserver(this);
2158+
}
21182159
consumer.onCompleted();
21192160
}
21202161

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,15 @@ interface ResultStreamConsumer {
153153
void onCompleted();
154154

155155
void onError(SpannerException e);
156+
157+
/**
158+
* Returns true if the stream should be cancelled when the Spanner client is closed. This
159+
* returns true for {@link com.google.cloud.spanner.BatchReadOnlyTransaction}, as these use a
160+
* non-pooled session. Pooled sessions are deleted when the Spanner client is closed, and this
161+
* automatically also cancels any query that uses the session, which means that we don't need to
162+
* explicitly cancel those queries when the Spanner client is closed.
163+
*/
164+
boolean cancelQueryWhenClientIsClosed();
156165
}
157166

158167
/** Handle for cancellation of a streaming read or query call. */

0 commit comments

Comments
 (0)