Added support for a new "on_close_connection" hook.

Author: karstenda

rust/loro-websocket-server/src/lib.rs

@@ -127,6 +127,18 @@ pub struct HandshakeAuthArgs<'a> {
 
 type HandshakeAuthFn = dyn Fn(HandshakeAuthArgs) -> bool + Send + Sync;
 
+/// Arguments provided to `on_close_connection`.
+pub struct CloseConnectionArgs {
+    pub workspace: String,
+    pub conn_id: u64,
+    pub rooms: Vec<(CrdtType, String)>,
+}
+
+type CloseConnectionFuture =
+    Pin<Box<dyn Future<Output = Result<(), String>> + Send + 'static>>;
+type CloseConnectionFn =
+    Arc<dyn Fn(CloseConnectionArgs) -> CloseConnectionFuture + Send + Sync>;
+
 #[derive(Clone)]
 pub struct ServerConfig<DocCtx = ()> {
     pub on_load_document: Option<LoadFn<DocCtx>>,
@@ -144,6 +156,9 @@ pub struct ServerConfig<DocCtx = ()> {
     ///
     /// Return true to accept, false to reject with 401.
     pub handshake_auth: Option<Arc<HandshakeAuthFn>>,
+    /// Optional hook invoked after a connection fully closes.
+    /// Receives the workspace id, connection id, and rooms the client had joined.
+    pub on_close_connection: Option<CloseConnectionFn>,
 }
 
 // CRDT document abstraction to reduce match-based branching
@@ -459,6 +474,7 @@ impl<DocCtx> Default for ServerConfig<DocCtx> {
             default_permission: Permission::Write,
             authenticate: None,
             handshake_auth: None,
+            on_close_connection: None,
         }
     }
 }
@@ -909,6 +925,7 @@ where
 
     // Capture config outside of non-async closure
     let handshake_auth = registry.config.handshake_auth.clone();
+    let close_connection = registry.config.on_close_connection.clone();
     let workspace_holder: Arc<std::sync::Mutex<Option<String>>> =
         Arc::new(std::sync::Mutex::new(None));
     let workspace_holder_c = workspace_holder.clone();
@@ -1422,6 +1439,11 @@ where
         }
     }
 
+    let rooms_for_hook: Vec<(CrdtType, String)> = joined_rooms
+        .into_iter()
+        .map(|RoomKey { crdt, room }| (crdt, room))
+        .collect();
+
     // cleanup
     {
         let mut h = hub.lock().await;
@@ -1430,6 +1452,18 @@ where
     // drop tx to stop writer
     drop(tx);
     let _ = sink_task.await;
+
+    if let Some(hook) = close_connection {
+        let args = CloseConnectionArgs {
+            workspace: workspace_id.clone(),
+            conn_id,
+            rooms: rooms_for_hook,
+        };
+        if let Err(e) = (hook)(args).await {
+            warn!(conn_id, %e, "on_close_connection hook failed");
+        }
+    }
+
     debug!(conn_id, "connection closed and cleaned up");
     Ok(())
 }