Fix deadlock between scheduler thread and websocket thread when server revokes auth token.

Also change assert() to FIREBASE_DEV_ASSERT() to compile out in release build.

PiperOrigin-RevId: 292447554

Author: chkuang-g

database/src/desktop/connection/connection.cc

@@ -14,10 +14,10 @@
 
 #include "database/src/desktop/connection/connection.h"
 
-#include <cassert>
 #include <cstdlib>
 #include <cstring>
 
+#include "app/src/assert.h"
 #include "app/src/log.h"
 #include "app/src/variant_util.h"
 #include "database/src/desktop/connection/util_connection.h"
@@ -62,17 +62,17 @@ Connection::Connection(scheduler::Scheduler* scheduler, const HostInfo& info,
       client_(nullptr),
       expected_incoming_frames_(0),
       logger_(logger) {
-  assert(scheduler);
-  assert(event_handler);
+  FIREBASE_DEV_ASSERT(scheduler);
+  FIREBASE_DEV_ASSERT(event_handler);
 
   // Create log id like "[conn_0]" for debugging
   std::stringstream log_id_stream;
   log_id_stream << "[conn_" << next_log_id_.fetch_add(1) << "]";
   log_id_ = log_id_stream.str();
 
   // Create web socket client regardless of its implementation
-  client_ =
-      CreateWebSocketClient(host_info_, this, opt_last_session_id, logger);
+  client_ = CreateWebSocketClient(host_info_, this, opt_last_session_id, logger,
+                                  scheduler);
 }
 
 Connection::~Connection() {
@@ -88,13 +88,13 @@ Connection::~Connection() {
   scheduler_->Schedule(new callback::CallbackValue1<scheduler::RequestHandle>(
       keep_alive_handler_, [](scheduler::RequestHandle handler) {
         if (handler.IsValid() && !handler.IsCancelled()) {
-          assert(handler.Cancel());
+          FIREBASE_DEV_ASSERT(handler.Cancel());
         }
       }));
 }
 
 void Connection::Open() {
-  assert(client_);
+  FIREBASE_DEV_ASSERT(client_);
   if (state_ != kStateNone) {
     logger_->LogError("%s Cannot open. Connection has be opened before",
                       log_id_.c_str());
@@ -107,7 +107,7 @@ void Connection::Open() {
 }
 
 void Connection::Close(DisconnectReason reason /* = kReasonOther */) {
-  assert(client_);
+  FIREBASE_DEV_ASSERT(client_);
 
   if (state_ == kStateDisconnected) {
     logger_->LogError("%s Cannot close. Connection has been closed.",
@@ -131,8 +131,8 @@ void Connection::Close(DisconnectReason reason /* = kReasonOther */) {
 }
 
 void Connection::Send(const Variant& message, bool is_sensitive) {
-  assert(client_);
-  assert(!message.is_null());
+  FIREBASE_DEV_ASSERT(client_);
+  FIREBASE_DEV_ASSERT(!message.is_null());
 
   if (state_ != kStateReady) {
     logger_->LogError("%s Tried to send on an unconnected connection",
@@ -174,67 +174,48 @@ void Connection::OnOpen() {
 
   logger_->LogDebug("%s websocket opened", log_id_.c_str());
 
-  scheduler_->Schedule(new callback::CallbackValue1<ConnectionRef>(
-      safe_this_, [](ConnectionRef conn_ref) {
-        ConnectionRefLock lock(&conn_ref);
-        auto connection = lock.GetReference();
-        if (connection != nullptr) {
-          assert(connection->state_ == kStateConnecting);
-
-          connection->ws_connected_ = true;
-
-          // Start periodic callback to keep the connection alive, by sending
-          // "0" to server
-          connection->keep_alive_handler_ = connection->scheduler_->Schedule(
-              new callback::CallbackValue1<ConnectionRef>(
-                  connection->safe_this_,
-                  [](ConnectionRef conn_ref) {
-                    ConnectionRefLock lock(&conn_ref);
-                    auto connection = lock.GetReference();
-                    if (connection != nullptr && connection->client_ &&
-                        connection->state_ == kStateReady) {
-                      connection->client_->Send("0");
-                    }
-                  }),
-              kKeepAliveTimeoutMs, kKeepAliveTimeoutMs);
-        }
-      }));
+  FIREBASE_DEV_ASSERT(state_ == kStateConnecting);
+
+  ws_connected_ = true;
+
+  // Start periodic callback to keep the connection alive, by sending
+  // "0" to server
+  keep_alive_handler_ = scheduler_->Schedule(
+      new callback::CallbackValue1<ConnectionRef>(
+          safe_this_,
+          [](ConnectionRef conn_ref) {
+            ConnectionRefLock lock(&conn_ref);
+            auto connection = lock.GetReference();
+            if (connection != nullptr && connection->client_ &&
+                connection->state_ == kStateReady) {
+              connection->client_->Send("0");
+            }
+          }),
+      kKeepAliveTimeoutMs, kKeepAliveTimeoutMs);
 }
 
 void Connection::OnMessage(const char* msg) {
   SAFE_REFERENCE_RETURN_VOID_IF_INVALID(ConnectionRefLock, lock, safe_this_);
 
   logger_->LogDebug("%s websocket message received", log_id_.c_str());
-  scheduler_->Schedule(new callback::CallbackValue1String1<ConnectionRef>(
-      safe_this_, msg, [](ConnectionRef conn_ref, const char* msg) {
-        ConnectionRefLock lock(&conn_ref);
-        auto connection = lock.GetReference();
-        if (connection != nullptr) {
-          connection->HandleIncomingFrame(msg);
-        }
-      }));
+
+  HandleIncomingFrame(msg);
 }
 
 void Connection::OnClose() {
   SAFE_REFERENCE_RETURN_VOID_IF_INVALID(ConnectionRefLock, lock, safe_this_);
 
   logger_->LogDebug("%s websocket closed", log_id_.c_str());
 
-  scheduler_->Schedule(new callback::CallbackValue1<ConnectionRef>(
-      safe_this_, [](ConnectionRef conn_ref) {
-        ConnectionRefLock lock(&conn_ref);
-        auto connection = lock.GetReference();
-        if (connection != nullptr && connection->state_ != kStateDisconnected) {
-          // No need to do anything if Close() has been called already.
-          // Otherwise, the cause could be either connection failure or
-          // connection lost, depending on whether the web socket has already
-          // been connected or not.
-          DisconnectReason reason = connection->ws_connected_
-                                        ? kDisconnectReasonConnectionLost
-                                        : kDisconnectReasonConnectionFailed;
-          connection->Close(reason);
-        }
-      }));
+  if (state_ != kStateDisconnected) {
+    // No need to do anything if Close() has been called already.
+    // Otherwise, the cause could be either connection failure or
+    // connection lost, depending on whether the web socket has already
+    // been connected or not.
+    DisconnectReason reason = ws_connected_ ? kDisconnectReasonConnectionLost
+                                            : kDisconnectReasonConnectionFailed;
+    Close(reason);
+  }
 }
 
 void Connection::OnError(const WebSocketClientErrorData& error_data) {
@@ -311,7 +292,7 @@ void Connection::ProcessMessage(const char* message) {
   logger_->LogDebug("%s ProcessMessage (length: %d)", log_id_.c_str(),
                     strlen(message));
 
-  assert(!message_data.is_null());
+  FIREBASE_DEV_ASSERT(!message_data.is_null());
 
   const auto& messageMap = message_data.map();
   auto itType = messageMap.find(kServerEnvelopeType);
@@ -357,7 +338,7 @@ void Connection::OnControlMessage(const Variant& data) {
   logger_->LogDebug("%s received control message: %s", log_id_.c_str(),
                     util::VariantToJson(data).c_str());
 
-  assert(!data.is_null());
+  FIREBASE_DEV_ASSERT(!data.is_null());
 
   const auto& data_map = data.map();
   auto itType = data_map.find(kServerControlMessageType);

database/src/desktop/connection/persistent_connection.cc

@@ -15,7 +15,6 @@
 #include "database/src/desktop/connection/persistent_connection.h"
 
 #include <algorithm>
-#include <cassert>
 #include <sstream>
 
 #include "app/src/app_common.h"
@@ -118,9 +117,9 @@ PersistentConnection::PersistentConnection(
       next_listen_id_(0),
       next_write_id_(0),
       logger_(logger) {
-  assert(app);
-  assert(scheduler);
-  assert(event_handler_);
+  FIREBASE_DEV_ASSERT(app);
+  FIREBASE_DEV_ASSERT(scheduler);
+  FIREBASE_DEV_ASSERT(event_handler_);
 
   // Create log id like "[pc_0]" for debugging
   std::stringstream log_id_stream;
@@ -208,7 +207,7 @@ void PersistentConnection::OnReady(int64_t timestamp,
 
   // Restore Auth
   logger_->LogDebug("%s calling restore state", log_id_.c_str());
-  assert(connection_state_ == kConnecting);
+  FIREBASE_DEV_ASSERT(connection_state_ == kConnecting);
 
   // Try to retrieve auth token synchronously when connection is ready.
   GetAuthToken(&auth_token_);
@@ -247,26 +246,26 @@ void PersistentConnection::HandleConnectStatsResponse(
 }
 
 void PersistentConnection::OnDataMessage(const Variant& message) {
-  assert(message.is_map());
+  FIREBASE_DEV_ASSERT(message.is_map());
 
   SAFE_REFERENCE_RETURN_VOID_IF_INVALID(ThisRefLock, lock, safe_this_);
 
   if (HasKey(message, kRequestNumber)) {
     auto it_request_number = message.map().find(kRequestNumber);
-    assert(it_request_number->second.is_numeric());
+    FIREBASE_DEV_ASSERT(it_request_number->second.is_numeric());
     uint64_t rn = it_request_number->second.int64_value();
 
     RequestDataPtr request_ptr;
     auto it_request = request_map_.find(rn);
-    assert(it_request != request_map_.end());
+    FIREBASE_DEV_ASSERT(it_request != request_map_.end());
     if (it_request != request_map_.end()) {
       request_ptr = Move(it_request->second);
       request_map_.erase(it_request);
     }
-    assert(request_ptr);
+    FIREBASE_DEV_ASSERT(request_ptr);
     if (request_ptr) {
       auto it_response_message = message.map().find(kResponseForRequest);
-      assert(it_response_message != message.map().end());
+      FIREBASE_DEV_ASSERT(it_response_message != message.map().end());
       if (it_response_message != message.map().end()) {
         logger_->LogDebug("%s Trigger handler for request %llu",
                           log_id_.c_str(), rn);
@@ -510,33 +509,41 @@ void PersistentConnection::TryScheduleReconnect() {
     return;
   }
 
-  assert(connection_state_ == kDisconnected);
+  FIREBASE_DEV_ASSERT(connection_state_ == kDisconnected);
   bool force_refresh = force_auth_refresh_;
   force_auth_refresh_ = false;
   logger_->LogDebug("%s Scheduling connection attempt", log_id_.c_str());
 
-  // TODO(chkuang): Implement Exponential Backoff Retry
-  connection_state_ = kGettingToken;
-  logger_->LogDebug("%s Trying to fetch auth token", log_id_.c_str());
-
-  // Get Token Asynchronously to make sure the token is not expired.
-  Future<std::string> future;
-  bool succeeded = app_->function_registry()->CallFunction(
-      ::firebase::internal::FnAuthGetTokenAsync, app_, &force_refresh, &future);
-  if (succeeded && future.status() != kFutureStatusInvalid) {
-    // Set pending future
-    MutexLock future_lock(pending_token_future_mutex_);
-    pending_token_future_ = future;
-    future.OnCompletion(OnTokenFutureComplete, this);
-  } else {
-    // Auth is not available now.  Start the connection anyway.
-    OpenNetworkConnection();
-  }
+  scheduler_->Schedule(new callback::CallbackValue2<ThisRef, bool>(
+      safe_this_, force_refresh, [](ThisRef ref, bool force_refresh) {
+        ThisRefLock lock(&ref);
+        auto* connection = lock.GetReference();
+        if (!connection) return;
+        // TODO(chkuang): Implement Exponential Backoff Retry
+        connection->connection_state_ = kGettingToken;
+        connection->logger_->LogDebug("%s Trying to fetch auth token",
+                                      connection->log_id_.c_str());
+
+        // Get Token Asynchronously to make sure the token is not expired.
+        Future<std::string> future;
+        bool succeeded = connection->app_->function_registry()->CallFunction(
+            ::firebase::internal::FnAuthGetTokenAsync, connection->app_,
+            &force_refresh, &future);
+        if (succeeded && future.status() != kFutureStatusInvalid) {
+          // Set pending future
+          MutexLock future_lock(connection->pending_token_future_mutex_);
+          connection->pending_token_future_ = future;
+          future.OnCompletion(OnTokenFutureComplete, connection);
+        } else {
+          // Auth is not available now.  Start the connection anyway.
+          connection->OpenNetworkConnection();
+        }
+      }));
 }
 
 void PersistentConnection::OnTokenFutureComplete(
     const Future<std::string>& result_data, void* user_data) {
-  assert(user_data);
+  FIREBASE_DEV_ASSERT(user_data);
   PersistentConnection* connection =
       static_cast<PersistentConnection*>(user_data);
   ThisRefLock lock(&connection->safe_this_);
@@ -568,7 +575,7 @@ void PersistentConnection::HandleTokenFuture(Future<std::string> future) {
       auth_token_ = *future.result();
       OpenNetworkConnection();
     } else {
-      assert(connection_state_ == kDisconnected);
+      FIREBASE_DEV_ASSERT(connection_state_ == kDisconnected);
       logger_->LogDebug(
           "%s Not opening connection after token refresh, because "
           "connection was set to disconnected",
@@ -583,7 +590,7 @@ void PersistentConnection::HandleTokenFuture(Future<std::string> future) {
 }
 
 void PersistentConnection::OpenNetworkConnection() {
-  assert(connection_state_ == kGettingToken);
+  FIREBASE_DEV_ASSERT(connection_state_ == kGettingToken);
 
   // User might have logged out. Positive auth status is handled after
   // authenticating with the server
@@ -823,10 +830,10 @@ void PersistentConnection::PutInternal(const char* action, const Path& path,
 }
 
 void PersistentConnection::SendPut(uint64_t write_id) {
-  assert(CanSendWrites());
+  FIREBASE_DEV_ASSERT(CanSendWrites());
 
   auto it_put = outstanding_puts_.find(write_id);
-  assert(it_put != outstanding_puts_.end());
+  FIREBASE_DEV_ASSERT(it_put != outstanding_puts_.end());
 
   it_put->second->MarkSent();
   SendSensitive(it_put->second->action.c_str(), false, it_put->second->data,
@@ -908,7 +915,7 @@ void PersistentConnection::SendSensitive(const char* action, bool sensitive,
                                          ResponsePtr response,
                                          ConnectionResponseHandler callback,
                                          uint64_t outstanding_id) {
-  assert(realtime_);
+  FIREBASE_DEV_ASSERT(realtime_);
 
   // Varient only accept int64_t
   int64_t rn = ++next_request_id_;
@@ -924,7 +931,7 @@ void PersistentConnection::SendSensitive(const char* action, bool sensitive,
 }
 
 void PersistentConnection::RestoreOutstandingRequests() {
-  assert(connection_state_ == kConnected);
+  FIREBASE_DEV_ASSERT(connection_state_ == kConnected);
 
   // Restore listens
   logger_->LogDebug("%s Restoring outstanding listens", log_id_.c_str());
@@ -950,7 +957,7 @@ void PersistentConnection::RestoreOutstandingRequests() {
 }
 
 void PersistentConnection::GetAuthToken(std::string* out) {
-  assert(out);
+  FIREBASE_DEV_ASSERT(out);
   app_->function_registry()->CallFunction(
       ::firebase::internal::FnAuthGetCurrentToken, app_, nullptr, out);
 }
@@ -1010,7 +1017,7 @@ void PersistentConnection::SendUnauth() {
 void PersistentConnection::HandleAuthTokenResponse(const Variant& message,
                                                    const ResponsePtr& response,
                                                    uint64_t outstanding_id) {
-  assert(response);
+  FIREBASE_DEV_ASSERT(response);
 
   connection_state_ = kConnected;
 

database/src/desktop/connection/util_connection.cc

@@ -23,11 +23,12 @@ namespace connection {
 
 UniquePtr<WebSocketClientInterface> CreateWebSocketClient(
     const HostInfo& info, WebSocketClientEventHandler* delegate,
-    const char* opt_last_session_id, Logger* logger) {
+    const char* opt_last_session_id, Logger* logger,
+    scheduler::Scheduler* scheduler) {
   // Currently we use uWebSockets implementation.
   std::string uri = info.GetConnectionUrl(opt_last_session_id);
   return MakeUnique<WebSocketClientImpl>(uri, info.user_agent(), logger,
-                                         delegate);
+                                         scheduler, delegate);
 }
 
 }  // namespace connection