Skip to main content

nautilus_interactive_brokers/common/
connection.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
16//! Connection management utilities for Interactive Brokers adapter.
17
18use std::{
19    fmt::Debug,
20    sync::{
21        Arc,
22        atomic::{AtomicBool, AtomicU32, Ordering},
23    },
24    time::Duration,
25};
26
27use anyhow::Context;
28use ibapi::client::Client;
29use nautilus_common::live::get_runtime;
30
31/// Connection manager for Interactive Brokers clients.
32///
33/// Handles automatic reconnection with exponential backoff, connection monitoring,
34/// and subscription resubscription on reconnect.
35#[derive(Debug, Clone)]
36pub struct ConnectionManager {
37    /// Host address for IB Gateway/TWS.
38    host: String,
39    /// Port for IB Gateway/TWS.
40    port: u16,
41    /// Client ID.
42    client_id: i32,
43    /// Current connection state.
44    is_connected: Arc<AtomicBool>,
45    /// Connection attempt counter.
46    attempt_count: Arc<AtomicU32>,
47    /// Maximum connection attempts (0 = infinite).
48    max_attempts: u32,
49    /// Whether to retry indefinitely.
50    retry_indefinitely: bool,
51    /// Current backoff duration.
52    current_backoff: Arc<parking_lot::Mutex<Duration>>,
53    /// Last disconnection time.
54    last_disconnection: Arc<parking_lot::Mutex<Option<tokio::time::Instant>>>,
55}
56
57impl ConnectionManager {
58    /// Create a new connection manager.
59    ///
60    /// # Arguments
61    ///
62    /// * `host` - Host address
63    /// * `port` - Port number
64    /// * `client_id` - Client ID
65    /// * `max_attempts` - Maximum connection attempts (0 = infinite)
66    pub fn new(host: String, port: u16, client_id: i32, max_attempts: u32) -> Self {
67        Self {
68            host,
69            port,
70            client_id,
71            is_connected: Arc::new(AtomicBool::new(false)),
72            attempt_count: Arc::new(AtomicU32::new(0)),
73            max_attempts,
74            retry_indefinitely: max_attempts == 0,
75            current_backoff: Arc::new(parking_lot::Mutex::new(Duration::from_secs(1))),
76            last_disconnection: Arc::new(parking_lot::Mutex::new(None)),
77        }
78    }
79
80    /// Connect to IB Gateway/TWS with automatic retry.
81    ///
82    /// # Returns
83    ///
84    /// Returns the connected client on success.
85    ///
86    /// # Errors
87    ///
88    /// Returns an error if connection fails after max attempts.
89    pub async fn connect_with_retry(&self) -> anyhow::Result<Arc<Client>> {
90        const MAX_BACKOFF: Duration = Duration::from_secs(60);
91        let mut attempt = 0;
92        let mut backoff = Duration::from_secs(1);
93
94        loop {
95            attempt += 1;
96            self.attempt_count.store(attempt, Ordering::Relaxed);
97
98            if !self.retry_indefinitely && attempt > self.max_attempts {
99                anyhow::bail!("Failed to connect after {} attempts", self.max_attempts);
100            }
101
102            tracing::debug!(
103                "Connection attempt {} to {}:{} (client_id: {})",
104                attempt,
105                self.host,
106                self.port,
107                self.client_id
108            );
109
110            let address = format!("{}:{}", self.host, self.port);
111            match Client::connect(&address, self.client_id).await {
112                Ok(client) => {
113                    tracing::info!(
114                        "Successfully connected to IB Gateway/TWS at {} (client_id: {})",
115                        address,
116                        self.client_id
117                    );
118
119                    self.is_connected.store(true, Ordering::Relaxed);
120                    self.attempt_count.store(0, Ordering::Relaxed);
121                    *self.current_backoff.lock() = Duration::from_secs(1);
122
123                    return Ok(Arc::new(client));
124                }
125                Err(e) => {
126                    tracing::warn!(
127                        "Connection attempt {} failed: {} (backoff: {:?})",
128                        attempt,
129                        e,
130                        backoff
131                    );
132
133                    if !self.retry_indefinitely && attempt >= self.max_attempts {
134                        return Err(e).context(format!(
135                            "Failed to connect after {} attempts",
136                            self.max_attempts
137                        ));
138                    }
139
140                    // Exponential backoff
141                    tokio::time::sleep(backoff).await;
142                    backoff = std::cmp::min(backoff * 2, MAX_BACKOFF);
143                    *self.current_backoff.lock() = backoff;
144                }
145            }
146        }
147    }
148
149    /// Check if currently connected.
150    pub fn is_connected(&self) -> bool {
151        self.is_connected.load(Ordering::Relaxed)
152    }
153
154    /// Mark connection as disconnected.
155    pub fn mark_disconnected(&self) {
156        self.is_connected.store(false, Ordering::Relaxed);
157        *self.last_disconnection.lock() = Some(tokio::time::Instant::now());
158    }
159
160    /// Get current attempt count.
161    pub fn attempt_count(&self) -> u32 {
162        self.attempt_count.load(Ordering::Relaxed)
163    }
164
165    /// Get current backoff duration.
166    pub fn current_backoff(&self) -> Duration {
167        *self.current_backoff.lock()
168    }
169
170    /// Get time since last disconnection.
171    pub fn time_since_disconnection(&self) -> Option<Duration> {
172        self.last_disconnection
173            .lock()
174            .map(|time: tokio::time::Instant| time.elapsed())
175    }
176}
177
178/// Connection watchdog for monitoring connection health.
179///
180/// Periodically checks connection status and triggers reconnection if needed.
181#[derive(Clone)]
182pub struct ConnectionWatchdog {
183    /// Connection manager.
184    manager: Arc<ConnectionManager>,
185    /// Check interval.
186    check_interval: Duration,
187    /// Client reference (for health checks).
188    client: Arc<parking_lot::Mutex<Option<Arc<Client>>>>,
189    /// Callback to call when reconnection is needed.
190    reconnect_callback: Arc<
191        dyn Fn() -> tokio::task::JoinHandle<anyhow::Result<Arc<Client>>> + Send + Sync + 'static,
192    >,
193    /// Whether the watchdog is running.
194    is_running: Arc<AtomicBool>,
195}
196
197impl Debug for ConnectionWatchdog {
198    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
199        f.debug_struct(stringify!(ConnectionWatchdog))
200            .field("check_interval", &self.check_interval)
201            .field("is_running", &self.is_running.load(Ordering::Relaxed))
202            .finish_non_exhaustive()
203    }
204}
205
206impl ConnectionWatchdog {
207    /// Create a new connection watchdog.
208    ///
209    /// # Arguments
210    ///
211    /// * `manager` - Connection manager
212    /// * `check_interval` - Interval between health checks
213    /// * `reconnect_callback` - Callback to trigger reconnection
214    pub fn new(
215        manager: Arc<ConnectionManager>,
216        check_interval: Duration,
217        reconnect_callback: Arc<
218            dyn Fn() -> tokio::task::JoinHandle<anyhow::Result<Arc<Client>>> + Send + Sync,
219        >,
220    ) -> Self {
221        Self {
222            manager,
223            check_interval,
224            client: Arc::new(parking_lot::Mutex::new(None)),
225            reconnect_callback,
226            is_running: Arc::new(AtomicBool::new(false)),
227        }
228    }
229
230    /// Set the client reference for health checks.
231    pub fn set_client(&self, client: Arc<Client>) {
232        *self.client.lock() = Some(client);
233    }
234
235    /// Start the watchdog.
236    pub fn start(&self) -> tokio::task::JoinHandle<()> {
237        let manager = Arc::clone(&self.manager);
238        let client = Arc::clone(&self.client);
239        let reconnect_callback = Arc::clone(&self.reconnect_callback);
240        let check_interval = self.check_interval;
241        let is_running = Arc::clone(&self.is_running);
242
243        is_running.store(true, Ordering::Relaxed);
244
245        get_runtime().spawn(async move {
246            tracing::debug!("Connection watchdog started");
247
248            while is_running.load(Ordering::Relaxed) {
249                tokio::time::sleep(check_interval).await;
250
251                if !manager.is_connected() {
252                    tracing::warn!(
253                        "Connection watchdog detected disconnection, triggering reconnection"
254                    );
255
256                    // Trigger reconnection
257                    let handle = reconnect_callback();
258
259                    // Wait for reconnection attempt
260                    match handle.await {
261                        Ok(Ok(new_client)) => {
262                            tracing::info!("Reconnection successful via watchdog");
263                            *client.lock() = Some(new_client);
264                            manager.is_connected.store(true, Ordering::Relaxed);
265                        }
266                        Ok(Err(e)) => {
267                            tracing::error!("Reconnection failed via watchdog: {}", e);
268                        }
269                        Err(e) => {
270                            tracing::error!("Reconnection task panicked: {}", e);
271                        }
272                    }
273                }
274            }
275
276            tracing::debug!("Connection watchdog stopped");
277        })
278    }
279
280    /// Stop the watchdog.
281    pub fn stop(&self) {
282        self.is_running.store(false, Ordering::Relaxed);
283    }
284}