receive revision from client

app_flowy/packages/flowy_sdk/lib/protobuf/flowy-ws/msg.pb.dart

@@ -9,16 +9,20 @@ import 'dart:core' as $core;
 
 import 'package:protobuf/protobuf.dart' as $pb;
 
+import 'msg.pbenum.dart';
+
+export 'msg.pbenum.dart';
+
 class WsMessage extends $pb.GeneratedMessage {
   static final $pb.BuilderInfo _i = $pb.BuilderInfo(const $core.bool.fromEnvironment('protobuf.omit_message_names') ? '' : 'WsMessage', createEmptyInstance: create)
-    ..aOS(1, const $core.bool.fromEnvironment('protobuf.omit_field_names') ? '' : 'source')
+    ..e<WsSource>(1, const $core.bool.fromEnvironment('protobuf.omit_field_names') ? '' : 'source', $pb.PbFieldType.OE, defaultOrMaker: WsSource.Doc, valueOf: WsSource.valueOf, enumValues: WsSource.values)
     ..a<$core.List<$core.int>>(2, const $core.bool.fromEnvironment('protobuf.omit_field_names') ? '' : 'data', $pb.PbFieldType.OY)
     ..hasRequiredFields = false
   ;
 
   WsMessage._() : super();
   factory WsMessage({
-    $core.String? source,
+    WsSource? source,
     $core.List<$core.int>? data,
   }) {
     final _result = create();
@@ -52,9 +56,9 @@ class WsMessage extends $pb.GeneratedMessage {
   static WsMessage? _defaultInstance;
 
   @$pb.TagNumber(1)
-  $core.String get source => $_getSZ(0);
+  WsSource get source => $_getN(0);
   @$pb.TagNumber(1)
-  set source($core.String v) { $_setString(0, v); }
+  set source(WsSource v) { setField(1, v); }
   @$pb.TagNumber(1)
   $core.bool hasSource() => $_has(0);
   @$pb.TagNumber(1)

app_flowy/packages/flowy_sdk/lib/protobuf/flowy-ws/msg.pbenum.dart

@@ -5,3 +5,20 @@
 // @dart = 2.12
 // ignore_for_file: annotate_overrides,camel_case_types,unnecessary_const,non_constant_identifier_names,library_prefixes,unused_import,unused_shown_name,return_of_invalid_type,unnecessary_this,prefer_final_fields
 
+// ignore_for_file: UNDEFINED_SHOWN_NAME
+import 'dart:core' as $core;
+import 'package:protobuf/protobuf.dart' as $pb;
+
+class WsSource extends $pb.ProtobufEnum {
+  static const WsSource Doc = WsSource._(0, const $core.bool.fromEnvironment('protobuf.omit_enum_names') ? '' : 'Doc');
+
+  static const $core.List<WsSource> values = <WsSource> [
+    Doc,
+  ];
+
+  static final $core.Map<$core.int, WsSource> _byValue = $pb.ProtobufEnum.initByValue(values);
+  static WsSource? valueOf($core.int value) => _byValue[value];
+
+  const WsSource._($core.int v, $core.String n) : super(v, n);
+}
+

app_flowy/packages/flowy_sdk/lib/protobuf/flowy-ws/msg.pbjson.dart

@@ -8,14 +8,24 @@
 import 'dart:core' as $core;
 import 'dart:convert' as $convert;
 import 'dart:typed_data' as $typed_data;
+@$core.Deprecated('Use wsSourceDescriptor instead')
+const WsSource$json = const {
+  '1': 'WsSource',
+  '2': const [
+    const {'1': 'Doc', '2': 0},
+  ],
+};
+
+/// Descriptor for `WsSource`. Decode as a `google.protobuf.EnumDescriptorProto`.
+final $typed_data.Uint8List wsSourceDescriptor = $convert.base64Decode('CghXc1NvdXJjZRIHCgNEb2MQAA==');
 @$core.Deprecated('Use wsMessageDescriptor instead')
 const WsMessage$json = const {
   '1': 'WsMessage',
   '2': const [
-    const {'1': 'source', '3': 1, '4': 1, '5': 9, '10': 'source'},
+    const {'1': 'source', '3': 1, '4': 1, '5': 14, '6': '.WsSource', '10': 'source'},
     const {'1': 'data', '3': 2, '4': 1, '5': 12, '10': 'data'},
   ],
 };
 
 /// Descriptor for `WsMessage`. Decode as a `google.protobuf.DescriptorProto`.
-final $typed_data.Uint8List wsMessageDescriptor = $convert.base64Decode('CglXc01lc3NhZ2USFgoGc291cmNlGAEgASgJUgZzb3VyY2USEgoEZGF0YRgCIAEoDFIEZGF0YQ==');
+final $typed_data.Uint8List wsMessageDescriptor = $convert.base64Decode('CglXc01lc3NhZ2USIQoGc291cmNlGAEgASgOMgkuV3NTb3VyY2VSBnNvdXJjZRISCgRkYXRhGAIgASgMUgRkYXRh');

backend/src/application.rs

@@ -16,6 +16,7 @@ use crate::{
     service::{
         app::router as app,
         doc::router as doc,
+        make_ws_biz_handlers,
         user::router as user,
         view::router as view,
         workspace::router as workspace,
@@ -23,6 +24,7 @@ use crate::{
         ws::WSServer,
     },
 };
+use flowy_ws::WsSource;
 
 pub struct Application {
     port: u16,
@@ -53,7 +55,7 @@ pub fn run(listener: TcpListener, app_ctx: AppContext) -> Result<Server, std::io
     let pg_pool = Data::new(pg_pool);
     let domain = domain();
     let secret: String = secret();
-
+    let ws_biz_handlers = Data::new(make_ws_biz_handlers());
     actix_rt::spawn(period_check(pg_pool.clone()));
 
     let server = HttpServer::new(move || {
@@ -67,6 +69,7 @@ pub fn run(listener: TcpListener, app_ctx: AppContext) -> Result<Server, std::io
             .service(user_scope())
             .app_data(ws_server.clone())
             .app_data(pg_pool.clone())
+            .app_data(ws_biz_handlers.clone())
     })
     .listen(listener)?
     .run();

backend/src/service/doc/mod.rs

@@ -1,6 +1,7 @@
 mod doc;
 pub mod router;
 mod sql_builder;
+pub mod ws_handler;
 
 pub(crate) use doc::*;
 pub use router::*;

backend/src/service/doc/ws_handler.rs

@@ -0,0 +1,17 @@
+use crate::service::{util::parse_from_bytes, ws::WsBizHandler};
+use bytes::Bytes;
+use flowy_document::protobuf::Revision;
+use protobuf::Message;
+
+pub struct DocWsBizHandler {}
+
+impl DocWsBizHandler {
+    pub fn new() -> Self { Self {} }
+}
+
+impl WsBizHandler for DocWsBizHandler {
+    fn receive_data(&self, data: Bytes) {
+        let revision: Revision = parse_from_bytes(&data).unwrap();
+        log::warn!("{:?}", revision);
+    }
+}

backend/src/service/mod.rs

@@ -1,3 +1,8 @@
+use crate::service::{doc::ws_handler::DocWsBizHandler, ws::WsBizHandlers};
+use flowy_ws::WsSource;
+use std::sync::Arc;
+use tokio::sync::RwLock;
+
 pub mod app;
 pub mod doc;
 pub(crate) mod log;
@@ -6,3 +11,16 @@ pub(crate) mod util;
 pub mod view;
 pub mod workspace;
 pub mod ws;
+
+pub fn make_ws_biz_handlers() -> WsBizHandlers {
+    let mut ws_biz_handlers = WsBizHandlers::new();
+
+    // doc
+    let doc_biz_handler = DocWsBizHandler::new();
+    ws_biz_handlers.register(WsSource::Doc, wrap(doc_biz_handler));
+
+    //
+    ws_biz_handlers
+}
+
+fn wrap<T>(val: T) -> Arc<RwLock<T>> { Arc::new(RwLock::new(val)) }

backend/src/service/ws/biz_handler.rs

@@ -0,0 +1,34 @@
+use bytes::Bytes;
+use dashmap::{mapref::one::Ref, DashMap};
+use flowy_ws::WsSource;
+use std::sync::Arc;
+use tokio::sync::RwLock;
+
+pub trait WsBizHandler: Send + Sync {
+    fn receive_data(&self, data: Bytes);
+}
+
+pub type BizHandler = Arc<RwLock<dyn WsBizHandler>>;
+
+pub struct WsBizHandlers {
+    inner: DashMap<WsSource, BizHandler>,
+}
+
+impl WsBizHandlers {
+    pub fn new() -> Self {
+        Self {
+            inner: DashMap::new(),
+        }
+    }
+
+    pub fn register(&self, source: WsSource, handler: BizHandler) {
+        self.inner.insert(source, handler);
+    }
+
+    pub fn get(&self, source: &WsSource) -> Option<BizHandler> {
+        match self.inner.get(source) {
+            None => None,
+            Some(handler) => Some(handler.clone()),
+        }
+    }
+}

backend/src/service/ws/mod.rs

@@ -1,7 +1,9 @@
+pub use biz_handler::*;
 pub use entities::message::*;
 pub use ws_client::*;
 pub use ws_server::*;
 
+mod biz_handler;
 pub(crate) mod entities;
 pub mod router;
 mod ws_client;

backend/src/service/ws/router.rs

@@ -1,4 +1,4 @@
-use crate::service::ws::{WSClient, WSServer};
+use crate::service::ws::{WSClient, WSServer, WsBizHandlers};
 use actix::Addr;
 
 use crate::service::user::LoggedUser;
@@ -17,10 +17,11 @@ pub async fn establish_ws_connection(
     payload: Payload,
     token: Path<String>,
     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());
+            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

@@ -5,41 +5,36 @@ use crate::{
         ClientMessage,
         MessageData,
         WSServer,
+        WsBizHandler,
+        WsBizHandlers,
     },
 };
-use actix::*;
+use actix::{fut::wrap_future, *};
+use actix_web::web::Data;
 use actix_web_actors::{ws, ws::Message::Text};
-use std::time::Instant;
+use bytes::Bytes;
+use flowy_ws::{WsMessage, WsSource};
+use std::{convert::TryFrom, pin::Pin, time::Instant};
+use tokio::sync::RwLock;
 
-//    Frontend          │                       Backend
-//
-//                      │
-// ┌──────────┐   WsMessage   ┌───────────┐  ClientMessage    ┌──────────┐
-// │  user 1  │─────────┼────▶│ws_client_1│──────────────────▶│ws_server │
-// └──────────┘               └───────────┘                   └──────────┘
-//                      │                                           │
-//                WsMessage                                         ▼
-// ┌──────────┐         │     ┌───────────┐    ClientMessage     Group
-// │  user 2  │◀──────────────│ws_client_2│◀───────┐        ┌───────────────┐
-// └──────────┘         │     └───────────┘        │        │  ws_user_1    │
-//                                                 │        │               │
-//                      │                          └────────│  ws_user_2    │
-// ┌──────────┐               ┌───────────┐                 │               │
-// │  user 3  │─────────┼────▶│ws_client_3│                 └───────────────┘
-// └──────────┘               └───────────┘
-//                      │
 pub struct WSClient {
     session_id: SessionId,
     server: Addr<WSServer>,
+    biz_handlers: Data<WsBizHandlers>,
     hb: Instant,
 }
 
 impl WSClient {
-    pub fn new<T: Into<SessionId>>(session_id: T, server: Addr<WSServer>) -> Self {
+    pub fn new<T: Into<SessionId>>(
+        session_id: T,
+        server: Addr<WSServer>,
+        biz_handlers: Data<WsBizHandlers>,
+    ) -> Self {
         Self {
             session_id: session_id.into(),
-            hb: Instant::now(),
             server,
+            biz_handlers,
+            hb: Instant::now(),
         }
     }
 
@@ -62,36 +57,16 @@ impl WSClient {
     }
 }
 
-impl Actor for WSClient {
+async fn handle_binary_message(biz_handlers: Data<WsBizHandlers>, bytes: Bytes) {
-    type Context = ws::WebsocketContext<Self>;
+    let message: WsMessage = WsMessage::try_from(bytes).unwrap();
-
+    match biz_handlers.get(&message.source) {
-    fn started(&mut self, ctx: &mut Self::Context) {
+        None => {
-        self.hb(ctx);
+            log::error!("Can't find the handler for {:?}", message.source);
-        let socket = ctx.address().recipient();
+        },
-        let connect = Connect {
+        Some(handler) => handler
-            socket,
+            .write()
-            sid: self.session_id.clone(),
+            .await
-        };
+            .receive_data(Bytes::from(message.data)),
-        self.server
-            .send(connect)
-            .into_actor(self)
-            .then(|res, _client, _ctx| {
-                match res {
-                    Ok(Ok(_)) => log::trace!("Send connect message to server success"),
-                    Ok(Err(e)) => log::error!("Send connect message to server failed: {:?}", e),
-                    Err(e) => log::error!("Send connect message to server failed: {:?}", e),
-                }
-                fut::ready(())
-            })
-            .wait(ctx);
-    }
-
-    fn stopping(&mut self, _: &mut Self::Context) -> Running {
-        self.server.do_send(Disconnect {
-            sid: self.session_id.clone(),
-        });
-
-        Running::Stop
     }
 }
 
@@ -106,9 +81,10 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WSClient {
                 // log::debug!("Receive {} pong {:?}", &self.session_id, &msg);
                 self.hb = Instant::now();
             },
-            Ok(ws::Message::Binary(bin)) => {
+            Ok(ws::Message::Binary(bytes)) => {
                 log::debug!(" Receive {} binary", &self.session_id);
-                self.send(MessageData::Binary(bin));
+                let biz_handlers = self.biz_handlers.clone();
+                ctx.spawn(wrap_future(handle_binary_message(biz_handlers, bytes)));
             },
             Ok(Text(_)) => {
                 log::warn!("Receive unexpected text message");
@@ -145,3 +121,36 @@ impl Handler<ClientMessage> for WSClient {
         }
     }
 }
+
+impl Actor for WSClient {
+    type Context = ws::WebsocketContext<Self>;
+
+    fn started(&mut self, ctx: &mut Self::Context) {
+        self.hb(ctx);
+        let socket = ctx.address().recipient();
+        let connect = Connect {
+            socket,
+            sid: self.session_id.clone(),
+        };
+        self.server
+            .send(connect)
+            .into_actor(self)
+            .then(|res, _client, _ctx| {
+                match res {
+                    Ok(Ok(_)) => log::trace!("Send connect message to server success"),
+                    Ok(Err(e)) => log::error!("Send connect message to server failed: {:?}", e),
+                    Err(e) => log::error!("Send connect message to server failed: {:?}", e),
+                }
+                fut::ready(())
+            })
+            .wait(ctx);
+    }
+
+    fn stopping(&mut self, _: &mut Self::Context) -> Running {
+        self.server.do_send(Disconnect {
+            sid: self.session_id.clone(),
+        });
+
+        Running::Stop
+    }
+}

rust-lib/flowy-document/src/protobuf/model/revision.rs

@@ -298,20 +298,20 @@ static file_descriptor_proto_data: &'static [u8] = b"\
     \n\x0erevision.proto\"i\n\x08Revision\x12\x1e\n\x0bbase_rev_id\x18\x01\
     \x20\x01(\x03R\tbaseRevId\x12\x15\n\x06rev_id\x18\x02\x20\x01(\x03R\x05r\
     evId\x12\x14\n\x05delta\x18\x03\x20\x01(\x0cR\x05delta\x12\x10\n\x03md5\
-    \x18\x04\x20\x01(\tR\x03md5J\x86\x02\n\x06\x12\x04\0\0\x06\x01\n\x08\n\
+    \x18\x04\x20\x01(\tR\x03md5J\x86\x02\n\x06\x12\x04\0\0\x07\x01\n\x08\n\
-    \x01\x0c\x12\x03\0\0\x12\n\n\n\x02\x04\0\x12\x04\x01\0\x06\x01\n\n\n\x03\
+    \x01\x0c\x12\x03\0\0\x12\n\n\n\x02\x04\0\x12\x04\x02\0\x07\x01\n\n\n\x03\
-    \x04\0\x01\x12\x03\x01\x08\x10\n\x0b\n\x04\x04\0\x02\0\x12\x03\x02\x04\
+    \x04\0\x01\x12\x03\x02\x08\x10\n\x0b\n\x04\x04\0\x02\0\x12\x03\x03\x04\
-    \x1a\n\x0c\n\x05\x04\0\x02\0\x05\x12\x03\x02\x04\t\n\x0c\n\x05\x04\0\x02\
+    \x1a\n\x0c\n\x05\x04\0\x02\0\x05\x12\x03\x03\x04\t\n\x0c\n\x05\x04\0\x02\
-    \0\x01\x12\x03\x02\n\x15\n\x0c\n\x05\x04\0\x02\0\x03\x12\x03\x02\x18\x19\
+    \0\x01\x12\x03\x03\n\x15\n\x0c\n\x05\x04\0\x02\0\x03\x12\x03\x03\x18\x19\
-    \n\x0b\n\x04\x04\0\x02\x01\x12\x03\x03\x04\x15\n\x0c\n\x05\x04\0\x02\x01\
+    \n\x0b\n\x04\x04\0\x02\x01\x12\x03\x04\x04\x15\n\x0c\n\x05\x04\0\x02\x01\
-    \x05\x12\x03\x03\x04\t\n\x0c\n\x05\x04\0\x02\x01\x01\x12\x03\x03\n\x10\n\
+    \x05\x12\x03\x04\x04\t\n\x0c\n\x05\x04\0\x02\x01\x01\x12\x03\x04\n\x10\n\
-    \x0c\n\x05\x04\0\x02\x01\x03\x12\x03\x03\x13\x14\n\x0b\n\x04\x04\0\x02\
+    \x0c\n\x05\x04\0\x02\x01\x03\x12\x03\x04\x13\x14\n\x0b\n\x04\x04\0\x02\
-    \x02\x12\x03\x04\x04\x14\n\x0c\n\x05\x04\0\x02\x02\x05\x12\x03\x04\x04\t\
+    \x02\x12\x03\x05\x04\x14\n\x0c\n\x05\x04\0\x02\x02\x05\x12\x03\x05\x04\t\
-    \n\x0c\n\x05\x04\0\x02\x02\x01\x12\x03\x04\n\x0f\n\x0c\n\x05\x04\0\x02\
+    \n\x0c\n\x05\x04\0\x02\x02\x01\x12\x03\x05\n\x0f\n\x0c\n\x05\x04\0\x02\
-    \x02\x03\x12\x03\x04\x12\x13\n\x0b\n\x04\x04\0\x02\x03\x12\x03\x05\x04\
+    \x02\x03\x12\x03\x05\x12\x13\n\x0b\n\x04\x04\0\x02\x03\x12\x03\x06\x04\
-    \x13\n\x0c\n\x05\x04\0\x02\x03\x05\x12\x03\x05\x04\n\n\x0c\n\x05\x04\0\
+    \x13\n\x0c\n\x05\x04\0\x02\x03\x05\x12\x03\x06\x04\n\n\x0c\n\x05\x04\0\
-    \x02\x03\x01\x12\x03\x05\x0b\x0e\n\x0c\n\x05\x04\0\x02\x03\x03\x12\x03\
+    \x02\x03\x01\x12\x03\x06\x0b\x0e\n\x0c\n\x05\x04\0\x02\x03\x03\x12\x03\
-    \x05\x11\x12b\x06proto3\
+    \x06\x11\x12b\x06proto3\
 ";
 
 static file_descriptor_proto_lazy: ::protobuf::rt::LazyV2<::protobuf::descriptor::FileDescriptorProto> = ::protobuf::rt::LazyV2::INIT;

rust-lib/flowy-document/src/protobuf/proto/revision.proto

@@ -1,4 +1,5 @@
 syntax = "proto3";
+
 message Revision {
     int64 base_rev_id = 1;
     int64 rev_id = 2;

rust-lib/flowy-document/src/services/cache.rs

@@ -8,23 +8,23 @@ use crate::{
 };
 
 pub(crate) struct DocCache {
-    doc_map: DashMap<DocId, Arc<EditDocContext>>,
+    inner: DashMap<DocId, Arc<EditDocContext>>,
 }
 
 impl DocCache {
-    pub(crate) fn new() -> Self { Self { doc_map: DashMap::new() } }
+    pub(crate) fn new() -> Self { Self { inner: DashMap::new() } }
 
     pub(crate) fn set(&self, doc: Arc<EditDocContext>) {
         let doc_id = doc.id.clone();
-        if self.doc_map.contains_key(&doc_id) {
+        if self.inner.contains_key(&doc_id) {
             log::warn!("Doc:{} already exists in cache", doc_id.as_ref());
         }
-        self.doc_map.insert(doc.id.clone(), doc);
+        self.inner.insert(doc.id.clone(), doc);
     }
 
     pub(crate) fn is_opened(&self, doc_id: &str) -> bool {
         let doc_id: DocId = doc_id.into();
-        self.doc_map.get(&doc_id).is_some()
+        self.inner.get(&doc_id).is_some()
     }
 
     pub(crate) fn get(&self, doc_id: &str) -> Result<Arc<EditDocContext>, DocError> {
@@ -32,13 +32,13 @@ impl DocCache {
             return Err(doc_not_found());
         }
         let doc_id: DocId = doc_id.into();
-        let opened_doc = self.doc_map.get(&doc_id).unwrap();
+        let opened_doc = self.inner.get(&doc_id).unwrap();
         Ok(opened_doc.clone())
     }
 
     pub(crate) fn remove(&self, id: &str) {
         let doc_id: DocId = id.into();
-        self.doc_map.remove(&doc_id);
+        self.inner.remove(&doc_id);
     }
 }
 

rust-lib/flowy-document/src/services/ws/ws_manager.rs

@@ -7,10 +7,6 @@ pub trait WsSender: Send + Sync {
     fn send_data(&self, data: Bytes) -> Result<(), DocError>;
 }
 
-lazy_static! {
-    pub static ref WS_ID: String = "Document".to_string();
-}
-
 pub struct WsManager {
     pub(crate) sender: Arc<dyn WsSender>,
     doc_handlers: HashMap<String, Arc<dyn WsHandler>>,

rust-lib/flowy-sdk/src/deps_resolve/document_deps.rs

@@ -2,11 +2,11 @@ use bytes::Bytes;
 use flowy_document::{
     errors::DocError,
     module::DocumentUser,
-    prelude::{WsManager, WsSender, WS_ID},
+    prelude::{WsManager, WsSender},
 };
 
 use flowy_user::{errors::ErrorCode, services::user::UserSession};
-use flowy_ws::{WsMessage, WsMessageHandler};
+use flowy_ws::{WsMessage, WsMessageHandler, WsSource};
 use parking_lot::RwLock;
 use std::{path::Path, sync::Arc};
 
@@ -73,7 +73,7 @@ struct WsSenderImpl {
 impl WsSender for WsSenderImpl {
     fn send_data(&self, data: Bytes) -> Result<(), DocError> {
         let msg = WsMessage {
-            source: WS_ID.clone(),
+            source: WsSource::Doc,
             data: data.to_vec(),
         };
         let _ = self.user.send_ws_msg(msg).map_err(|e| DocError::internal().context(e))?;
@@ -86,7 +86,7 @@ struct WsDocumentResolver {
 }
 
 impl WsMessageHandler for WsDocumentResolver {
-    fn source(&self) -> String { WS_ID.clone() }
+    fn source(&self) -> WsSource { WsSource::Doc }
 
     fn receive_message(&self, msg: WsMessage) {
         let data = Bytes::from(msg.data);

rust-lib/flowy-ws/Cargo.toml

@@ -14,7 +14,7 @@ futures-util = "0.3.17"
 futures-channel = "0.3.17"
 tokio = {version = "1", features = ["full"]}
 futures = "0.3.17"
-bytes = "0.5"
+bytes = "1.0"
 pin-project = "1.0.0"
 futures-core = { version = "0.3", default-features = false }
 paste = "1"

rust-lib/flowy-ws/src/msg.rs

@@ -1,34 +1,51 @@
 use bytes::Bytes;
-use flowy_derive::ProtoBuf;
+use flowy_derive::{ProtoBuf, ProtoBuf_Enum};
 use std::convert::{TryFrom, TryInto};
-use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::tungstenite::Message as TokioMessage;
 
 #[derive(ProtoBuf, Debug, Clone, Default)]
 pub struct WsMessage {
     #[pb(index = 1)]
-    pub source: String,
+    pub source: WsSource,
 
     #[pb(index = 2)]
     pub data: Vec<u8>,
 }
 
-impl std::convert::Into<Message> for WsMessage {
+#[derive(ProtoBuf_Enum, Debug, Clone, Eq, PartialEq, Hash)]
-    fn into(self) -> Message {
+pub enum WsSource {
+    Doc = 0,
+}
+
+impl std::default::Default for WsSource {
+    fn default() -> Self { WsSource::Doc }
+}
+
+impl ToString for WsSource {
+    fn to_string(&self) -> String {
+        match self {
+            WsSource::Doc => "0".to_string(),
+        }
+    }
+}
+
+impl std::convert::Into<TokioMessage> for WsMessage {
+    fn into(self) -> TokioMessage {
         let result: Result<Bytes, ::protobuf::ProtobufError> = self.try_into();
         match result {
-            Ok(bytes) => Message::Binary(bytes.to_vec()),
+            Ok(bytes) => TokioMessage::Binary(bytes.to_vec()),
             Err(e) => {
                 log::error!("WsMessage serialize error: {:?}", e);
-                Message::Binary(vec![])
+                TokioMessage::Binary(vec![])
             },
         }
     }
 }
 
-impl std::convert::From<Message> for WsMessage {
+impl std::convert::From<TokioMessage> for WsMessage {
-    fn from(value: Message) -> Self {
+    fn from(value: TokioMessage) -> Self {
         match value {
-            Message::Binary(bytes) => WsMessage::try_from(Bytes::from(bytes)).unwrap(),
+            TokioMessage::Binary(bytes) => WsMessage::try_from(Bytes::from(bytes)).unwrap(),
             _ => {
                 log::error!("WsMessage deserialize failed. Unsupported message");
                 WsMessage::default()

rust-lib/flowy-ws/src/protobuf/model/msg.rs

@@ -26,7 +26,7 @@
 #[derive(PartialEq,Clone,Default)]
 pub struct WsMessage {
     // message fields
-    pub source: ::std::string::String,
+    pub source: WsSource,
     pub data: ::std::vec::Vec<u8>,
     // special fields
     pub unknown_fields: ::protobuf::UnknownFields,
@@ -44,32 +44,21 @@ impl WsMessage {
         ::std::default::Default::default()
     }
 
-    // string source = 1;
+    // .WsSource source = 1;
 
 
-    pub fn get_source(&self) -> &str {
+    pub fn get_source(&self) -> WsSource {
-        &self.source
+        self.source
     }
     pub fn clear_source(&mut self) {
-        self.source.clear();
+        self.source = WsSource::Doc;
     }
 
     // Param is passed by value, moved
-    pub fn set_source(&mut self, v: ::std::string::String) {
+    pub fn set_source(&mut self, v: WsSource) {
         self.source = v;
     }
 
-    // Mutable pointer to the field.
-    // If field is not initialized, it is initialized with default value first.
-    pub fn mut_source(&mut self) -> &mut ::std::string::String {
-        &mut self.source
-    }
-
-    // Take field
-    pub fn take_source(&mut self) -> ::std::string::String {
-        ::std::mem::replace(&mut self.source, ::std::string::String::new())
-    }
-
     // bytes data = 2;
 
 
@@ -107,7 +96,7 @@ impl ::protobuf::Message for WsMessage {
             let (field_number, wire_type) = is.read_tag_unpack()?;
             match field_number {
                 1 => {
-                    ::protobuf::rt::read_singular_proto3_string_into(wire_type, is, &mut self.source)?;
+                    ::protobuf::rt::read_proto3_enum_with_unknown_fields_into(wire_type, is, &mut self.source, 1, &mut self.unknown_fields)?
                 },
                 2 => {
                     ::protobuf::rt::read_singular_proto3_bytes_into(wire_type, is, &mut self.data)?;
@@ -124,8 +113,8 @@ impl ::protobuf::Message for WsMessage {
     #[allow(unused_variables)]
     fn compute_size(&self) -> u32 {
         let mut my_size = 0;
-        if !self.source.is_empty() {
+        if self.source != WsSource::Doc {
-            my_size += ::protobuf::rt::string_size(1, &self.source);
+            my_size += ::protobuf::rt::enum_size(1, self.source);
         }
         if !self.data.is_empty() {
             my_size += ::protobuf::rt::bytes_size(2, &self.data);
@@ -136,8 +125,8 @@ impl ::protobuf::Message for WsMessage {
     }
 
     fn write_to_with_cached_sizes(&self, os: &mut ::protobuf::CodedOutputStream<'_>) -> ::protobuf::ProtobufResult<()> {
-        if !self.source.is_empty() {
+        if self.source != WsSource::Doc {
-            os.write_string(1, &self.source)?;
+            os.write_enum(1, ::protobuf::ProtobufEnum::value(&self.source))?;
         }
         if !self.data.is_empty() {
             os.write_bytes(2, &self.data)?;
@@ -180,7 +169,7 @@ impl ::protobuf::Message for WsMessage {
         static descriptor: ::protobuf::rt::LazyV2<::protobuf::reflect::MessageDescriptor> = ::protobuf::rt::LazyV2::INIT;
         descriptor.get(|| {
             let mut fields = ::std::vec::Vec::new();
-            fields.push(::protobuf::reflect::accessor::make_simple_field_accessor::<_, ::protobuf::types::ProtobufTypeString>(
+            fields.push(::protobuf::reflect::accessor::make_simple_field_accessor::<_, ::protobuf::types::ProtobufTypeEnum<WsSource>>(
                 "source",
                 |m: &WsMessage| { &m.source },
                 |m: &mut WsMessage| { &mut m.source },
@@ -206,7 +195,7 @@ impl ::protobuf::Message for WsMessage {
 
 impl ::protobuf::Clear for WsMessage {
     fn clear(&mut self) {
-        self.source.clear();
+        self.source = WsSource::Doc;
         self.data.clear();
         self.unknown_fields.clear();
     }
@@ -224,17 +213,68 @@ impl ::protobuf::reflect::ProtobufValue for WsMessage {
     }
 }
 
+#[derive(Clone,PartialEq,Eq,Debug,Hash)]
+pub enum WsSource {
+    Doc = 0,
+}
+
+impl ::protobuf::ProtobufEnum for WsSource {
+    fn value(&self) -> i32 {
+        *self as i32
+    }
+
+    fn from_i32(value: i32) -> ::std::option::Option<WsSource> {
+        match value {
+            0 => ::std::option::Option::Some(WsSource::Doc),
+            _ => ::std::option::Option::None
+        }
+    }
+
+    fn values() -> &'static [Self] {
+        static values: &'static [WsSource] = &[
+            WsSource::Doc,
+        ];
+        values
+    }
+
+    fn enum_descriptor_static() -> &'static ::protobuf::reflect::EnumDescriptor {
+        static descriptor: ::protobuf::rt::LazyV2<::protobuf::reflect::EnumDescriptor> = ::protobuf::rt::LazyV2::INIT;
+        descriptor.get(|| {
+            ::protobuf::reflect::EnumDescriptor::new_pb_name::<WsSource>("WsSource", file_descriptor_proto())
+        })
+    }
+}
+
+impl ::std::marker::Copy for WsSource {
+}
+
+impl ::std::default::Default for WsSource {
+    fn default() -> Self {
+        WsSource::Doc
+    }
+}
+
+impl ::protobuf::reflect::ProtobufValue for WsSource {
+    fn as_ref(&self) -> ::protobuf::reflect::ReflectValueRef {
+        ::protobuf::reflect::ReflectValueRef::Enum(::protobuf::ProtobufEnum::descriptor(self))
+    }
+}
+
 static file_descriptor_proto_data: &'static [u8] = b"\
-    \n\tmsg.proto\"7\n\tWsMessage\x12\x16\n\x06source\x18\x01\x20\x01(\tR\
+    \n\tmsg.proto\"B\n\tWsMessage\x12!\n\x06source\x18\x01\x20\x01(\x0e2\t.W\
-    \x06source\x12\x12\n\x04data\x18\x02\x20\x01(\x0cR\x04dataJ\x98\x01\n\
+    sSourceR\x06source\x12\x12\n\x04data\x18\x02\x20\x01(\x0cR\x04data*\x13\
-    \x06\x12\x04\0\0\x05\x01\n\x08\n\x01\x0c\x12\x03\0\0\x12\n\n\n\x02\x04\0\
+    \n\x08WsSource\x12\x07\n\x03Doc\x10\0J\xd9\x01\n\x06\x12\x04\0\0\x08\x01\
-    \x12\x04\x02\0\x05\x01\n\n\n\x03\x04\0\x01\x12\x03\x02\x08\x11\n\x0b\n\
+    \n\x08\n\x01\x0c\x12\x03\0\0\x12\n\n\n\x02\x04\0\x12\x04\x02\0\x05\x01\n\
-    \x04\x04\0\x02\0\x12\x03\x03\x04\x16\n\x0c\n\x05\x04\0\x02\0\x05\x12\x03\
+    \n\n\x03\x04\0\x01\x12\x03\x02\x08\x11\n\x0b\n\x04\x04\0\x02\0\x12\x03\
-    \x03\x04\n\n\x0c\n\x05\x04\0\x02\0\x01\x12\x03\x03\x0b\x11\n\x0c\n\x05\
+    \x03\x04\x18\n\x0c\n\x05\x04\0\x02\0\x06\x12\x03\x03\x04\x0c\n\x0c\n\x05\
-    \x04\0\x02\0\x03\x12\x03\x03\x14\x15\n\x0b\n\x04\x04\0\x02\x01\x12\x03\
+    \x04\0\x02\0\x01\x12\x03\x03\r\x13\n\x0c\n\x05\x04\0\x02\0\x03\x12\x03\
-    \x04\x04\x13\n\x0c\n\x05\x04\0\x02\x01\x05\x12\x03\x04\x04\t\n\x0c\n\x05\
+    \x03\x16\x17\n\x0b\n\x04\x04\0\x02\x01\x12\x03\x04\x04\x13\n\x0c\n\x05\
-    \x04\0\x02\x01\x01\x12\x03\x04\n\x0e\n\x0c\n\x05\x04\0\x02\x01\x03\x12\
+    \x04\0\x02\x01\x05\x12\x03\x04\x04\t\n\x0c\n\x05\x04\0\x02\x01\x01\x12\
-    \x03\x04\x11\x12b\x06proto3\
+    \x03\x04\n\x0e\n\x0c\n\x05\x04\0\x02\x01\x03\x12\x03\x04\x11\x12\n\n\n\
+    \x02\x05\0\x12\x04\x06\0\x08\x01\n\n\n\x03\x05\0\x01\x12\x03\x06\x05\r\n\
+    \x0b\n\x04\x05\0\x02\0\x12\x03\x07\x04\x0c\n\x0c\n\x05\x05\0\x02\0\x01\
+    \x12\x03\x07\x04\x07\n\x0c\n\x05\x05\0\x02\0\x02\x12\x03\x07\n\x0bb\x06p\
+    roto3\
 ";
 
 static file_descriptor_proto_lazy: ::protobuf::rt::LazyV2<::protobuf::descriptor::FileDescriptorProto> = ::protobuf::rt::LazyV2::INIT;

rust-lib/flowy-ws/src/protobuf/proto/msg.proto

@@ -1,6 +1,9 @@
 syntax = "proto3";
 
 message WsMessage {
-    string source = 1;
+    WsSource source = 1;
     bytes data = 2;
 }
+enum WsSource {
+    Doc = 0;
+}

rust-lib/flowy-ws/src/ws.rs

@@ -2,6 +2,7 @@ use crate::{
     connect::{Retry, WsConnectionFuture},
     errors::WsError,
     WsMessage,
+    WsSource,
 };
 use flowy_net::errors::ServerError;
 use futures_channel::mpsc::{UnboundedReceiver, UnboundedSender};
@@ -23,7 +24,7 @@ use tokio_tungstenite::tungstenite::{
 pub type MsgReceiver = UnboundedReceiver<Message>;
 pub type MsgSender = UnboundedSender<Message>;
 pub trait WsMessageHandler: Sync + Send + 'static {
-    fn source(&self) -> String;
+    fn source(&self) -> WsSource;
     fn receive_message(&self, msg: WsMessage);
 }
 
@@ -50,7 +51,7 @@ pub enum WsState {
 }
 
 pub struct WsController {
-    handlers: HashMap<String, Arc<dyn WsMessageHandler>>,
+    handlers: HashMap<WsSource, Arc<dyn WsMessageHandler>>,
     state_notify: Arc<RwLock<WsStateNotify>>,
     #[allow(dead_code)]
     addr: Option<String>,
@@ -83,7 +84,7 @@ impl WsController {
     pub fn add_handler(&mut self, handler: Arc<dyn WsMessageHandler>) -> Result<(), WsError> {
         let source = handler.source();
         if self.handlers.contains_key(&source) {
-            log::error!("{} source is already registered", source);
+            log::error!("WsSource's {:?} is already registered", source);
         }
         self.handlers.insert(source, handler);
         Ok(())
@@ -163,11 +164,11 @@ impl WsController {
 pub struct WsHandlerFuture {
     #[pin]
     msg_rx: MsgReceiver,
-    handlers: HashMap<String, Arc<dyn WsMessageHandler>>,
+    handlers: HashMap<WsSource, Arc<dyn WsMessageHandler>>,
 }
 
 impl WsHandlerFuture {
-    fn new(handlers: HashMap<String, Arc<dyn WsMessageHandler>>, msg_rx: MsgReceiver) -> Self { Self { msg_rx, handlers } }
+    fn new(handlers: HashMap<WsSource, Arc<dyn WsMessageHandler>>, msg_rx: MsgReceiver) -> Self { Self { msg_rx, handlers } }
 }
 
 impl Future for WsHandlerFuture {
@@ -203,19 +204,16 @@ impl WsSender {
         Ok(())
     }
 
-    pub fn send_text(&self, source: &str, text: &str) -> Result<(), WsError> {
+    pub fn send_text(&self, source: WsSource, text: &str) -> Result<(), WsError> {
         let msg = WsMessage {
-            source: source.to_string(),
+            source,
             data: text.as_bytes().to_vec(),
         };
         self.send_msg(msg)
     }
 
-    pub fn send_binary(&self, source: &str, bytes: Vec<u8>) -> Result<(), WsError> {
+    pub fn send_binary(&self, source: WsSource, bytes: Vec<u8>) -> Result<(), WsError> {
-        let msg = WsMessage {
+        let msg = WsMessage { source, data: bytes };
-            source: source.to_string(),
-            data: bytes,
-        };
         self.send_msg(msg)
     }