nautilus_interactive_brokers/common/
connection.rs1use 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#[derive(Debug, Clone)]
36pub struct ConnectionManager {
37 host: String,
39 port: u16,
41 client_id: i32,
43 is_connected: Arc<AtomicBool>,
45 attempt_count: Arc<AtomicU32>,
47 max_attempts: u32,
49 retry_indefinitely: bool,
51 current_backoff: Arc<parking_lot::Mutex<Duration>>,
53 last_disconnection: Arc<parking_lot::Mutex<Option<tokio::time::Instant>>>,
55}
56
57impl ConnectionManager {
58 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 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 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 pub fn is_connected(&self) -> bool {
151 self.is_connected.load(Ordering::Relaxed)
152 }
153
154 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 pub fn attempt_count(&self) -> u32 {
162 self.attempt_count.load(Ordering::Relaxed)
163 }
164
165 pub fn current_backoff(&self) -> Duration {
167 *self.current_backoff.lock()
168 }
169
170 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#[derive(Clone)]
182pub struct ConnectionWatchdog {
183 manager: Arc<ConnectionManager>,
185 check_interval: Duration,
187 client: Arc<parking_lot::Mutex<Option<Arc<Client>>>>,
189 reconnect_callback: Arc<
191 dyn Fn() -> tokio::task::JoinHandle<anyhow::Result<Arc<Client>>> + Send + Sync + 'static,
192 >,
193 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 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 pub fn set_client(&self, client: Arc<Client>) {
232 *self.client.lock() = Some(client);
233 }
234
235 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 let handle = reconnect_callback();
258
259 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 pub fn stop(&self) {
282 self.is_running.store(false, Ordering::Relaxed);
283 }
284}