chore: use rayon to write snapshot updates (#1423)

* chore: use rayon to write snapshot updates

* chore: use rayon

* chore: rayon deps

* chore: client api check

* fix: encode version

* chore: rename function, add verbose logs

* chore: use default sv

* chore: logs

* chore: replaying small collab on current thread

Author: appflowy

Cargo.lock

@@ -256,6 +256,7 @@ impl WorkspaceControllerActor {
     Ok(())
   }
 
+  #[instrument(level = "trace", skip_all)]
   pub(crate) fn set_connection_status(&self, status: ConnectionStatus) {
     sync_trace!("set connection status: {:?}", status);
     self.status_tx.send_replace(status);
@@ -305,6 +306,7 @@ impl WorkspaceControllerActor {
     }
   }
 
+  #[instrument(level = "trace", skip_all)]
   async fn handle_action(actor: &Arc<Self>, action: WorkspaceAction) {
     let id = actor.db.client_id();
     sync_trace!("[{}] action {:?}", id, action);
@@ -358,6 +360,7 @@ impl WorkspaceControllerActor {
     }
   }
 
+  #[instrument(level = "trace", skip_all)]
   async fn handle_send(&self, msg: ClientMessage, source: ActionSource) -> anyhow::Result<()> {
     if let ClientMessage::Update {
       object_id,
@@ -422,7 +425,7 @@ impl WorkspaceControllerActor {
         .borrow()
         .disconnected_reason()
         .as_ref()
-        .map(|v| v.retriable() || matches!(v, DisconnectedReason::ReachMaximumRetry))
+        .map(|v| v.retriable_when_editing())
         .unwrap_or(false);
 
       if should_retry {
@@ -441,6 +444,7 @@ impl WorkspaceControllerActor {
     }
   }
 
+  #[instrument(level = "trace", skip_all)]
   pub(crate) async fn handle_connect(
     actor: &Arc<Self>,
     access_token: String,
@@ -469,7 +473,12 @@ impl WorkspaceControllerActor {
 
     match result {
       None => {
-        actor.set_connection_status(ConnectionStatus::Disconnected { reason: None });
+        sync_info!("[{}] connection established failed", actor.db.client_id());
+        actor.set_connection_status(ConnectionStatus::Disconnected {
+          reason: Some(DisconnectedReason::Unexpected(
+            "Establish connect failed".into(),
+          )),
+        });
       },
       Some(connection) => {
         sync_debug!("[{}] connected to {}", client_id, actor.options.url);
@@ -502,6 +511,7 @@ impl WorkspaceControllerActor {
     let reason = Self::remote_receiver_loop(weak_actor.clone(), stream, cancel.clone())
       .await
       .err();
+
     if let Some(actor) = weak_actor.upgrade() {
       if reason.is_some() {
         sync_error!("failed to receive messages from server: {:?}", reason);
@@ -521,7 +531,7 @@ impl WorkspaceControllerActor {
     while let Some(res) = stream.next().await {
       if cancel.is_cancelled() {
         sync_trace!("remote receiver loop cancelled");
-        break;
+        return Err(DisconnectedReason::UserDisconnect("User disconnect".into()));
       }
       let actor = match weak_actor.upgrade() {
         Some(inner) => inner,
@@ -580,8 +590,7 @@ impl WorkspaceControllerActor {
               frame.reason
             ),
           }
-          cancel.cancel();
-          break;
+          return Err(DisconnectedReason::ServerForceClose);
         },
       }
     }
@@ -1112,17 +1121,17 @@ impl WorkspaceControllerActor {
       res = fut => {
         match res {
           Ok((stream, _resp)) => {
-            info!("establishing WebSocket successfully");
+            sync_info!("establishing WebSocket successfully");
             Ok(Some(stream))
           },
           Err(err) => {
-            error!("establishing WebSocket failed");
+            sync_error!("establishing WebSocket failed");
             Err(AppError::from(err).into())
           }
         }
       }
       _ = cancel.cancelled() => {
-        tracing::debug!("connection cancelled");
+        tracing::debug!("establishing connection cancelled");
         Ok(None)
       }
     }

libs/client-api/src/v2/actor.rs

@@ -22,6 +22,7 @@ use tokio_stream::StreamExt;
 use tokio_tungstenite::tungstenite::error::ProtocolError;
 use tokio_tungstenite::tungstenite::{Error, Message};
 use tokio_util::sync::CancellationToken;
+use tracing::instrument;
 use yrs::block::ClientID;
 
 #[derive(Clone)]
@@ -188,6 +189,7 @@ pub enum DisconnectedReason {
   MessageLoopEnd(Arc<str>),
   CannotHandleReceiveMessage(Arc<str>),
   UserDisconnect(Arc<str>),
+  ServerForceClose,
   Unauthorized(Arc<str>),
 }
 
@@ -206,6 +208,7 @@ impl Display for DisconnectedReason {
       },
       DisconnectedReason::UserDisconnect(reason) => write!(f, "user disconnect: {}", reason),
       DisconnectedReason::Unauthorized(reason) => write!(f, "unauthorized: {}", reason),
+      DisconnectedReason::ServerForceClose => write!(f, "server force close"),
     }
   }
 }
@@ -247,6 +250,16 @@ impl DisconnectedReason {
       Self::Unexpected(..) | Self::ResetWithoutClosingHandshake
     )
   }
+
+  pub fn retriable_when_editing(&self) -> bool {
+    matches!(
+      self,
+      Self::Unexpected(..)
+        | Self::ResetWithoutClosingHandshake
+        | DisconnectedReason::Unauthorized(_)
+        | DisconnectedReason::ReachMaximumRetry
+    )
+  }
 }
 
 #[derive(Debug, Clone)]
@@ -345,6 +358,7 @@ impl ReconnectTarget for WorkspaceControllerActor {
     WorkspaceControllerActor::handle_connect(&self, token).await
   }
 
+  #[instrument(level = "trace", skip_all)]
   fn set_disconnected(&self, reason: DisconnectedReason) {
     self.set_connection_status(ConnectionStatus::Disconnected {
       reason: Some(reason),

libs/client-api/src/v2/controller.rs

@@ -35,6 +35,7 @@ impl CollabUpdateSink {
       self.collab_type,
       &msg.sender,
       msg.data.clone(),
+      msg.flags,
     )
     .query_async(&mut *lock)
     .await?;

libs/collab-stream/src/collab_update_sink.rs

@@ -492,6 +492,7 @@ impl UpdateStreamMessage {
     collab_type: CollabType,
     sender: &CollabOrigin,
     update: Vec<u8>,
+    flag: UpdateFlags,
   ) -> Cmd {
     let mut cmd = cmd("XADD");
     cmd
@@ -504,7 +505,9 @@ impl UpdateStreamMessage {
       .arg("sender")
       .arg(sender.to_string())
       .arg("data")
-      .arg(update.to_vec());
+      .arg(update.to_vec())
+      .arg("flags")
+      .arg(flag);
     cmd
   }
 }

libs/collab-stream/src/model.rs

@@ -22,7 +22,7 @@ uuid.workspace = true
 async-trait.workspace = true
 serde_json.workspace = true
 anyhow.workspace = true
-infra.workspace = true
+infra = { workspace = true }
 prometheus-client = "0.22.3"
 chrono = "0.4.39"
 serde.workspace = true

libs/indexer/Cargo.toml

@@ -4,6 +4,5 @@ pub mod error;
 pub mod metrics;
 pub mod queue;
 pub mod scheduler;
-pub mod thread_pool;
 mod unindexed_workspace;
 pub mod vector;

libs/indexer/src/lib.rs

@@ -16,7 +16,9 @@ tokio = { workspace = true, optional = true }
 pin-project.workspace = true
 futures = "0.3.30"
 validator = { workspace = true, features = ["validator_derive", "derive"] }
+thiserror = { workspace = true }
+rayon = { workspace = true }
 
 [features]
 file_util = ["tokio/fs"]
-request_util = ["reqwest"]
+request_util = ["reqwest"]

libs/infra/Cargo.toml

@@ -4,4 +4,6 @@ pub mod env_util;
 pub mod file_util;
 #[cfg(feature = "request_util")]
 pub mod reqwest;
+
+pub mod thread_pool;
 pub mod validate;