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

wsio_client/
builder.rs
