ri/ws/
client.rs

1//! Copyright 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18//! # WebSocket Client Implementation
19
20use super::*;
21use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
22
23#[derive(Debug, Clone)]
24#[cfg_attr(feature = "pyo3", pyclass)]
25pub struct RiWSClientConfig {
26    pub heartbeat_interval: u64,
27    pub heartbeat_timeout: u64,
28    pub max_message_size: usize,
29    pub connect_timeout: u64,
30    pub auto_reconnect: bool,
31    pub reconnect_interval: u64,
32}
33
34#[cfg(feature = "pyo3")]
35#[pymethods]
36impl RiWSClientConfig {
37    #[new]
38    fn new() -> Self {
39        Self::default()
40    }
41    
42    #[getter]
43    fn get_heartbeat_interval(&self) -> u64 {
44        self.heartbeat_interval
45    }
46    
47    #[setter]
48    fn set_heartbeat_interval(&mut self, val: u64) {
49        self.heartbeat_interval = val;
50    }
51    
52    #[getter]
53    fn get_heartbeat_timeout(&self) -> u64 {
54        self.heartbeat_timeout
55    }
56    
57    #[setter]
58    fn set_heartbeat_timeout(&mut self, val: u64) {
59        self.heartbeat_timeout = val;
60    }
61    
62    #[getter]
63    fn get_max_message_size(&self) -> usize {
64        self.max_message_size
65    }
66    
67    #[setter]
68    fn set_max_message_size(&mut self, val: usize) {
69        self.max_message_size = val;
70    }
71    
72    #[getter]
73    fn get_connect_timeout(&self) -> u64 {
74        self.connect_timeout
75    }
76    
77    #[setter]
78    fn set_connect_timeout(&mut self, val: u64) {
79        self.connect_timeout = val;
80    }
81    
82    #[getter]
83    fn get_auto_reconnect(&self) -> bool {
84        self.auto_reconnect
85    }
86    
87    #[setter]
88    fn set_auto_reconnect(&mut self, val: bool) {
89        self.auto_reconnect = val;
90    }
91    
92    #[getter]
93    fn get_reconnect_interval(&self) -> u64 {
94        self.reconnect_interval
95    }
96    
97    #[setter]
98    fn set_reconnect_interval(&mut self, val: u64) {
99        self.reconnect_interval = val;
100    }
101}
102
103impl Default for RiWSClientConfig {
104    fn default() -> Self {
105        Self {
106            heartbeat_interval: 30,
107            heartbeat_timeout: 60,
108            max_message_size: 65536,
109            connect_timeout: 10,
110            auto_reconnect: false,
111            reconnect_interval: 5,
112        }
113    }
114}
115
116#[derive(Debug, Clone)]
117#[cfg_attr(feature = "pyo3", pyclass)]
118pub struct RiWSClientStats {
119    pub total_connections: u64,
120    pub total_messages_sent: u64,
121    pub total_messages_received: u64,
122    pub total_bytes_sent: u64,
123    pub total_bytes_received: u64,
124    pub connection_errors: u64,
125    pub message_errors: u64,
126    pub last_connected_at: Option<u64>,
127}
128
129#[cfg(feature = "pyo3")]
130#[pymethods]
131impl RiWSClientStats {
132    #[getter]
133    fn get_total_connections(&self) -> u64 {
134        self.total_connections
135    }
136    
137    #[getter]
138    fn get_total_messages_sent(&self) -> u64 {
139        self.total_messages_sent
140    }
141    
142    #[getter]
143    fn get_total_messages_received(&self) -> u64 {
144        self.total_messages_received
145    }
146    
147    #[getter]
148    fn get_total_bytes_sent(&self) -> u64 {
149        self.total_bytes_sent
150    }
151    
152    #[getter]
153    fn get_total_bytes_received(&self) -> u64 {
154        self.total_bytes_received
155    }
156    
157    #[getter]
158    fn get_connection_errors(&self) -> u64 {
159        self.connection_errors
160    }
161    
162    #[getter]
163    fn get_message_errors(&self) -> u64 {
164        self.message_errors
165    }
166    
167    #[getter]
168    fn get_last_connected_at(&self) -> Option<u64> {
169        self.last_connected_at
170    }
171}
172
173impl RiWSClientStats {
174    pub fn new() -> Self {
175        Self {
176            total_connections: 0,
177            total_messages_sent: 0,
178            total_messages_received: 0,
179            total_bytes_sent: 0,
180            total_bytes_received: 0,
181            connection_errors: 0,
182            message_errors: 0,
183            last_connected_at: None,
184        }
185    }
186
187    fn record_connection(&mut self) {
188        self.total_connections += 1;
189        self.last_connected_at = Some(chrono::Utc::now().timestamp() as u64);
190    }
191
192    #[allow(dead_code)]
193    fn record_message_sent(&mut self, size: usize) {
194        self.total_messages_sent += 1;
195        self.total_bytes_sent += size as u64;
196    }
197
198    #[allow(dead_code)]
199    fn record_message_received(&mut self, size: usize) {
200        self.total_messages_received += 1;
201        self.total_bytes_received += size as u64;
202    }
203
204    #[allow(dead_code)]
205    fn record_connection_error(&mut self) {
206        self.connection_errors += 1;
207    }
208
209    #[allow(dead_code)]
210    fn record_message_error(&mut self) {
211        self.message_errors += 1;
212    }
213}
214
215impl Default for RiWSClientStats {
216    fn default() -> Self {
217        Self::new()
218    }
219}
220
221pub struct RiWSClient {
222    config: RiWSClientConfig,
223    stats: Arc<RwLock<RiWSClientStats>>,
224    connected: Arc<RwLock<bool>>,
225    server_url: String,
226}
227
228#[cfg(feature = "pyo3")]
229#[pyclass]
230pub struct RiWSClientPy {
231    inner: RiWSClient,
232}
233
234#[cfg(feature = "pyo3")]
235#[pymethods]
236impl RiWSClientPy {
237    #[new]
238    fn new(server_url: String) -> Self {
239        Self {
240            inner: RiWSClient::new(server_url),
241        }
242    }
243
244    #[staticmethod]
245    fn with_config(server_url: String, config: RiWSClientConfig) -> Self {
246        Self {
247            inner: RiWSClient::with_config(server_url, config),
248        }
249    }
250
251    fn get_stats(&self) -> RiWSClientStats {
252        self.inner.get_stats()
253    }
254
255    fn is_connected(&self) -> bool {
256        tokio::runtime::Handle::try_current()
257            .map(|handle| handle.block_on(async { self.inner.is_connected().await }))
258            .unwrap_or(false)
259    }
260
261    fn connect(&mut self) -> PyResult<()> {
262        let rt = tokio::runtime::Runtime::new()
263            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
264        
265        rt.block_on(async {
266            self.inner.connect().await
267        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
268    }
269
270    fn send(&self, data: Vec<u8>) -> PyResult<()> {
271        let rt = tokio::runtime::Runtime::new()
272            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
273        
274        rt.block_on(async {
275            self.inner.send(&data).await
276        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
277    }
278
279    fn send_text(&self, text: String) -> PyResult<()> {
280        let rt = tokio::runtime::Runtime::new()
281            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
282        
283        rt.block_on(async {
284            self.inner.send_text(&text).await
285        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
286    }
287
288    fn close(&mut self) -> PyResult<()> {
289        let rt = tokio::runtime::Runtime::new()
290            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
291        
292        rt.block_on(async {
293            self.inner.close().await
294        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
295    }
296}
297
298impl RiWSClient {
299    pub fn new(server_url: String) -> Self {
300        Self::with_config(server_url, RiWSClientConfig::default())
301    }
302
303    pub fn with_config(server_url: String, config: RiWSClientConfig) -> Self {
304        Self {
305            config,
306            stats: Arc::new(RwLock::new(RiWSClientStats::new())),
307            connected: Arc::new(RwLock::new(false)),
308            server_url,
309        }
310    }
311
312    pub fn get_stats(&self) -> RiWSClientStats {
313        self.stats.try_read()
314            .map(|guard| guard.clone())
315            .unwrap_or_else(|_| RiWSClientStats::new())
316    }
317
318    pub async fn is_connected(&self) -> bool {
319        *self.connected.read().await
320    }
321
322    pub async fn connect(&mut self) -> RiResult<()> {
323        if *self.connected.read().await {
324            return Ok(());
325        }
326
327        let url = self.server_url.parse::<http::Uri>().map_err(|e| WSError::Connection {
328            message: format!("Invalid WebSocket URL: {}", e)
329        })?;
330
331        let ws_config = WebSocketConfig {
332            max_message_size: Some(self.config.max_message_size),
333            max_frame_size: Some(self.config.max_message_size),
334            max_write_buffer_size: self.config.max_message_size,
335            ..Default::default()
336        };
337
338        let (_ws_stream, _response) = tokio_tungstenite::connect_async_with_config(
339            &url,
340            Some(ws_config),
341            true,
342        )
343        .await
344        .map_err(|e| WSError::Connection {
345            message: format!("Failed to connect to WebSocket server: {}", e)
346        })?;
347
348        *self.connected.write().await = true;
349        self.stats.write().await.record_connection();
350
351        tracing::info!("WebSocket client connected to {}", self.server_url);
352        Ok(())
353    }
354
355    pub async fn send(&self, _data: &[u8]) -> RiResult<()> {
356        if !*self.connected.read().await {
357            return Err(WSError::Connection {
358                message: "Not connected to WebSocket server".to_string()
359            }.into());
360        }
361        Ok(())
362    }
363
364    pub async fn send_text(&self, _text: &str) -> RiResult<()> {
365        if !*self.connected.read().await {
366            return Err(WSError::Connection {
367                message: "Not connected to WebSocket server".to_string()
368            }.into());
369        }
370        Ok(())
371    }
372
373    pub async fn close(&mut self) -> RiResult<()> {
374        *self.connected.write().await = false;
375        tracing::info!("WebSocket client disconnected from {}", self.server_url);
376        Ok(())
377    }
378
379    pub async fn disconnect(&mut self) {
380        let _ = self.close().await;
381    }
382}
383
384impl Clone for RiWSClient {
385    fn clone(&self) -> Self {
386        Self {
387            config: self.config.clone(),
388            stats: self.stats.clone(),
389            connected: self.connected.clone(),
390            server_url: self.server_url.clone(),
391        }
392    }
393}
394
395impl Default for RiWSClient {
396    fn default() -> Self {
397        Self::new("ws://127.0.0.1:8080".to_string())
398    }
399}
ri/ws/client.rs

ri/ws/
client.rs
