Don't invalidate token if we were able to receive data (#8766)

Firestore/core/src/remote/stream.cc

@@ -53,6 +53,8 @@ const AsyncQueue::Milliseconds kBackoffInitialDelay{std::chrono::seconds(1)};
 const AsyncQueue::Milliseconds kBackoffMaxDelay{std::chrono::seconds(60)};
 /** The time a stream stays open after it is marked idle. */
 const AsyncQueue::Milliseconds kIdleTimeout{std::chrono::seconds(60)};
+/** The time a stream stays open until we consider it healthy. */
+const AsyncQueue::Milliseconds kHealthyTimeout{std::chrono::seconds(10)};
 
 }  // namespace
 
@@ -60,20 +62,22 @@ Stream::Stream(const std::shared_ptr<AsyncQueue>& worker_queue,
                std::shared_ptr<AuthCredentialsProvider> credentials_provider,
                GrpcConnection* grpc_connection,
                TimerId backoff_timer_id,
-               TimerId idle_timer_id)
+               TimerId idle_timer_id,
+               TimerId health_check_timer_id)
     : backoff_{worker_queue, backoff_timer_id, kBackoffFactor,
                kBackoffInitialDelay, kBackoffMaxDelay},
       credentials_provider_{std::move(credentials_provider)},
       worker_queue_{worker_queue},
       grpc_connection_{grpc_connection},
-      idle_timer_id_{idle_timer_id} {
+      idle_timer_id_{idle_timer_id},
+      health_check_timer_id_{health_check_timer_id} {
 }
 
 // Check state
 
 bool Stream::IsOpen() const {
   EnsureOnQueue();
-  return state_ == State::Open;
+  return state_ == State::Open || state_ == State::Healthy;
 }
 
 bool Stream::IsStarted() const {
@@ -147,6 +151,15 @@ void Stream::OnStreamStart() {
 
   state_ = State::Open;
   NotifyStreamOpen();
+
+  health_check_ = worker_queue_->EnqueueAfterDelay(
+      kHealthyTimeout, health_check_timer_id_, [this] {
+        {
+          if (IsOpen()) {
+            state_ = State::Healthy;
+          }
+        }
+      });
 }
 
 // Backoff
@@ -249,6 +262,7 @@ void Stream::Close(const Status& status) {
   // execute).
   CancelIdleCheck();
   backoff_.Cancel();
+  health_check_.Cancel();
 
   // Step 3 (both): increment close count, which invalidates long-lived
   // callbacks, guaranteeing they won't execute against a new instance of the
@@ -288,9 +302,14 @@ void Stream::HandleErrorStatus(const Status& status) {
         "%s Using maximum backoff delay to prevent overloading the backend.",
         GetDebugDescription());
     backoff_.ResetToMax();
-  } else if (status.code() == Error::kErrorUnauthenticated) {
-    // "unauthenticated" error means the token was rejected. Try force
-    // refreshing it in case it just expired.
+  } else if (status.code() == Error::kErrorUnauthenticated &&
+             state_ != State::Healthy) {
+    // "unauthenticated" error means the token was rejected. This should rarely
+    // happen since both Auth and AppCheck ensure a sufficient TTL when we
+    // request a token. If a user manually resets their system clock this can
+    // fail, however. In this case, we should get a kErrorUnauthenticated error
+    // before we received the first message and we need to invalidate the token
+    // to ensure that we fetch a new token.
     credentials_provider_->InvalidateToken();
   }
 }

Firestore/core/src/remote/stream.h

@@ -104,6 +104,13 @@ class Stream : public GrpcStreamObserver,
      */
     Open,
 
+    /**
+     * The stream is healthy and has been connected for more than 10 seconds. We
+     * therefore assume that the credentials we passed were valid.
+     * Both `IsStarted` and `IsOpen` will return true.
+     */
+    Healthy,
+
     /**
      * The stream encountered an error. The next start attempt will back off.
      * While in this state, `IsStarted` will return false.
@@ -124,7 +131,8 @@ class Stream : public GrpcStreamObserver,
              credentials_provider,
          GrpcConnection* grpc_connection,
          util::TimerId backoff_timer_id,
-         util::TimerId idle_timer_id);
+         util::TimerId idle_timer_id,
+         util::TimerId health_check_timer_id);
 
   /**
    * Starts the stream. Only allowed if `IsStarted` returns false. The stream is
@@ -233,7 +241,9 @@ class Stream : public GrpcStreamObserver,
   GrpcConnection* grpc_connection_ = nullptr;
 
   util::TimerId idle_timer_id_{};
+  util::TimerId health_check_timer_id_{};
   util::DelayedOperation idleness_timer_;
+  util::DelayedOperation health_check_;
 
   // Used to prevent auth if the stream happens to be restarted before token is
   // received.

Firestore/core/src/remote/watch_stream.cc

@@ -45,8 +45,12 @@ WatchStream::WatchStream(
     Serializer serializer,
     GrpcConnection* grpc_connection,
     WatchStreamCallback* callback)
-    : Stream{async_queue, std::move(credentials_provider), grpc_connection,
-             TimerId::ListenStreamConnectionBackoff, TimerId::ListenStreamIdle},
+    : Stream{async_queue,
+             std::move(credentials_provider),
+             grpc_connection,
+             TimerId::ListenStreamConnectionBackoff,
+             TimerId::ListenStreamIdle,
+             TimerId::HealthCheckTimeout},
       watch_serializer_{std::move(serializer)},
       callback_{NOT_NULL(callback)} {
 }

Firestore/core/src/remote/write_stream.cc

@@ -46,8 +46,12 @@ WriteStream::WriteStream(
     Serializer serializer,
     GrpcConnection* grpc_connection,
     WriteStreamCallback* callback)
-    : Stream{async_queue, std::move(credentials_provider), grpc_connection,
-             TimerId::WriteStreamConnectionBackoff, TimerId::WriteStreamIdle},
+    : Stream{async_queue,
+             std::move(credentials_provider),
+             grpc_connection,
+             TimerId::WriteStreamConnectionBackoff,
+             TimerId::WriteStreamIdle,
+             TimerId::HealthCheckTimeout},
       write_serializer_{std::move(serializer)},
       callback_{NOT_NULL(callback)} {
 }

Firestore/core/src/util/async_queue.h

@@ -40,15 +40,17 @@ enum class TimerId {
   All,
 
   /**
-   * The following 4 timers are used in `Stream` for the listen and write
+   * The following 5 timers are used in `Stream` for the listen and write
    * streams. The "Idle" timer is used to close the stream due to inactivity.
    * The "ConnectionBackoff" timer is used to restart a stream once the
-   * appropriate backoff delay has elapsed.
+   * appropriate backoff delay has elapsed. The health check is used to mark
+   * a stream healthy if it has not received an error during its initial setup.
    */
   ListenStreamIdle,
   ListenStreamConnectionBackoff,
   WriteStreamIdle,
   WriteStreamConnectionBackoff,
+  HealthCheckTimeout,
 
   /**
    * A timer used in `OnlineStateTracker` to transition from

Firestore/core/test/unit/remote/stream_test.cc

@@ -55,14 +55,19 @@ using Type = GrpcCompletion::Type;
 
 const auto kIdleTimerId = TimerId::ListenStreamIdle;
 const auto kBackoffTimerId = TimerId::ListenStreamConnectionBackoff;
+const auto kHealthCheckTimerId = TimerId::HealthCheckTimeout;
 
 class TestStream : public Stream {
  public:
   TestStream(const std::shared_ptr<AsyncQueue>& worker_queue,
              GrpcStreamTester* tester,
              std::shared_ptr<AuthCredentialsProvider> credentials_provider)
-      : Stream{worker_queue, credentials_provider,
-               /*GrpcConnection=*/nullptr, kBackoffTimerId, kIdleTimerId},
+      : Stream{worker_queue,
+               credentials_provider,
+               /*GrpcConnection=*/nullptr,
+               kBackoffTimerId,
+               kIdleTimerId,
+               kHealthCheckTimerId},
         tester_{tester} {
   }
 
@@ -519,6 +524,16 @@ TEST_F(StreamTest, RefreshesTokenUponExpiration) {
             States({"GetToken", "InvalidateToken", "GetToken"}));
 }
 
+TEST_F(StreamTest, TokenIsNotInvalidatedOnceStreamIsHealthy) {
+  StartStream();
+  worker_queue->RunScheduledOperationsUntil(kHealthCheckTimerId);
+  ForceFinish({{Type::Read, CompletionResult::Error},
+               {Type::Finish, grpc::Status{grpc::UNAUTHENTICATED, ""}}});
+  // Error "Unauthenticated" on a healthy connection should not invalidate the
+  // token.
+  EXPECT_EQ(credentials->observed_states(), States({"GetToken"}));
+}
+
 }  // namespace remote
 }  // namespace firestore
 }  // namespace firebase