Skip to main content

wsio_client/
builder.rs

1use std::{
2    sync::Arc,
3    time::Duration,
4};
5
6use anyhow::{
7    Result,
8    bail,
9};
10use serde::{
11    Serialize,
12    de::DeserializeOwned,
13};
14use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
15use url::Url;
16
17use crate::{
18    WsIoClient,
19    config::WsIoClientConfig,
20    core::packet::codecs::WsIoPacketCodec,
21    runtime::WsIoClientRuntime,
22    session::WsIoClientSession,
23};
24
25// Structs
26pub struct WsIoClientBuilder {
27    config: WsIoClientConfig,
28    connect_url: Url,
29}
30
31impl WsIoClientBuilder {
32    pub(crate) fn new(mut url: Url) -> Result<Self> {
33        if !matches!(url.scheme(), "ws" | "wss") {
34            bail!("Invalid URL scheme: {}", url.scheme());
35        }
36
37        let mut query_pairs = url.query_pairs().collect::<Vec<_>>();
38        query_pairs.retain(|(k, _)| k != "namespace");
39        query_pairs.push(("namespace".into(), Self::normalize_url_path(url.path()).into()));
40        let query = query_pairs
41            .iter()
42            .map(|(k, v)| format!("{k}={v}"))
43            .collect::<Vec<_>>()
44            .join("&");
45
46        url.set_query(Some(&query));
47        url.set_path("ws.io");
48        Ok(Self {
49            config: WsIoClientConfig {
50                init_handler: None,
51                init_handler_timeout: Duration::from_secs(3),
52                init_packet_timeout: Duration::from_secs(5),
53                on_session_close_handler: None,
54                on_session_close_handler_timeout: Duration::from_secs(2),
55                on_session_ready_handler: None,
56                packet_codec: WsIoPacketCodec::SerdeJson,
57                ping_interval: Duration::from_secs(25),
58                ready_packet_timeout: Duration::from_secs(5),
59                reconnect_delay: Duration::from_secs(1),
60                websocket_config: WebSocketConfig::default()
61                    .max_frame_size(Some(8 * 1024 * 1024))
62                    .max_message_size(Some(16 * 1024 * 1024))
63                    .max_write_buffer_size(2 * 1024 * 1024)
64                    .read_buffer_size(8 * 1024)
65                    .write_buffer_size(8 * 1024),
66            },
67            connect_url: url,
68        })
69    }
70
71    // Private methods
72    fn normalize_url_path(path: &str) -> String {
73        format!(
74            "/{}",
75            path.split('/').filter(|s| !s.is_empty()).collect::<Vec<_>>().join("/")
76        )
77    }
78
79    // Public methods
80    pub fn build(self) -> WsIoClient {
81        WsIoClient(WsIoClientRuntime::new(self.config, self.connect_url))
82    }
83
84    pub fn init_handler_timeout(mut self, duration: Duration) -> Self {
85        self.config.init_handler_timeout = duration;
86        self
87    }
88
89    pub fn init_packet_timeout(mut self, duration: Duration) -> Self {
90        self.config.init_packet_timeout = duration;
91        self
92    }
93
94    pub fn on_session_close<H, Fut>(mut self, handler: H) -> Self
95    where
96        H: Fn(Arc<WsIoClientSession>) -> Fut + Send + Sync + 'static,
97        Fut: Future<Output = Result<()>> + Send + 'static,
98    {
99        self.config.on_session_close_handler = Some(Box::new(move |session| Box::pin(handler(session))));
100        self
101    }
102
103    pub fn on_session_close_handler_timeout(mut self, duration: Duration) -> Self {
104        self.config.on_session_close_handler_timeout = duration;
105        self
106    }
107
108    pub fn on_session_ready<H, Fut>(mut self, handler: H) -> Self
109    where
110        H: Fn(Arc<WsIoClientSession>) -> Fut + Send + Sync + 'static,
111        Fut: Future<Output = Result<()>> + Send + 'static,
112    {
113        self.config.on_session_ready_handler = Some(Arc::new(move |session| Box::pin(handler(session))));
114        self
115    }
116
117    pub fn packet_codec(mut self, packet_codec: WsIoPacketCodec) -> Self {
118        self.config.packet_codec = packet_codec;
119        self
120    }
121
122    pub fn ping_interval(mut self, duration: Duration) -> Self {
123        self.config.ping_interval = duration;
124        self
125    }
126
127    pub fn ready_packet_timeout(mut self, duration: Duration) -> Self {
128        self.config.ready_packet_timeout = duration;
129        self
130    }
131
132    pub fn reconnect_delay(mut self, delay: Duration) -> Self {
133        self.config.reconnect_delay = delay;
134        self
135    }
136
137    pub fn request_path(mut self, request_path: impl AsRef<str>) -> Self {
138        self.connect_url
139            .set_path(&Self::normalize_url_path(request_path.as_ref()));
140
141        self
142    }
143
144    pub fn websocket_config(mut self, websocket_config: WebSocketConfig) -> Self {
145        self.config.websocket_config = websocket_config;
146        self
147    }
148
149    pub fn websocket_config_mut<F: FnOnce(&mut WebSocketConfig)>(mut self, f: F) -> Self {
150        f(&mut self.config.websocket_config);
151        self
152    }
153
154    pub fn with_init_handler<H, Fut, D, R>(mut self, handler: H) -> WsIoClientBuilder
155    where
156        H: Fn(Arc<WsIoClientSession>, Option<D>) -> Fut + Send + Sync + 'static,
157        Fut: Future<Output = Result<Option<R>>> + Send + 'static,
158        D: DeserializeOwned + Send + 'static,
159        R: Serialize + Send + 'static,
160    {
161        let handler = Arc::new(handler);
162        self.config.init_handler = Some(Box::new(move |session, bytes, packet_codec| {
163            let handler = handler.clone();
164            Box::pin(async move {
165                handler(session, bytes.map(|bytes| packet_codec.decode_data(bytes)).transpose()?)
166                    .await?
167                    .map(|data| packet_codec.encode_data(&data))
168                    .transpose()
169            })
170        }));
171
172        self
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    #[test]
181    fn test_builder_new_valid_ws_url() {
182        let result = WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap());
183        assert!(result.is_ok());
184    }
185
186    #[test]
187    fn test_builder_new_valid_wss_url() {
188        let result = WsIoClientBuilder::new(Url::parse("wss://localhost:8080/socket").unwrap());
189        assert!(result.is_ok());
190    }
191
192    #[test]
193    fn test_builder_new_invalid_scheme() {
194        let result = WsIoClientBuilder::new(Url::parse("http://localhost:8080/socket").unwrap());
195        assert!(result.is_err());
196        if let Err(e) = result {
197            let err_msg = format!("{e}");
198            assert!(err_msg.contains("Invalid URL scheme"));
199        }
200    }
201
202    #[test]
203    fn test_builder_configuration_chaining() {
204        WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap())
205            .unwrap()
206            .init_handler_timeout(Duration::from_secs(10))
207            .init_packet_timeout(Duration::from_secs(15))
208            .on_session_close_handler_timeout(Duration::from_secs(5))
209            .packet_codec(WsIoPacketCodec::SerdeJson)
210            .ping_interval(Duration::from_secs(30))
211            .ready_packet_timeout(Duration::from_secs(10))
212            .reconnect_delay(Duration::from_secs(5))
213            .request_path("/custom/path")
214            .build();
215    }
216
217    #[test]
218    fn test_builder_request_path_normalizes() {
219        // Test that request_path properly normalizes paths
220        WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap())
221            .unwrap()
222            .request_path("/multiple//slashes///path/")
223            .build();
224    }
225
226    #[test]
227    fn test_builder_websocket_config_override() {
228        WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap())
229            .unwrap()
230            .websocket_config_mut(|config| {
231                *config = config.max_frame_size(Some(1024 * 1024));
232            })
233            .build();
234    }
235
236    #[test]
237    fn test_builder_all_timeout_configurations() {
238        WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap())
239            .unwrap()
240            .init_handler_timeout(Duration::from_secs(1))
241            .init_packet_timeout(Duration::from_secs(2))
242            .on_session_close_handler_timeout(Duration::from_secs(3))
243            .ready_packet_timeout(Duration::from_secs(4))
244            .build();
245    }
246
247    #[test]
248    fn test_builder_reconnect_delay_configuration() {
249        WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap())
250            .unwrap()
251            .reconnect_delay(Duration::from_millis(500))
252            .build();
253    }
254}