Skip to content

Commit 7096899

Browse files
sakthivelmaniisakthivelmanii
committed
fix: Avoid blocking thread usage in AsyncResultSet
1 parent 11ead4e commit 7096899

File tree

13 files changed

+232
-77
lines changed

13 files changed

+232
-77
lines changed

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -768,9 +768,11 @@ ResultSet executeQueryInternalWithOptions(
768768
rpc.getExecuteQueryRetrySettings(),
769769
rpc.getExecuteQueryRetryableCodes()) {
770770
@Override
771-
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
771+
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken,
772+
AsyncResultSet.StreamMessageListener streamListener) {
772773
GrpcStreamIterator stream =
773774
new GrpcStreamIterator(statement, prefetchChunks, cancelQueryWhenClientIsClosed);
775+
stream.registerListener(streamListener);
774776
if (partitionToken != null) {
775777
request.setPartitionToken(partitionToken);
776778
}
@@ -791,8 +793,8 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
791793
getTransactionChannelHint(),
792794
isRouteToLeader());
793795
session.markUsed(clock.instant());
794-
call.request(prefetchChunks);
795796
stream.setCall(call, request.getTransaction().hasBegin());
797+
call.request(prefetchChunks);
796798
return stream;
797799
}
798800

@@ -959,9 +961,11 @@ ResultSet readInternalWithOptions(
959961
rpc.getReadRetrySettings(),
960962
rpc.getReadRetryableCodes()) {
961963
@Override
962-
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
964+
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken,
965+
AsyncResultSet.StreamMessageListener streamListener) {
963966
GrpcStreamIterator stream =
964967
new GrpcStreamIterator(prefetchChunks, cancelQueryWhenClientIsClosed);
968+
stream.registerListener(streamListener);
965969
TransactionSelector selector = null;
966970
if (resumeToken != null) {
967971
builder.setResumeToken(resumeToken);
@@ -980,8 +984,8 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
980984
getTransactionChannelHint(),
981985
isRouteToLeader());
982986
session.markUsed(clock.instant());
983-
call.request(prefetchChunks);
984987
stream.setCall(call, /* withBeginTransaction = */ builder.getTransaction().hasBegin());
988+
call.request(prefetchChunks);
985989
return stream;
986990
}
987991

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,10 @@ interface CloseableIterator<T> extends Iterator<T> {
150150
void close(@Nullable String message);
151151

152152
boolean isWithBeginTransaction();
153+
154+
default boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) {
155+
return false;
156+
}
153157
}
154158

155159
static double valueProtoToFloat64(com.google.protobuf.Value proto) {

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818

1919
import com.google.api.core.ApiFuture;
2020
import com.google.common.base.Function;
21+
import com.google.spanner.v1.PartialResultSet;
22+
2123
import java.util.List;
2224
import java.util.concurrent.ExecutionException;
2325
import java.util.concurrent.Executor;
@@ -223,4 +225,20 @@ interface ReadyCallback {
223225
* @param transformer function which will be used to transform the row. It should not return null.
224226
*/
225227
<T> List<T> toList(Function<StructReader, T> transformer) throws SpannerException;
228+
229+
/**
230+
* An interface to register the listener for streaming gRPC request. It will be called when a chunk is received
231+
* from gRPC streaming call.
232+
*/
233+
interface StreamMessageListener {
234+
void onStreamMessage(PartialResultSet partialResultSet, int prefetchChunks, int currentBufferSize, StreamMessageRequestor streamMessageRequestor);
235+
}
236+
237+
/**
238+
* An interface to request more messages from the gRPC streaming call. It will be implemented by the class which has access
239+
* to SpannerRpc.StreamingCall object
240+
*/
241+
interface StreamMessageRequestor {
242+
void requestMessages(int numOfMessages);
243+
}
226244
}

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

Lines changed: 70 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818

1919
import com.google.api.core.ApiFuture;
2020
import com.google.api.core.ApiFutures;
21-
import com.google.api.core.ListenableFutureToApiFuture;
2221
import com.google.api.core.SettableApiFuture;
2322
import com.google.api.gax.core.ExecutorProvider;
2423
import com.google.cloud.spanner.AbstractReadContext.ListenableAsyncResultSet;
@@ -29,28 +28,26 @@
2928
import com.google.common.collect.ImmutableList;
3029
import com.google.common.util.concurrent.ListeningScheduledExecutorService;
3130
import com.google.common.util.concurrent.MoreExecutors;
31+
import com.google.spanner.v1.PartialResultSet;
3232
import com.google.spanner.v1.ResultSetMetadata;
3333
import com.google.spanner.v1.ResultSetStats;
34+
3435
import java.util.Collection;
3536
import java.util.LinkedList;
3637
import java.util.List;
37-
import java.util.concurrent.BlockingDeque;
38-
import java.util.concurrent.Callable;
39-
import java.util.concurrent.CountDownLatch;
40-
import java.util.concurrent.ExecutionException;
41-
import java.util.concurrent.Executor;
42-
import java.util.concurrent.Future;
43-
import java.util.concurrent.LinkedBlockingDeque;
38+
import java.util.concurrent.*;
4439
import java.util.logging.Level;
4540
import java.util.logging.Logger;
4641

4742
/** Default implementation for {@link AsyncResultSet}. */
48-
class AsyncResultSetImpl extends ForwardingStructReader implements ListenableAsyncResultSet {
43+
class AsyncResultSetImpl extends ForwardingStructReader implements ListenableAsyncResultSet, AsyncResultSet.StreamMessageListener {
4944
private static final Logger log = Logger.getLogger(AsyncResultSetImpl.class.getName());
5045

46+
5147
/** State of an {@link AsyncResultSetImpl}. */
5248
private enum State {
5349
INITIALIZED,
50+
IN_PROGRESS,
5451
/** SYNC indicates that the {@link ResultSet} is used in sync pattern. */
5552
SYNC,
5653
CONSUMING,
@@ -120,7 +117,7 @@ private enum State {
120117
*/
121118
private volatile boolean finished;
122119

123-
private volatile ApiFuture<Void> result;
120+
private volatile SettableApiFuture<Void> result;
124121

125122
/**
126123
* This variable indicates whether {@link #tryNext()} has returned {@link CursorState#DONE} or a
@@ -329,12 +326,12 @@ public void run() {
329326
private final CallbackRunnable callbackRunnable = new CallbackRunnable();
330327

331328
/**
332-
* {@link ProduceRowsCallable} reads data from the underlying {@link ResultSet}, places these in
329+
* {@link ProduceRowsRunnable} reads data from the underlying {@link ResultSet}, places these in
333330
* the buffer and dispatches the {@link CallbackRunnable} when data is ready to be consumed.
334331
*/
335-
private class ProduceRowsCallable implements Callable<Void> {
332+
private class ProduceRowsRunnable implements Runnable {
336333
@Override
337-
public Void call() throws Exception {
334+
public void run() {
338335
boolean stop = false;
339336
boolean hasNext = false;
340337
try {
@@ -393,12 +390,17 @@ public Void call() throws Exception {
393390
}
394391
// Call the callback if there are still rows in the buffer that need to be processed.
395392
while (!stop) {
396-
waitIfPaused();
397-
startCallbackIfNecessary();
398-
// Make sure we wait until the callback runner has actually finished.
399-
consumingLatch.await();
400-
synchronized (monitor) {
401-
stop = cursorReturnedDoneOrException;
393+
try {
394+
waitIfPaused();
395+
startCallbackIfNecessary();
396+
// Make sure we wait until the callback runner has actually finished.
397+
consumingLatch.await();
398+
synchronized (monitor) {
399+
stop = cursorReturnedDoneOrException;
400+
}
401+
} catch (InterruptedException e) {
402+
result.setException(e);
403+
return;
402404
}
403405
}
404406
} finally {
@@ -410,14 +412,16 @@ public Void call() throws Exception {
410412
}
411413
synchronized (monitor) {
412414
if (executionException != null) {
415+
result.setException(executionException);
413416
throw executionException;
414417
}
415418
if (state == State.CANCELLED) {
419+
result.setException(CANCELLED_EXCEPTION);
416420
throw CANCELLED_EXCEPTION;
417421
}
418422
}
419423
}
420-
return null;
424+
result.set(null);
421425
}
422426

423427
private void waitIfPaused() throws InterruptedException {
@@ -449,6 +453,21 @@ private void startCallbackWithBufferLatchIfNecessary(int bufferLatch) {
449453
}
450454
}
451455

456+
private class InitiateStreamingRunnable implements Runnable {
457+
458+
@Override
459+
public void run() {
460+
try {
461+
if(!initiateStreaming(AsyncResultSetImpl.this)) {
462+
initiateProduceRows();
463+
}
464+
} catch (SpannerException e) {
465+
executionException = e;
466+
initiateProduceRows();
467+
}
468+
}
469+
}
470+
452471
/** Sets the callback for this {@link AsyncResultSet}. */
453472
@Override
454473
public ApiFuture<Void> setCallback(Executor exec, ReadyCallback cb) {
@@ -458,16 +477,21 @@ public ApiFuture<Void> setCallback(Executor exec, ReadyCallback cb) {
458477
this.state == State.INITIALIZED, "callback may not be set multiple times");
459478

460479
// Start to fetch data and buffer these.
461-
this.result =
462-
new ListenableFutureToApiFuture<>(this.service.submit(new ProduceRowsCallable()));
480+
this.result = SettableApiFuture.create();
481+
this.state = State.IN_PROGRESS;
482+
this.service.execute(new InitiateStreamingRunnable());
463483
this.executor = MoreExecutors.newSequentialExecutor(Preconditions.checkNotNull(exec));
464484
this.callback = Preconditions.checkNotNull(cb);
465-
this.state = State.RUNNING;
466485
pausedLatch.countDown();
467486
return result;
468487
}
469488
}
470489

490+
private void initiateProduceRows() {
491+
this.service.execute(new ProduceRowsRunnable());
492+
this.state = State.RUNNING;
493+
}
494+
471495
Future<Void> getResult() {
472496
return result;
473497
}
@@ -480,6 +504,7 @@ public void cancel() {
480504
"cannot cancel a result set without a callback");
481505
state = State.CANCELLED;
482506
pausedLatch.countDown();
507+
this.result.setException(CANCELLED_EXCEPTION);
483508
}
484509
}
485510

@@ -578,6 +603,11 @@ public ResultSetMetadata getMetadata() {
578603
return delegateResultSet.get().getMetadata();
579604
}
580605

606+
@Override
607+
public boolean initiateStreaming(StreamMessageListener streamMessageListener) {
608+
return delegateResultSet.get().initiateStreaming(streamMessageListener);
609+
}
610+
581611
@Override
582612
protected void checkValidState() {
583613
synchronized (monitor) {
@@ -593,4 +623,21 @@ public Struct getCurrentRowAsStruct() {
593623
checkValidState();
594624
return currentRow;
595625
}
626+
627+
@Override
628+
public void onStreamMessage(PartialResultSet partialResultSet, int prefetchChunks, int currentBufferSize, StreamMessageRequestor streamMessageRequestor) {
629+
synchronized (monitor) {
630+
if (state == State.IN_PROGRESS) {
631+
// if PartialResultSet contains resume token or buffer size is more than configured size or we have reached
632+
// end of stream, we can start the thread
633+
boolean startJobThread = !partialResultSet.getResumeToken().isEmpty()
634+
|| currentBufferSize > prefetchChunks || partialResultSet == GrpcStreamIterator.END_OF_STREAM;
635+
if (startJobThread){
636+
initiateProduceRows();
637+
} else {
638+
streamMessageRequestor.requestMessages(1);
639+
}
640+
}
641+
}
642+
}
596643
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,4 +102,9 @@ public ResultSetStats getStats() {
102102
public ResultSetMetadata getMetadata() {
103103
return delegate.get().getMetadata();
104104
}
105+
106+
@Override
107+
public boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) {
108+
return delegate.get().initiateStreaming(streamMessageListener);
109+
}
105110
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,11 @@ public ResultSetMetadata getMetadata() {
123123
return metadata;
124124
}
125125

126+
@Override
127+
public boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) {
128+
return iterator.initiateStreaming(streamMessageListener);
129+
}
130+
126131
@Override
127132
public void close() {
128133
synchronized (this) {

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

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,24 +23,28 @@
2323
import com.google.common.collect.AbstractIterator;
2424
import com.google.common.util.concurrent.Uninterruptibles;
2525
import com.google.spanner.v1.PartialResultSet;
26+
import org.threeten.bp.Duration;
27+
28+
import javax.annotation.Nullable;
29+
import java.util.Optional;
2630
import java.util.concurrent.BlockingQueue;
2731
import java.util.concurrent.LinkedBlockingQueue;
2832
import java.util.concurrent.TimeUnit;
2933
import java.util.logging.Level;
3034
import java.util.logging.Logger;
31-
import javax.annotation.Nullable;
32-
import org.threeten.bp.Duration;
3335

3436
/** Adapts a streaming read/query call into an iterator over partial result sets. */
3537
@VisibleForTesting
3638
class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
37-
implements CloseableIterator<PartialResultSet> {
39+
implements CloseableIterator<PartialResultSet>, AsyncResultSet.StreamMessageRequestor {
3840
private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName());
39-
private static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build();
41+
public static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build();
42+
private AsyncResultSet.StreamMessageListener streamMessageListener;
4043

4144
private final ConsumerImpl consumer;
4245
private final BlockingQueue<PartialResultSet> stream;
4346
private final Statement statement;
47+
private final int prefetchChunks;
4448

4549
private SpannerRpc.StreamingCall call;
4650
private volatile boolean withBeginTransaction;
@@ -57,15 +61,20 @@ class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
5761
GrpcStreamIterator(
5862
Statement statement, int prefetchChunks, boolean cancelQueryWhenClientIsClosed) {
5963
this.statement = statement;
64+
this.prefetchChunks = prefetchChunks;
6065
this.consumer = new ConsumerImpl(cancelQueryWhenClientIsClosed);
6166
// One extra to allow for END_OF_STREAM message.
62-
this.stream = new LinkedBlockingQueue<>(prefetchChunks + 1);
67+
this.stream = new LinkedBlockingQueue<>((prefetchChunks * 2) + 1);
6368
}
6469

6570
protected final SpannerRpc.ResultStreamConsumer consumer() {
6671
return consumer;
6772
}
6873

74+
public void registerListener(AsyncResultSet.StreamMessageListener streamMessageListener) {
75+
this.streamMessageListener = streamMessageListener;
76+
}
77+
6978
public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) {
7079
this.call = call;
7180
this.withBeginTransaction = withBeginTransaction;
@@ -135,6 +144,12 @@ protected final PartialResultSet computeNext() {
135144
private void addToStream(PartialResultSet results) {
136145
// We assume that nothing from the user will interrupt gRPC event threads.
137146
Uninterruptibles.putUninterruptibly(stream, results);
147+
onStreamMessage(results);
148+
}
149+
150+
@Override
151+
public void requestMessages(int numOfMessages) {
152+
call.request(numOfMessages);
138153
}
139154

140155
private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer {
@@ -182,4 +197,9 @@ public boolean cancelQueryWhenClientIsClosed() {
182197
return this.cancelQueryWhenClientIsClosed;
183198
}
184199
}
200+
201+
private void onStreamMessage(PartialResultSet partialResultSet) {
202+
Optional.ofNullable(streamMessageListener).ifPresent(
203+
sl -> sl.onStreamMessage(partialResultSet, prefetchChunks, stream.size(), this));
204+
}
185205
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,10 @@ ResultSetStats getStats() {
127127
return statistics;
128128
}
129129

130+
boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) {
131+
return stream.initiateStreaming(streamMessageListener);
132+
}
133+
130134
Type type() {
131135
checkState(type != null, "metadata has not been received");
132136
return type;

0 commit comments

Comments
 (0)