Prepare shutdown to be a public api. (#589)

* Deregister firestore when shutdown + app.delete calls shutdown

* synchronized scheduling + soft shutdown

* Add comments and more tests

* addressing comments #1

* fixing format in tools

* Stop leakages

* better comments

* Handle app delete shutdown differently

* No more error-prone hack

* Change how to interact with MultiDbComponent.

* add missing sync keyword

Author: Hui-Wu

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

@@ -31,9 +31,16 @@ public static FirebaseFirestore newFirebaseFirestore(
       String persistenceKey,
       CredentialsProvider credentialsProvider,
       AsyncQueue asyncQueue,
-      FirebaseApp firebaseApp) {
+      FirebaseApp firebaseApp,
+      FirebaseFirestore.InstanceRegistry instanceRegistry) {
     return new FirebaseFirestore(
-        context, databaseId, persistenceKey, credentialsProvider, asyncQueue, firebaseApp);
+        context,
+        databaseId,
+        persistenceKey,
+        credentialsProvider,
+        asyncQueue,
+        firebaseApp,
+        instanceRegistry);
   }
 
   /** Makes the shutdown method accessible. */

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

@@ -20,6 +20,7 @@
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testCollection;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testCollectionWithDocs;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testDocument;
+import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testFirebaseApp;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testFirestore;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.waitFor;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.waitForException;
@@ -30,6 +31,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
@@ -38,6 +40,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
+import com.google.firebase.FirebaseApp;
 import com.google.firebase.Timestamp;
 import com.google.firebase.firestore.FirebaseFirestoreException.Code;
 import com.google.firebase.firestore.Query.Direction;
@@ -1023,4 +1026,74 @@ public void testClearPersistenceWhileRunningFails() {
     FirebaseFirestoreException firestoreException = (FirebaseFirestoreException) e;
     assertEquals(Code.FAILED_PRECONDITION, firestoreException.getCode());
   }
+
+  @Test
+  public void testRestartFirestoreLeadsToNewInstance() {
+    FirebaseApp app = testFirebaseApp();
+    FirebaseFirestore instance = FirebaseFirestore.getInstance(app);
+    FirebaseFirestore sameInstance = FirebaseFirestore.getInstance(app);
+
+    assertSame(instance, sameInstance);
+    waitFor(instance.document("abc/123").set(Collections.singletonMap("field", 100L)));
+
+    instance.shutdown();
+    FirebaseFirestore newInstance = FirebaseFirestore.getInstance(app);
+
+    // Verify new instance works.
+    DocumentSnapshot doc = waitFor(newInstance.document("abc/123").get());
+    assertEquals(doc.get("field"), 100L);
+    waitFor(newInstance.document("abc/123").delete());
+
+    // Verify it is different instance.
+    assertNotSame(instance, newInstance);
+  }
+
+  @Test
+  public void testAppDeleteLeadsToFirestoreShutdown() {
+    FirebaseApp app = testFirebaseApp();
+    FirebaseFirestore instance = FirebaseFirestore.getInstance(app);
+    waitFor(instance.document("abc/123").set(Collections.singletonMap("Field", 100)));
+
+    app.delete();
+
+    assertTrue(instance.getClient().isShutdown());
+  }
+
+  @Test
+  public void testNewOperationThrowsAfterFirestoreShutdown() {
+    FirebaseFirestore instance = testFirestore();
+    DocumentReference reference = instance.document("abc/123");
+    waitFor(reference.set(Collections.singletonMap("Field", 100)));
+
+    instance.shutdown();
+
+    final String expectedMessage = "The client has already been shutdown";
+    expectError(() -> waitFor(reference.get()), expectedMessage);
+    expectError(() -> waitFor(reference.update("Field", 1)), expectedMessage);
+    expectError(
+        () -> waitFor(reference.set(Collections.singletonMap("Field", 1))), expectedMessage);
+    expectError(
+        () -> waitFor(instance.runBatch((batch) -> batch.update(reference, "Field", 1))),
+        expectedMessage);
+    expectError(
+        () -> waitFor(instance.runTransaction(transaction -> transaction.get(reference))),
+        expectedMessage);
+  }
+
+  @Test
+  public void testShutdownCalledMultipleTimes() {
+    FirebaseFirestore instance = testFirestore();
+    DocumentReference reference = instance.document("abc/123");
+    waitFor(reference.set(Collections.singletonMap("Field", 100)));
+
+    instance.shutdown();
+
+    final String expectedMessage = "The client has already been shutdown";
+    expectError(() -> waitFor(reference.get()), expectedMessage);
+
+    // Calling a second time should go through and change nothing.
+    instance.shutdown();
+
+    expectError(() -> waitFor(reference.get()), expectedMessage);
+  }
 }

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

@@ -25,6 +25,7 @@
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
+import com.google.firebase.FirebaseApp;
 import com.google.firebase.firestore.AccessHelper;
 import com.google.firebase.firestore.BuildConfig;
 import com.google.firebase.firestore.CollectionReference;
@@ -145,6 +146,10 @@ protected String checkAuthority(String authority) {
     return settings.build();
   }
 
+  public static FirebaseApp testFirebaseApp() {
+    return FirebaseApp.initializeApp(ApplicationProvider.getApplicationContext());
+  }
+
   /** Initializes a new Firestore instance that uses the default project. */
   public static FirebaseFirestore testFirestore() {
     return testFirestore(newTestSettings());
@@ -257,7 +262,8 @@ public static FirebaseFirestore testFirestore(
             persistenceKey,
             new EmptyCredentialsProvider(),
             asyncQueue,
-            /*firebaseApp=*/ null);
+            /*firebaseApp=*/ null,
+            /*instanceRegistry=*/ (dbId) -> {});
     waitFor(AccessHelper.clearPersistence(firestore));
     firestore.setFirestoreSettings(settings);
     firestoreStatus.put(firestore, true);

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

@@ -51,6 +51,12 @@
 @PublicApi
 public class FirebaseFirestore {
 
+  /** Provides a registry management interface for {@code FirebaseFirestore} instances. */
+  public interface InstanceRegistry {
+    /** Removes the Firestore instance with given name from registry. */
+    void remove(@NonNull String databaseId);
+  }
+
   private static final String TAG = "FirebaseFirestore";
   private final Context context;
   // This is also used as private lock object for this instance. There is nothing inherent about
@@ -61,6 +67,9 @@ public class FirebaseFirestore {
   private final AsyncQueue asyncQueue;
   private final FirebaseApp firebaseApp;
   private final UserDataConverter dataConverter;
+  // When user requests to shutdown, use this to notify `FirestoreMultiDbComponent` to deregister
+  // this instance.
+  private final InstanceRegistry instanceRegistry;
   private FirebaseFirestoreSettings settings;
   private volatile FirestoreClient client;
 
@@ -94,7 +103,8 @@ static FirebaseFirestore newInstance(
       @NonNull Context context,
       @NonNull FirebaseApp app,
       @Nullable InternalAuthProvider authProvider,
-      @NonNull String database) {
+      @NonNull String database,
+      @NonNull InstanceRegistry instanceRegistry) {
     String projectId = app.getOptions().getProjectId();
     if (projectId == null) {
       throw new IllegalArgumentException("FirebaseOptions.getProjectId() cannot be null");
@@ -117,7 +127,10 @@ static FirebaseFirestore newInstance(
     // so there is no need to include it in the persistence key.
     String persistenceKey = app.getName();
 
-    return new FirebaseFirestore(context, databaseId, persistenceKey, provider, queue, app);
+    FirebaseFirestore firestore =
+        new FirebaseFirestore(
+            context, databaseId, persistenceKey, provider, queue, app, instanceRegistry);
+    return firestore;
   }
 
   @VisibleForTesting
@@ -127,7 +140,8 @@ static FirebaseFirestore newInstance(
       String persistenceKey,
       CredentialsProvider credentialsProvider,
       AsyncQueue asyncQueue,
-      @Nullable FirebaseApp firebaseApp) {
+      @Nullable FirebaseApp firebaseApp,
+      InstanceRegistry instanceRegistry) {
     this.context = checkNotNull(context);
     this.databaseId = checkNotNull(checkNotNull(databaseId));
     this.dataConverter = new UserDataConverter(databaseId);
@@ -136,6 +150,7 @@ static FirebaseFirestore newInstance(
     this.asyncQueue = checkNotNull(asyncQueue);
     // NOTE: We allow firebaseApp to be null in tests only.
     this.firebaseApp = firebaseApp;
+    this.instanceRegistry = instanceRegistry;
 
     settings = new FirebaseFirestoreSettings.Builder().build();
   }
@@ -329,13 +344,39 @@ public Task<Void> runBatch(@NonNull WriteBatch.Function batchFunction) {
     return batch.commit();
   }
 
-  @VisibleForTesting
-  Task<Void> shutdown() {
+  Task<Void> shutdownInternal() {
     // The client must be initialized to ensure that all subsequent API usage throws an exception.
     this.ensureClientConfigured();
     return client.shutdown();
   }
 
+  /**
+   * Shuts down this FirebaseFirestore instance.
+   *
+   * <p>After shutdown only the {@link #clearPersistence()} method may be used. Any other method
+   * will throw an {@link IllegalStateException}.
+   *
+   * <p>To restart after shutdown, simply create a new instance of FirebaseFirestore with {@link
+   * #getInstance()} or {@link #getInstance(FirebaseApp)}.
+   *
+   * <p>Shutdown does not cancel any pending writes and any tasks that are awaiting a response from
+   * the server will not be resolved. The next time you start this instance, it will resume
+   * attempting to send these writes to the server.
+   *
+   * <p>Note: Under normal circumstances, calling <code>shutdown()</code> is not required. This
+   * method is useful only when you want to force this instance to release all of its resources or
+   * in combination with {@link #clearPersistence} to ensure that all local state is destroyed
+   * between test runs.
+   *
+   * @return A <code>Task</code> that is resolved when the instance has been successfully shut down.
+   */
+  @VisibleForTesting
+  // TODO(b/135755126): Make this public and remove @VisibleForTesting
+  Task<Void> shutdown() {
+    instanceRegistry.remove(this.getDatabaseId().getDatabaseId());
+    return shutdownInternal();
+  }
+
   @VisibleForTesting
   AsyncQueue getAsyncQueue() {
     return asyncQueue;
@@ -396,7 +437,7 @@ public static void setLoggingEnabled(boolean loggingEnabled) {
   @PublicApi
   public Task<Void> clearPersistence() {
     final TaskCompletionSource<Void> source = new TaskCompletionSource<>();
-    asyncQueue.enqueueAndForget(
+    asyncQueue.enqueueAndForgetEvenAfterShutdown(
         () -> {
           try {
             if (client != null && !client.isShutdown()) {

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

@@ -18,12 +18,16 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import com.google.firebase.FirebaseApp;
+import com.google.firebase.FirebaseAppLifecycleListener;
+import com.google.firebase.FirebaseOptions;
 import com.google.firebase.auth.internal.InternalAuthProvider;
 import java.util.HashMap;
 import java.util.Map;
 
 /** Multi-resource container for Firestore. */
-class FirestoreMultiDbComponent {
+class FirestoreMultiDbComponent
+    implements FirebaseAppLifecycleListener, FirebaseFirestore.InstanceRegistry {
+
   /**
    * A static map from instance key to FirebaseFirestore instances. Instance keys are database
    * names.
@@ -41,16 +45,38 @@ class FirestoreMultiDbComponent {
     this.context = context;
     this.app = app;
     this.authProvider = authProvider;
+    this.app.addLifecycleEventListener(this);
   }
 
-  /** Provides instances of Firestore for given database names. */
+  /** Provides instances of Firestore for given database IDs. */
   @NonNull
-  synchronized FirebaseFirestore get(@NonNull String databaseName) {
-    FirebaseFirestore firestore = instances.get(databaseName);
+  synchronized FirebaseFirestore get(@NonNull String databaseId) {
+    FirebaseFirestore firestore = instances.get(databaseId);
     if (firestore == null) {
-      firestore = FirebaseFirestore.newInstance(context, app, authProvider, databaseName);
-      instances.put(databaseName, firestore);
+      firestore = FirebaseFirestore.newInstance(context, app, authProvider, databaseId, this);
+      instances.put(databaseId, firestore);
     }
     return firestore;
   }
+
+  /**
+   * Remove the instance of a given database ID from this component, such that if {@link
+   * FirestoreMultiDbComponent#get(String)} is called again with the same name, a new instance of
+   * {@link FirebaseFirestore} is created.
+   *
+   * <p>It is a no-op if there is no instance associated with the given database name.
+   */
+  @Override
+  public synchronized void remove(@NonNull String databaseId) {
+    instances.remove(databaseId);
+  }
+
+  @Override
+  public synchronized void onDeleted(String firebaseAppName, FirebaseOptions options) {
+    // Shuts down all database instances and remove them from registry map when App is deleted.
+    for (Map.Entry<String, FirebaseFirestore> entry : instances.entrySet()) {
+      entry.getValue().shutdownInternal();
+      instances.remove(entry.getKey());
+    }
+  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/core/FirestoreClient.java

@@ -75,7 +75,6 @@ public final class FirestoreClient implements RemoteStore.RemoteStoreCallback {
   private RemoteStore remoteStore;
   private SyncEngine syncEngine;
   private EventManager eventManager;
-  private volatile boolean clientShutdown = false;
 
   // LRU-related
   @Nullable private LruGarbageCollector.Scheduler lruScheduler;
@@ -138,21 +137,21 @@ public Task<Void> enableNetwork() {
   /** Shuts down this client, cancels all writes / listeners, and releases all resources. */
   public Task<Void> shutdown() {
     credentialsProvider.removeChangeListener();
-    return asyncQueue.enqueue(
+    return asyncQueue.enqueueAndInitiateShutdown(
         () -> {
-          if (!this.clientShutdown) {
-            remoteStore.shutdown();
-            persistence.shutdown();
-            if (lruScheduler != null) {
-              lruScheduler.stop();
-            }
-            this.clientShutdown = true;
+          remoteStore.shutdown();
+          persistence.shutdown();
+          if (lruScheduler != null) {
+            lruScheduler.stop();
           }
         });
   }
 
+  /** Returns true if this client has been shutdown. */
   public boolean isShutdown() {
-    return this.clientShutdown;
+    // Technically, the asyncQueue is still running, but only accepting tasks related to shutdown
+    // or supposed to be run after shutdown. It is effectively shut down to the eyes of users.
+    return this.asyncQueue.isShuttingDown();
   }
 
   /** Starts listening to a query. */
@@ -272,8 +271,8 @@ private void initialize(Context context, User user, boolean usePersistence, long
   }
 
   private void verifyNotShutdown() {
-    if (this.clientShutdown) {
-      throw new IllegalArgumentException("The client has already been shutdown");
+    if (this.isShutdown()) {
+      throw new IllegalStateException("The client has already been shutdown");
     }
   }