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 = crate::py_runtime::py_runtime();
263        
264        rt.block_on(async {
265            self.inner.connect().await
266        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
267    }
268
269    fn send(&self, data: Vec<u8>) -> PyResult<()> {
270        let rt = crate::py_runtime::py_runtime();
271        
272        rt.block_on(async {
273            self.inner.send(&data).await
274        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
275    }
276
277    fn send_text(&self, text: String) -> PyResult<()> {
278        let rt = crate::py_runtime::py_runtime();
279        
280        rt.block_on(async {
281            self.inner.send_text(&text).await
282        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
283    }
284
285    fn close(&mut self) -> PyResult<()> {
286        let rt = crate::py_runtime::py_runtime();
287        
288        rt.block_on(async {
289            self.inner.close().await
290        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
291    }
292}
293
294impl RiWSClient {
295    pub fn new(server_url: String) -> Self {
296        Self::with_config(server_url, RiWSClientConfig::default())
297    }
298
299    pub fn with_config(server_url: String, config: RiWSClientConfig) -> Self {
300        Self {
301            config,
302            stats: Arc::new(RwLock::new(RiWSClientStats::new())),
303            connected: Arc::new(RwLock::new(false)),
304            server_url,
305        }
306    }
307
308    pub fn get_stats(&self) -> RiWSClientStats {
309        self.stats.try_read()
310            .map(|guard| guard.clone())
311            .unwrap_or_else(|_| RiWSClientStats::new())
312    }
313
314    pub async fn is_connected(&self) -> bool {
315        *self.connected.read().await
316    }
317
318    pub async fn connect(&mut self) -> RiResult<()> {
319        if *self.connected.read().await {
320            return Ok(());
321        }
322
323        let url = self.server_url.parse::<http::Uri>().map_err(|e| WSError::Connection {
324            message: format!("Invalid WebSocket URL: {}", e)
325        })?;
326
327        let ws_config = WebSocketConfig {
328            max_message_size: Some(self.config.max_message_size),
329            max_frame_size: Some(self.config.max_message_size),
330            max_write_buffer_size: self.config.max_message_size,
331            ..Default::default()
332        };
333
334        let (_ws_stream, _response) = tokio_tungstenite::connect_async_with_config(
335            &url,
336            Some(ws_config),
337            true,
338        )
339        .await
340        .map_err(|e| WSError::Connection {
341            message: format!("Failed to connect to WebSocket server: {}", e)
342        })?;
343
344        *self.connected.write().await = true;
345        self.stats.write().await.record_connection();
346
347        tracing::info!("WebSocket client connected to {}", self.server_url);
348        Ok(())
349    }
350
351    pub async fn send(&self, _data: &[u8]) -> RiResult<()> {
352        if !*self.connected.read().await {
353            return Err(WSError::Connection {
354                message: "Not connected to WebSocket server".to_string()
355            }.into());
356        }
357        Ok(())
358    }
359
360    pub async fn send_text(&self, _text: &str) -> RiResult<()> {
361        if !*self.connected.read().await {
362            return Err(WSError::Connection {
363                message: "Not connected to WebSocket server".to_string()
364            }.into());
365        }
366        Ok(())
367    }
368
369    pub async fn close(&mut self) -> RiResult<()> {
370        *self.connected.write().await = false;
371        tracing::info!("WebSocket client disconnected from {}", self.server_url);
372        Ok(())
373    }
374
375    pub async fn disconnect(&mut self) {
376        let _ = self.close().await;
377    }
378}
379
380impl Clone for RiWSClient {
381    fn clone(&self) -> Self {
382        Self {
383            config: self.config.clone(),
384            stats: self.stats.clone(),
385            connected: self.connected.clone(),
386            server_url: self.server_url.clone(),
387        }
388    }
389}
390
391impl Default for RiWSClient {
392    fn default() -> Self {
393        Self::new("ws://127.0.0.1:8080".to_string())
394    }
395}
ri/ws/client.rs

ri/ws/
client.rs
