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