Fix Android Connectivity Monitor (v2)

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

@@ -174,6 +174,9 @@ public void run() {
   /** The time a stream stays open after it is marked idle. */
   private static final long IDLE_TIMEOUT_MS = TimeUnit.MINUTES.toMillis(1);
 
+  /** Maximum backoff time when reconnecting. */
+  private static final long RECONNECT_BACKOFF_MAX_DELAY_MS = TimeUnit.SECONDS.toMillis(10);
+
   @Nullable private DelayedTask idleTimer;
 
   private final FirestoreChannel firestoreChannel;
@@ -263,8 +266,9 @@ public void start() {
    *
    * @param finalState the intended state of the stream after closing.
    * @param status the status to emit to the listener.
+   * @param forceNewConnection whether to use a new connection the next time we open the stream
    */
-  private void close(State finalState, Status status) {
+  protected void close(State finalState, Status status, boolean forceNewConnection) {
     hardAssert(isStarted(), "Only started streams should be closed.");
     hardAssert(
         finalState == State.Error || status.equals(Status.OK),
@@ -290,18 +294,26 @@ private void close(State finalState, Status status) {
     if (code == Code.OK) {
       // If this is an intentional close ensure we don't delay our next connection attempt.
       backoff.reset();
-
     } else if (code == Code.RESOURCE_EXHAUSTED) {
       Logger.debug(
           getClass().getSimpleName(),
           "(%x) Using maximum backoff delay to prevent overloading the backend.",
           System.identityHashCode(this));
       backoff.resetToMax();
-
     } else if (code == Code.UNAUTHENTICATED) {
       // "unauthenticated" error means the token was rejected. Try force refreshing it in case it
       // just expired.
       firestoreChannel.invalidateToken();
+    } else if (code == Code.UNAVAILABLE) {
+      // These exceptions are thrown when the gRPC stream is closed with an connection error. For
+      // these cases, we need to use a new connection for the next connection attempt, which is
+      // done by marking the underlying channel as idle.
+      if (forceNewConnection
+          || status.getCause() instanceof java.net.ConnectException
+          || status.getCause() instanceof java.net.UnknownHostException) {
+        backoff.setTemporaryMaxDelay(RECONNECT_BACKOFF_MAX_DELAY_MS);
+        firestoreChannel.markChannelIdle();
+      }
     }
 
     if (finalState != State.Error) {
@@ -333,6 +345,10 @@ private void close(State finalState, Status status) {
     listener.onClose(status);
   }
 
+  protected void close(State finalState, Status status) {
+    close(finalState, status, false);
+  }
+
   /**
    * Can be overridden to perform additional cleanup before the stream is closed. Calling
    * super.tearDown() is not required.

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

@@ -26,6 +26,8 @@
 import com.google.firebase.firestore.core.DatabaseInfo;
 import com.google.firebase.firestore.model.DatabaseId;
 import com.google.firebase.firestore.util.AsyncQueue;
+import com.google.firebase.firestore.util.AsyncQueue.DelayedTask;
+import com.google.firebase.firestore.util.AsyncQueue.TimerId;
 import com.google.firebase.firestore.util.Util;
 import io.grpc.ClientCall;
 import io.grpc.ForwardingClientCall;
@@ -53,6 +55,11 @@ class FirestoreChannel {
   private static final String X_GOOG_API_CLIENT_VALUE =
       "gl-java/ fire/" + BuildConfig.VERSION_NAME + " grpc/";
 
+  // This timeout is used when attempting to establish a connection. If a connection attempt
+  // does not succeed in time, we close the stream and restart the connection, rather than having
+  // it hang indefinitely.
+  private static final int CONNECTIVITY_ATTEMPT_TIMEOUT_MS = 15 * 1000;
+
   /** The async worker queue that is used to dispatch events. */
   private final AsyncQueue asyncQueue;
 
@@ -92,6 +99,15 @@ public void shutdown() {
     callProvider.shutdown();
   }
 
+  /**
+   * Marks the underlying gRPC channel as idle. This allows on-going RPCs to continue, but the next
+   * RPC on the channel will trigger the creation of a new connection. This method is primarily used
+   * to reset the underlying connection.
+   */
+  public void markChannelIdle() {
+    this.callProvider.markChannelIdle();
+  }
+
   /**
    * Creates and starts a new bi-directional streaming RPC. The stream cannot accept message before
    * the observer's `onOpen()` callback is invoked.
@@ -154,6 +170,18 @@ public void onReady() {
           call[0].request(1);
         });
 
+    DelayedTask connectivityAttemptTimer =
+        asyncQueue.enqueueAfterDelay(
+            TimerId.CONNECTIVITY_ATTEMPT_TIMER,
+            CONNECTIVITY_ATTEMPT_TIMEOUT_MS,
+            () -> {
+              // Reset the underlying connection and restart the stream.
+              callProvider.clearConnectivityTimer();
+              markChannelIdle();
+              runBidiStreamingRpc(method, observer);
+            });
+    callProvider.setConnectivityAttemptTimer(connectivityAttemptTimer);
+
     return new ForwardingClientCall<ReqT, RespT>() {
       @Override
       protected ClientCall<ReqT, RespT> delegate() {

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

@@ -23,13 +23,15 @@
 import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.firestore.core.DatabaseInfo;
 import com.google.firebase.firestore.util.AsyncQueue;
+import com.google.firebase.firestore.util.AsyncQueue.DelayedTask;
 import com.google.firebase.firestore.util.Executors;
 import com.google.firebase.firestore.util.Logger;
 import com.google.firebase.firestore.util.Supplier;
 import com.google.firestore.v1.FirestoreGrpc;
 import io.grpc.CallCredentials;
 import io.grpc.CallOptions;
 import io.grpc.ClientCall;
+import io.grpc.ConnectivityState;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.MethodDescriptor;
@@ -50,6 +52,10 @@ public class GrpcCallProvider {
 
   private CallOptions callOptions;
 
+  // A timer that elapses after CONNECTIVTY_ATTEMPT_TIMEOUT_MS, at which point we close the
+  // stream, reset the underlying connection, and try connecting again.
+  private DelayedTask connectivityAttemptTimer;
+
   /**
    * Helper function to globally override the channel that RPCs use. Useful for testing when you
    * want to bypass SSL certificate checking.
@@ -89,6 +95,17 @@ public static void overrideChannelBuilder(
             });
   }
 
+  /**
+   * Marks the underlying gRPC channel as idle. This allows on-going RPCs to continue, but the next
+   * RPC on the channel will trigger the creation of a new connection.
+   */
+  void markChannelIdle() {
+    ManagedChannel channel = this.channelTask.getResult();
+    if (channel != null) {
+      this.channelTask.getResult().enterIdle();
+    }
+  }
+
   /** Sets up the SSL provider and configures the gRPC channel. */
   private ManagedChannel initChannel(Context context, DatabaseInfo databaseInfo) {
     try {
@@ -198,4 +215,48 @@ void shutdown() {
       Thread.currentThread().interrupt();
     }
   }
+
+  /**
+   * Sets the connectivity attempt timer to track. Existing timers must complete or be cancelled
+   * before a new timer can be set.
+   */
+  void setConnectivityAttemptTimer(DelayedTask connectivityAttemptTimer) {
+    // If an existing timer is already running, we want to continue using that one.
+    if (this.connectivityAttemptTimer != null) {
+      connectivityAttemptTimer.cancel();
+    } else {
+      try {
+        ManagedChannel channel = this.channelTask.getResult();
+        if (channel != null) {
+          this.connectivityAttemptTimer = connectivityAttemptTimer;
+          listenToConnectivityState(channel);
+        }
+      } catch (IllegalStateException e) {
+        // When the ManagedChannel is first initialized, channelTask could still be running.
+        connectivityAttemptTimer.cancel();
+      }
+    }
+  }
+
+  /**
+   * Cancels the connectivityStateTimer if the new state indicates grpc is online. Otherwise, we
+   * reset the listener for the next state change.
+   */
+  private void listenToConnectivityState(ManagedChannel channel) {
+    ConnectivityState newState = channel.getState(false);
+    // Check that the new state is online, then cancel timer.
+    if (newState == ConnectivityState.READY) {
+      clearConnectivityTimer();
+    } else {
+      channel.notifyWhenStateChanged(newState, () -> listenToConnectivityState(channel));
+    }
+  }
+
+  /** Clears the connectivity timer if it exists. */
+  void clearConnectivityTimer() {
+    if (connectivityAttemptTimer != null) {
+      connectivityAttemptTimer.cancel();
+      connectivityAttemptTimer = null;
+    }
+  }
 }

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

@@ -143,8 +143,8 @@ void handleWatchStreamFailure(Status status) {
         logClientOfflineWarningIfNecessary(
             String.format(
                 Locale.ENGLISH,
-                "Connection failed %d times. Most recent error: %s",
-                MAX_WATCH_STREAM_FAILURES,
+                "Backend didn't respond within %d seconds. Most recent error: %s\n",
+                ONLINE_STATE_TIMEOUT_MS / 1000,
                 status));
         setAndBroadcastState(OnlineState.OFFLINE);
       }

firebase-firestore/src/main/java/com/google/firebase/firestore/util/AsyncQueue.java

@@ -74,7 +74,13 @@ public enum TimerId {
      * A timer used to retry transactions. Since there can be multiple concurrent transactions,
      * multiple of these may be in the queue at a given time.
      */
-    RETRY_TRANSACTION
+    RETRY_TRANSACTION,
+    /**
+     * A timer used in RemoteStore to monitor when a connection attempt is unsuccessful and retry
+     * accordingly. While `ONLINE_STATE_TIMEOUT` is used to transition from UNKNOWN to OFFLINE,
+     * `CONNECTIVITY_ATTEMPT_TIMER` is used for tracking and retrying connectivity attempts.
+     */
+    CONNECTIVITY_ATTEMPT_TIMER
   }
 
   /**

firebase-firestore/src/main/java/com/google/firebase/firestore/util/ExponentialBackoff.java

@@ -36,8 +36,8 @@ public class ExponentialBackoff {
   private final TimerId timerId;
   private final long initialDelayMs;
   private final double backoffFactor;
-  private final long maxDelayMs;
 
+  private long maxDelayMs;
   private long currentBaseMs;
   private long lastAttemptTime;
   private DelayedTask timerTask;
@@ -103,6 +103,16 @@ public void resetToMax() {
     currentBaseMs = maxDelayMs;
   }
 
+  /**
+   * Set the backoff's maximum delay for only the next call to backoffAndRun, after which the delay
+   * will be reset to maxDelayMs.
+   *
+   * @param newMax The temporary maximum delay to set.
+   */
+  public void setTemporaryMaxDelay(long newMax) {
+    maxDelayMs = newMax;
+  }
+
   /**
    * Waits for currentDelayMs, increases the delay and runs the specified task. If there was a
    * pending backoff task waiting to run already, it will be canceled.
@@ -151,6 +161,9 @@ public void backoffAndRun(Runnable task) {
     } else if (currentBaseMs > maxDelayMs) {
       currentBaseMs = maxDelayMs;
     }
+
+    // Reset max delay to the default.
+    maxDelayMs = DEFAULT_BACKOFF_MAX_DELAY_MS;
   }
 
   public void cancel() {

firebase-firestore/src/test/java/com/google/firebase/firestore/remote/MockDatastore.java

@@ -67,6 +67,12 @@ public void stop() {
       open = false;
     }
 
+    @Override
+    protected void close(State finalState, Status status, boolean forceNewConnection) {
+      open = false;
+      listener.onClose(status);
+    }
+
     @Override
     public boolean isStarted() {
       return open;