fix: internalStream should handle duplicate onCompletes. (#1523)

googleapis · Jan 11, 2024 · b3067d7 · b3067d7
1 parent 4dd720a
commit b3067d7
Showing 1 changed file with 8 additions and 11 deletions.
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java
@@ -1477,8 +1477,6 @@ public void stream(@Nonnull final ApiStreamObserver<DocumentSnapshot> responseOb
 
     internalStream(
         new QuerySnapshotObserver() {
-          boolean hasCompleted = false;
-
           @Override
           public void onNext(QueryDocumentSnapshot documentSnapshot) {
             responseObserver.onNext(documentSnapshot);
@@ -1491,8 +1489,6 @@ public void onError(Throwable throwable) {
 
           @Override
           public void onCompleted() {
-            if (hasCompleted) return;
-            hasCompleted = true;
             responseObserver.onCompleted();
           }
         },
@@ -1660,6 +1656,10 @@ private void internalStream(
           boolean firstResponse;
           int numDocuments;
 
+          // The stream's `onComplete()` could be called more than once,
+          // this flag makes sure only the first one is actually processed.
+          boolean hasCompleted = false;
+
           @Override
           public void onStart(StreamController streamController) {}
 
@@ -1695,7 +1695,7 @@ public void onResponse(RunQueryResponse response) {
                       "Firestore.Query: Completed",
                       ImmutableMap.of(
                           "numDocuments", AttributeValue.longAttributeValue(numDocuments)));
-              documentObserver.onCompleted(readTime);
+              onComplete();
             }
           }
 
@@ -1723,6 +1723,9 @@ public void onError(Throwable throwable) {
 
           @Override
           public void onComplete() {
+            if (hasCompleted) return;
+            hasCompleted = true;
+
             Tracing.getTracer()
                 .getCurrentSpan()
                 .addAnnotation(
@@ -1788,9 +1791,6 @@ ApiFuture<QuerySnapshot> get(@Nullable ByteString transactionId) {
     internalStream(
         new QuerySnapshotObserver() {
           final List<QueryDocumentSnapshot> documentSnapshots = new ArrayList<>();
-          // The stream's onCompleted could be called more than once,
-          // this flag makes sure only the first one is actually processed.
-          boolean hasCompleted = false;
 
           @Override
           public void onNext(QueryDocumentSnapshot documentSnapshot) {
@@ -1804,9 +1804,6 @@ public void onError(Throwable throwable) {
 
           @Override
           public void onCompleted() {
-            if (hasCompleted) return;
-            hasCompleted = true;
-
             // The results for limitToLast queries need to be flipped since we reversed the
             // ordering constraints before sending the query to the backend.
             List<QueryDocumentSnapshot> resultView =