SDK support for logical termination in Firestore BatchGetDocuments (#3622)

* SDK support for logical termination in Firestore BatchGetDocuments

Author: cherylEnkidu

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/Datastore.java

@@ -14,10 +14,13 @@
 
 package com.google.firebase.firestore.remote;
 
+import static com.google.firebase.firestore.util.Util.exceptionFromStatus;
+
 import android.content.Context;
 import android.os.Build;
 import androidx.annotation.Nullable;
 import com.google.android.gms.tasks.Task;
+import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.firebase.firestore.FirebaseFirestoreException;
 import com.google.firebase.firestore.auth.CredentialsProvider;
 import com.google.firebase.firestore.auth.User;
@@ -36,6 +39,7 @@
 import io.grpc.Status;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -170,21 +174,18 @@ public Task<List<MutableDocument>> lookup(List<DocumentKey> keys) {
     for (DocumentKey key : keys) {
       builder.addDocuments(serializer.encodeKey(key));
     }
-    return channel
-        .runStreamingResponseRpc(FirestoreGrpc.getBatchGetDocumentsMethod(), builder.build())
-        .continueWith(
-            workerQueue.getExecutor(),
-            task -> {
-              if (!task.isSuccessful()) {
-                if (task.getException() instanceof FirebaseFirestoreException
-                    && ((FirebaseFirestoreException) task.getException()).getCode()
-                        == FirebaseFirestoreException.Code.UNAUTHENTICATED) {
-                  channel.invalidateToken();
-                }
-              }
+    List<BatchGetDocumentsResponse> responses = new ArrayList<>();
+    TaskCompletionSource<List<MutableDocument>> completionSource = new TaskCompletionSource<>();
 
+    channel.runStreamingResponseRpc(
+        FirestoreGrpc.getBatchGetDocumentsMethod(),
+        builder.build(),
+        new FirestoreChannel.StreamingListener<BatchGetDocumentsResponse>() {
+          @Override
+          public void onMessage(BatchGetDocumentsResponse message) {
+            responses.add(message);
+            if (responses.size() == keys.size()) {
               Map<DocumentKey, MutableDocument> resultMap = new HashMap<>();
-              List<BatchGetDocumentsResponse> responses = task.getResult();
               for (BatchGetDocumentsResponse response : responses) {
                 MutableDocument doc = serializer.decodeMaybeDocument(response);
                 resultMap.put(doc.getKey(), doc);
@@ -193,8 +194,25 @@ public Task<List<MutableDocument>> lookup(List<DocumentKey> keys) {
               for (DocumentKey key : keys) {
                 results.add(resultMap.get(key));
               }
-              return results;
-            });
+              completionSource.trySetResult(results);
+            }
+          }
+
+          @Override
+          public void onClose(Status status) {
+            if (status.isOk()) {
+              completionSource.trySetResult(Collections.emptyList());
+            } else {
+              FirebaseFirestoreException exception = exceptionFromStatus(status);
+              if (exception.getCode() == FirebaseFirestoreException.Code.UNAUTHENTICATED) {
+                channel.invalidateToken();
+              }
+              completionSource.trySetException(exception);
+            }
+          }
+        });
+
+    return completionSource.getTask();
   }
 
   /**

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/FirestoreChannel.java

@@ -33,8 +33,6 @@
 import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
 import io.grpc.Status;
-import java.util.ArrayList;
-import java.util.List;
 
 /**
  * Wrapper class around io.grpc.Channel that adds headers, exception handling and simplifies
@@ -44,6 +42,15 @@
  */
 public class FirestoreChannel {
 
+  /** Listen to changes inside runStreamingResponseRpc */
+  public abstract static class StreamingListener<T> {
+    public StreamingListener() {}
+
+    public void onMessage(T message) {}
+
+    public void onClose(Status status) {}
+  }
+
   private static final Metadata.Key<String> X_GOOG_API_CLIENT_HEADER =
       Metadata.Key.of("x-goog-api-client", Metadata.ASCII_STRING_MARSHALLER);
 
@@ -184,36 +191,28 @@ public void halfClose() {
   }
 
   /** Creates and starts a streaming response RPC. */
-  <ReqT, RespT> Task<List<RespT>> runStreamingResponseRpc(
-      MethodDescriptor<ReqT, RespT> method, ReqT request) {
-    TaskCompletionSource<List<RespT>> tcs = new TaskCompletionSource<>();
-
+  <ReqT, RespT> void runStreamingResponseRpc(
+      MethodDescriptor<ReqT, RespT> method,
+      ReqT request,
+      FirestoreChannel.StreamingListener<RespT> callback) {
     callProvider
         .createClientCall(method)
         .addOnCompleteListener(
             asyncQueue.getExecutor(),
             result -> {
               ClientCall<ReqT, RespT> call = result.getResult();
-
-              List<RespT> results = new ArrayList<>();
-
               call.start(
                   new ClientCall.Listener<RespT>() {
                     @Override
                     public void onMessage(RespT message) {
-                      results.add(message);
-
+                      callback.onMessage(message);
                       // Make sure next message can be delivered
                       call.request(1);
                     }
 
                     @Override
                     public void onClose(Status status, Metadata trailers) {
-                      if (status.isOk()) {
-                        tcs.setResult(results);
-                      } else {
-                        tcs.setException(exceptionFromStatus(status));
-                      }
+                      callback.onClose(status);
                     }
                   },
                   requestHeaders());
@@ -224,8 +223,6 @@ public void onClose(Status status, Metadata trailers) {
               call.sendMessage(request);
               call.halfClose();
             });
-
-    return tcs.getTask();
   }
 
   /** Creates and starts a single response RPC. */