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::{
15    http::Request,
16    protocol::WebSocketConfig,
17};
18use url::Url;
19
20use crate::{
21    WsIoClient,
22    config::WsIoClientConfig,
23    core::packet::codecs::WsIoPacketCodec,
24    runtime::WsIoClientRuntime,
25    session::WsIoClientSession,
26};
27
28// Structs
29
30/// Builder for configuring and creating a [`WsIoClient`].
31///
32/// The URL passed to the client constructor selects the namespace from its path,
33/// while the actual WebSocket request path defaults to `/ws.io`.
34pub struct WsIoClientBuilder {
35    config: WsIoClientConfig,
36    connect_url: Url,
37}
38
39impl WsIoClientBuilder {
40    pub(crate) fn new(mut url: Url) -> Result<Self> {
41        if !matches!(url.scheme(), "ws" | "wss") {
42            bail!("Invalid URL scheme: {}", url.scheme());
43        }
44
45        let mut query_pairs = url.query_pairs().collect::<Vec<_>>();
46        query_pairs.retain(|(k, _)| k != "namespace");
47        query_pairs.push(("namespace".into(), Self::normalize_url_path(url.path()).into()));
48        let query = query_pairs
49            .iter()
50            .map(|(k, v)| format!("{k}={v}"))
51            .collect::<Vec<_>>()
52            .join("&");
53
54        url.set_query(Some(&query));
55        url.set_path("ws.io");
56        Ok(Self {
57            config: WsIoClientConfig {
58                disconnect_timeout: Duration::from_secs(5),
59                init_handler: None,
60                init_handler_timeout: Duration::from_secs(3),
61                init_packet_timeout: Duration::from_secs(5),
62                on_session_close_handler: None,
63                on_session_close_handler_timeout: Duration::from_secs(2),
64                on_session_ready_handler: None,
65                packet_codec: WsIoPacketCodec::SerdeJson,
66                ping_interval: Duration::from_secs(25),
67                ready_packet_timeout: Duration::from_secs(5),
68                reconnect_delay: Duration::from_secs(1),
69                request_modifier: None,
70                websocket_config: WebSocketConfig::default()
71                    .max_frame_size(Some(8 * 1024 * 1024))
72                    .max_message_size(Some(16 * 1024 * 1024))
73                    .max_write_buffer_size(2 * 1024 * 1024)
74                    .read_buffer_size(8 * 1024)
75                    .write_buffer_size(8 * 1024),
76            },
77            connect_url: url,
78        })
79    }
80
81    // Private methods
82    fn normalize_url_path(path: &str) -> String {
83        format!(
84            "/{}",
85            path.split('/').filter(|s| !s.is_empty()).collect::<Vec<_>>().join("/")
86        )
87    }
88
89    // Public methods
90
91    /// Builds a [`WsIoClient`] with the accumulated configuration.
92    pub fn build(self) -> WsIoClient {
93        WsIoClient(WsIoClientRuntime::new(self.config, self.connect_url))
94    }
95
96    /// Sets how long `disconnect().await` waits for graceful WebSocket
97    /// shutdown before aborting the connection read/write tasks.
98    pub fn disconnect_timeout(mut self, duration: Duration) -> Self {
99        self.config.disconnect_timeout = duration;
100        self
101    }
102
103    /// Sets the maximum duration allowed for the init handler to run.
104    ///
105    /// The init handler is registered with [`Self::with_init_handler`] and is
106    /// invoked after the server sends the init packet.
107    pub fn init_handler_timeout(mut self, duration: Duration) -> Self {
108        self.config.init_handler_timeout = duration;
109        self
110    }
111
112    /// Sets how long the client waits for the server init packet after the
113    /// WebSocket connection is established.
114    ///
115    /// If the init packet is not received before this timeout, the session is
116    /// closed and the runtime may reconnect according to [`Self::reconnect_delay`].
117    pub fn init_packet_timeout(mut self, duration: Duration) -> Self {
118        self.config.init_packet_timeout = duration;
119        self
120    }
121
122    /// Registers a handler that runs when a session closes.
123    ///
124    /// The handler is awaited during session cleanup and is bounded by
125    /// [`Self::on_session_close_handler_timeout`].
126    pub fn on_session_close<H, Fut>(mut self, handler: H) -> Self
127    where
128        H: Fn(Arc<WsIoClientSession>) -> Fut + Send + Sync + 'static,
129        Fut: Future<Output = Result<()>> + Send + 'static,
130    {
131        self.config.on_session_close_handler = Some(Box::new(move |session| Box::pin(handler(session))));
132        self
133    }
134
135    /// Sets the maximum duration allowed for the session-close handler to run.
136    pub fn on_session_close_handler_timeout(mut self, duration: Duration) -> Self {
137        self.config.on_session_close_handler_timeout = duration;
138        self
139    }
140
141    /// Registers a handler that runs after a session becomes ready.
142    ///
143    /// The handler is spawned asynchronously after the ready packet is received,
144    /// so it does not block the connection handshake.
145    pub fn on_session_ready<H, Fut>(mut self, handler: H) -> Self
146    where
147        H: Fn(Arc<WsIoClientSession>) -> Fut + Send + Sync + 'static,
148        Fut: Future<Output = Result<()>> + Send + 'static,
149    {
150        self.config.on_session_ready_handler = Some(Arc::new(move |session| Box::pin(handler(session))));
151        self
152    }
153
154    /// Sets the packet codec used to encode and decode ws.io protocol packets.
155    ///
156    /// This must match the server namespace codec.
157    pub fn packet_codec(mut self, packet_codec: WsIoPacketCodec) -> Self {
158        self.config.packet_codec = packet_codec;
159        self
160    }
161
162    /// Sets the interval for client heartbeat frames.
163    ///
164    /// After session initialization starts, the client periodically sends a
165    /// one-byte binary WebSocket frame. The server treats single-byte binary
166    /// frames as heartbeats and ignores them before packet decoding.
167    pub fn ping_interval(mut self, duration: Duration) -> Self {
168        self.config.ping_interval = duration;
169        self
170    }
171
172    /// Sets how long the client waits for the server ready packet.
173    ///
174    /// The ready timeout starts after the client handles the server init packet
175    /// and sends its init response.
176    pub fn ready_packet_timeout(mut self, duration: Duration) -> Self {
177        self.config.ready_packet_timeout = duration;
178        self
179    }
180
181    /// Sets the delay before the runtime attempts another connection.
182    ///
183    /// This delay is used after a connection attempt/session ends while the client
184    /// runtime is still running.
185    pub fn reconnect_delay(mut self, delay: Duration) -> Self {
186        self.config.reconnect_delay = delay;
187        self
188    }
189
190    /// Registers an async modifier for the WebSocket HTTP request.
191    ///
192    /// Use this to add headers or adjust request metadata before
193    /// `connect_async_with_config` is called.
194    pub fn request_modifier<M, Fut>(mut self, modifier: M) -> Self
195    where
196        M: Fn(Request<()>) -> Fut + Send + Sync + 'static,
197        Fut: Future<Output = Result<Request<()>>> + Send + 'static,
198    {
199        self.config.request_modifier = Some(Box::new(move |request| Box::pin(modifier(request))));
200        self
201    }
202
203    /// Sets the WebSocket HTTP request path.
204    ///
205    /// Paths are normalized to a single leading slash with empty path segments
206    /// removed. This controls the request URI path, not the namespace query value
207    /// inferred from the original URL passed to the builder.
208    pub fn request_path(mut self, request_path: impl AsRef<str>) -> Self {
209        self.connect_url
210            .set_path(&Self::normalize_url_path(request_path.as_ref()));
211
212        self
213    }
214
215    /// Replaces the full Tungstenite WebSocket configuration.
216    ///
217    /// This controls transport limits and buffer sizes passed to the WebSocket
218    /// connection. It is also used to derive internal channel capacity from the
219    /// configured max-write/write-buffer ratio.
220    pub fn websocket_config(mut self, websocket_config: WebSocketConfig) -> Self {
221        self.config.websocket_config = websocket_config;
222        self
223    }
224
225    /// Mutates the current Tungstenite WebSocket configuration in place.
226    ///
227    /// Prefer this when you want to adjust one or two fields while keeping the
228    /// builder defaults for the rest.
229    pub fn websocket_config_mut<F: FnOnce(&mut WebSocketConfig)>(mut self, f: F) -> Self {
230        f(&mut self.config.websocket_config);
231        self
232    }
233
234    /// Registers the client-side init handler.
235    ///
236    /// The handler receives the session and the optional server init payload
237    /// decoded as `D`. Its optional return value is encoded as `R` and sent back
238    /// to the server as the client init response.
239    pub fn with_init_handler<H, Fut, D, R>(mut self, handler: H) -> WsIoClientBuilder
240    where
241        H: Fn(Arc<WsIoClientSession>, Option<D>) -> Fut + Send + Sync + 'static,
242        Fut: Future<Output = Result<Option<R>>> + Send + 'static,
243        D: DeserializeOwned + Send + 'static,
244        R: Serialize + Send + 'static,
245    {
246        let handler = Arc::new(handler);
247        self.config.init_handler = Some(Box::new(move |session, bytes, packet_codec| {
248            let handler = handler.clone();
249            Box::pin(async move {
250                handler(session, bytes.map(|bytes| packet_codec.decode_data(bytes)).transpose()?)
251                    .await?
252                    .map(|data| packet_codec.encode_data(&data))
253                    .transpose()
254            })
255        }));
256
257        self
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use tokio_tungstenite::tungstenite::http::HeaderValue;
264
265    use super::*;
266
267    const TEST_URL: &str = "ws://localhost:8080/socket";
268
269    fn test_builder() -> WsIoClientBuilder {
270        WsIoClientBuilder::new(Url::parse(TEST_URL).unwrap()).unwrap()
271    }
272
273    #[test]
274    fn test_builder_new_valid_ws_url_sets_default_request_path_and_namespace_query() {
275        let builder = test_builder();
276
277        assert_eq!(builder.connect_url.path(), "/ws.io");
278        assert_eq!(
279            builder
280                .connect_url
281                .query_pairs()
282                .find(|(key, _)| key == "namespace")
283                .map(|(_, value)| value.into_owned()),
284            Some("/socket".into())
285        );
286    }
287
288    #[test]
289    fn test_builder_new_valid_wss_url() {
290        let result = WsIoClientBuilder::new(Url::parse("wss://localhost:8080/socket").unwrap());
291        assert!(result.is_ok());
292    }
293
294    #[test]
295    fn test_builder_new_invalid_scheme() {
296        let result = WsIoClientBuilder::new(Url::parse("http://localhost:8080/socket").unwrap());
297        assert!(result.is_err());
298        if let Err(e) = result {
299            let err_msg = format!("{e}");
300            assert!(err_msg.contains("Invalid URL scheme"));
301        }
302    }
303
304    #[test]
305    fn test_builder_configuration_chaining_updates_runtime_config() {
306        let builder = test_builder()
307            .disconnect_timeout(Duration::from_secs(20))
308            .init_handler_timeout(Duration::from_secs(10))
309            .init_packet_timeout(Duration::from_secs(15))
310            .on_session_close_handler_timeout(Duration::from_secs(5))
311            .packet_codec(WsIoPacketCodec::SerdeJson)
312            .ping_interval(Duration::from_secs(30))
313            .ready_packet_timeout(Duration::from_secs(10))
314            .reconnect_delay(Duration::from_secs(5))
315            .request_path("/custom/path");
316
317        assert_eq!(builder.connect_url.path(), "/custom/path");
318
319        let client = builder.build();
320
321        let config = &client.0.config;
322        assert_eq!(config.disconnect_timeout, Duration::from_secs(20));
323        assert_eq!(config.init_handler_timeout, Duration::from_secs(10));
324        assert_eq!(config.init_packet_timeout, Duration::from_secs(15));
325        assert_eq!(config.on_session_close_handler_timeout, Duration::from_secs(5));
326        assert!(matches!(config.packet_codec, WsIoPacketCodec::SerdeJson));
327        assert_eq!(config.ping_interval, Duration::from_secs(30));
328        assert_eq!(config.ready_packet_timeout, Duration::from_secs(10));
329        assert_eq!(config.reconnect_delay, Duration::from_secs(5));
330    }
331
332    #[test]
333    fn test_builder_request_path_normalizes() {
334        let builder = test_builder().request_path("/multiple//slashes///path/");
335
336        assert_eq!(builder.connect_url.path(), "/multiple/slashes/path");
337    }
338
339    #[test]
340    fn test_builder_websocket_config_override() {
341        let client = test_builder()
342            .websocket_config_mut(|config| {
343                *config = config.max_frame_size(Some(1024 * 1024));
344            })
345            .build();
346
347        assert_eq!(client.0.config.websocket_config.max_frame_size, Some(1024 * 1024));
348    }
349
350    #[test]
351    fn test_builder_websocket_config_replaces_defaults() {
352        let config = WebSocketConfig::default().max_frame_size(Some(42));
353        let client = test_builder().websocket_config(config).build();
354
355        assert_eq!(client.0.config.websocket_config.max_frame_size, Some(42));
356    }
357
358    #[test]
359    fn test_builder_with_init_and_session_handlers_registers_callbacks() {
360        let client = test_builder()
361            .with_init_handler(|_session, _data: Option<String>| async { Ok(Some("response".to_string())) })
362            .on_session_ready(|_session| async { Ok(()) })
363            .on_session_close(|_session| async { Ok(()) })
364            .build();
365
366        assert!(client.0.config.init_handler.is_some());
367        assert!(client.0.config.on_session_ready_handler.is_some());
368        assert!(client.0.config.on_session_close_handler.is_some());
369    }
370
371    #[test]
372    fn test_builder_request_modifier_registers_async_callback() {
373        let client = test_builder()
374            .request_modifier(|mut request| async move {
375                request
376                    .headers_mut()
377                    .insert("x-wsio-test", HeaderValue::from_static("enabled"));
378
379                Ok(request)
380            })
381            .build();
382
383        assert!(client.0.config.request_modifier.is_some());
384    }
385
386    #[test]
387    fn test_builder_all_timeout_configurations() {
388        let client = test_builder()
389            .disconnect_timeout(Duration::from_millis(500))
390            .init_handler_timeout(Duration::from_secs(1))
391            .init_packet_timeout(Duration::from_secs(2))
392            .on_session_close_handler_timeout(Duration::from_secs(3))
393            .ready_packet_timeout(Duration::from_secs(4))
394            .build();
395
396        assert_eq!(client.0.config.disconnect_timeout, Duration::from_millis(500));
397        assert_eq!(client.0.config.init_handler_timeout, Duration::from_secs(1));
398        assert_eq!(client.0.config.init_packet_timeout, Duration::from_secs(2));
399        assert_eq!(client.0.config.on_session_close_handler_timeout, Duration::from_secs(3));
400        assert_eq!(client.0.config.ready_packet_timeout, Duration::from_secs(4));
401    }
402
403    #[test]
404    fn test_builder_reconnect_delay_configuration() {
405        let client = test_builder().reconnect_delay(Duration::from_millis(500)).build();
406
407        assert_eq!(client.0.config.reconnect_delay, Duration::from_millis(500));
408    }
409}
wsio_client/builder.rs

wsio_client/
builder.rs
