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
13pub struct WsIoServerBuilder {
14    config: WsIoServerConfig,
15}
16
17impl WsIoServerBuilder {
18    pub(crate) fn new() -> Self {
19        Self {
20            config: WsIoServerConfig {
21                broadcast_concurrency_limit: 512,
22                http_request_upgrade_timeout: Duration::from_secs(3),
23                init_request_handler_timeout: Duration::from_secs(3),
24                init_response_handler_timeout: Duration::from_secs(3),
25                init_response_timeout: Duration::from_secs(5),
26                middleware_execution_timeout: Duration::from_secs(2),
27                on_close_handler_timeout: Duration::from_secs(2),
28                on_connect_handler_timeout: Duration::from_secs(3),
29                packet_codec: WsIoPacketCodec::SerdeJson,
30                request_path: "/ws.io".into(),
31                websocket_config: WebSocketConfig::default()
32                    .max_frame_size(Some(8 * 1024 * 1024))
33                    .max_message_size(Some(16 * 1024 * 1024))
34                    .max_write_buffer_size(2 * 1024 * 1024)
35                    .read_buffer_size(8 * 1024)
36                    .write_buffer_size(8 * 1024),
37            },
38        }
39    }
40
41    // Public methods
42    pub fn broadcast_concurrency_limit(mut self, broadcast_concurrency_limit: usize) -> Self {
43        self.config.broadcast_concurrency_limit = broadcast_concurrency_limit;
44        self
45    }
46
47    pub fn build(self) -> WsIoServer {
48        WsIoServer(WsIoServerRuntime::new(self.config))
49    }
50
51    pub fn http_request_upgrade_timeout(mut self, duration: Duration) -> Self {
52        self.config.http_request_upgrade_timeout = duration;
53        self
54    }
55
56    pub fn init_request_handler_timeout(mut self, duration: Duration) -> Self {
57        self.config.init_request_handler_timeout = duration;
58        self
59    }
60
61    pub fn init_response_handler_timeout(mut self, duration: Duration) -> Self {
62        self.config.init_response_handler_timeout = duration;
63        self
64    }
65
66    pub fn init_response_timeout(mut self, duration: Duration) -> Self {
67        self.config.init_response_timeout = duration;
68        self
69    }
70
71    pub fn middleware_execution_timeout(mut self, duration: Duration) -> Self {
72        self.config.middleware_execution_timeout = duration;
73        self
74    }
75
76    pub fn on_close_handler_timeout(mut self, duration: Duration) -> Self {
77        self.config.on_close_handler_timeout = duration;
78        self
79    }
80
81    pub fn on_connect_handler_timeout(mut self, duration: Duration) -> Self {
82        self.config.on_connect_handler_timeout = duration;
83        self
84    }
85
86    pub fn packet_codec(mut self, packet_codec: WsIoPacketCodec) -> Self {
87        self.config.packet_codec = packet_codec;
88        self
89    }
90
91    pub fn request_path(mut self, request_path: impl AsRef<str>) -> Self {
92        self.config.request_path = request_path.as_ref().into();
93        self
94    }
95
96    pub fn websocket_config(mut self, websocket_config: WebSocketConfig) -> Self {
97        self.config.websocket_config = websocket_config;
98        self
99    }
100
101    pub fn websocket_config_mut<F: FnOnce(&mut WebSocketConfig)>(mut self, f: F) -> Self {
102        f(&mut self.config.websocket_config);
103        self
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use std::time::Duration;
110
111    use super::*;
112
113    #[test]
114    fn test_builder_configuration_chaining() {
115        let server = WsIoServer::builder()
116            .broadcast_concurrency_limit(1024)
117            .init_request_handler_timeout(Duration::from_secs(1))
118            .init_response_handler_timeout(Duration::from_secs(2))
119            .init_response_timeout(Duration::from_secs(3))
120            .middleware_execution_timeout(Duration::from_secs(4))
121            .on_close_handler_timeout(Duration::from_secs(5))
122            .on_connect_handler_timeout(Duration::from_secs(6))
123            .packet_codec(WsIoPacketCodec::Msgpack)
124            .request_path("/custom")
125            .websocket_config_mut(|config| {
126                *config = config.max_frame_size(Some(999));
127            })
128            .build();
129
130        // Access internal config through the built runtime
131        let config = &server.0.config;
132        assert_eq!(config.broadcast_concurrency_limit, 1024);
133        assert_eq!(config.init_request_handler_timeout, Duration::from_secs(1));
134        assert_eq!(config.init_response_handler_timeout, Duration::from_secs(2));
135        assert_eq!(config.init_response_timeout, Duration::from_secs(3));
136        assert_eq!(config.middleware_execution_timeout, Duration::from_secs(4));
137        assert_eq!(config.on_close_handler_timeout, Duration::from_secs(5));
138        assert_eq!(config.on_connect_handler_timeout, Duration::from_secs(6));
139        assert!(matches!(config.packet_codec, WsIoPacketCodec::Msgpack));
140        assert_eq!(config.request_path, "/custom");
141        assert_eq!(config.websocket_config.max_frame_size, Some(999));
142    }
143
144    #[test]
145    fn test_builder_websocket_config_override() {
146        let config = WebSocketConfig::default().max_frame_size(Some(42));
147        let server = WsIoServer::builder().websocket_config(config).build();
148        assert_eq!(server.0.config.websocket_config.max_frame_size, Some(42));
149    }
150}