From c8569d341d3dcf2fd3d9280082ebf2c2abfd2b73 Mon Sep 17 00:00:00 2001 From: Liam Date: Fri, 11 Sep 2026 12:58:02 +0200 Subject: [PATCH] refactor(core): move remote request handling into core behind remote-adapter feature Moves the duplicated remote request handling from the redis, postgres and mongodb adapters into socketioxide-core behind the existing remote-adapter feature flag. - RemoteRequestHandler trait with default implementations of recv_req and all the recv_* handlers (broadcast, broadcast_with_ack, rooms, add/del sockets, fetch sockets, disconnect sockets). Adapters only provide their local adapter, send_res and an optional recv_heartbeat hook. - insert_response_handler/wait_responses helpers covering the get_res channel setup, response type filter, take/timeout and handler cleanup. Some tiny differences were normalized on purpose: send_res now takes (req_id, req_origin) everywhere with the strictest bounds, a request without options warns and is skipped for all adapters, the loopback check is done in the shared dispatch and logs use the same style. The postgres adapter get_res now also wraps the response stream in a DropStream, fixing a leak of response handler entries for rooms and fetch sockets requests. No behavior or wire-format change. Closes #774 --- crates/socketioxide-core/src/adapter/mod.rs | 16 +- .../socketioxide-core/src/adapter/request.rs | 356 ++++++++++++++++++ .../socketioxide-core/src/adapter/stream.rs | 40 +- crates/socketioxide-mongodb/src/lib.rs | 225 +++-------- crates/socketioxide-postgres/src/lib.rs | 269 ++++--------- crates/socketioxide-redis/src/lib.rs | 230 +++-------- 6 files changed, 605 insertions(+), 531 deletions(-) create mode 100644 crates/socketioxide-core/src/adapter/request.rs diff --git a/crates/socketioxide-core/src/adapter/mod.rs b/crates/socketioxide-core/src/adapter/mod.rs index 8e3f6089..c6408b4d 100644 --- a/crates/socketioxide-core/src/adapter/mod.rs +++ b/crates/socketioxide-core/src/adapter/mod.rs @@ -30,6 +30,8 @@ pub mod heartbeat; #[cfg(feature = "remote-adapter")] pub mod remote_packet; #[cfg(feature = "remote-adapter")] +pub mod request; +#[cfg(feature = "remote-adapter")] pub mod stream; /// A room identifier @@ -726,7 +728,7 @@ pub struct RemoteSocketData { } #[cfg(test)] -mod test { +pub(crate) mod test { use smallvec::smallvec; use std::{ @@ -737,12 +739,12 @@ mod test { use super::*; - struct StubSockets { + pub(crate) struct StubSockets { sockets: HashSet, path: Str, } impl StubSockets { - fn new(sockets: &[Sid]) -> Self { + pub(crate) fn new(sockets: &[Sid]) -> Self { let sockets = HashSet::from_iter(sockets.iter().copied()); Self { sockets, @@ -751,7 +753,7 @@ mod test { } } - struct StubAckStream; + pub(crate) struct StubAckStream; impl Stream for StubAckStream { type Item = (Sid, Result); fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { @@ -764,7 +766,7 @@ mod test { } } #[derive(Debug, Serialize, Deserialize)] - struct StubError; + pub(crate) struct StubError; impl std::fmt::Display for StubError { fn fmt(&self, _: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { Ok(()) @@ -825,7 +827,9 @@ mod test { } } - fn create_adapter(sockets: [Sid; S]) -> CoreLocalAdapter { + pub(crate) fn create_adapter( + sockets: [Sid; S], + ) -> CoreLocalAdapter { CoreLocalAdapter::new(StubSockets::new(&sockets)) } diff --git a/crates/socketioxide-core/src/adapter/request.rs b/crates/socketioxide-core/src/adapter/request.rs new file mode 100644 index 00000000..b834fcae --- /dev/null +++ b/crates/socketioxide-core/src/adapter/request.rs @@ -0,0 +1,356 @@ +//! Shared remote request handling for remote adapters. +//! +//! The [`RemoteRequestHandler`] trait provides default implementations for the +//! requests a node can receive from another node. Adapters only have to provide +//! access to their local adapter and a way to send responses back. + +use std::{fmt, future::Future, sync::Arc}; + +use futures_util::StreamExt; +use serde::Serialize; + +use crate::{ + Sid, Uid, + adapter::{ + BroadcastOptions, CoreLocalAdapter, Room, SocketEmitter, + errors::AdapterError, + remote_packet::{RequestIn, RequestTypeIn, Response, ResponseType}, + }, + packet::Packet, +}; + +/// Handle requests received from other nodes of the cluster. +/// +/// Every remote adapter implements this trait. The default methods dispatch the +/// request to the right handler and apply it on the local adapter. +pub trait RemoteRequestHandler: Send + Sync + 'static { + /// The error returned when sending a response. + type Error: std::error::Error + Into + Send + 'static; + + /// The local adapter, used to apply the requests on the local sockets. + fn local(&self) -> &CoreLocalAdapter; + + /// Send a response to the node that sent the request. + fn send_res( + &self, + req_id: Sid, + req_origin: Uid, + res: Response, + ) -> impl Future> + Send + 'static; + + /// Handle a heartbeat request received from a remote node. + /// + /// It is a no-op by default as some adapters do not rely on heartbeats to + /// track the liveness of the other nodes. + fn recv_heartbeat(self: &Arc, req_type: RequestTypeIn, origin: Uid) { + let _ = (self, req_type, origin); + } + + /// Handle a request received from another node. + fn recv_req(self: &Arc, req: RequestIn) { + let RequestIn { + node_id, + id, + r#type, + opts, + } = req; + + // Ignore loopback requests. + if node_id == self.local().server_id() { + return; + } + + tracing::trace!(?r#type, ?id, ?node_id, "incoming request"); + match (r#type, opts) { + (req_type @ (RequestTypeIn::Heartbeat | RequestTypeIn::InitHeartbeat), _) => { + Self::recv_heartbeat(self, req_type, node_id); + } + (r#type, Some(opts)) => match r#type { + RequestTypeIn::Broadcast(p) => self.recv_broadcast(opts, p), + RequestTypeIn::BroadcastWithAck(p) => { + self.clone().recv_broadcast_with_ack(node_id, id, p, opts) + } + RequestTypeIn::DisconnectSockets => self.recv_disconnect_sockets(opts), + RequestTypeIn::AllRooms => self.recv_rooms(node_id, id, opts), + RequestTypeIn::AddSockets(rooms) => self.recv_add_sockets(opts, rooms), + RequestTypeIn::DelSockets(rooms) => self.recv_del_sockets(opts, rooms), + RequestTypeIn::FetchSockets => self.recv_fetch_sockets(node_id, id, opts), + RequestTypeIn::Heartbeat | RequestTypeIn::InitHeartbeat => unreachable!(), + }, + (r#type, None) => { + tracing::warn!(?node_id, ?r#type, "request is missing options"); + } + } + } + + /// Broadcast a packet to the local sockets matching the options. + fn recv_broadcast(&self, opts: BroadcastOptions, packet: Packet) { + if let Err(e) = self.local().broadcast(packet, opts) { + let ns = self.local().path(); + let node_id = self.local().server_id(); + tracing::warn!( + %node_id, + ?ns, + "remote request broadcast handler: {:?}", + e + ); + } + } + + /// Disconnect the local sockets matching the options. + fn recv_disconnect_sockets(&self, opts: BroadcastOptions) { + if let Err(e) = self.local().disconnect_socket(opts) { + let ns = self.local().path(); + let node_id = self.local().server_id(); + tracing::warn!( + %node_id, + ?ns, + "remote request disconnect sockets handler: {:?}", + e + ); + } + } + + /// Broadcast a packet to the local sockets matching the options and send back + /// the expected ack count and the acks as they are received. + fn recv_broadcast_with_ack( + self: Arc, + origin: Uid, + req_id: Sid, + packet: Packet, + opts: BroadcastOptions, + ) { + let node_id = self.local().server_id(); + let (stream, count) = self.local().broadcast_with_ack(packet, opts, None); + tokio::spawn(async move { + let on_err = |err| { + let ns = self.local().path(); + let node_id = self.local().server_id(); + tracing::warn!( + %node_id, + ?ns, + "remote request broadcast with ack handler errors: {:?}", + err + ); + }; + // First send the count of expected acks to the server that sent the request. + // This is used to keep track of the number of expected acks. + let res = Response { + r#type: ResponseType::<()>::BroadcastAckCount(count), + node_id, + }; + if let Err(err) = self.send_res(req_id, origin, res).await { + on_err(err); + return; + } + + // Then send the acks as they are received. + futures_util::pin_mut!(stream); + while let Some(ack) = stream.next().await { + let res = Response { + r#type: ResponseType::BroadcastAck(ack), + node_id, + }; + if let Err(err) = self.send_res(req_id, origin, res).await { + on_err(err); + return; + } + } + }); + } + + /// Send back all the local room names. + fn recv_rooms(&self, origin: Uid, req_id: Sid, opts: BroadcastOptions) { + let rooms = self.local().rooms(opts); + let res = Response { + r#type: ResponseType::<()>::AllRooms(rooms), + node_id: self.local().server_id(), + }; + let fut = self.send_res(req_id, origin, res); + let ns = self.local().path().clone(); + let uid = self.local().server_id(); + tokio::spawn(async move { + if let Err(err) = fut.await { + tracing::warn!(?uid, ?ns, "remote request rooms handler: {:?}", err); + } + }); + } + + /// Add the local sockets matching the options to the rooms. + fn recv_add_sockets(&self, opts: BroadcastOptions, rooms: Vec) { + self.local().add_sockets(opts, rooms); + } + + /// Remove the local sockets matching the options from the rooms. + fn recv_del_sockets(&self, opts: BroadcastOptions, rooms: Vec) { + self.local().del_sockets(opts, rooms); + } + + /// Send back the data of all the local sockets matching the options. + fn recv_fetch_sockets(&self, origin: Uid, req_id: Sid, opts: BroadcastOptions) { + let sockets = self.local().fetch_sockets(opts); + let res = Response { + node_id: self.local().server_id(), + r#type: ResponseType::FetchSockets(sockets), + }; + let fut = self.send_res(req_id, origin, res); + let ns = self.local().path().clone(); + let uid = self.local().server_id(); + tokio::spawn(async move { + if let Err(err) = fut.await { + tracing::warn!(?uid, ?ns, "remote request fetch sockets handler: {:?}", err); + } + }); + } +} + +#[cfg(test)] +mod tests { + use std::{ + fmt, + sync::atomic::{AtomicUsize, Ordering}, + time::Duration, + }; + + use tokio::sync::mpsc; + + use super::*; + use crate::{ + Value, + adapter::{ + errors::AdapterError, + test::{StubSockets, create_adapter}, + }, + }; + + struct TestHandler { + local: CoreLocalAdapter, + responses: mpsc::UnboundedSender, + heartbeats: AtomicUsize, + } + + impl RemoteRequestHandler for TestHandler { + type Error = AdapterError; + + fn local(&self) -> &CoreLocalAdapter { + &self.local + } + + fn send_res( + &self, + _req_id: Sid, + _req_origin: Uid, + res: Response, + ) -> impl Future> + Send + 'static { + let tx = self.responses.clone(); + async move { + tx.send(format!("{res:?}")).ok(); + Ok(()) + } + } + + fn recv_heartbeat(self: &Arc, _req_type: RequestTypeIn, _origin: Uid) { + self.heartbeats.fetch_add(1, Ordering::SeqCst); + } + } + + fn handler( + sockets: [Sid; S], + ) -> (Arc, mpsc::UnboundedReceiver) { + let (tx, rx) = mpsc::unbounded_channel(); + let handler = TestHandler { + local: create_adapter(sockets), + responses: tx, + heartbeats: AtomicUsize::new(0), + }; + (Arc::new(handler), rx) + } + + fn req(r#type: RequestTypeIn, opts: Option) -> RequestIn { + RequestIn { + node_id: Uid::new(), + id: Sid::new(), + r#type, + opts, + } + } + + #[test] + fn heartbeat_hook_is_called() { + let (handler, _rx) = handler([Sid::new()]); + handler.recv_req(req(RequestTypeIn::InitHeartbeat, None)); + assert_eq!(handler.heartbeats.load(Ordering::SeqCst), 1); + } + + #[test] + fn loopback_request_is_ignored() { + let (handler, _rx) = handler([Sid::new()]); + let mut req = req(RequestTypeIn::InitHeartbeat, None); + req.node_id = Uid::ZERO; + handler.recv_req(req); + assert_eq!(handler.heartbeats.load(Ordering::SeqCst), 0); + } + + #[test] + fn add_and_del_sockets() { + let sid = Sid::new(); + let (handler, _rx) = handler([sid]); + let opts = BroadcastOptions::new(sid); + + handler.recv_req(req( + RequestTypeIn::AddSockets(vec!["room1".into()]), + Some(opts.clone()), + )); + assert!(handler.local().socket_rooms(sid).contains("room1")); + + handler.recv_req(req( + RequestTypeIn::DelSockets(vec!["room1".into()]), + Some(opts), + )); + assert!(handler.local().socket_rooms(sid).is_empty()); + } + + #[test] + fn request_without_opts_is_skipped() { + let (handler, mut rx) = handler([Sid::new()]); + handler.recv_req(req(RequestTypeIn::AllRooms, None)); + assert!(rx.try_recv().is_err()); + } + + #[tokio::test] + async fn rooms_response_is_sent() { + let sid = Sid::new(); + let (handler, mut rx) = handler([sid]); + handler.local().add_all(sid, ["room1"]); + handler.recv_req(req( + RequestTypeIn::AllRooms, + Some(BroadcastOptions::new(sid)), + )); + + let res = tokio::time::timeout(Duration::from_secs(1), rx.recv()) + .await + .unwrap() + .unwrap(); + assert!(res.contains("AllRooms"), "unexpected response: {res}"); + } + + #[tokio::test] + async fn broadcast_with_ack_count_is_sent() { + let sid = Sid::new(); + let (handler, mut rx) = handler([sid]); + let packet = Packet::event("/", Value::Str("test".into(), None)); + handler.recv_req(req( + RequestTypeIn::BroadcastWithAck(packet), + Some(BroadcastOptions::new(sid)), + )); + + let res = tokio::time::timeout(Duration::from_secs(1), rx.recv()) + .await + .unwrap() + .unwrap(); + assert!( + res.contains("BroadcastAckCount"), + "unexpected response: {res}" + ); + } +} diff --git a/crates/socketioxide-core/src/adapter/stream.rs b/crates/socketioxide-core/src/adapter/stream.rs index 17d5391c..5c9adfa3 100644 --- a/crates/socketioxide-core/src/adapter/stream.rs +++ b/crates/socketioxide-core/src/adapter/stream.rs @@ -11,14 +11,14 @@ use std::{ }; use futures_core::{FusedStream, Stream}; -use futures_util::{StreamExt, stream::TakeUntil}; +use futures_util::{StreamExt, future, stream::TakeUntil}; use pin_project_lite::pin_project; use serde::de::DeserializeOwned; use tokio::{sync::mpsc, time}; use crate::Sid; use crate::adapter::AckStreamItem; -use crate::adapter::remote_packet::{Response, ResponseType}; +use crate::adapter::remote_packet::{Response, ResponseType, ResponseTypeId}; use crate::adapter::SocketEmitter; @@ -217,6 +217,42 @@ impl fmt::Debug for AckStream { } } +/// Register the response channel of a request so responses can be routed to it. +/// +/// The receiver is returned to be consumed by [`wait_responses`]. The channel is +/// always big enough to hold at least one response. +pub fn insert_response_handler( + handlers: &Arc>>, + req_id: Sid, + serv_cnt: usize, +) -> mpsc::Receiver { + let (tx, rx) = mpsc::channel(std::cmp::max(serv_cnt, 1)); + handlers.lock().unwrap().insert(req_id, tx); + rx +} + +/// Filter the responses of a request by type, wait for at most `serv_cnt` of them +/// and stop after `timeout`. +/// +/// The response handler is removed from the map when the returned stream is dropped. +pub fn wait_responses( + stream: S, + handlers: Arc>>, + req_id: Sid, + response_type: ResponseTypeId, + serv_cnt: usize, + timeout: Duration, +) -> impl Stream> +where + S: Stream>, +{ + let stream = stream + .filter(move |item| future::ready(ResponseTypeId::from(&item.r#type) == response_type)) + .take(serv_cnt) + .take_until(time::sleep(timeout)); + DropStream::new(stream, handlers, req_id) +} + #[cfg(test)] mod tests { use std::sync::{Arc, Mutex}; diff --git a/crates/socketioxide-mongodb/src/lib.rs b/crates/socketioxide-mongodb/src/lib.rs index 9d253ee2..5d598a07 100644 --- a/crates/socketioxide-mongodb/src/lib.rs +++ b/crates/socketioxide-mongodb/src/lib.rs @@ -102,8 +102,9 @@ use futures_core::{Stream, future::Future}; use futures_util::{StreamExt, future::Either}; use serde::{Serialize, de::DeserializeOwned}; use socketioxide_core::adapter::remote_packet::{ - RequestIn, RequestOut, RequestTypeIn, RequestTypeOut, Response, ResponseType, ResponseTypeId, + RequestIn, RequestOut, RequestTypeIn, RequestTypeOut, Response, ResponseTypeId, }; +use socketioxide_core::adapter::request::RemoteRequestHandler; use socketioxide_core::{ Sid, Uid, adapter::errors::{AdapterError, BroadcastError}, @@ -111,7 +112,9 @@ use socketioxide_core::{ BroadcastOptions, CoreAdapter, CoreLocalAdapter, DefinedAdapter, RemoteSocketData, Room, RoomParam, SocketEmitter, Spawnable, heartbeat::{HeartbeatSender, HeartbeatTracker, heartbeat_loop}, - stream::{AckStream, ChanStream, DropStream, ResponseHandlers}, + stream::{ + AckStream, ChanStream, ResponseHandlers, insert_response_handler, wait_responses, + }, }, packet::Packet, }; @@ -561,7 +564,7 @@ impl CustomMongoDbAdapter { .. }) if target.is_none_or(|id| id == self.uid) => { tracing::debug!(?target, "request header"); - if let Err(e) = self.recv_req(data).await { + if let Err(e) = self.recv_item(data) { tracing::warn!("error receiving request from driver: {e}"); } } @@ -597,153 +600,13 @@ impl CustomMongoDbAdapter { } } - async fn recv_req(self: &Arc, req: Vec) -> Result<(), Error> { + /// Decode a request received from the driver and handle it. + fn recv_item(self: &Arc, req: Vec) -> Result<(), Error> { let req = rmp_serde::from_slice::(&req)?; - tracing::trace!(?req, "incoming request"); - match (req.r#type, req.opts) { - (RequestTypeIn::Broadcast(p), Some(opts)) => self.recv_broadcast(opts, p), - (RequestTypeIn::BroadcastWithAck(p), Some(opts)) => self - .clone() - .recv_broadcast_with_ack(req.node_id, req.id, p, opts), - (RequestTypeIn::DisconnectSockets, Some(opts)) => self.recv_disconnect_sockets(opts), - (RequestTypeIn::AllRooms, Some(opts)) => self.recv_rooms(req.node_id, req.id, opts), - (RequestTypeIn::AddSockets(rooms), Some(opts)) => self.recv_add_sockets(opts, rooms), - (RequestTypeIn::DelSockets(rooms), Some(opts)) => self.recv_del_sockets(opts, rooms), - (RequestTypeIn::FetchSockets, Some(opts)) => { - self.recv_fetch_sockets(req.node_id, req.id, opts) - } - req_type @ (RequestTypeIn::Heartbeat | RequestTypeIn::InitHeartbeat, _) => { - self.recv_heartbeat(&self.heartbeat, req_type.0, req.node_id) - } - _ => (), - } + >::recv_req(self, req); Ok(()) } - fn recv_broadcast(&self, opts: BroadcastOptions, packet: Packet) { - tracing::trace!(?opts, "incoming broadcast"); - if let Err(e) = self.local.broadcast(packet, opts) { - let ns = self.local.path(); - tracing::warn!(?self.uid, ?ns, "remote request broadcast handler: {:?}", e); - } - } - - fn recv_disconnect_sockets(&self, opts: BroadcastOptions) { - if let Err(e) = self.local.disconnect_socket(opts) { - let ns = self.local.path(); - tracing::warn!( - ?self.uid, - ?ns, - "remote request disconnect sockets handler: {:?}", - e - ); - } - } - - fn recv_broadcast_with_ack( - self: Arc, - origin: Uid, - req_id: Sid, - packet: Packet, - opts: BroadcastOptions, - ) { - let (stream, count) = self.local.broadcast_with_ack(packet, opts, None); - tokio::spawn(async move { - let on_err = |err| { - let ns = self.local.path(); - tracing::warn!( - ?self.uid, - ?ns, - "remote request broadcast with ack handler errors: {:?}", - err - ); - }; - // First send the count of expected acks to the server that sent the request. - // This is used to keep track of the number of expected acks. - let res = Response { - r#type: ResponseType::<()>::BroadcastAckCount(count), - node_id: self.uid, - }; - if let Err(err) = self.send_res(req_id, origin, res).await { - on_err(err); - return; - } - - // Then send the acks as they are received. - futures_util::pin_mut!(stream); - while let Some(ack) = stream.next().await { - let res = Response { - r#type: ResponseType::BroadcastAck(ack), - node_id: self.uid, - }; - if let Err(err) = self.send_res(req_id, origin, res).await { - on_err(err); - return; - } - } - }); - } - - fn recv_rooms(&self, origin: Uid, req_id: Sid, opts: BroadcastOptions) { - let rooms = self.local.rooms(opts); - let res = Response { - r#type: ResponseType::<()>::AllRooms(rooms), - node_id: self.uid, - }; - let fut = self.send_res(req_id, origin, res); - let ns = self.local.path().clone(); - let uid = self.uid; - tokio::spawn(async move { - if let Err(err) = fut.await { - tracing::warn!(?uid, ?ns, "remote request rooms handler: {:?}", err); - } - }); - } - - fn recv_add_sockets(&self, opts: BroadcastOptions, rooms: Vec) { - self.local.add_sockets(opts, rooms); - } - - fn recv_del_sockets(&self, opts: BroadcastOptions, rooms: Vec) { - self.local.del_sockets(opts, rooms); - } - fn recv_fetch_sockets(&self, origin: Uid, req_id: Sid, opts: BroadcastOptions) { - let sockets = self.local.fetch_sockets(opts); - let res = Response { - node_id: self.uid, - r#type: ResponseType::FetchSockets(sockets), - }; - let fut = self.send_res(req_id, origin, res); - let ns = self.local.path().clone(); - let uid = self.uid; - tokio::spawn(async move { - if let Err(err) = fut.await { - tracing::warn!(?uid, ?ns, "remote request fetch sockets handler: {:?}", err); - } - }); - } - - /// Send a response to the node that sent the request. - fn send_res( - &self, - req_id: Sid, - req_origin: Uid, - res: Response, - ) -> impl Future>> + Send + 'static { - tracing::trace!(?res, "sending response for {req_id} req to {req_origin}"); - let driver = self.driver.clone(); - let head = ItemHeader::Res { - request: req_id, - target: req_origin, - }; - let res = self.new_packet(head, &res); - - async move { - driver.emit(&res?).await.map_err(Error::from_driver)?; - Ok(()) - } - } - /// Await for all the responses from the remote servers. /// If the target node is specified, only await for the response from that node. async fn get_res( @@ -758,24 +621,25 @@ impl CustomMongoDbAdapter { } else { 1 }; - let (tx, rx) = mpsc::channel(std::cmp::max(remote_serv_cnt, 1)); - self.responses.lock().unwrap().insert(req_id, tx); - let stream = ChanStream::new(rx) - .filter_map(|Item { header, data, .. }| { - let data = match rmp_serde::from_slice::>(&data) { - Ok(data) => Some(data), - Err(e) => { - tracing::warn!(header = ?header, "error decoding response: {e}"); - None - } - }; - future::ready(data) - }) - .filter(move |item| future::ready(ResponseTypeId::from(&item.r#type) == response_type)) - .take(remote_serv_cnt) - .take_until(tokio::time::sleep(self.config.request_timeout)); - let stream = DropStream::new(stream, self.responses.clone(), req_id); - Ok(stream) + let rx = insert_response_handler(&self.responses, req_id, remote_serv_cnt); + let stream = ChanStream::new(rx).filter_map(|Item { header, data, .. }| { + let data = match rmp_serde::from_slice::>(&data) { + Ok(data) => Some(data), + Err(e) => { + tracing::warn!(header = ?header, "error decoding response: {e}"); + None + } + }; + future::ready(data) + }); + Ok(wait_responses( + stream, + self.responses.clone(), + req_id, + response_type, + remote_serv_cnt, + self.config.request_timeout, + )) } fn new_packet(&self, head: ItemHeader, data: &impl Serialize) -> Result> { @@ -806,6 +670,39 @@ impl HeartbeatSender for CustomMongoDbAdapter } } +impl RemoteRequestHandler for CustomMongoDbAdapter { + type Error = Error; + + fn local(&self) -> &CoreLocalAdapter { + &self.local + } + + /// Send a response to the node that sent the request. + fn send_res( + &self, + req_id: Sid, + req_origin: Uid, + res: Response, + ) -> impl Future>> + Send + 'static { + tracing::trace!(?res, "sending response for {req_id} req to {req_origin}"); + let driver = self.driver.clone(); + let head = ItemHeader::Res { + request: req_id, + target: req_origin, + }; + let res = self.new_packet(head, &res); + + async move { + driver.emit(&res?).await.map_err(Error::from_driver)?; + Ok(()) + } + } + + fn recv_heartbeat(self: &Arc, req_type: RequestTypeIn, origin: Uid) { + HeartbeatSender::recv_heartbeat(self, &self.heartbeat, req_type, origin); + } +} + /// The result of the init future. #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct InitRes(futures_core::future::BoxFuture<'static, Result<(), D::Error>>); diff --git a/crates/socketioxide-postgres/src/lib.rs b/crates/socketioxide-postgres/src/lib.rs index f956e421..f61419df 100644 --- a/crates/socketioxide-postgres/src/lib.rs +++ b/crates/socketioxide-postgres/src/lib.rs @@ -65,10 +65,12 @@ use socketioxide_core::{ errors::{AdapterError, BroadcastError}, heartbeat::{HeartbeatSender, HeartbeatTracker, heartbeat_loop}, remote_packet::{ - RequestIn, RequestOut, RequestTypeIn, RequestTypeOut, Response, ResponseType, - ResponseTypeId, + RequestIn, RequestOut, RequestTypeIn, RequestTypeOut, Response, ResponseTypeId, + }, + request::RemoteRequestHandler, + stream::{ + AckStream, ChanStream, ResponseHandlers, insert_response_handler, wait_responses, }, - stream::{AckStream, ChanStream, ResponseHandlers}, }, packet::Packet, }; @@ -755,184 +757,6 @@ impl CustomPostgresAdapter { } } - fn recv_req(self: &Arc, req: RequestIn) { - tracing::trace!(?req, "incoming request"); - match (req.r#type, req.opts) { - (RequestTypeIn::Broadcast(p), Some(opts)) => self.recv_broadcast(opts, p), - (RequestTypeIn::BroadcastWithAck(p), Some(opts)) => self - .clone() - .recv_broadcast_with_ack(req.node_id, req.id, p, opts), - (RequestTypeIn::DisconnectSockets, Some(opts)) => self.recv_disconnect_sockets(opts), - (RequestTypeIn::AllRooms, Some(opts)) => self.recv_rooms(req.node_id, req.id, opts), - (RequestTypeIn::AddSockets(rooms), Some(opts)) => self.recv_add_sockets(opts, rooms), - (RequestTypeIn::DelSockets(rooms), Some(opts)) => self.recv_del_sockets(opts, rooms), - (RequestTypeIn::FetchSockets, Some(opts)) => { - self.recv_fetch_sockets(req.node_id, req.id, opts) - } - req_type @ (RequestTypeIn::Heartbeat | RequestTypeIn::InitHeartbeat, _) => { - self.recv_heartbeat(&self.heartbeat, req_type.0, req.node_id) - } - _ => (), - } - } - - fn recv_broadcast(&self, opts: BroadcastOptions, packet: Packet) { - tracing::trace!(?opts, "incoming broadcast"); - if let Err(e) = self.local.broadcast(packet, opts) { - let ns = self.local.path(); - tracing::warn!(node_id = %self.local.server_id(), ?ns, "remote request broadcast handler: {:?}", e); - } - } - - fn recv_disconnect_sockets(&self, opts: BroadcastOptions) { - if let Err(e) = self.local.disconnect_socket(opts) { - let ns = self.local.path(); - tracing::warn!( - node_id = %self.local.server_id(), - %ns, - "remote request disconnect sockets handler: {:?}", - e - ); - } - } - - fn recv_broadcast_with_ack( - self: Arc, - origin: Uid, - req_id: Sid, - packet: Packet, - opts: BroadcastOptions, - ) { - let (stream, count) = self.local.broadcast_with_ack(packet, opts, None); - tokio::spawn(async move { - let on_err = |err| { - let ns = self.local.path(); - tracing::warn!( - node_id = %self.local.server_id(), - %ns, - "remote request broadcast with ack handler errors: {:?}", - err - ); - }; - // First send the count of expected acks to the server that sent the request. - // This is used to keep track of the number of expected acks. - let res = Response { - r#type: ResponseType::<()>::BroadcastAckCount(count), - node_id: self.local.server_id(), - }; - if let Err(err) = self.send_res(req_id, origin, res).await { - on_err(err); - return; - } - - // Then send the acks as they are received. - futures_util::pin_mut!(stream); - while let Some(ack) = stream.next().await { - let res = Response { - r#type: ResponseType::BroadcastAck(ack), - node_id: self.local.server_id(), - }; - if let Err(err) = self.send_res(req_id, origin, res).await { - on_err(err); - return; - } - } - }); - } - - fn recv_rooms(&self, origin: Uid, req_id: Sid, opts: BroadcastOptions) { - let rooms = self.local.rooms(opts); - let res = Response { - r#type: ResponseType::<()>::AllRooms(rooms), - node_id: self.local.server_id(), - }; - let fut = self.send_res(req_id, origin, res); - let ns = self.local.path().clone(); - let uid = self.local.server_id(); - tokio::spawn(async move { - if let Err(err) = fut.await { - tracing::warn!(?uid, ?ns, "remote request rooms handler: {:?}", err); - } - }); - } - - fn recv_add_sockets(&self, opts: BroadcastOptions, rooms: Vec) { - self.local.add_sockets(opts, rooms); - } - - fn recv_del_sockets(&self, opts: BroadcastOptions, rooms: Vec) { - self.local.del_sockets(opts, rooms); - } - fn recv_fetch_sockets(&self, origin: Uid, req_id: Sid, opts: BroadcastOptions) { - let sockets = self.local.fetch_sockets(opts); - let res = Response { - node_id: self.local.server_id(), - r#type: ResponseType::FetchSockets(sockets), - }; - let fut = self.send_res(req_id, origin, res); - let ns = self.local.path().clone(); - let uid = self.local.server_id(); - tokio::spawn(async move { - if let Err(err) = fut.await { - tracing::warn!(?uid, ?ns, "remote request fetch sockets handler: {:?}", err); - } - }); - } - - /// Send a response to the node that sent the request. - /// - /// When the serialized response exceeds [`PostgresAdapterConfig::payload_threshold`], it is - /// stored in the attachment table and only the row id travels over NOTIFY as a - /// [`ResponsePayload::Attachment`]. - fn send_res( - &self, - req_id: Sid, - req_origin: Uid, - res: Response, - ) -> impl Future>> + 'static { - tracing::trace!(?res, "sending response for {req_id} req to {req_origin}"); - let driver = self.driver.clone(); - let chan = self.get_response_chan(req_origin); - let table = self.config.table_name.clone(); - let threshold = self.config.payload_threshold; - let node_id = self.local.server_id(); - let is_binary = res.is_binary(); - async move { - let body = if is_binary { - rmp_serde::to_vec(&res)? - } else { - serde_json::to_vec(&res)? - }; - - let payload = if body.len() >= threshold || is_binary { - let id = driver - .push_attachment(&table, &body) - .await - .map_err(Error::Driver)?; - ResponsePayload::Attachment { id, is_binary } - } else { - assert!( - !is_binary, - "binary packets should be stored in attachment table and serialized in msgpack" - ); - let body = unsafe { String::from_utf8_unchecked(body) }; - ResponsePayload::Data(RawValue::from_string(body)?) - }; - - let message = serde_json::to_string(&ResponsePacket { - req_id, - node_id, - payload, - })?; - - driver - .notify(&chan, &message) - .await - .map_err(Error::Driver)?; - Ok(()) - } - } - /// Await for all the responses from the remote servers. /// If the target node is specified, only await for the response from that node. async fn get_res( @@ -948,8 +772,7 @@ impl CustomPostgresAdapter { 1 }; - let (tx, rx) = mpsc::channel(std::cmp::max(remote_serv_cnt, 1)); - self.responses.lock().unwrap().insert(req_id, tx); + let rx = insert_response_handler(&self.responses, req_id, remote_serv_cnt); let stream = ChanStream::new(rx); // Overlap attachment fetches across servers while preserving arrival order so that @@ -968,12 +791,16 @@ impl CustomPostgresAdapter { } }) .buffered(concurrency) - .filter_map(future::ready) - .filter(move |item| future::ready(ResponseTypeId::from(&item.r#type) == response_type)) - .take(remote_serv_cnt) - .take_until(tokio::time::sleep(self.config.request_timeout)); + .filter_map(future::ready); - Ok(stream) + Ok(wait_responses( + stream, + self.responses.clone(), + req_id, + response_type, + remote_serv_cnt, + self.config.request_timeout, + )) } // == All channels are hashed to avoid thresspassing the 63 bytes limit on postgres channel == @@ -1059,6 +886,72 @@ impl HeartbeatSender for CustomPostgresAdapter RemoteRequestHandler for CustomPostgresAdapter { + type Error = Error; + + fn local(&self) -> &CoreLocalAdapter { + &self.local + } + + /// Send a response to the node that sent the request. + /// + /// When the serialized response exceeds [`PostgresAdapterConfig::payload_threshold`], it is + /// stored in the attachment table and only the row id travels over NOTIFY as a + /// [`ResponsePayload::Attachment`]. + fn send_res( + &self, + req_id: Sid, + req_origin: Uid, + res: Response, + ) -> impl Future>> + Send + 'static { + tracing::trace!(?res, "sending response for {req_id} req to {req_origin}"); + let driver = self.driver.clone(); + let chan = self.get_response_chan(req_origin); + let table = self.config.table_name.clone(); + let threshold = self.config.payload_threshold; + let node_id = self.local.server_id(); + let is_binary = res.is_binary(); + async move { + let body = if is_binary { + rmp_serde::to_vec(&res)? + } else { + serde_json::to_vec(&res)? + }; + + let payload = if body.len() >= threshold || is_binary { + let id = driver + .push_attachment(&table, &body) + .await + .map_err(Error::Driver)?; + ResponsePayload::Attachment { id, is_binary } + } else { + assert!( + !is_binary, + "binary packets should be stored in attachment table and serialized in msgpack" + ); + let body = unsafe { String::from_utf8_unchecked(body) }; + ResponsePayload::Data(RawValue::from_string(body)?) + }; + + let message = serde_json::to_string(&ResponsePacket { + req_id, + node_id, + payload, + })?; + + driver + .notify(&chan, &message) + .await + .map_err(Error::Driver)?; + Ok(()) + } + } + + fn recv_heartbeat(self: &Arc, req_type: RequestTypeIn, origin: Uid) { + HeartbeatSender::recv_heartbeat(self, &self.heartbeat, req_type, origin); + } +} + fn hash_chan(chan: &str) -> String { let hash = xxhash_rust::xxh3::xxh3_64(chan.as_bytes()); format!("ch_{:x}", hash) diff --git a/crates/socketioxide-redis/src/lib.rs b/crates/socketioxide-redis/src/lib.rs index 4930ea2a..4b361c35 100644 --- a/crates/socketioxide-redis/src/lib.rs +++ b/crates/socketioxide-redis/src/lib.rs @@ -154,9 +154,12 @@ use futures_core::Stream; use futures_util::{StreamExt, future::Either}; use serde::{Serialize, de::DeserializeOwned}; use socketioxide_core::adapter::remote_packet::{ - RequestIn, RequestOut, RequestTypeIn, RequestTypeOut, Response, ResponseType, ResponseTypeId, + RequestIn, RequestOut, RequestTypeOut, Response, ResponseTypeId, +}; +use socketioxide_core::adapter::request::RemoteRequestHandler; +use socketioxide_core::adapter::stream::{ + AckStream, ResponseHandlers, insert_response_handler, wait_responses, }; -use socketioxide_core::adapter::stream::{AckStream, DropStream, ResponseHandlers}; use socketioxide_core::{ Sid, Uid, adapter::errors::{AdapterError, BroadcastError}, @@ -166,7 +169,7 @@ use socketioxide_core::{ }, packet::Packet, }; -use tokio::{sync::mpsc, time}; +use tokio::sync::mpsc; /// Drivers are an abstraction over the pub/sub backend used by the adapter. /// You can use the provided implementation or implement your own. @@ -697,7 +700,7 @@ impl CustomRedisAdapter { ) { while let Some((chan, item)) = stream.next().await { if chan.starts_with(&self.req_chan) { - if let Err(e) = self.recv_req(item) { + if let Err(e) = self.recv_item(item) { let ns = self.local.path(); let uid = self.uid; tracing::warn!(?uid, ?ns, "request handler error: {e}"); @@ -719,137 +722,13 @@ impl CustomRedisAdapter { } } - /// Handle a generic request received from the request channel. - fn recv_req(self: &Arc, item: Vec) -> Result<(), Error> { + /// Decode a request received from the request channel and handle it. + fn recv_item(self: &Arc, item: Vec) -> Result<(), Error> { let req: RequestIn = rmp_serde::from_slice(&item)?; - if req.node_id == self.uid { - return Ok(()); - } - - tracing::trace!(?req, "handling request"); - let Some(opts) = req.opts else { - tracing::warn!(?req, "request is missing options"); - return Ok(()); - }; - - match req.r#type { - RequestTypeIn::Broadcast(p) => self.recv_broadcast(opts, p), - RequestTypeIn::BroadcastWithAck(p) => { - self.clone() - .recv_broadcast_with_ack(req.node_id, req.id, p, opts) - } - RequestTypeIn::DisconnectSockets => self.recv_disconnect_sockets(opts), - RequestTypeIn::AllRooms => self.recv_rooms(req.node_id, req.id, opts), - RequestTypeIn::AddSockets(rooms) => self.recv_add_sockets(opts, rooms), - RequestTypeIn::DelSockets(rooms) => self.recv_del_sockets(opts, rooms), - RequestTypeIn::FetchSockets => self.recv_fetch_sockets(req.node_id, req.id, opts), - _ => (), - }; + >::recv_req(self, req); Ok(()) } - fn recv_broadcast(&self, opts: BroadcastOptions, packet: Packet) { - if let Err(e) = self.local.broadcast(packet, opts) { - let ns = self.local.path(); - tracing::warn!(?self.uid, ?ns, "remote request broadcast handler: {:?}", e); - } - } - - fn recv_disconnect_sockets(&self, opts: BroadcastOptions) { - if let Err(e) = self.local.disconnect_socket(opts) { - let ns = self.local.path(); - tracing::warn!( - ?self.uid, - ?ns, - "remote request disconnect sockets handler: {:?}", - e - ); - } - } - - fn recv_broadcast_with_ack( - self: Arc, - origin: Uid, - req_id: Sid, - packet: Packet, - opts: BroadcastOptions, - ) { - let (stream, count) = self.local.broadcast_with_ack(packet, opts, None); - tokio::spawn(async move { - let on_err = |err| { - let ns = self.local.path(); - tracing::warn!( - ?origin, - ?ns, - "remote request broadcast with ack handler errors: {:?}", - err - ); - }; - // First send the count of expected acks to the server that sent the request. - // This is used to keep track of the number of expected acks. - let res = Response { - r#type: ResponseType::<()>::BroadcastAckCount(count), - node_id: self.uid, - }; - if let Err(err) = self.send_res(origin, req_id, res).await { - on_err(err); - return; - } - - // Then send the acks as they are received. - futures_util::pin_mut!(stream); - while let Some(ack) = stream.next().await { - let res = Response { - r#type: ResponseType::BroadcastAck(ack), - node_id: self.uid, - }; - if let Err(err) = self.send_res(origin, req_id, res).await { - on_err(err); - return; - } - } - }); - } - - fn recv_rooms(&self, origin: Uid, req_id: Sid, opts: BroadcastOptions) { - let rooms = self.local.rooms(opts); - let res = Response { - r#type: ResponseType::<()>::AllRooms(rooms), - node_id: self.uid, - }; - let fut = self.send_res(origin, req_id, res); - let ns = self.local.path().clone(); - let uid = self.uid; - tokio::spawn(async move { - if let Err(err) = fut.await { - tracing::warn!(?uid, ?ns, "remote request rooms handler: {:?}", err); - } - }); - } - - fn recv_add_sockets(&self, opts: BroadcastOptions, rooms: Vec) { - self.local.add_sockets(opts, rooms); - } - - fn recv_del_sockets(&self, opts: BroadcastOptions, rooms: Vec) { - self.local.del_sockets(opts, rooms); - } - fn recv_fetch_sockets(&self, origin: Uid, req_id: Sid, opts: BroadcastOptions) { - let sockets = self.local.fetch_sockets(opts); - let res = Response { - node_id: self.uid, - r#type: ResponseType::FetchSockets(sockets), - }; - let fut = self.send_res(origin, req_id, res); - let ns = self.local.path().clone(); - let uid = self.uid; - tokio::spawn(async move { - if let Err(err) = fut.await { - tracing::warn!(?uid, ?ns, "remote request fetch sockets handler: {:?}", err); - } - }); - } - async fn send_req(&self, req: RequestOut<'_>, target_uid: Option) -> Result<(), Error> { tracing::trace!(?req, "sending request"); let req = rmp_serde::to_vec(&req)?; @@ -862,28 +741,6 @@ impl CustomRedisAdapter { Ok(()) } - fn send_res( - &self, - req_node_id: Uid, - req_id: Sid, - res: Response, - ) -> impl Future>> + Send + 'static { - let chan = self.get_res_chan(req_node_id); - tracing::trace!(?res, "sending response to {}", &chan); - // We send the req_id separated from the response object. - // This allows to partially decode the response and route by the req_id - // before fully deserializing it. - let res = rmp_serde::to_vec(&(req_id, res)); - let driver = self.driver.clone(); - async move { - driver - .publish(chan, res?) - .await - .map_err(Error::from_driver)?; - Ok(()) - } - } - /// Await for all the responses from the remote servers. async fn get_res( &self, @@ -897,24 +754,25 @@ impl CustomRedisAdapter { } else { 1 }; - let (tx, rx) = mpsc::channel(std::cmp::max(remote_serv_cnt, 1)); - self.responses.lock().unwrap().insert(req_id, tx); - let stream = MessageStream::new(rx) - .filter_map(|item| { - let data = match rmp_serde::from_slice::<(Sid, Response)>(&item) { - Ok((_, data)) => Some(data), - Err(e) => { - tracing::warn!("error decoding response: {e}"); - None - } - }; - future::ready(data) - }) - .filter(move |item| future::ready(ResponseTypeId::from(&item.r#type) == response_type)) - .take(remote_serv_cnt) - .take_until(time::sleep(self.config.request_timeout)); - let stream = DropStream::new(stream, self.responses.clone(), req_id); - Ok(stream) + let rx = insert_response_handler(&self.responses, req_id, remote_serv_cnt); + let stream = MessageStream::new(rx).filter_map(|item| { + let data = match rmp_serde::from_slice::<(Sid, Response)>(&item) { + Ok((_, data)) => Some(data), + Err(e) => { + tracing::warn!("error decoding response: {e}"); + None + } + }; + future::ready(data) + }); + Ok(wait_responses( + stream, + self.responses.clone(), + req_id, + response_type, + remote_serv_cnt, + self.config.request_timeout, + )) } /// Little wrapper to map the error type. @@ -928,6 +786,36 @@ impl CustomRedisAdapter { } } +impl RemoteRequestHandler for CustomRedisAdapter { + type Error = Error; + + fn local(&self) -> &CoreLocalAdapter { + &self.local + } + + fn send_res( + &self, + req_id: Sid, + req_origin: Uid, + res: Response, + ) -> impl Future>> + Send + 'static { + let chan = self.get_res_chan(req_origin); + tracing::trace!(?res, "sending response to {}", &chan); + // We send the req_id separated from the response object. + // This allows to partially decode the response and route by the req_id + // before fully deserializing it. + let res = rmp_serde::to_vec(&(req_id, res)); + let driver = self.driver.clone(); + async move { + driver + .publish(chan, res?) + .await + .map_err(Error::from_driver)?; + Ok(()) + } + } +} + /// Checks if the namespace path is valid /// Panics if the path is empty or contains a `#` fn check_ns(path: &str) -> Result<(), InitError> {