nautilus_interactive_brokers/common/
shared_client.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//! Shared IB API client connection per (host, port, client_id).
17//!
18//! Data, execution, and historical clients use a single TCP connection per logical
19//! connection to avoid client ID conflicts and redundant connections (parity with
20//! Python's get_cached_ib_client).
21
22use std::{
23    collections::HashMap,
24    fmt::Debug,
25    future::Future,
26    ops::Deref,
27    sync::{Arc, LazyLock},
28    time::Duration,
29};
30
31use anyhow::Context;
32use ibapi::{Error, client::Client};
33use parking_lot::Mutex;
34
35/// Key for the connection registry: (host, port, client_id).
36#[derive(Clone, Debug, Eq, Hash, PartialEq)]
37struct ConnectionKey(String, u16, i32);
38
39enum RegistryEntry<T> {
40    Connected {
41        client: Arc<T>,
42        ref_count: u32,
43    },
44    Connecting {
45        flight: Arc<InFlight<T>>,
46        ref_count: u32,
47    },
48}
49
50struct InFlight<T> {
51    result: Mutex<Option<Result<Arc<T>, String>>>,
52    notify: tokio::sync::Notify,
53}
54
55impl<T> Default for InFlight<T> {
56    fn default() -> Self {
57        Self {
58            result: Mutex::new(None),
59            notify: tokio::sync::Notify::new(),
60        }
61    }
62}
63
64impl<T> InFlight<T> {
65    fn complete(&self, result: Result<Arc<T>, String>) {
66        *self.result.lock() = Some(result);
67        self.notify.notify_waiters();
68    }
69
70    async fn wait(&self) -> anyhow::Result<Arc<T>> {
71        loop {
72            let notified = self.notify.notified();
73
74            if let Some(result) = self.result.lock().clone() {
75                return result.map_err(anyhow::Error::msg);
76            }
77
78            notified.await;
79        }
80    }
81}
82
83struct SharedRegistry<T> {
84    entries: Mutex<HashMap<ConnectionKey, RegistryEntry<T>>>,
85}
86
87impl<T> Default for SharedRegistry<T> {
88    fn default() -> Self {
89        Self {
90            entries: Mutex::new(HashMap::new()),
91        }
92    }
93}
94
95impl<T> SharedRegistry<T> {
96    async fn acquire<Connect, ConnectFuture, IsConnected>(
97        &self,
98        key: ConnectionKey,
99        is_connected: IsConnected,
100        connect: Connect,
101    ) -> anyhow::Result<Arc<T>>
102    where
103        Connect: FnOnce() -> ConnectFuture,
104        ConnectFuture: Future<Output = anyhow::Result<T>>,
105        IsConnected: Fn(&T) -> bool,
106    {
107        let (flight, should_connect) = {
108            let mut entries = self.entries.lock();
109
110            match entries.get_mut(&key) {
111                Some(RegistryEntry::Connected { client, ref_count }) if is_connected(client) => {
112                    *ref_count = ref_count
113                        .checked_add(1)
114                        .context("Shared IB client reference count overflow")?;
115                    return Ok(Arc::clone(client));
116                }
117                Some(RegistryEntry::Connecting { flight, .. }) => (Arc::clone(flight), false),
118                Some(RegistryEntry::Connected { .. }) | None => {
119                    let flight = Arc::new(InFlight::default());
120                    entries.insert(
121                        key.clone(),
122                        RegistryEntry::Connecting {
123                            flight: Arc::clone(&flight),
124                            ref_count: 1,
125                        },
126                    );
127
128                    (flight, true)
129                }
130            }
131        };
132
133        if !should_connect {
134            let client = flight.wait().await?;
135            let mut entries = self.entries.lock();
136            match entries.get_mut(&key) {
137                Some(RegistryEntry::Connected {
138                    client: current,
139                    ref_count,
140                }) if Arc::ptr_eq(current, &client) => {
141                    *ref_count = ref_count
142                        .checked_add(1)
143                        .context("Shared IB client reference count overflow")?;
144                }
145                None => {
146                    entries.insert(
147                        key,
148                        RegistryEntry::Connected {
149                            client: Arc::clone(&client),
150                            ref_count: 1,
151                        },
152                    );
153                }
154                _ => anyhow::bail!("Shared IB client was replaced while acquisition was waiting"),
155            }
156
157            return Ok(client);
158        }
159
160        let mut guard = ConnectGuard {
161            registry: self,
162            key: key.clone(),
163            flight: Arc::clone(&flight),
164            armed: true,
165        };
166
167        let result = connect().await.map(Arc::new).map_err(|e| format!("{e:#}"));
168        {
169            let mut entries = self.entries.lock();
170
171            let ref_count = match entries.get(&key) {
172                Some(RegistryEntry::Connecting {
173                    flight: current,
174                    ref_count,
175                }) if Arc::ptr_eq(current, &flight) => *ref_count,
176                _ => anyhow::bail!("Shared IB client in-flight registry entry was replaced"),
177            };
178
179            match &result {
180                Ok(client) => {
181                    entries.insert(
182                        key.clone(),
183                        RegistryEntry::Connected {
184                            client: Arc::clone(client),
185                            ref_count,
186                        },
187                    );
188                }
189                Err(_) => {
190                    entries.remove(&key);
191                }
192            }
193        }
194
195        flight.complete(result.clone());
196        guard.armed = false;
197        result.map_err(anyhow::Error::msg)
198    }
199
200    fn release(&self, key: &ConnectionKey, client: &Arc<T>) {
201        let mut entries = self.entries.lock();
202
203        let Some(RegistryEntry::Connected {
204            client: registered,
205            ref_count,
206        }) = entries.get_mut(key)
207        else {
208            return;
209        };
210
211        if !Arc::ptr_eq(registered, client) {
212            return;
213        }
214
215        *ref_count = ref_count.saturating_sub(1);
216        if *ref_count == 0 {
217            entries.remove(key);
218        }
219    }
220
221    fn ref_count(&self, key: &ConnectionKey) -> Option<u32> {
222        match self.entries.lock().get(key) {
223            Some(
224                RegistryEntry::Connected { ref_count, .. }
225                | RegistryEntry::Connecting { ref_count, .. },
226            ) => Some(*ref_count),
227            None => None,
228        }
229    }
230}
231
232/// Fails the in-flight connect when the winning `acquire` future is dropped mid-connect,
233/// so waiters error out and the key can be reacquired instead of hanging forever.
234struct ConnectGuard<'a, T> {
235    registry: &'a SharedRegistry<T>,
236    key: ConnectionKey,
237    flight: Arc<InFlight<T>>,
238    armed: bool,
239}
240
241impl<T> Drop for ConnectGuard<'_, T> {
242    fn drop(&mut self) {
243        if !self.armed {
244            return;
245        }
246
247        let mut entries = self.registry.entries.lock();
248        if let Some(RegistryEntry::Connecting { flight, .. }) = entries.get(&self.key)
249            && Arc::ptr_eq(flight, &self.flight)
250        {
251            entries.remove(&self.key);
252        }
253
254        drop(entries);
255        self.flight
256            .complete(Err(String::from("Shared IB client connect was cancelled")));
257    }
258}
259
260/// Global registry: one shared client per (host, port, client_id) with ref count.
261static REGISTRY: LazyLock<Arc<SharedRegistry<Client>>> =
262    LazyLock::new(|| Arc::new(SharedRegistry::default()));
263
264/// Handle to a shared IB client; when dropped, ref count is decremented and the
265/// connection is removed from the registry when the count reaches zero.
266pub struct SharedClientHandle {
267    client: Arc<Client>,
268    registry: Arc<SharedRegistry<Client>>,
269    key: ConnectionKey,
270}
271
272impl Debug for SharedClientHandle {
273    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
274        f.debug_struct(stringify!(SharedClientHandle))
275            .field("key", &self.key)
276            .finish_non_exhaustive()
277    }
278}
279
280impl SharedClientHandle {
281    fn new(client: Arc<Client>, registry: Arc<SharedRegistry<Client>>, key: ConnectionKey) -> Self {
282        Self {
283            client,
284            registry,
285            key,
286        }
287    }
288
289    /// Returns a reference to the underlying `Arc<Client>` for call sites that need it.
290    pub fn as_arc(&self) -> &Arc<Client> {
291        &self.client
292    }
293}
294
295impl Deref for SharedClientHandle {
296    type Target = Client;
297
298    fn deref(&self) -> &Self::Target {
299        self.client.as_ref()
300    }
301}
302
303impl Drop for SharedClientHandle {
304    fn drop(&mut self) {
305        self.registry.release(&self.key, &self.client);
306    }
307}
308
309/// Returns a handle to the shared IB client for the given (host, port, client_id).
310/// If a connection already exists, its ref count is incremented and the same client
311/// is returned. Otherwise a new connection is established and registered.
312///
313/// # Errors
314///
315/// Returns an error if connecting to IB Gateway/TWS fails.
316pub async fn get_or_connect(
317    host: &str,
318    port: u16,
319    client_id: i32,
320    connection_timeout_secs: u64,
321) -> anyhow::Result<SharedClientHandle> {
322    let key = ConnectionKey(host.to_string(), port, client_id);
323    let registry = Arc::clone(&REGISTRY);
324
325    log::debug!(
326        "Acquiring shared IB client (host={host}, port={port}, client_id={client_id}, timeout_secs={connection_timeout_secs})"
327    );
328
329    let address = format!("{host}:{port}");
330    let connect_timeout = Duration::from_secs(connection_timeout_secs);
331
332    let client = registry
333        .acquire(key.clone(), Client::is_connected, || async move {
334            log::debug!(
335                "Establishing shared IB connection to {address} with timeout {connect_timeout:?}"
336            );
337            tokio::time::timeout(connect_timeout, Client::connect(&address, client_id))
338                .await
339                .map_err(|_| {
340                    anyhow::anyhow!(
341                        "Timed out connecting to IB Gateway/TWS after {connection_timeout_secs}s; verify socket API settings and use a TWS/IB Gateway version supported by the IB integration guide"
342                    )
343                })?
344                .map_err(connection_error)
345                .context("Failed to connect to IB Gateway/TWS")
346        })
347        .await?;
348
349    if registry.ref_count(&key).is_some_and(|count| count > 1) {
350        log::debug!("Reusing shared IB client (host={host}, port={port}, client_id={client_id})");
351    }
352
353    Ok(SharedClientHandle::new(client, registry, key))
354}
355
356fn connection_error(e: Error) -> anyhow::Error {
357    match e {
358        Error::ServerVersion(required, actual, _) => anyhow::anyhow!(
359            "TWS/IB Gateway server protocol {actual} is unsupported; server protocol {required} or newer is required. Upgrade TWS/IB Gateway"
360        ),
361        e => anyhow::Error::new(e),
362    }
363}
364
365#[cfg(test)]
366mod tests {
367    use std::{
368        sync::atomic::{AtomicUsize, Ordering},
369        task::{Context as TaskContext, Waker},
370    };
371
372    use rstest::rstest;
373    use tokio::io::{AsyncReadExt, AsyncWriteExt};
374
375    use super::*;
376
377    #[derive(Debug)]
378    struct TestClient;
379
380    #[rstest]
381    #[tokio::test]
382    async fn concurrent_acquire_uses_one_connection_and_reserves_both_references() {
383        let registry = Arc::new(SharedRegistry::default());
384        let key = ConnectionKey(String::from("127.0.0.1"), 7_497, 42);
385        let connect_count = Arc::new(AtomicUsize::new(0));
386        let connect_started = Arc::new(tokio::sync::Notify::new());
387        let release_connect = Arc::new(tokio::sync::Notify::new());
388
389        let first = {
390            let registry = Arc::clone(&registry);
391            let key = key.clone();
392            let connect_count = Arc::clone(&connect_count);
393            let connect_started = Arc::clone(&connect_started);
394            let release_connect = Arc::clone(&release_connect);
395
396            tokio::spawn(async move {
397                registry
398                    .acquire(
399                        key,
400                        |_| true,
401                        || async move {
402                            connect_count.fetch_add(1, Ordering::SeqCst);
403                            connect_started.notify_one();
404                            release_connect.notified().await;
405                            Ok::<_, anyhow::Error>(TestClient)
406                        },
407                    )
408                    .await
409            })
410        };
411
412        connect_started.notified().await;
413
414        let second = registry.acquire(
415            key.clone(),
416            |_| true,
417            || async {
418                connect_count.fetch_add(1, Ordering::SeqCst);
419                anyhow::bail!("second connector must not run")
420            },
421        );
422
423        tokio::pin!(second);
424        assert!(
425            second
426                .as_mut()
427                .poll(&mut TaskContext::from_waker(Waker::noop()))
428                .is_pending()
429        );
430        assert_eq!(registry.ref_count(&key), Some(1));
431
432        release_connect.notify_one();
433
434        let first_client = first.await.unwrap().unwrap();
435        let second_client = second.await.unwrap();
436
437        assert!(Arc::ptr_eq(&first_client, &second_client));
438        assert_eq!(connect_count.load(Ordering::SeqCst), 1);
439        assert_eq!(registry.ref_count(&key), Some(2));
440
441        registry.release(&key, &first_client);
442        assert_eq!(registry.ref_count(&key), Some(1));
443        registry.release(&key, &second_client);
444        assert_eq!(registry.ref_count(&key), None);
445    }
446
447    #[rstest]
448    #[tokio::test]
449    async fn cancelled_connect_releases_key_and_fails_waiters() {
450        let registry = Arc::new(SharedRegistry::default());
451        let key = ConnectionKey(String::from("127.0.0.1"), 7_497, 43);
452        let connect_started = Arc::new(tokio::sync::Notify::new());
453
454        let winner = {
455            let registry = Arc::clone(&registry);
456            let key = key.clone();
457            let connect_started = Arc::clone(&connect_started);
458
459            tokio::spawn(async move {
460                registry
461                    .acquire(
462                        key,
463                        |_| true,
464                        || async move {
465                            connect_started.notify_one();
466                            std::future::pending::<()>().await;
467                            Ok::<_, anyhow::Error>(TestClient)
468                        },
469                    )
470                    .await
471            })
472        };
473
474        connect_started.notified().await;
475
476        let waiter = registry.acquire(
477            key.clone(),
478            |_| true,
479            || async { anyhow::bail!("waiter must not connect while flight is active") },
480        );
481        tokio::pin!(waiter);
482        assert!(
483            waiter
484                .as_mut()
485                .poll(&mut TaskContext::from_waker(Waker::noop()))
486                .is_pending()
487        );
488
489        winner.abort();
490        assert!(winner.await.unwrap_err().is_cancelled());
491
492        let waiter_error = tokio::time::timeout(Duration::from_secs(1), waiter)
493            .await
494            .expect("waiter must fail fast after cancellation")
495            .unwrap_err();
496        assert_eq!(
497            waiter_error.to_string(),
498            "Shared IB client connect was cancelled"
499        );
500        assert_eq!(registry.ref_count(&key), None);
501
502        let client = registry
503            .acquire(
504                key.clone(),
505                |_| true,
506                || async move { Ok::<_, anyhow::Error>(TestClient) },
507            )
508            .await
509            .unwrap();
510        assert_eq!(registry.ref_count(&key), Some(1));
511        registry.release(&key, &client);
512        assert_eq!(registry.ref_count(&key), None);
513    }
514    #[rstest]
515    #[case(1, false)]
516    #[case(3, false)]
517    #[case(1, true)]
518    #[tokio::test]
519    async fn cancelled_waiters_do_not_retain_connection(
520        #[case] waiter_count: usize,
521        #[case] cancel_after_connect: bool,
522    ) {
523        let registry = SharedRegistry::default();
524        let key = ConnectionKey(String::from("127.0.0.1"), 7_497, 44);
525        let release_connect = tokio::sync::Notify::new();
526
527        let winner = registry.acquire(
528            key.clone(),
529            |_| true,
530            || async {
531                release_connect.notified().await;
532                Ok(TestClient)
533            },
534        );
535
536        tokio::pin!(winner);
537        assert!(
538            winner
539                .as_mut()
540                .poll(&mut TaskContext::from_waker(Waker::noop()))
541                .is_pending()
542        );
543
544        let mut waiters = (0..waiter_count)
545            .map(|_| {
546                Box::pin(registry.acquire(
547                    key.clone(),
548                    |_| true,
549                    || async { anyhow::bail!("waiter must not connect") },
550                ))
551            })
552            .collect::<Vec<_>>();
553
554        for waiter in &mut waiters {
555            assert!(
556                waiter
557                    .as_mut()
558                    .poll(&mut TaskContext::from_waker(Waker::noop()))
559                    .is_pending()
560            );
561        }
562
563        if !cancel_after_connect {
564            waiters.clear();
565        }
566
567        release_connect.notify_one();
568        let client = winner.await.unwrap();
569        waiters.clear();
570        assert_eq!(registry.ref_count(&key), Some(1));
571        registry.release(&key, &client);
572        assert_eq!(registry.ref_count(&key), None);
573    }
574
575    #[rstest]
576    #[tokio::test]
577    async fn cancelled_old_waiter_preserves_replacement_connection() {
578        let registry = SharedRegistry::default();
579        let key = ConnectionKey(String::from("127.0.0.1"), 7_497, 45);
580        let mut winner = Box::pin(registry.acquire(
581            key.clone(),
582            |_| true,
583            || async { std::future::pending::<anyhow::Result<TestClient>>().await },
584        ));
585        assert!(
586            winner
587                .as_mut()
588                .poll(&mut TaskContext::from_waker(Waker::noop()))
589                .is_pending()
590        );
591        let mut waiter = Box::pin(registry.acquire(
592            key.clone(),
593            |_| true,
594            || async { anyhow::bail!("waiter must not connect") },
595        ));
596        assert!(
597            waiter
598                .as_mut()
599                .poll(&mut TaskContext::from_waker(Waker::noop()))
600                .is_pending()
601        );
602        drop(winner);
603        let replacement = registry
604            .acquire(key.clone(), |_| true, || async { Ok(TestClient) })
605            .await
606            .unwrap();
607        drop(waiter);
608        assert_eq!(registry.ref_count(&key), Some(1));
609        registry.release(&key, &replacement);
610        assert_eq!(registry.ref_count(&key), None);
611    }
612    #[rstest]
613    #[case(187)]
614    #[case(212)]
615    #[tokio::test]
616    async fn unsupported_server_is_rejected_before_start_api(#[case] server_version: i32) {
617        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
618        let port = listener.local_addr().unwrap().port();
619
620        let server = tokio::spawn(async move {
621            for _ in 0..2 {
622                let (mut socket, _) = listener.accept().await.unwrap();
623                let mut magic = [0_u8; 4];
624                socket.read_exact(&mut magic).await.unwrap();
625                let size = socket.read_u32().await.unwrap();
626                let mut version = vec![0_u8; size as usize];
627                socket.read_exact(&mut version).await.unwrap();
628                let version = String::from_utf8(version).unwrap();
629                let response = format!("{server_version}\020260911 12:00:00 UTC\0");
630                socket.write_u32(response.len() as u32).await.unwrap();
631                socket.write_all(response.as_bytes()).await.unwrap();
632                let mut extra = [0_u8; 1];
633                let read = tokio::time::timeout(Duration::from_secs(2), socket.read(&mut extra))
634                    .await
635                    .unwrap()
636                    .unwrap();
637                assert_eq!(magic, *b"API\0");
638                assert_eq!(version.split("..").next(), Some("v213"));
639                assert_eq!(
640                    read, 0,
641                    "unsupported servers must receive no StartApi message"
642                );
643            }
644        });
645
646        let mut errors = Vec::new();
647        for _ in 0..2 {
648            errors.push(
649                get_or_connect("127.0.0.1", port, 917, 1)
650                    .await
651                    .unwrap_err()
652                    .to_string(),
653            );
654        }
655
656        server.await.unwrap();
657        let expected = format!(
658            "Failed to connect to IB Gateway/TWS: TWS/IB Gateway server protocol {server_version} is unsupported; server protocol 213 or newer is required. Upgrade TWS/IB Gateway"
659        );
660        assert_eq!(errors, vec![expected.clone(), expected]);
661        assert_eq!(
662            REGISTRY.ref_count(&ConnectionKey("127.0.0.1".to_string(), port, 917)),
663            None
664        );
665    }
666}
nautilus_interactive_brokers/common/shared_client.rs

nautilus_interactive_brokers/common/
shared_client.rs
