Skip to content

fix(spanner): prevent stream leaks and race conditions on result set close #14665

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
olavloite wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(spanner): prevent stream leaks and race conditions on result set close #14665

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
3 changes: 2 additions & 1 deletion java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,8 +140,9 @@ public void close() {
}
closed = true;
}
listener.onDone(iterator.isWithBeginTransaction());
boolean withBeginTransaction = iterator.isWithBeginTransaction();
iterator.close("ResultSet closed");
listener.onDone(withBeginTransaction);
}

@Override
Expand Down
178 changes: 142 additions & 36 deletions ...anner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import com.google.common.util.concurrent.Uninterruptibles;
import com.google.spanner.v1.PartialResultSet;
import java.time.Duration;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
Expand All @@ -38,6 +39,40 @@
@VisibleForTesting
class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
implements CloseableIterator<PartialResultSet> {

@VisibleForTesting
static final class Timeout {
final long duration;
final TimeUnit unit;

Timeout(long duration, TimeUnit unit) {
this.duration = duration;
this.unit = Preconditions.checkNotNull(unit);
}

@Override
public boolean equals(Object o) {
if (this == o) {
return true;
}
if (!(o instanceof Timeout)) {
return false;
}
Timeout that = (Timeout) o;
return this.duration == that.duration && this.unit == that.unit;
}

@Override
public int hashCode() {
return Objects.hash(duration, unit);
}

@Override
public String toString() {
return duration + " " + unit;
}
}

private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName());
static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build();
private final int prefetchChunks;
Expand All @@ -47,13 +82,15 @@ class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
private final BlockingQueue<PartialResultSet> stream;
private final Statement statement;

private SpannerRpc.StreamingCall call;
private volatile SpannerRpc.StreamingCall call;
private boolean callCleared;
private boolean closed;
@Nullable private String closeMessage;
private volatile boolean withBeginTransaction;
private final boolean lastStatement;
private TimeUnit streamWaitTimeoutUnit;
private long streamWaitTimeoutValue;
private SpannerException error;
private boolean done;
@Nullable private volatile Timeout streamWaitTimeout;
@Nullable private volatile SpannerException error;
private volatile boolean done;

@VisibleForTesting
GrpcStreamIterator(
Expand Down Expand Up @@ -84,8 +121,28 @@ void registerListener(AsyncResultSet.StreamMessageListener streamMessageListener
}

public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) {
this.call = call;
this.withBeginTransaction = withBeginTransaction;
Timeout timeout = extractTimeout(call);
boolean shouldCancel;
String messageToCancel;
synchronized (this) {
this.withBeginTransaction = withBeginTransaction;
shouldCancel = this.closed && !this.callCleared;
if (!this.callCleared && !this.closed) {
this.call = call;
}
if (timeout != null) {
this.streamWaitTimeout = timeout;
}
messageToCancel = this.closeMessage;
}
if (shouldCancel) {
call.cancel(messageToCancel);
}
}

@VisibleForTesting
@Nullable
static Timeout extractTimeout(SpannerRpc.StreamingCall call) {
ApiCallContext callContext = call.getCallContext();
Duration streamWaitTimeout =
callContext == null ? null : callContext.getStreamWaitTimeoutDuration();
Expand All @@ -94,29 +151,52 @@ public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction)
// value is more than 1 second, which is lower than the precision that would normally be
// used by the stream watchdog (which uses a precision of 10 seconds by default).
if (streamWaitTimeout.getSeconds() > 0L) {
streamWaitTimeoutValue = streamWaitTimeout.getSeconds();
streamWaitTimeoutUnit = TimeUnit.SECONDS;
return new Timeout(streamWaitTimeout.getSeconds(), TimeUnit.SECONDS);
} else if (streamWaitTimeout.getNano() > 0) {
streamWaitTimeoutValue = streamWaitTimeout.getNano();
streamWaitTimeoutUnit = TimeUnit.NANOSECONDS;
return new Timeout(streamWaitTimeout.getNano(), TimeUnit.NANOSECONDS);
}
// Note that if the stream-wait-timeout is zero, we won't set a timeout at all.
// That is consistent with ApiCallContext#withStreamWaitTimeout(Duration.ZERO).
}
return null;
}

@Override
public void close(@Nullable String message) {
if (call != null) {
call.cancel(message);
SpannerRpc.StreamingCall callToCancel = null;
synchronized (this) {
if (!closed) {
closed = true;
closeMessage = message;
if (!callCleared) {
callToCancel = this.call;
this.call = null;
}
}
}
if (callToCancel != null) {
callToCancel.cancel(message);
}
}

@Override
@InternalApi
public void requestPrefetchChunks() {
Preconditions.checkState(call != null, "The StreamingCall object is not initialized");
call.request(prefetchChunks);
SpannerRpc.StreamingCall callToRequest = null;
synchronized (this) {
if (this.call != null) {
callToRequest = this.call;
} else if (!this.closed && !this.callCleared) {
// This should never happen during normal operation because startStream(...) always
// initializes the call via setCall(...) before returning the GrpcStreamIterator and
// calling requestPrefetchChunks(). Reaching this point means requestPrefetchChunks()
// was invoked before setCall(...).
throw new IllegalStateException("The StreamingCall object is not initialized");
}
}
if (callToRequest != null) {
callToRequest.request(prefetchChunks);
}
}

@Override
Expand All @@ -131,35 +211,51 @@ public boolean isLastStatement() {

@Override
protected final PartialResultSet computeNext() {
PartialResultSet next;
PartialResultSet next = takeNextFromStream();
if (next != END_OF_STREAM) {
SpannerRpc.StreamingCall currentCall = this.call;
if (currentCall != null) {
currentCall.request(1);
}
if (next.getLast()) {
clearCall();
}
return next;
}

clearCall();
if (error != null) {
throw SpannerExceptionFactory.asSpannerException(error);
}
endOfData();
return null;
}

private PartialResultSet takeNextFromStream() {
try {
if (streamWaitTimeoutUnit != null) {
next = stream.poll(streamWaitTimeoutValue, streamWaitTimeoutUnit);
Timeout timeout = this.streamWaitTimeout;
if (timeout != null) {
PartialResultSet next = stream.poll(timeout.duration, timeout.unit);
if (next == null) {
throw SpannerExceptionFactory.newSpannerException(
ErrorCode.DEADLINE_EXCEEDED, "stream wait timeout");
}
return next;
} else {
next = stream.take();
return stream.take();
}
} catch (InterruptedException e) {
// Treat interrupt as a request to cancel the read.
throw SpannerExceptionFactory.propagateInterrupt(e);
}
if (next != END_OF_STREAM) {
call.request(1);
return next;
}
}

private void clearCall() {
// All done - close() no longer needs to cancel the call.
call = null;

if (error != null) {
throw SpannerExceptionFactory.asSpannerException(error);
synchronized (this) {
callCleared = true;
call = null;
}

endOfData();
return null;
}

private void addToStream(PartialResultSet results) {
Expand All @@ -179,16 +275,13 @@ private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer {
public void onPartialResultSet(PartialResultSet results) {
addToStream(results);
if (results.getLast()) {
done = true;
addToStream(END_OF_STREAM);
addEndOfStream();
}
}

@Override
public void onCompleted() {
if (!done) {
addToStream(END_OF_STREAM);
}
addEndOfStream();
}

@Override
Expand All @@ -203,7 +296,20 @@ public void onError(SpannerException exception) {
}
}
error = exception;
addToStream(END_OF_STREAM);
addEndOfStream();
}

private void addEndOfStream() {
boolean shouldAddEndOfStream = false;
synchronized (GrpcStreamIterator.this) {
if (!done) {
done = true;
shouldAddEndOfStream = true;
}
}
if (shouldAddEndOfStream) {
addToStream(END_OF_STREAM);
}
}

@Override
Expand Down
73 changes: 70 additions & 3 deletions ...oud-spanner/src/test/java/com/google/cloud/spanner/CloseSpannerWithOpenResultSetTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,14 @@
public class CloseSpannerWithOpenResultSetTest extends AbstractMockServerTest {

Spanner createSpanner() {
return spannerOptionsBuilder().build().getService();
}

Spanner createSpanner(int prefetchChunks) {
return spannerOptionsBuilder().setPrefetchChunks(prefetchChunks).build().getService();
}

private SpannerOptions.Builder spannerOptionsBuilder() {
return SpannerOptions.newBuilder()
.setProjectId("p")
.setHost(String.format("http://localhost:%d", getPort()))
Expand All @@ -54,9 +62,7 @@ Spanner createSpanner() {
.setSessionPoolOption(
SessionPoolOptions.newBuilder()
.setWaitForMinSessionsDuration(Duration.ofSeconds(5L))
.build())
.build()
.getService();
.build());
}

@BeforeClass
Expand Down Expand Up @@ -122,4 +128,65 @@ public void testStreamsAreCleanedUp() throws Exception {
0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers());
}
}

@Test(timeout = 30000)
public void testQueryWithPrefetchChunksOne() {
try (Spanner spanner = createSpanner(1)) {
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
for (int i = 0; i < 10; i++) {
try (ResultSet resultSet = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) {
int count = 0;
while (resultSet.next()) {
assertNotNull(resultSet.getCurrentRowAsStruct());
count++;
}
assertTrue(count > 0);
}
}
assertEquals(
0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers());
}
}

@Test(timeout = 30000)
public void testQueryWithPrefetchChunksOneAndError() {
String errorSql = "select * from will_fail";
Statement errorStatement = Statement.of(errorSql);
mockSpanner.putStatementResult(
StatementResult.exception(
errorStatement,
Status.INVALID_ARGUMENT.withDescription("Invalid syntax").asRuntimeException()));
try (Spanner spanner = createSpanner(1)) {
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
try (ResultSet resultSet = client.singleUse().executeQuery(errorStatement)) {
SpannerException exception = assertThrows(SpannerException.class, resultSet::next);
assertEquals(ErrorCode.INVALID_ARGUMENT, exception.getErrorCode());
}
assertEquals(
0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers());
}
}

@Test(timeout = 30000)
public void testCloseResultSetCleansUpResponseObservers() {
try (Spanner spanner = createSpanner()) {
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
// 1. Partial consumption then close:
try (ResultSet resultSet = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) {
assertTrue(resultSet.next());
assertNotNull(resultSet.getCurrentRowAsStruct());
}
assertEquals(
0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers());

// 2. Full consumption then close:
try (ResultSet resultSet = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) {
while (resultSet.next()) {
assertNotNull(resultSet.getCurrentRowAsStruct());
}
}
assertEquals(
0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers());
}
}
}
Loading
Loading