Skip to main content

wsio_server/
builder.rs

1use std::time::Duration;
2
3use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
4
5use crate::{
6    WsIoServer,
7    config::WsIoServerConfig,
8    core::packet::codecs::WsIoPacketCodec,
9    runtime::WsIoServerRuntime,
10};
11
12// Structs
13
14/// Builder for configuring and creating a [`WsIoServer`].
15///
16/// Server-level settings become the defaults inherited by namespaces created
17/// from the server. Namespace builders may override most of these values per
18/// namespace.
19#[derive(Debug)]
20pub struct WsIoServerBuilder {
21    config: WsIoServerConfig,
22}
23
24impl WsIoServerBuilder {
25    pub(crate) fn new() -> Self {
26        Self {
27            config: WsIoServerConfig {
28                broadcast_concurrency_limit: 512,
29                http_request_upgrade_timeout: Duration::from_secs(3),
30                init_request_handler_timeout: Duration::from_secs(3),
31                init_response_handler_timeout: Duration::from_secs(3),
32                init_response_timeout: Duration::from_secs(5),
33                middleware_execution_timeout: Duration::from_secs(2),
34                on_close_handler_timeout: Duration::from_secs(2),
35                on_connect_handler_timeout: Duration::from_secs(3),
36                packet_codec: WsIoPacketCodec::SerdeJson,
37                request_path: "/ws.io".to_owned(),
38                websocket_config: WebSocketConfig::default()
39                    .max_frame_size(Some(8 * 1024 * 1024))
40                    .max_message_size(Some(16 * 1024 * 1024))
41                    .max_write_buffer_size(2 * 1024 * 1024)
42                    .read_buffer_size(8 * 1024)
43                    .write_buffer_size(8 * 1024),
44            },
45        }
46    }
47
48    // Public methods
49    /// Sets the default maximum number of broadcast send operations to run at
50    /// once.
51    ///
52    /// This value is inherited by namespace builders and passed to
53    /// `StreamExt::for_each_concurrent`; `0` is treated as no concurrency limit.
54    pub fn broadcast_concurrency_limit(mut self, broadcast_concurrency_limit: usize) -> Self {
55        self.config.broadcast_concurrency_limit = broadcast_concurrency_limit;
56        self
57    }
58
59    /// Builds a [`WsIoServer`] with the accumulated configuration.
60    pub fn build(self) -> WsIoServer {
61        WsIoServer(WsIoServerRuntime::new(self.config))
62    }
63
64    /// Sets the default timeout for a matched HTTP request to finish the
65    /// WebSocket upgrade.
66    ///
67    /// The timeout wraps the HTTP adapter's upgrade future. Namespace builders
68    /// inherit this value and may override it.
69    pub fn http_request_upgrade_timeout(mut self, duration: Duration) -> Self {
70        self.config.http_request_upgrade_timeout = duration;
71        self
72    }
73
74    /// Sets the default maximum duration allowed for init-request handlers.
75    ///
76    /// Init-request handlers are registered per namespace with
77    /// `WsIoServerNamespaceBuilder::with_init_request`.
78    pub fn init_request_handler_timeout(mut self, duration: Duration) -> Self {
79        self.config.init_request_handler_timeout = duration;
80        self
81    }
82
83    /// Sets the default maximum duration allowed for init-response handlers.
84    ///
85    /// Init-response handlers are registered per namespace with
86    /// `WsIoServerNamespaceBuilder::with_init_response`.
87    pub fn init_response_handler_timeout(mut self, duration: Duration) -> Self {
88        self.config.init_response_handler_timeout = duration;
89        self
90    }
91
92    /// Sets the default timeout for waiting on a client init-response packet.
93    ///
94    /// This timeout starts after the server sends its init packet. Namespace
95    /// builders inherit this value and may override it.
96    pub fn init_response_timeout(mut self, duration: Duration) -> Self {
97        self.config.init_response_timeout = duration;
98        self
99    }
100
101    /// Sets the default maximum duration allowed for namespace middleware.
102    ///
103    /// Middleware is registered per namespace with
104    /// `WsIoServerNamespaceBuilder::with_middleware`.
105    pub fn middleware_execution_timeout(mut self, duration: Duration) -> Self {
106        self.config.middleware_execution_timeout = duration;
107        self
108    }
109
110    /// Sets the default maximum duration allowed for per-connection close
111    /// handlers.
112    ///
113    /// This applies to handlers registered through `WsIoServerConnection::on_close`.
114    pub fn on_close_handler_timeout(mut self, duration: Duration) -> Self {
115        self.config.on_close_handler_timeout = duration;
116        self
117    }
118
119    /// Sets the default maximum duration allowed for namespace on-connect
120    /// handlers.
121    ///
122    /// On-connect handlers are registered per namespace with
123    /// `WsIoServerNamespaceBuilder::on_connect`.
124    pub fn on_connect_handler_timeout(mut self, duration: Duration) -> Self {
125        self.config.on_connect_handler_timeout = duration;
126        self
127    }
128
129    /// Sets the default packet codec for namespaces.
130    ///
131    /// The codec is used for ws.io protocol packets and init payload data.
132    /// Namespace builders inherit this value and may override it.
133    pub fn packet_codec(mut self, packet_codec: WsIoPacketCodec) -> Self {
134        self.config.packet_codec = packet_codec;
135        self
136    }
137
138    /// Sets the HTTP request path handled by the server adapter.
139    ///
140    /// Requests whose URI path does not match this value pass through to the
141    /// wrapped service. Client namespace routing is carried separately in the
142    /// `namespace` query parameter.
143    pub fn request_path(mut self, request_path: impl AsRef<str>) -> Self {
144        self.config.request_path = request_path.as_ref().to_owned();
145        self
146    }
147
148    /// Replaces the default Tungstenite WebSocket configuration.
149    ///
150    /// Namespace builders inherit this value. It controls transport limits and
151    /// buffer sizes and is also used to derive internal channel capacity from the
152    /// configured max-write/write-buffer ratio.
153    pub fn websocket_config(mut self, websocket_config: WebSocketConfig) -> Self {
154        self.config.websocket_config = websocket_config;
155        self
156    }
157
158    /// Mutates the current default Tungstenite WebSocket configuration in place.
159    ///
160    /// Prefer this when you want to adjust one or two fields while keeping the
161    /// builder defaults for the rest.
162    pub fn websocket_config_mut<F: FnOnce(&mut WebSocketConfig)>(mut self, f: F) -> Self {
163        f(&mut self.config.websocket_config);
164        self
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use std::time::Duration;
171
172    use super::*;
173
174    #[test]
175    fn test_builder_configuration_chaining() {
176        let server = WsIoServer::builder()
177            .broadcast_concurrency_limit(1024)
178            .http_request_upgrade_timeout(Duration::from_millis(750))
179            .init_request_handler_timeout(Duration::from_secs(1))
180            .init_response_handler_timeout(Duration::from_secs(2))
181            .init_response_timeout(Duration::from_secs(3))
182            .middleware_execution_timeout(Duration::from_secs(4))
183            .on_close_handler_timeout(Duration::from_secs(5))
184            .on_connect_handler_timeout(Duration::from_secs(6))
185            .packet_codec(WsIoPacketCodec::Msgpack)
186            .request_path("/custom")
187            .websocket_config_mut(|config| {
188                *config = config.max_frame_size(Some(999));
189            })
190            .build();
191
192        // Access internal config through the built runtime
193        let config = &server.0.config;
194        assert_eq!(config.broadcast_concurrency_limit, 1024);
195        assert_eq!(config.http_request_upgrade_timeout, Duration::from_millis(750));
196        assert_eq!(config.init_request_handler_timeout, Duration::from_secs(1));
197        assert_eq!(config.init_response_handler_timeout, Duration::from_secs(2));
198        assert_eq!(config.init_response_timeout, Duration::from_secs(3));
199        assert_eq!(config.middleware_execution_timeout, Duration::from_secs(4));
200        assert_eq!(config.on_close_handler_timeout, Duration::from_secs(5));
201        assert_eq!(config.on_connect_handler_timeout, Duration::from_secs(6));
202        assert!(matches!(config.packet_codec, WsIoPacketCodec::Msgpack));
203        assert_eq!(config.request_path, "/custom");
204        assert_eq!(config.websocket_config.max_frame_size, Some(999));
205    }
206
207    #[test]
208    fn test_builder_websocket_config_override() {
209        let config = WebSocketConfig::default().max_frame_size(Some(42));
210        let server = WsIoServer::builder().websocket_config(config).build();
211        assert_eq!(server.0.config.websocket_config.max_frame_size, Some(42));
212    }
213}