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
12pub 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 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 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}