nautilus_network/python/
socket.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use std::{sync::atomic::Ordering, time::Duration};
17
18use nautilus_core::python::{clone_py_object, to_pyruntime_err, to_pyvalue_err};
19use pyo3::{Py, prelude::*};
20use tokio_tungstenite::tungstenite::stream::Mode;
21
22use crate::{
23    mode::ConnectionMode,
24    socket::{SocketClient, SocketConfig, TcpMessageHandler, WriterCommand},
25};
26
27#[pymethods]
28#[pyo3_stub_gen::derive::gen_stub_pymethods]
29impl SocketConfig {
30    /// Configuration for TCP socket connection.
31    #[new]
32    #[expect(clippy::too_many_arguments, clippy::needless_pass_by_value)]
33    #[pyo3(signature = (url, ssl, suffix, handler, heartbeat=None, reconnect_timeout_ms=10_000, reconnect_delay_initial_ms=2_000, reconnect_delay_max_ms=30_000, reconnect_backoff_factor=1.5, reconnect_jitter_ms=100, connection_max_retries=5, reconnect_max_attempts=None, idle_timeout_ms=None, certs_dir=None))]
34    fn py_new(
35        url: String,
36        ssl: bool,
37        suffix: Vec<u8>,
38        handler: Py<PyAny>,
39        heartbeat: Option<(u64, Vec<u8>)>,
40        reconnect_timeout_ms: Option<u64>,
41        reconnect_delay_initial_ms: Option<u64>,
42        reconnect_delay_max_ms: Option<u64>,
43        reconnect_backoff_factor: Option<f64>,
44        reconnect_jitter_ms: Option<u64>,
45        connection_max_retries: Option<u32>,
46        reconnect_max_attempts: Option<u32>,
47        idle_timeout_ms: Option<u64>,
48        certs_dir: Option<String>,
49    ) -> PyResult<Self> {
50        let mode = if ssl { Mode::Tls } else { Mode::Plain };
51
52        // Create function pointer that calls Python handler
53        let handler_clone = clone_py_object(&handler);
54        let message_handler: TcpMessageHandler = std::sync::Arc::new(move |data: &[u8]| {
55            Python::attach(|py| {
56                if let Err(e) = handler_clone.call1(py, (data,)) {
57                    log::error!("Error calling Python message handler: {e}");
58                }
59            });
60        });
61
62        let config = Self {
63            url,
64            mode,
65            suffix,
66            message_handler: Some(message_handler),
67            heartbeat,
68            reconnect_timeout_ms,
69            reconnect_delay_initial_ms,
70            reconnect_delay_max_ms,
71            reconnect_backoff_factor,
72            reconnect_jitter_ms,
73            connection_max_retries,
74            reconnect_max_attempts,
75            idle_timeout_ms,
76            certs_dir,
77        };
78        config.validate().map_err(to_pyvalue_err)?;
79        Ok(config)
80    }
81}
82
83#[pymethods]
84#[pyo3_stub_gen::derive::gen_stub_pymethods]
85impl SocketClient {
86    /// Connect to the server.
87    #[staticmethod]
88    #[pyo3(name = "connect")]
89    #[pyo3(signature = (config, post_connection=None, post_reconnection=None, post_disconnection=None))]
90    fn py_connect(
91        config: SocketConfig,
92        post_connection: Option<Py<PyAny>>,
93        post_reconnection: Option<Py<PyAny>>,
94        post_disconnection: Option<Py<PyAny>>,
95        py: Python<'_>,
96    ) -> PyResult<Bound<'_, PyAny>> {
97        // Convert Python callbacks to function pointers
98        let post_connection_fn = post_connection.map(|callback| {
99            let callback_clone = clone_py_object(&callback);
100            std::sync::Arc::new(move || {
101                Python::attach(|py| {
102                    if let Err(e) = callback_clone.call0(py) {
103                        log::error!("Error calling post_connection handler: {e}");
104                    }
105                });
106            }) as std::sync::Arc<dyn Fn() + Send + Sync>
107        });
108
109        let post_reconnection_fn = post_reconnection.map(|callback| {
110            let callback_clone = clone_py_object(&callback);
111            std::sync::Arc::new(move || {
112                Python::attach(|py| {
113                    if let Err(e) = callback_clone.call0(py) {
114                        log::error!("Error calling post_reconnection handler: {e}");
115                    }
116                });
117            }) as std::sync::Arc<dyn Fn() + Send + Sync>
118        });
119
120        let post_disconnection_fn = post_disconnection.map(|callback| {
121            let callback_clone = clone_py_object(&callback);
122            std::sync::Arc::new(move || {
123                Python::attach(|py| {
124                    if let Err(e) = callback_clone.call0(py) {
125                        log::error!("Error calling post_disconnection handler: {e}");
126                    }
127                });
128            }) as std::sync::Arc<dyn Fn() + Send + Sync>
129        });
130
131        pyo3_async_runtimes::tokio::future_into_py(py, async move {
132            Self::connect(
133                config,
134                post_connection_fn,
135                post_reconnection_fn,
136                post_disconnection_fn,
137            )
138            .await
139            .map_err(to_pyruntime_err)
140        })
141    }
142
143    /// Check if the client connection is active.
144    ///
145    /// Returns `true` if the client is connected and has not been signalled to disconnect.
146    /// The client will automatically retry connection based on its configuration.
147    #[pyo3(name = "is_active")]
148    #[expect(clippy::needless_pass_by_value)]
149    fn py_is_active(slf: PyRef<'_, Self>) -> bool {
150        slf.is_active()
151    }
152
153    /// Check if the client is reconnecting.
154    ///
155    /// Returns `true` if the client lost connection and is attempting to reestablish it.
156    /// The client will automatically retry connection based on its configuration.
157    #[pyo3(name = "is_reconnecting")]
158    #[expect(clippy::needless_pass_by_value)]
159    fn py_is_reconnecting(slf: PyRef<'_, Self>) -> bool {
160        slf.is_reconnecting()
161    }
162
163    /// Check if the client is disconnecting.
164    ///
165    /// Returns `true` if the client is in disconnect mode.
166    #[pyo3(name = "is_disconnecting")]
167    #[expect(clippy::needless_pass_by_value)]
168    fn py_is_disconnecting(slf: PyRef<'_, Self>) -> bool {
169        slf.is_disconnecting()
170    }
171
172    /// Check if the client is closed.
173    ///
174    /// Returns `true` if the client has been explicitly disconnected or reached
175    /// maximum reconnection attempts. In this state, the client cannot be reused
176    /// and a new client must be created for further connections.
177    #[pyo3(name = "is_closed")]
178    #[expect(clippy::needless_pass_by_value)]
179    fn py_is_closed(slf: PyRef<'_, Self>) -> bool {
180        slf.is_closed()
181    }
182
183    #[pyo3(name = "mode")]
184    #[expect(clippy::needless_pass_by_value)]
185    fn py_mode(slf: PyRef<'_, Self>) -> String {
186        slf.connection_mode().to_string()
187    }
188
189    /// Reconnect the client.
190    #[pyo3(name = "reconnect")]
191    #[expect(clippy::needless_pass_by_value)]
192    fn py_reconnect<'py>(slf: PyRef<'_, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
193        let connection_mode = slf.connection_mode.clone();
194        let state_notify = slf.state_notify.clone();
195        let mode_str = ConnectionMode::from_atomic(&connection_mode).to_string();
196        log::debug!("Reconnect from mode {mode_str}");
197
198        pyo3_async_runtimes::tokio::future_into_py(py, async move {
199            match ConnectionMode::from_atomic(&connection_mode) {
200                ConnectionMode::Reconnect => {
201                    log::warn!("Cannot reconnect - socket already reconnecting");
202                }
203                ConnectionMode::Disconnect => {
204                    log::warn!("Cannot reconnect - socket disconnecting");
205                }
206                ConnectionMode::Closed => {
207                    log::warn!("Cannot reconnect - socket closed");
208                }
209                ConnectionMode::Active => {
210                    // CAS so a concurrent close cannot be overwritten back to Reconnect
211                    if !ConnectionMode::request_reconnect(&connection_mode) {
212                        log::warn!("Cannot reconnect - socket no longer active");
213                        return Ok(());
214                    }
215                    state_notify.notify_one();
216
217                    let fallback_interval = Duration::from_millis(100);
218                    let timeout = tokio::time::timeout(Duration::from_secs(30), async {
219                        loop {
220                            let notified = state_notify.notified();
221
222                            let current = ConnectionMode::from_atomic(&connection_mode);
223                            if current.is_active() {
224                                return Ok(());
225                            }
226
227                            if current.is_closed() || current.is_disconnect() {
228                                return Err("Connection closed during reconnect");
229                            }
230
231                            tokio::select! {
232                                biased;
233                                () = notified => {}
234                                () = tokio::time::sleep(fallback_interval) => {}
235                            }
236                        }
237                    })
238                    .await;
239
240                    match timeout {
241                        Ok(Ok(())) => log::debug!("Reconnected successfully"),
242                        Ok(Err(e)) => log::warn!("Reconnect aborted: {e}"),
243                        Err(_) => log::warn!("Reconnect timed out after 30s"),
244                    }
245                }
246            }
247
248            Ok(())
249        })
250    }
251
252    /// Close the client.
253    ///
254    /// Controller task will periodically check the disconnect mode
255    /// and shutdown the client if it is not alive.
256    #[pyo3(name = "close")]
257    #[expect(clippy::needless_pass_by_value)]
258    fn py_close<'py>(slf: PyRef<'_, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
259        let connection_mode = slf.connection_mode.clone();
260        let state_notify = slf.state_notify.clone();
261        let mode_str = ConnectionMode::from_atomic(&connection_mode).to_string();
262        log::debug!("Close from mode {mode_str}");
263
264        pyo3_async_runtimes::tokio::future_into_py(py, async move {
265            match ConnectionMode::from_atomic(&connection_mode) {
266                ConnectionMode::Closed => {
267                    log::debug!("Socket already closed");
268                }
269                ConnectionMode::Disconnect => {
270                    log::debug!("Socket already disconnecting");
271                }
272                _ => {
273                    // Preserve a CLOSED terminal state reached concurrently
274                    ConnectionMode::request_disconnect(&connection_mode);
275                    state_notify.notify_one();
276
277                    let timeout = tokio::time::timeout(Duration::from_secs(5), async {
278                        while !ConnectionMode::from_atomic(&connection_mode).is_closed() {
279                            tokio::time::sleep(Duration::from_millis(10)).await;
280                        }
281                    })
282                    .await;
283
284                    if timeout.is_err() {
285                        log::warn!("Timeout waiting for socket to close, forcing closed state");
286                        connection_mode.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
287                    }
288                }
289            }
290
291            Ok(())
292        })
293    }
294
295    /// Send bytes data to the connection.
296    ///
297    /// # Errors
298    ///
299    /// - Throws an Exception if it is not able to send data.
300    #[pyo3(name = "send")]
301    #[expect(clippy::needless_pass_by_value)]
302    fn py_send<'py>(
303        slf: PyRef<'_, Self>,
304        data: Vec<u8>,
305        py: Python<'py>,
306    ) -> PyResult<Bound<'py, PyAny>> {
307        log::trace!("Sending {}", String::from_utf8_lossy(&data));
308
309        let connection_mode = slf.connection_mode.clone();
310        let state_notify = slf.state_notify.clone();
311        let writer_tx = slf.writer_tx.clone();
312
313        pyo3_async_runtimes::tokio::future_into_py(py, async move {
314            match ConnectionMode::from_atomic(&connection_mode) {
315                ConnectionMode::Disconnect | ConnectionMode::Closed => {
316                    let msg = format!(
317                        "Cannot send data ({}): socket closed",
318                        String::from_utf8_lossy(&data)
319                    );
320
321                    let io_err = std::io::Error::new(std::io::ErrorKind::NotConnected, msg);
322                    return Err(to_pyruntime_err(io_err));
323                }
324                mode if !mode.is_active() => {
325                    let timeout = Duration::from_secs(2);
326                    let fallback_interval = Duration::from_millis(100);
327
328                    log::debug!("Waiting for client to become ACTIVE before sending (2s)...");
329
330                    match tokio::time::timeout(timeout, async {
331                        loop {
332                            let notified = state_notify.notified();
333
334                            let mode = ConnectionMode::from_atomic(&connection_mode);
335                            if mode.is_active() {
336                                return Ok(());
337                            }
338
339                            if matches!(mode, ConnectionMode::Disconnect | ConnectionMode::Closed) {
340                                return Err("Client disconnected waiting to send");
341                            }
342
343                            tokio::select! {
344                                biased;
345                                () = notified => {}
346                                () = tokio::time::sleep(fallback_interval) => {}
347                            }
348                        }
349                    })
350                    .await
351                    {
352                        Ok(Ok(())) => log::debug!("Client now active"),
353                        Ok(Err(e)) => {
354                            let err_msg = format!(
355                                "Failed sending data ({}): {e}",
356                                String::from_utf8_lossy(&data)
357                            );
358
359                            let io_err =
360                                std::io::Error::new(std::io::ErrorKind::NotConnected, err_msg);
361                            return Err(to_pyruntime_err(io_err));
362                        }
363                        Err(_) => {
364                            let err_msg = format!(
365                                "Failed sending data ({}): timeout waiting to become ACTIVE",
366                                String::from_utf8_lossy(&data)
367                            );
368
369                            let io_err = std::io::Error::new(std::io::ErrorKind::TimedOut, err_msg);
370                            return Err(to_pyruntime_err(io_err));
371                        }
372                    }
373                }
374                _ => {}
375            }
376
377            let msg = WriterCommand::Send(data.into());
378            writer_tx.send(msg).map_err(to_pyruntime_err)
379        })
380    }
381}
nautilus_network/python/socket.rs

nautilus_network/python/
socket.rs
