chore: fix user awareness fetch (#5048)

* chore: fix user awareness fetch

* chore: update

* chore: update

frontend/rust-lib/flowy-user/src/event_handler.rs

@@ -7,7 +7,7 @@ use lib_infra::box_any::BoxAny;
 use serde_json::Value;
 use std::sync::Weak;
 use std::{convert::TryInto, sync::Arc};
-use tracing::event;
+use tracing::{event, trace};
 
 use crate::entities::*;
 use crate::notification::{send_notification, UserNotification};
@@ -562,6 +562,8 @@ pub async fn get_all_reminder_event_handler(
     .into_iter()
     .map(ReminderPB::from)
     .collect::<Vec<_>>();
+
+  trace!("number of reminders: {}", reminders.len());
   data_result_ok(reminders.into())
 }
 

frontend/rust-lib/flowy-user/src/services/data_import/appflowy_data_import.rs

@@ -26,7 +26,7 @@ use flowy_folder_pub::entities::{AppFlowyData, ImportData};
 use flowy_folder_pub::folder_builder::{ParentChildViews, ViewBuilder};
 use flowy_sqlite::kv::StorePreferences;
 use flowy_user_pub::cloud::{UserCloudService, UserCollabParams};
-use flowy_user_pub::entities::{awareness_oid_from_user_uuid, Authenticator};
+use flowy_user_pub::entities::{user_awareness_object_id, Authenticator};
 use flowy_user_pub::session::Session;
 use parking_lot::{Mutex, RwLock};
 use std::collections::{HashMap, HashSet};
@@ -128,8 +128,13 @@ pub(crate) fn import_appflowy_data_folder(
     all_imported_object_ids.retain(|id| id != &imported_session.user_workspace.id);
     all_imported_object_ids
       .retain(|id| id != &imported_session.user_workspace.workspace_database_object_id);
-    all_imported_object_ids
-      .retain(|id| id != &awareness_oid_from_user_uuid(&imported_session.user_uuid).to_string());
+    all_imported_object_ids.retain(|id| {
+      id != &user_awareness_object_id(
+        &imported_session.user_uuid,
+        &imported_session.user_workspace.id,
+      )
+      .to_string()
+    });
 
     // import database view tracker
     migrate_database_view_tracker(

frontend/rust-lib/flowy-user/src/user_manager/manager.rs

@@ -13,7 +13,7 @@ use flowy_user_pub::entities::*;
 use flowy_user_pub::workspace_service::UserWorkspaceService;
 use serde_json::Value;
 use std::string::ToString;
-use std::sync::atomic::{AtomicI64, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
 use std::sync::{Arc, Weak};
 use tokio::sync::{Mutex, RwLock};
 use tokio_stream::StreamExt;
@@ -37,7 +37,6 @@ use crate::services::data_import::importer::import_data;
 use crate::services::data_import::ImportContext;
 
 use crate::services::sqlite_sql::user_sql::{select_user_profile, UserTable, UserTableChangeset};
-use crate::user_manager::manager_user_awareness::UserAwarenessDataSource;
 use crate::user_manager::manager_user_encryption::validate_encryption_sign;
 use crate::user_manager::manager_user_workspace::save_user_workspaces;
 use crate::user_manager::user_login_state::UserAuthProcess;
@@ -55,6 +54,7 @@ pub struct UserManager {
   auth_process: Mutex<Option<UserAuthProcess>>,
   pub(crate) authenticate_user: Arc<AuthenticateUser>,
   refresh_user_profile_since: AtomicI64,
+  pub(crate) is_loading_awareness: Arc<AtomicBool>,
 }
 
 impl UserManager {
@@ -80,6 +80,7 @@ impl UserManager {
       authenticate_user,
       refresh_user_profile_since,
       user_workspace_service,
+      is_loading_awareness: Arc::new(AtomicBool::new(false)),
     });
 
     let weak_user_manager = Arc::downgrade(&user_manager);
@@ -252,9 +253,7 @@ impl UserManager {
       self.authenticate_user.vacuum_database_if_need();
       let cloud_config = get_cloud_config(session.user_id, &self.store_preferences);
       // Init the user awareness
-      self
-        .initialize_user_awareness(&session, UserAwarenessDataSource::Local)
-        .await;
+      self.initialize_user_awareness(&session).await;
 
       user_status_callback
         .did_init(
@@ -324,10 +323,7 @@ impl UserManager {
       .save_auth_data(&response, &authenticator, &session)
       .await?;
 
-    let _ = self
-      .initialize_user_awareness(&session, UserAwarenessDataSource::Remote)
-      .await;
-
+    let _ = self.initialize_user_awareness(&session).await;
     self
       .user_status_callback
       .read()
@@ -412,12 +408,6 @@ impl UserManager {
   ) -> FlowyResult<()> {
     let new_session = Session::from(&response);
     self.prepare_user(&new_session).await;
-
-    let user_awareness_source = if response.is_new_user {
-      UserAwarenessDataSource::Local
-    } else {
-      UserAwarenessDataSource::Remote
-    };
     self
       .save_auth_data(&response, authenticator, &new_session)
       .await?;
@@ -433,10 +423,6 @@ impl UserManager {
       )
       .await?;
 
-    self
-      .initialize_user_awareness(&new_session, user_awareness_source)
-      .await;
-
     if response.is_new_user {
       if let Some(old_user) = migration_user {
         event!(

frontend/rust-lib/flowy-user/src/user_manager/manager_user_awareness.rs

@@ -1,16 +1,17 @@
+use std::sync::atomic::Ordering;
 use std::sync::{Arc, Weak};
 
 use anyhow::Context;
 use collab::core::collab::{DocStateSource, MutexCollab};
 use collab_entity::reminder::Reminder;
 use collab_entity::CollabType;
-use collab_integrate::collab_builder::CollabBuilderConfig;
+use collab_integrate::collab_builder::{AppFlowyCollabBuilder, CollabBuilderConfig};
 use collab_user::core::{MutexUserAwareness, UserAwareness};
-use tracing::{error, instrument, trace};
+use tracing::{error, info, instrument, trace};
 
 use collab_integrate::CollabKVDB;
 use flowy_error::{ErrorCode, FlowyError, FlowyResult};
-use flowy_user_pub::entities::awareness_oid_from_user_uuid;
+use flowy_user_pub::entities::user_awareness_object_id;
 
 use crate::entities::ReminderPB;
 use crate::user_manager::UserManager;
@@ -99,12 +100,8 @@ impl UserManager {
       .await
   }
 
-  pub async fn initialize_user_awareness(
-    &self,
-    session: &Session,
-    source: UserAwarenessDataSource,
-  ) {
-    match self.try_initial_user_awareness(session, source).await {
+  pub async fn initialize_user_awareness(&self, session: &Session) {
+    match self.try_initial_user_awareness(session).await {
       Ok(_) => trace!("User awareness initialized"),
       Err(e) => error!("Failed to initialize user awareness: {:?}", e),
     }
@@ -123,35 +120,77 @@ impl UserManager {
   /// # Returns
   /// - Returns `Ok(())` if the user's awareness is successfully initialized.
   /// - May return errors of type `FlowyError` if any issues arise during the initialization.
-  #[instrument(level = "info", skip(self, session, source), err)]
-  async fn try_initial_user_awareness(
-    &self,
-    session: &Session,
-    source: UserAwarenessDataSource,
-  ) -> FlowyResult<()> {
-    trace!("Initializing user awareness from {:?}", source);
+  #[instrument(level = "info", skip(self, session), err)]
+  async fn try_initial_user_awareness(&self, session: &Session) -> FlowyResult<()> {
+    if self.is_loading_awareness.load(Ordering::SeqCst) {
+      return Ok(());
+    }
+    self.is_loading_awareness.store(true, Ordering::SeqCst);
+
+    let object_id =
+      user_awareness_object_id(&session.user_uuid, &session.user_workspace.id).to_string();
+    trace!("Initializing user awareness {}", object_id);
     let collab_db = self.get_collab_db(session.user_id)?;
-    let user_awareness = match source {
-      UserAwarenessDataSource::Local => {
-        let collab = self
-          .collab_for_user_awareness(session, collab_db, vec![])
-          .await?;
-        MutexUserAwareness::new(UserAwareness::create(collab, None))
-      },
-      UserAwarenessDataSource::Remote => {
-        let data = self
-          .cloud_services
+    let weak_cloud_services = Arc::downgrade(&self.cloud_services);
+    let weak_user_awareness = Arc::downgrade(&self.user_awareness);
+    let weak_builder = self.collab_builder.clone();
+    let weak_is_loading_awareness = Arc::downgrade(&self.is_loading_awareness);
+    let session = session.clone();
+    tokio::spawn(async move {
+      if let (Some(cloud_services), Some(user_awareness)) =
+        (weak_cloud_services.upgrade(), weak_user_awareness.upgrade())
+      {
+        let result = cloud_services
           .get_user_service()?
-          .get_user_awareness_doc_state(session.user_id)
-          .await?;
-        trace!("Get user awareness collab: {}", data.len());
-        let collab = self
-          .collab_for_user_awareness(session, collab_db, data)
-          .await?;
-        MutexUserAwareness::new(UserAwareness::create(collab, None))
-      },
-    };
-    self.user_awareness.lock().await.replace(user_awareness);
+          .get_user_awareness_doc_state(session.user_id, &session.user_workspace.id, &object_id)
+          .await;
+
+        if let Some(is_loading_awareness) = weak_is_loading_awareness.upgrade() {
+          is_loading_awareness.store(false, Ordering::SeqCst);
+        }
+
+        let awareness = match result {
+          Ok(data) => {
+            trace!("Get user awareness collab from remote: {}", data.len());
+            let collab = Self::collab_for_user_awareness(
+              &weak_builder,
+              session.user_id,
+              &object_id,
+              collab_db,
+              DocStateSource::FromDocState(data),
+            )
+            .await?;
+            MutexUserAwareness::new(UserAwareness::create(collab, None))
+          },
+          Err(err) => {
+            if err.is_record_not_found() {
+              info!("User awareness not found, creating new");
+              let collab = Self::collab_for_user_awareness(
+                &weak_builder,
+                session.user_id,
+                &object_id,
+                collab_db,
+                DocStateSource::FromDisk,
+              )
+              .await?;
+              MutexUserAwareness::new(UserAwareness::create(collab, None))
+            } else {
+              error!("Failed to fetch user awareness: {:?}", err);
+              return Err(err);
+            }
+          },
+        };
+
+        user_awareness.lock().await.replace(awareness);
+        Ok(())
+      } else {
+        if let Some(is_loading_awareness) = weak_is_loading_awareness.upgrade() {
+          is_loading_awareness.store(false, Ordering::SeqCst);
+        }
+        Ok(())
+      }
+    });
+
     Ok(())
   }
 
@@ -161,22 +200,22 @@ impl UserManager {
   /// using a collaboration builder. This instance is specifically geared towards handling
   /// user awareness.
   async fn collab_for_user_awareness(
-    &self,
-    session: &Session,
+    collab_builder: &Weak<AppFlowyCollabBuilder>,
+    uid: i64,
+    object_id: &str,
     collab_db: Weak<CollabKVDB>,
-    doc_state: Vec<u8>,
+    doc_state: DocStateSource,
   ) -> Result<Arc<MutexCollab>, FlowyError> {
-    let collab_builder = self.collab_builder.upgrade().ok_or(FlowyError::new(
+    let collab_builder = collab_builder.upgrade().ok_or(FlowyError::new(
       ErrorCode::Internal,
       "Unexpected error: collab builder is not available",
     ))?;
-    let user_awareness_id = awareness_oid_from_user_uuid(&session.user_uuid);
     let collab = collab_builder
       .build(
-        session.user_id,
-        &user_awareness_id.to_string(),
+        uid,
+        object_id,
         CollabType::UserAwareness,
-        DocStateSource::FromDocState(doc_state),
+        doc_state,
         collab_db,
         CollabBuilderConfig::default().sync_enable(true),
       )
@@ -204,9 +243,7 @@ impl UserManager {
     match &*user_awareness {
       None => {
         if let Ok(session) = self.get_session() {
-          self
-            .initialize_user_awareness(&session, UserAwarenessDataSource::Remote)
-            .await;
+          self.initialize_user_awareness(&session).await;
         }
         default_value
       },
@@ -214,12 +251,3 @@ impl UserManager {
     }
   }
 }
-
-/// Indicate using which data source to initialize the user awareness
-/// If the user is not a new user, the local data source is used. Otherwise, the remote data source is used.
-/// When using the remote data source, the user awareness will be initialized from the remote server.
-#[derive(Debug)]
-pub enum UserAwarenessDataSource {
-  Local,
-  Remote,
-}