Add FirestoreClientProvider

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/AccessHelper.java

@@ -33,7 +33,6 @@ public static FirebaseFirestore newFirebaseFirestore(
       String persistenceKey,
       CredentialsProvider<User> authProvider,
       CredentialsProvider<String> appCheckProvider,
-      AsyncQueue asyncQueue,
       Function<FirebaseFirestoreSettings, ComponentProvider> componentProviderFactory,
       FirebaseApp firebaseApp,
       FirebaseFirestore.InstanceRegistry instanceRegistry) {
@@ -43,14 +42,13 @@ public static FirebaseFirestore newFirebaseFirestore(
         persistenceKey,
         authProvider,
         appCheckProvider,
-        asyncQueue,
         componentProviderFactory,
         firebaseApp,
         instanceRegistry,
         null);
   }
 
   public static AsyncQueue getAsyncQueue(FirebaseFirestore firestore) {
-    return firestore.getAsyncQueue();
+    return firestore.clientProvider.getAsyncQueue();
   }
 }

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/CompositeIndexQueryTest.java

@@ -742,7 +742,7 @@ public void testMultipleInequalityReadFromCacheWhenOffline() {
     assertEquals(2L, snapshot1.size());
     assertFalse(snapshot1.getMetadata().isFromCache());
 
-    waitFor(collection.firestore.getClient().disableNetwork());
+    waitFor(collection.firestore.disableNetwork());
 
     QuerySnapshot snapshot2 = waitFor(query.get());
     assertEquals(2L, snapshot2.size());

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/FirestoreTest.java

@@ -49,6 +49,7 @@
 import com.google.firebase.firestore.FirebaseFirestoreException.Code;
 import com.google.firebase.firestore.Query.Direction;
 import com.google.firebase.firestore.auth.User;
+import com.google.firebase.firestore.core.FirestoreClient;
 import com.google.firebase.firestore.model.DatabaseId;
 import com.google.firebase.firestore.testutil.EventAccumulator;
 import com.google.firebase.firestore.testutil.IntegrationTestUtil;
@@ -1110,7 +1111,7 @@ public void testRestartFirestoreLeadsToNewInstance() {
     assertSame(instance, sameInstance);
     waitFor(instance.document("abc/123").set(Collections.singletonMap("field", 100L)));
 
-    instance.terminate();
+    waitFor(instance.terminate());
     FirebaseFirestore newInstance = FirebaseFirestore.getInstance(app);
     newInstance.setFirestoreSettings(newTestSettings());
 
@@ -1132,7 +1133,7 @@ public void testAppDeleteLeadsToFirestoreTerminate() {
 
     app.delete();
 
-    assertTrue(instance.getClient().isTerminated());
+    assertTrue(instance.callClient(FirestoreClient::isTerminated));
   }
 
   @Test

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/QueryTest.java

@@ -536,12 +536,12 @@ public void testQueriesFireFromCacheWhenOffline() {
     assertFalse(querySnapshot.getMetadata().isFromCache());
 
     // offline event with fromCache=true
-    waitFor(collection.firestore.getClient().disableNetwork());
+    waitFor(collection.firestore.disableNetwork());
     querySnapshot = accum.await();
     assertTrue(querySnapshot.getMetadata().isFromCache());
 
     // back online event with fromCache=false
-    waitFor(collection.firestore.getClient().enableNetwork());
+    waitFor(collection.firestore.enableNetwork());
     querySnapshot = accum.await();
     assertFalse(querySnapshot.getMetadata().isFromCache());
 

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/ServerTimestampTest.java

@@ -214,7 +214,7 @@ public void testServerTimestampsCanReturnPreviousValueOfDifferentType() {
   @Test
   public void testServerTimestampsCanRetainPreviousValueThroughConsecutiveUpdates() {
     writeInitialData();
-    waitFor(docRef.getFirestore().getClient().disableNetwork());
+    waitFor(docRef.getFirestore().disableNetwork());
     accumulator.awaitRemoteEvent();
 
     docRef.update("a", FieldValue.serverTimestamp());
@@ -226,7 +226,7 @@ public void testServerTimestampsCanRetainPreviousValueThroughConsecutiveUpdates(
     localSnapshot = accumulator.awaitLocalEvent();
     assertEquals(42L, localSnapshot.get("a", ServerTimestampBehavior.PREVIOUS));
 
-    waitFor(docRef.getFirestore().getClient().enableNetwork());
+    waitFor(docRef.getFirestore().enableNetwork());
 
     DocumentSnapshot remoteSnapshot = accumulator.awaitRemoteEvent();
     assertThat(remoteSnapshot.get("a")).isInstanceOf(Timestamp.class);
@@ -235,7 +235,7 @@ public void testServerTimestampsCanRetainPreviousValueThroughConsecutiveUpdates(
   @Test
   public void testServerTimestampsUsesPreviousValueFromLocalMutation() {
     writeInitialData();
-    waitFor(docRef.getFirestore().getClient().disableNetwork());
+    waitFor(docRef.getFirestore().disableNetwork());
     accumulator.awaitRemoteEvent();
 
     docRef.update("a", FieldValue.serverTimestamp());
@@ -249,7 +249,7 @@ public void testServerTimestampsUsesPreviousValueFromLocalMutation() {
     localSnapshot = accumulator.awaitLocalEvent();
     assertEquals(1337L, localSnapshot.get("a", ServerTimestampBehavior.PREVIOUS));
 
-    waitFor(docRef.getFirestore().getClient().enableNetwork());
+    waitFor(docRef.getFirestore().enableNetwork());
 
     DocumentSnapshot remoteSnapshot = accumulator.awaitRemoteEvent();
     assertThat(remoteSnapshot.get("a")).isInstanceOf(Timestamp.class);

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/TransactionTest.java

@@ -371,7 +371,7 @@ public void testIncrementTransactionally() {
     AtomicInteger started = new AtomicInteger(0);
 
     FirebaseFirestore firestore = testFirestore();
-    firestore.getAsyncQueue().skipDelaysForTimerId(TimerId.RETRY_TRANSACTION);
+    AccessHelper.getAsyncQueue(firestore).skipDelaysForTimerId(TimerId.RETRY_TRANSACTION);
     DocumentReference doc = firestore.collection("counters").document();
     waitFor(doc.set(map("count", 5.0)));
 
@@ -437,7 +437,7 @@ public void testUpdateTransactionally() {
     AtomicInteger counter = new AtomicInteger(0);
 
     FirebaseFirestore firestore = testFirestore();
-    firestore.getAsyncQueue().skipDelaysForTimerId(TimerId.RETRY_TRANSACTION);
+    AccessHelper.getAsyncQueue(firestore).skipDelaysForTimerId(TimerId.RETRY_TRANSACTION);
     DocumentReference doc = firestore.collection("counters").document();
     waitFor(doc.set(map("count", 5.0, "other", "yes")));
 
@@ -532,7 +532,7 @@ public void testUpdatePOJOTransactionally() {
     AtomicInteger started = new AtomicInteger(0);
 
     FirebaseFirestore firestore = testFirestore();
-    firestore.getAsyncQueue().skipDelaysForTimerId(TimerId.RETRY_TRANSACTION);
+    AccessHelper.getAsyncQueue(firestore).skipDelaysForTimerId(TimerId.RETRY_TRANSACTION);
     DocumentReference doc = firestore.collection("counters").document();
     waitFor(doc.set(new POJO(5.0, "no", "clean")));
 
@@ -601,7 +601,7 @@ public void testRetriesWhenDocumentThatWasReadWithoutBeingWrittenChanges() {
   @Test
   public void testReadingADocTwiceWithDifferentVersions() {
     FirebaseFirestore firestore = testFirestore();
-    firestore.getAsyncQueue().skipDelaysForTimerId(TimerId.RETRY_TRANSACTION);
+    AccessHelper.getAsyncQueue(firestore).skipDelaysForTimerId(TimerId.RETRY_TRANSACTION);
     DocumentReference doc = firestore.collection("counters").document();
     waitFor(doc.set(map("count", 15.0)));
     AtomicInteger counter = new AtomicInteger(0);

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/ValidationTest.java

@@ -453,7 +453,7 @@ public void queriesCannotBeSortedByAnUncommittedServerTimestamp() {
     CollectionReference collection = testCollection();
 
     // Ensure the server timestamp stays uncommitted for the first half of the test
-    waitFor(collection.firestore.getClient().disableNetwork());
+    waitFor(collection.firestore.disableNetwork());
 
     TaskCompletionSource<Void> offlineCallbackDone = new TaskCompletionSource<>();
     TaskCompletionSource<Void> onlineCallbackDone = new TaskCompletionSource<>();
@@ -497,7 +497,7 @@ public void queriesCannotBeSortedByAnUncommittedServerTimestamp() {
     document.set(map("timestamp", FieldValue.serverTimestamp()));
     waitFor(offlineCallbackDone.getTask());
 
-    waitFor(collection.firestore.getClient().enableNetwork());
+    waitFor(collection.firestore.enableNetwork());
     waitFor(onlineCallbackDone.getTask());
 
     listenerRegistration.remove();

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/testutil/IntegrationTestUtil.java

@@ -320,7 +320,6 @@ public static FirebaseFirestore testFirestore(
             persistenceKey,
             MockCredentialsProvider.instance(),
             new EmptyAppCheckTokenProvider(),
-            asyncQueue,
             ComponentProvider::defaultFactory,
             /* firebaseApp= */ null,
             /* instanceRegistry= */ (dbId) -> {});

firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateQuery.java

@@ -66,8 +66,7 @@ public Task<AggregateQuerySnapshot> get(@NonNull AggregateSource source) {
     TaskCompletionSource<AggregateQuerySnapshot> tcs = new TaskCompletionSource<>();
     query
         .firestore
-        .getClient()
-        .runAggregateQuery(query.query, aggregateFieldList)
+        .callClient(client -> client.runAggregateQuery(query.query, aggregateFieldList))
         .continueWith(
             Executors.DIRECT_EXECUTOR,
             (task) -> {

firebase-firestore/src/main/java/com/google/firebase/firestore/DocumentReference.java

@@ -29,7 +29,6 @@
 import com.google.firebase.firestore.core.ActivityScope;
 import com.google.firebase.firestore.core.AsyncEventListener;
 import com.google.firebase.firestore.core.EventManager.ListenOptions;
-import com.google.firebase.firestore.core.ListenerRegistrationImpl;
 import com.google.firebase.firestore.core.QueryListener;
 import com.google.firebase.firestore.core.UserData.ParsedSetData;
 import com.google.firebase.firestore.core.UserData.ParsedUpdateData;
@@ -38,11 +37,12 @@
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.mutation.DeleteMutation;
+import com.google.firebase.firestore.model.mutation.Mutation;
 import com.google.firebase.firestore.model.mutation.Precondition;
 import com.google.firebase.firestore.util.Assert;
 import com.google.firebase.firestore.util.Executors;
 import com.google.firebase.firestore.util.Util;
-import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
@@ -165,9 +165,9 @@ public Task<Void> set(@NonNull Object data, @NonNull SetOptions options) {
         options.isMerge()
             ? firestore.getUserDataReader().parseMergeData(data, options.getFieldMask())
             : firestore.getUserDataReader().parseSetData(data);
+    List<Mutation> mutations = singletonList(parsed.toMutation(key, Precondition.NONE));
     return firestore
-        .getClient()
-        .write(Collections.singletonList(parsed.toMutation(key, Precondition.NONE)))
+        .callClient(client -> client.write(mutations))
         .continueWith(Executors.DIRECT_EXECUTOR, voidErrorTransformer());
   }
 
@@ -229,9 +229,9 @@ public Task<Void> update(
   }
 
   private Task<Void> update(@NonNull ParsedUpdateData parsedData) {
+    List<Mutation> mutations = singletonList(parsedData.toMutation(key, Precondition.exists(true)));
     return firestore
-        .getClient()
-        .write(Collections.singletonList(parsedData.toMutation(key, Precondition.exists(true))))
+        .callClient(client -> client.write(mutations))
         .continueWith(Executors.DIRECT_EXECUTOR, voidErrorTransformer());
   }
 
@@ -242,9 +242,9 @@ private Task<Void> update(@NonNull ParsedUpdateData parsedData) {
    */
   @NonNull
   public Task<Void> delete() {
+    List<Mutation> mutations = singletonList(new DeleteMutation(key, Precondition.NONE));
     return firestore
-        .getClient()
-        .write(singletonList(new DeleteMutation(key, Precondition.NONE)))
+        .callClient(client -> client.write(mutations))
         .continueWith(Executors.DIRECT_EXECUTOR, voidErrorTransformer());
   }
 
@@ -274,8 +274,7 @@ public Task<DocumentSnapshot> get() {
   public Task<DocumentSnapshot> get(@NonNull Source source) {
     if (source == Source.CACHE) {
       return firestore
-          .getClient()
-          .getDocumentFromLocalCache(key)
+          .callClient(client -> client.getDocumentFromLocalCache(key))
           .continueWith(
               Executors.DIRECT_EXECUTOR,
               (Task<Document> task) -> {
@@ -531,11 +530,17 @@ private ListenerRegistration addSnapshotListenerInternal(
         new AsyncEventListener<>(userExecutor, viewListener);
 
     com.google.firebase.firestore.core.Query query = asQuery();
-    QueryListener queryListener = firestore.getClient().listen(query, options, asyncListener);
 
-    return ActivityScope.bind(
-        activity,
-        new ListenerRegistrationImpl(firestore.getClient(), queryListener, asyncListener));
+    return firestore.callClient(
+        client -> {
+          QueryListener queryListener = client.listen(query, options, asyncListener);
+          return ActivityScope.bind(
+              activity,
+              () -> {
+                asyncListener.mute();
+                client.stopListening(queryListener);
+              });
+        });
   }
 
   @Override