send acked to client using ws

backend/Cargo.toml

@@ -57,6 +57,7 @@ parking_lot = "0.11"
 md5 = "0.7.0"
 futures-core = { version = "0.3", default-features = false }
 pin-project = "1.0.0"
+byteorder = {version = "1.3.4"}
 
 flowy-user = { path = "../rust-lib/flowy-user" }
 flowy-workspace = { path = "../rust-lib/flowy-workspace" }

backend/src/application.rs

@@ -21,7 +21,7 @@ use crate::{
         view::router as view,
         workspace::router as workspace,
         ws,
-        ws::WSServer,
+        ws::WsServer,
     },
 };
 
@@ -142,7 +142,7 @@ async fn init_app_context(configuration: &Settings) -> AppContext {
             configuration.database
         ));
 
-    let ws_server = WSServer::new().start();
+    let ws_server = WsServer::new().start();
 
     AppContext::new(ws_server, pg_pool)
 }

backend/src/context.rs

@@ -1,15 +1,15 @@
-use crate::service::ws::WSServer;
+use crate::service::ws::WsServer;
 use actix::Addr;
 
 use sqlx::PgPool;
 
 pub struct AppContext {
-    pub ws_server: Addr<WSServer>,
+    pub ws_server: Addr<WsServer>,
     pub pg_pool: PgPool,
 }
 
 impl AppContext {
-    pub fn new(ws_server: Addr<WSServer>, db_pool: PgPool) -> Self {
+    pub fn new(ws_server: Addr<WsServer>, db_pool: PgPool) -> Self {
         AppContext {
             ws_server,
             pg_pool: db_pool,

backend/src/service/doc/edit_doc.rs

@@ -1,14 +1,22 @@
-use crate::service::doc::update_doc;
+use crate::service::{
+    doc::update_doc,
+    ws::{entities::Socket, WsClientData, WsMessageAdaptor},
+};
 use actix_web::web::Data;
+use byteorder::{BigEndian, ByteOrder, WriteBytesExt};
+use bytes::Bytes;
 use flowy_document::{
+    entities::ws::{WsDataType, WsDocumentData},
     protobuf::{Doc, Revision, UpdateDocParams},
     services::doc::Document,
 };
 use flowy_net::errors::{internal_error, ServerError};
 use flowy_ot::core::Delta;
+use flowy_ws::{protobuf::WsModule, WsMessage};
 use parking_lot::RwLock;
+use protobuf::Message;
 use sqlx::PgPool;
-use std::{sync::Arc, time::Duration};
+use std::{convert::TryInto, sync::Arc, time::Duration};
 
 pub(crate) struct EditDoc {
     doc_id: String,
@@ -27,15 +35,31 @@ impl EditDoc {
         })
     }
 
-    #[tracing::instrument(level = "debug", skip(self, revision))]
-    pub(crate) async fn apply_revision(&self, revision: Revision) -> Result<(), ServerError> {
+    #[tracing::instrument(level = "debug", skip(self, socket, revision))]
+    pub(crate) async fn apply_revision(
+        &self,
+        socket: Socket,
+        revision: Revision,
+    ) -> Result<(), ServerError> {
         let delta = Delta::from_bytes(revision.delta).map_err(internal_error)?;
         match self.document.try_write_for(Duration::from_millis(300)) {
             None => {
                 log::error!("Failed to acquire write lock of document");
             },
-            Some(mut w) => {
-                let _ = w.apply_delta(delta).map_err(internal_error)?;
+            Some(mut write_guard) => {
+                let _ = write_guard.apply_delta(delta).map_err(internal_error)?;
+                let mut wtr = vec![];
+                let _ = wtr.write_i64::<BigEndian>(revision.rev_id);
+
+                let data = WsDocumentData {
+                    id: self.doc_id.clone(),
+                    ty: WsDataType::Acked,
+                    data: wtr,
+                };
+
+                let msg: WsMessage = data.into();
+                let bytes: Bytes = msg.try_into().unwrap();
+                socket.do_send(WsMessageAdaptor(bytes));
             },
         }
 

backend/src/service/doc/ws_handler.rs

@@ -1,5 +1,9 @@
 use super::edit_doc::EditDoc;
-use crate::service::{doc::read_doc, util::parse_from_bytes, ws::WsBizHandler};
+use crate::service::{
+    doc::read_doc,
+    util::parse_from_bytes,
+    ws::{WsBizHandler, WsClientData},
+};
 use actix_web::web::Data;
 use bytes::Bytes;
 use flowy_document::{
@@ -13,22 +17,22 @@ use sqlx::PgPool;
 use std::{collections::HashMap, sync::Arc};
 
 pub struct DocWsBizHandler {
-    inner: Arc<Inner>,
+    doc_manager: Arc<EditDocManager>,
 }
 
 impl DocWsBizHandler {
     pub fn new(pg_pool: Data<PgPool>) -> Self {
         Self {
-            inner: Arc::new(Inner::new(pg_pool)),
+            doc_manager: Arc::new(EditDocManager::new(pg_pool)),
         }
     }
 }
 
 impl WsBizHandler for DocWsBizHandler {
-    fn receive_data(&self, data: Bytes) {
-        let inner = self.inner.clone();
+    fn receive_data(&self, client_data: WsClientData) {
+        let doc_manager = self.doc_manager.clone();
         actix_rt::spawn(async move {
-            let result = inner.handle(data).await;
+            let result = doc_manager.handle(client_data).await;
             match result {
                 Ok(_) => {},
                 Err(e) => log::error!("WsBizHandler handle data error: {:?}", e),
@@ -37,12 +41,12 @@ impl WsBizHandler for DocWsBizHandler {
     }
 }
 
-struct Inner {
+struct EditDocManager {
     pg_pool: Data<PgPool>,
     edit_docs: RwLock<HashMap<String, Arc<EditDoc>>>,
 }
 
-impl Inner {
+impl EditDocManager {
     fn new(pg_pool: Data<PgPool>) -> Self {
         Self {
             pg_pool,
@@ -50,16 +54,22 @@ impl Inner {
         }
     }
 
-    async fn handle(&self, data: Bytes) -> Result<(), ServerError> {
-        let document_data: WsDocumentData = parse_from_bytes(&data)?;
+    async fn handle(&self, client_data: WsClientData) -> Result<(), ServerError> {
+        let document_data: WsDocumentData = parse_from_bytes(&client_data.data)?;
 
         match document_data.ty {
-            WsDataType::Command => {},
+            WsDataType::Acked => {},
             WsDataType::Delta => {
                 let revision: Revision = parse_from_bytes(&document_data.data)?;
                 let edited_doc = self.get_edit_doc(&revision.doc_id).await?;
                 tokio::spawn(async move {
-                    edited_doc.apply_revision(revision).await.unwrap();
+                    match edited_doc
+                        .apply_revision(client_data.socket, revision)
+                        .await
+                    {
+                        Ok(_) => {},
+                        Err(e) => log::error!("Doc apply revision failed: {:?}", e),
+                    }
                 });
             },
         }

backend/src/service/ws/biz_handler.rs

@@ -1,9 +1,10 @@
+use crate::service::ws::WsClientData;
 use bytes::Bytes;
 use flowy_ws::WsModule;
 use std::{collections::HashMap, sync::Arc};
 
 pub trait WsBizHandler: Send + Sync {
-    fn receive_data(&self, data: Bytes);
+    fn receive_data(&self, client_data: WsClientData);
 }
 
 pub type BizHandler = Arc<dyn WsBizHandler>;

backend/src/service/ws/entities/connect.rs

@@ -1,10 +1,10 @@
-use crate::service::ws::ClientMessage;
+use crate::service::ws::WsMessageAdaptor;
 use actix::{Message, Recipient};
 use flowy_net::errors::ServerError;
 use serde::{Deserialize, Serialize};
 use std::fmt::Formatter;
 
-pub type Socket = Recipient<ClientMessage>;
+pub type Socket = Recipient<WsMessageAdaptor>;
 
 #[derive(Serialize, Deserialize, Debug, Clone, Hash, PartialEq, Eq)]
 pub struct SessionId(pub String);

backend/src/service/ws/entities/message.rs

@@ -3,38 +3,12 @@ use actix::Message;
 use bytes::Bytes;
 use std::fmt::Formatter;
 
-#[derive(Debug, Clone)]
-pub enum MessageData {
-    Binary(Bytes),
-    Connect(SessionId),
-    Disconnect(SessionId),
-}
-
 #[derive(Debug, Message, Clone)]
 #[rtype(result = "()")]
-pub struct ClientMessage {
-    pub session_id: SessionId,
-    pub data: MessageData,
-}
+pub struct WsMessageAdaptor(pub Bytes);
 
-impl ClientMessage {
-    pub fn new<T: Into<SessionId>>(session_id: T, data: MessageData) -> Self {
-        ClientMessage {
-            session_id: session_id.into(),
-            data,
-        }
-    }
-}
+impl std::ops::Deref for WsMessageAdaptor {
+    type Target = Bytes;
 
-impl std::fmt::Display for ClientMessage {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        let content = match &self.data {
-            MessageData::Binary(_) => "[Binary]".to_owned(),
-            MessageData::Connect(_) => "[Connect]".to_owned(),
-            MessageData::Disconnect(_) => "[Disconnect]".to_owned(),
-        };
-
-        let desc = format!("{}:{}", &self.session_id, content);
-        f.write_str(&desc)
-    }
+    fn deref(&self) -> &Self::Target { &self.0 }
 }

backend/src/service/ws/router.rs

@@ -1,4 +1,4 @@
-use crate::service::ws::{WSClient, WSServer, WsBizHandlers};
+use crate::service::ws::{WsBizHandlers, WsClient, WsServer};
 use actix::Addr;
 
 use crate::service::user::LoggedUser;
@@ -16,12 +16,12 @@ pub async fn establish_ws_connection(
     request: HttpRequest,
     payload: Payload,
     token: Path<String>,
-    server: Data<Addr<WSServer>>,
+    server: Data<Addr<WsServer>>,
     biz_handlers: Data<WsBizHandlers>,
 ) -> Result<HttpResponse, Error> {
     match LoggedUser::from_token(token.clone()) {
         Ok(user) => {
-            let client = WSClient::new(&user.user_id, server.get_ref().clone(), biz_handlers);
+            let client = WsClient::new(&user.user_id, server.get_ref().clone(), biz_handlers);
             let result = ws::start(client, &request, payload);
             match result {
                 Ok(response) => Ok(response.into()),

backend/src/service/ws/ws_client.rs

@@ -1,12 +1,11 @@
 use crate::{
     config::{HEARTBEAT_INTERVAL, PING_TIMEOUT},
     service::ws::{
-        entities::{Connect, Disconnect, SessionId},
-        ClientMessage,
-        MessageData,
-        WSServer,
+        entities::{Connect, Disconnect, SessionId, Socket},
         WsBizHandler,
         WsBizHandlers,
+        WsMessageAdaptor,
+        WsServer,
     },
 };
 use actix::*;
@@ -16,17 +15,22 @@ use bytes::Bytes;
 use flowy_ws::WsMessage;
 use std::{convert::TryFrom, time::Instant};
 
-pub struct WSClient {
+pub struct WsClientData {
+    pub(crate) socket: Socket,
+    pub(crate) data: Bytes,
+}
+
+pub struct WsClient {
     session_id: SessionId,
-    server: Addr<WSServer>,
+    server: Addr<WsServer>,
     biz_handlers: Data<WsBizHandlers>,
     hb: Instant,
 }
 
-impl WSClient {
+impl WsClient {
     pub fn new<T: Into<SessionId>>(
         session_id: T,
-        server: Addr<WSServer>,
+        server: Addr<WsServer>,
         biz_handlers: Data<WsBizHandlers>,
     ) -> Self {
         Self {
@@ -50,24 +54,25 @@ impl WSClient {
         });
     }
 
-    fn send(&self, data: MessageData) {
-        let msg = ClientMessage::new(self.session_id.clone(), data);
-        self.server.do_send(msg);
-    }
-
-    fn handle_binary_message(&self, bytes: Bytes) {
+    fn handle_binary_message(&self, bytes: Bytes, socket: Socket) {
         // TODO: ok to unwrap?
         let message: WsMessage = WsMessage::try_from(bytes).unwrap();
         match self.biz_handlers.get(&message.module) {
             None => {
                 log::error!("Can't find the handler for {:?}", message.module);
             },
-            Some(handler) => handler.receive_data(Bytes::from(message.data)),
+            Some(handler) => {
+                let client_data = WsClientData {
+                    socket,
+                    data: Bytes::from(message.data),
+                };
+                handler.receive_data(client_data)
+            },
         }
     }
 }
 
-impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WSClient {
+impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WsClient {
     fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
         match msg {
             Ok(ws::Message::Ping(msg)) => {
@@ -80,13 +85,13 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WSClient {
             },
             Ok(ws::Message::Binary(bytes)) => {
                 log::debug!(" Receive {} binary", &self.session_id);
-                self.handle_binary_message(bytes);
+                let socket = ctx.address().recipient();
+                self.handle_binary_message(bytes, socket);
             },
             Ok(Text(_)) => {
                 log::warn!("Receive unexpected text message");
             },
             Ok(ws::Message::Close(reason)) => {
-                self.send(MessageData::Disconnect(self.session_id.clone()));
                 ctx.close(reason);
                 ctx.stop();
             },
@@ -104,21 +109,13 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WSClient {
     }
 }
 
-impl Handler<ClientMessage> for WSClient {
+impl Handler<WsMessageAdaptor> for WsClient {
     type Result = ();
 
-    fn handle(&mut self, msg: ClientMessage, ctx: &mut Self::Context) {
-        match msg.data {
-            MessageData::Binary(binary) => {
-                ctx.binary(binary);
-            },
-            MessageData::Connect(_) => {},
-            MessageData::Disconnect(_) => {},
-        }
-    }
+    fn handle(&mut self, msg: WsMessageAdaptor, ctx: &mut Self::Context) { ctx.binary(msg.0); }
 }
 
-impl Actor for WSClient {
+impl Actor for WsClient {
     type Context = ws::WebsocketContext<Self>;
 
     fn started(&mut self, ctx: &mut Self::Context) {

backend/src/service/ws/ws_server.rs

@@ -1,32 +1,31 @@
 use crate::service::ws::{
     entities::{Connect, Disconnect, Session, SessionId},
-    ClientMessage,
-    MessageData,
+    WsMessageAdaptor,
 };
 use actix::{Actor, Context, Handler};
 use dashmap::DashMap;
 use flowy_net::errors::ServerError;
 
-pub struct WSServer {
+pub struct WsServer {
     sessions: DashMap<SessionId, Session>,
 }
 
-impl WSServer {
+impl WsServer {
     pub fn new() -> Self {
         Self {
             sessions: DashMap::new(),
         }
     }
 
-    pub fn send(&self, _msg: ClientMessage) { unimplemented!() }
+    pub fn send(&self, _msg: WsMessageAdaptor) { unimplemented!() }
 }
 
-impl Actor for WSServer {
+impl Actor for WsServer {
     type Context = Context<Self>;
     fn started(&mut self, _ctx: &mut Self::Context) {}
 }
 
-impl Handler<Connect> for WSServer {
+impl Handler<Connect> for WsServer {
     type Result = Result<(), ServerError>;
     fn handle(&mut self, msg: Connect, _ctx: &mut Context<Self>) -> Self::Result {
         let session: Session = msg.into();
@@ -36,7 +35,7 @@ impl Handler<Connect> for WSServer {
     }
 }
 
-impl Handler<Disconnect> for WSServer {
+impl Handler<Disconnect> for WsServer {
     type Result = Result<(), ServerError>;
     fn handle(&mut self, msg: Disconnect, _: &mut Context<Self>) -> Self::Result {
         self.sessions.remove(&msg.sid);
@@ -44,20 +43,16 @@ impl Handler<Disconnect> for WSServer {
     }
 }
 
-impl Handler<ClientMessage> for WSServer {
+impl Handler<WsMessageAdaptor> for WsServer {
     type Result = ();
 
-    fn handle(&mut self, msg: ClientMessage, _ctx: &mut Context<Self>) -> Self::Result {
-        match msg.data {
-            MessageData::Binary(_) => {},
-            MessageData::Connect(_) => {},
-            MessageData::Disconnect(_) => {},
-        }
+    fn handle(&mut self, _msg: WsMessageAdaptor, _ctx: &mut Context<Self>) -> Self::Result {
+        unimplemented!()
     }
 }
 
-impl actix::Supervised for WSServer {
-    fn restarting(&mut self, _ctx: &mut Context<WSServer>) {
+impl actix::Supervised for WsServer {
+    fn restarting(&mut self, _ctx: &mut Context<WsServer>) {
         log::warn!("restarting");
     }
 }