firebase-ios-sdk/1499: Restart streams on any token change. (#222)

[Port of firebase/firebase-js-sdk#1120]

See firebase/firebase-ios-sdk#1499

This reworks our "user listener" into a "credential change listener" that
fires on any token change, even if the User did not change. This allows
us to restart our streams (but not switch mutation queues, etc.) on token
changes.

Author: mikelehen

firebase-firestore/src/main/java/com/google/firebase/firestore/auth/CredentialsProvider.java

@@ -27,17 +27,17 @@ public abstract class CredentialsProvider {
   public abstract Task<String> getToken();
 
   /**
-   * Marks the last retrieved token as invalid, making the next {@link #GetToken} request force
+   * Marks the last retrieved token as invalid, making the next {@link #getToken} request force
    * refresh the token.
    */
   public abstract void invalidateToken();
 
   /**
-   * Sets the listener to be notified of user changes (sign-in / sign-out). It is immediately called
-   * once with the initial user.
+   * Sets the listener to be notified of credential changes (sign-in / sign-out, token changes). It
+   * is immediately called once with the initial user.
    */
-  public abstract void setUserChangeListener(Listener<User> listener);
+  public abstract void setChangeListener(Listener<User> changeListener);
 
-  /** Removes the listener set with {@link #setUserChangeListener}. */
-  public abstract void removeUserChangeListener();
+  /** Removes the listener set with {@link #setChangeListener}. */
+  public abstract void removeChangeListener();
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/auth/EmptyCredentialsProvider.java

@@ -32,10 +32,10 @@ public Task<String> getToken() {
   public void invalidateToken() {}
 
   @Override
-  public void setUserChangeListener(Listener<User> listener) {
-    listener.onValue(User.UNAUTHENTICATED);
+  public void setChangeListener(Listener<User> changeListener) {
+    changeListener.onValue(User.UNAUTHENTICATED);
   }
 
   @Override
-  public void removeUserChangeListener() {}
+  public void removeChangeListener() {}
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/auth/FirebaseAuthCredentialsProvider.java

@@ -29,13 +29,12 @@
 /**
  * FirebaseAuthCredentialsProvider uses Firebase Auth via {@link FirebaseApp} to get an auth token.
  *
- * <p>NOTE: To simplify the implementation, it requires that you call {@link #setUserChangeListener}
- * no more than once and don't call {@link #getToken} after calling {@link
- * #removeUserChangeListener}.
+ * <p>NOTE: To simplify the implementation, it requires that you call {@link #setChangeListener} no
+ * more than once and don't call {@link #getToken} after calling {@link #removeChangeListener}.
  *
  * <p>This class must be implemented to be thread-safe since getToken() and
- * set/removeUserChangeListener() are called from the Firestore worker thread, but the getToken()
- * Task callbacks and user change notifications will be executed on arbitrary different threads.
+ * set/removeChangeListener() are called from the Firestore worker thread, but the getToken() Task
+ * callbacks and user change notifications will be executed on arbitrary different threads.
  */
 public final class FirebaseAuthCredentialsProvider extends CredentialsProvider {
 
@@ -45,18 +44,18 @@ public final class FirebaseAuthCredentialsProvider extends CredentialsProvider {
 
   /**
    * The listener registered with FirebaseApp; used to stop receiving auth changes once
-   * userChangeListener is removed.
+   * changeListener is removed.
    */
   private final FirebaseApp.IdTokenListener idTokenListener;
 
-  /** The listener to be notified of user changes (sign-in / sign-out). */
-  @Nullable private Listener<User> userChangeListener;
+  /** The listener to be notified of credential changes (sign-in / sign-out, token changes). */
+  @Nullable private Listener<User> changeListener;
 
   /** The current user as reported to us via our IdTokenListener. */
   private User currentUser;
 
-  /** Counter used to detect if the user changed while a getToken request was outstanding. */
-  private int userCounter;
+  /** Counter used to detect if the token changed while a getToken request was outstanding. */
+  private int tokenCounter;
 
   private boolean forceRefresh;
 
@@ -72,20 +71,17 @@ public FirebaseAuthCredentialsProvider(FirebaseApp app) {
     idTokenListener =
         token -> {
           synchronized (this) {
-            User user = getUser(app);
-            if (!currentUser.equals(user)) {
-              currentUser = user;
-              userCounter++;
+            currentUser = getUser(app);
+            tokenCounter++;
 
-              if (userChangeListener != null) {
-                userChangeListener.onValue(currentUser);
-              }
+            if (changeListener != null) {
+              changeListener.onValue(currentUser);
             }
           }
         };
 
     currentUser = getUser(app);
-    userCounter = 0;
+    tokenCounter = 0;
 
     FirebaseAppHelper.addIdTokenListener(app, idTokenListener);
   }
@@ -96,17 +92,17 @@ public synchronized Task<String> getToken() {
     forceRefresh = false;
     Task<GetTokenResult> res = FirebaseAppHelper.getToken(app, doForceRefresh);
 
-    // Take note of the current value of the userCounter so that this method can fail (with a
-    // FirebaseFirestoreException) if there is a user change while the request is outstanding.
-    final int savedCounter = userCounter;
+    // Take note of the current value of the tokenCounter so that this method can fail (with a
+    // FirebaseFirestoreException) if there is a token change while the request is outstanding.
+    final int savedCounter = tokenCounter;
     return res.continueWith(
         task -> {
           synchronized (this) {
-            // Cancel the request since the user changed while the request was outstanding so the
-            // response is likely for a previous user (which user, we can't be sure).
-            if (savedCounter != userCounter) {
+            // Cancel the request since the token changed while the request was outstanding so the
+            // response is potentially for a previous user (which user, we can't be sure).
+            if (savedCounter != tokenCounter) {
               throw new FirebaseFirestoreException(
-                  "getToken aborted due to user change", Code.ABORTED);
+                  "getToken aborted due to token change", Code.ABORTED);
             }
             if (!task.isSuccessful()) {
               Exception exception = task.getException();
@@ -131,16 +127,16 @@ public synchronized void invalidateToken() {
   }
 
   @Override
-  public synchronized void setUserChangeListener(@NonNull Listener<User> listener) {
-    userChangeListener = listener;
+  public synchronized void setChangeListener(@NonNull Listener<User> changeListener) {
+    this.changeListener = changeListener;
 
     // Fire the initial event.
-    listener.onValue(currentUser);
+    changeListener.onValue(currentUser);
   }
 
   @Override
-  public synchronized void removeUserChangeListener() {
-    userChangeListener = null;
+  public synchronized void removeChangeListener() {
+    changeListener = null;
     FirebaseAppHelper.removeIdTokenListener(app, idTokenListener);
   }
 

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

@@ -49,6 +49,7 @@
 import com.google.firebase.firestore.remote.RemoteSerializer;
 import com.google.firebase.firestore.remote.RemoteStore;
 import com.google.firebase.firestore.util.AsyncQueue;
+import com.google.firebase.firestore.util.Logger;
 import io.grpc.Status;
 import java.util.Collections;
 import java.util.List;
@@ -61,6 +62,8 @@
  */
 public final class FirestoreClient implements RemoteStore.RemoteStoreCallback {
 
+  private static final String LOG_TAG = "FirestoreClient";
+
   private final DatabaseInfo databaseInfo;
   private final CredentialsProvider credentialsProvider;
   private final AsyncQueue asyncQueue;
@@ -83,17 +86,21 @@ public FirestoreClient(
 
     TaskCompletionSource<User> firstUser = new TaskCompletionSource<>();
     final AtomicBoolean initialized = new AtomicBoolean(false);
-    credentialsProvider.setUserChangeListener(
+    credentialsProvider.setChangeListener(
         (User user) -> {
           if (initialized.compareAndSet(false, true)) {
             hardAssert(!firstUser.getTask().isComplete(), "Already fulfilled first user task");
             firstUser.setResult(user);
           } else {
-            asyncQueue.enqueueAndForget(() -> syncEngine.handleUserChange(user));
+            asyncQueue.enqueueAndForget(
+                () -> {
+                  Logger.debug(LOG_TAG, "Credential changed. Current user: %s", user.getUid());
+                  syncEngine.handleCredentialChange(user);
+                });
           }
         });
 
-    // Defer initialization until we get the current user from the userChangeListener. This is
+    // Defer initialization until we get the current user from the changeListener. This is
     // guaranteed to be synchronously dispatched onto our worker queue, so we will be initialized
     // before any subsequently queued work runs.
     asyncQueue.enqueueAndForget(
@@ -118,7 +125,7 @@ public Task<Void> enableNetwork() {
 
   /** Shuts down this client, cancels all writes / listeners, and releases all resources. */
   public Task<Void> shutdown() {
-    credentialsProvider.removeUserChangeListener();
+    credentialsProvider.removeChangeListener();
     return asyncQueue.enqueue(
         () -> {
           remoteStore.shutdown();
@@ -197,6 +204,7 @@ private void initialize(Context context, User user, boolean usePersistence) {
     // Note: The initialization work must all be synchronous (we can't dispatch more work) since
     // external write/listen operations could get queued to run before that subsequent work
     // completes.
+    Logger.debug(LOG_TAG, "Initializing. user=%s", user.getUid());
 
     GarbageCollector garbageCollector;
     if (usePersistence) {

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

@@ -539,14 +539,17 @@ public Map<DocumentKey, Integer> getCurrentLimboDocuments() {
     return new HashMap<>(limboTargetsByKey);
   }
 
-  public void handleUserChange(User user) {
+  public void handleCredentialChange(User user) {
+    boolean userChanged = !currentUser.equals(user);
     currentUser = user;
 
-    // Notify local store and emit any resulting events from swapping out the mutation queue.
-    ImmutableSortedMap<DocumentKey, MaybeDocument> changes = localStore.handleUserChange(user);
-    emitNewSnapshot(changes, /*remoteEvent=*/ null);
+    if (userChanged) {
+      // Notify local store and emit any resulting events from swapping out the mutation queue.
+      ImmutableSortedMap<DocumentKey, MaybeDocument> changes = localStore.handleUserChange(user);
+      emitNewSnapshot(changes, /*remoteEvent=*/ null);
+    }
 
     // Notify remote store so it can restart its streams.
-    remoteStore.handleUserChange(user);
+    remoteStore.handleCredentialChange();
   }
 }

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

@@ -18,7 +18,6 @@
 
 import android.support.annotation.Nullable;
 import com.google.firebase.database.collection.ImmutableSortedSet;
-import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.core.OnlineState;
 import com.google.firebase.firestore.core.Transaction;
 import com.google.firebase.firestore.local.LocalStore;
@@ -237,13 +236,13 @@ public void shutdown() {
    * <p>In response the remote store tears down streams and clears up any tracked operations that
    * should not persist across users. Restarts the streams if appropriate.
    */
-  public void handleUserChange(User user) {
-    Logger.debug(LOG_TAG, "Changing user to " + user);
+  public void handleCredentialChange() {
     // If the network has been explicitly disabled, make sure we don't accidentally re-enable it.
     if (isNetworkEnabled()) {
       // Tear down and re-create our network streams. This will ensure we get a fresh auth token
       // for the new user and re-fill the write pipeline with new mutations from the LocalStore
       // (since mutations are per-user).
+      Logger.debug(LOG_TAG, "Restarting streams for new credential.");
       disableNetworkInternal();
       onlineStateTracker.updateState(OnlineState.UNKNOWN);
       enableNetwork();

firebase-firestore/src/test/java/com/google/firebase/firestore/spec/SpecTestCase.java

@@ -731,7 +731,7 @@ private void doChangeUser(@Nullable String uid) throws Exception {
     currentUser = new User(uid);
     queue.runSync(
         () -> {
-          syncEngine.handleUserChange(currentUser);
+          syncEngine.handleCredentialChange(currentUser);
         });
   }