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