fix(spanner): unflake unit tests

googleapis · Jan 6, 2025 · b32f148 · b32f148
1 parent 6225efa
commit b32f148
Showing 1 changed file with 13 additions and 10 deletions.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java
@@ -65,10 +65,11 @@ abstract class ResumableStreamIterator extends AbstractIterator<PartialResultSet
   private static final Logger logger = Logger.getLogger(ResumableStreamIterator.class.getName());
   private BackOff backOff;
   private final LinkedList<PartialResultSet> buffer = new LinkedList<>();
+  private final Object monitor = new Object();
   private final int maxBufferSize;
   private final ISpan span;
   private final TraceWrapper tracer;
-  private CloseableIterator<PartialResultSet> stream;
+  private volatile CloseableIterator<PartialResultSet> stream;
   private ByteString resumeToken;
   private boolean finished;
   /**
@@ -317,15 +318,17 @@ && prepareIteratorForRetryOnDifferentGrpcChannel()) {
   }
 
   private void startGrpcStreaming() {
-    if (stream == null) {
-      span.addAnnotation(
-          "Starting/Resuming stream",
-          "ResumeToken",
-          resumeToken == null ? "null" : resumeToken.toStringUtf8());
-      try (IScope scope = tracer.withSpan(span)) {
-        // When start a new stream set the Span as current to make the gRPC Span a child of
-        // this Span.
-        stream = checkNotNull(startStream(resumeToken, streamMessageListener));
+    synchronized (monitor) {
+      if (stream == null) {
+        span.addAnnotation(
+            "Starting/Resuming stream",
+            "ResumeToken",
+            resumeToken == null ? "null" : resumeToken.toStringUtf8());
+        try (IScope scope = tracer.withSpan(span)) {
+          // When start a new stream set the Span as current to make the gRPC Span a child of
+          // this Span.
+          stream = checkNotNull(startStream(resumeToken, streamMessageListener));
+        }
       }
     }
   }