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    max_reconnect_attempts: Option<u32>,
322) -> anyhow::Result<SharedClientHandle> {
323    let key = ConnectionKey(host.to_string(), port, client_id);
324    let registry = Arc::clone(&REGISTRY);
325
326    log::debug!(
327        "Acquiring shared IB client (host={host}, port={port}, client_id={client_id}, timeout_secs={connection_timeout_secs})"
328    );
329
330    let address = format!("{host}:{port}");
331    let connect_timeout = Duration::from_secs(connection_timeout_secs);
332
333    let client = registry
334        .acquire(key.clone(), Client::is_connected, || async move {
335            log::debug!(
336                "Establishing shared IB connection to {address} with timeout {connect_timeout:?}"
337            );
338            let mut builder = Client::builder().address(address).client_id(client_id);
339            if let Some(attempts) = max_reconnect_attempts {
340                builder = builder.max_reconnect_attempts(attempts);
341            }
342            tokio::time::timeout(connect_timeout, builder.connect())
343                .await
344                .map_err(|_| {
345                    anyhow::anyhow!(
346                        "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"
347                    )
348                })?
349                .map_err(connection_error)
350                .context("Failed to connect to IB Gateway/TWS")
351        })
352        .await?;
353
354    if registry.ref_count(&key).is_some_and(|count| count > 1) {
355        log::debug!("Reusing shared IB client (host={host}, port={port}, client_id={client_id})");
356    }
357
358    Ok(SharedClientHandle::new(client, registry, key))
359}
360
361fn connection_error(e: Error) -> anyhow::Error {
362    match e {
363        Error::ServerVersion(required, actual, _) => anyhow::anyhow!(
364            "TWS/IB Gateway server protocol {actual} is unsupported; server protocol {required} or newer is required. Upgrade TWS/IB Gateway"
365        ),
366        e => anyhow::Error::new(e),
367    }
368}
369
370#[cfg(test)]
371mod tests {
372    use std::{
373        sync::atomic::{AtomicUsize, Ordering},
374        task::{Context as TaskContext, Waker},
375    };
376
377    use rstest::rstest;
378    use tokio::io::{AsyncReadExt, AsyncWriteExt};
379
380    use super::*;
381
382    #[derive(Debug)]
383    struct TestClient;
384
385    #[rstest]
386    #[tokio::test]
387    async fn concurrent_acquire_uses_one_connection_and_reserves_both_references() {
388        let registry = Arc::new(SharedRegistry::default());
389        let key = ConnectionKey(String::from("127.0.0.1"), 7_497, 42);
390        let connect_count = Arc::new(AtomicUsize::new(0));
391        let connect_started = Arc::new(tokio::sync::Notify::new());
392        let release_connect = Arc::new(tokio::sync::Notify::new());
393
394        let first = {
395            let registry = Arc::clone(&registry);
396            let key = key.clone();
397            let connect_count = Arc::clone(&connect_count);
398            let connect_started = Arc::clone(&connect_started);
399            let release_connect = Arc::clone(&release_connect);
400
401            tokio::spawn(async move {
402                registry
403                    .acquire(
404                        key,
405                        |_| true,
406                        || async move {
407                            connect_count.fetch_add(1, Ordering::SeqCst);
408                            connect_started.notify_one();
409                            release_connect.notified().await;
410                            Ok::<_, anyhow::Error>(TestClient)
411                        },
412                    )
413                    .await
414            })
415        };
416
417        connect_started.notified().await;
418
419        let second = registry.acquire(
420            key.clone(),
421            |_| true,
422            || async {
423                connect_count.fetch_add(1, Ordering::SeqCst);
424                anyhow::bail!("second connector must not run")
425            },
426        );
427
428        tokio::pin!(second);
429        assert!(
430            second
431                .as_mut()
432                .poll(&mut TaskContext::from_waker(Waker::noop()))
433                .is_pending()
434        );
435        assert_eq!(registry.ref_count(&key), Some(1));
436
437        release_connect.notify_one();
438
439        let first_client = first.await.unwrap().unwrap();
440        let second_client = second.await.unwrap();
441
442        assert!(Arc::ptr_eq(&first_client, &second_client));
443        assert_eq!(connect_count.load(Ordering::SeqCst), 1);
444        assert_eq!(registry.ref_count(&key), Some(2));
445
446        registry.release(&key, &first_client);
447        assert_eq!(registry.ref_count(&key), Some(1));
448        registry.release(&key, &second_client);
449        assert_eq!(registry.ref_count(&key), None);
450    }
451
452    #[rstest]
453    #[tokio::test]
454    async fn cancelled_connect_releases_key_and_fails_waiters() {
455        let registry = Arc::new(SharedRegistry::default());
456        let key = ConnectionKey(String::from("127.0.0.1"), 7_497, 43);
457        let connect_started = Arc::new(tokio::sync::Notify::new());
458
459        let winner = {
460            let registry = Arc::clone(&registry);
461            let key = key.clone();
462            let connect_started = Arc::clone(&connect_started);
463
464            tokio::spawn(async move {
465                registry
466                    .acquire(
467                        key,
468                        |_| true,
469                        || async move {
470                            connect_started.notify_one();
471                            std::future::pending::<()>().await;
472                            Ok::<_, anyhow::Error>(TestClient)
473                        },
474                    )
475                    .await
476            })
477        };
478
479        connect_started.notified().await;
480
481        let waiter = registry.acquire(
482            key.clone(),
483            |_| true,
484            || async { anyhow::bail!("waiter must not connect while flight is active") },
485        );
486        tokio::pin!(waiter);
487        assert!(
488            waiter
489                .as_mut()
490                .poll(&mut TaskContext::from_waker(Waker::noop()))
491                .is_pending()
492        );
493
494        winner.abort();
495        assert!(winner.await.unwrap_err().is_cancelled());
496
497        let waiter_error = tokio::time::timeout(Duration::from_secs(1), waiter)
498            .await
499            .expect("waiter must fail fast after cancellation")
500            .unwrap_err();
501        assert_eq!(
502            waiter_error.to_string(),
503            "Shared IB client connect was cancelled"
504        );
505        assert_eq!(registry.ref_count(&key), None);
506
507        let client = registry
508            .acquire(
509                key.clone(),
510                |_| true,
511                || async move { Ok::<_, anyhow::Error>(TestClient) },
512            )
513            .await
514            .unwrap();
515        assert_eq!(registry.ref_count(&key), Some(1));
516        registry.release(&key, &client);
517        assert_eq!(registry.ref_count(&key), None);
518    }
519    #[rstest]
520    #[case(1, false)]
521    #[case(3, false)]
522    #[case(1, true)]
523    #[tokio::test]
524    async fn cancelled_waiters_do_not_retain_connection(
525        #[case] waiter_count: usize,
526        #[case] cancel_after_connect: bool,
527    ) {
528        let registry = SharedRegistry::default();
529        let key = ConnectionKey(String::from("127.0.0.1"), 7_497, 44);
530        let release_connect = tokio::sync::Notify::new();
531
532        let winner = registry.acquire(
533            key.clone(),
534            |_| true,
535            || async {
536                release_connect.notified().await;
537                Ok(TestClient)
538            },
539        );
540
541        tokio::pin!(winner);
542        assert!(
543            winner
544                .as_mut()
545                .poll(&mut TaskContext::from_waker(Waker::noop()))
546                .is_pending()
547        );
548
549        let mut waiters = (0..waiter_count)
550            .map(|_| {
551                Box::pin(registry.acquire(
552                    key.clone(),
553                    |_| true,
554                    || async { anyhow::bail!("waiter must not connect") },
555                ))
556            })
557            .collect::<Vec<_>>();
558
559        for waiter in &mut waiters {
560            assert!(
561                waiter
562                    .as_mut()
563                    .poll(&mut TaskContext::from_waker(Waker::noop()))
564                    .is_pending()
565            );
566        }
567
568        if !cancel_after_connect {
569            waiters.clear();
570        }
571
572        release_connect.notify_one();
573        let client = winner.await.unwrap();
574        waiters.clear();
575        assert_eq!(registry.ref_count(&key), Some(1));
576        registry.release(&key, &client);
577        assert_eq!(registry.ref_count(&key), None);
578    }
579
580    #[rstest]
581    #[tokio::test]
582    async fn cancelled_old_waiter_preserves_replacement_connection() {
583        let registry = SharedRegistry::default();
584        let key = ConnectionKey(String::from("127.0.0.1"), 7_497, 45);
585        let mut winner = Box::pin(registry.acquire(
586            key.clone(),
587            |_| true,
588            || async { std::future::pending::<anyhow::Result<TestClient>>().await },
589        ));
590        assert!(
591            winner
592                .as_mut()
593                .poll(&mut TaskContext::from_waker(Waker::noop()))
594                .is_pending()
595        );
596        let mut waiter = Box::pin(registry.acquire(
597            key.clone(),
598            |_| true,
599            || async { anyhow::bail!("waiter must not connect") },
600        ));
601        assert!(
602            waiter
603                .as_mut()
604                .poll(&mut TaskContext::from_waker(Waker::noop()))
605                .is_pending()
606        );
607        drop(winner);
608        let replacement = registry
609            .acquire(key.clone(), |_| true, || async { Ok(TestClient) })
610            .await
611            .unwrap();
612        drop(waiter);
613        assert_eq!(registry.ref_count(&key), Some(1));
614        registry.release(&key, &replacement);
615        assert_eq!(registry.ref_count(&key), None);
616    }
617    #[rstest]
618    #[case(187)]
619    #[case(212)]
620    #[tokio::test]
621    async fn unsupported_server_is_rejected_before_start_api(#[case] server_version: i32) {
622        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
623        let port = listener.local_addr().unwrap().port();
624
625        let server = tokio::spawn(async move {
626            for _ in 0..2 {
627                let (mut socket, _) = listener.accept().await.unwrap();
628                let mut magic = [0_u8; 4];
629                socket.read_exact(&mut magic).await.unwrap();
630                let size = socket.read_u32().await.unwrap();
631                let mut version = vec![0_u8; size as usize];
632                socket.read_exact(&mut version).await.unwrap();
633                let version = String::from_utf8(version).unwrap();
634                let response = format!("{server_version}\020260911 12:00:00 UTC\0");
635                socket.write_u32(response.len() as u32).await.unwrap();
636                socket.write_all(response.as_bytes()).await.unwrap();
637                let mut extra = [0_u8; 1];
638                let read = tokio::time::timeout(Duration::from_secs(2), socket.read(&mut extra))
639                    .await
640                    .unwrap()
641                    .unwrap();
642                assert_eq!(magic, *b"API\0");
643                assert_eq!(version.split("..").next(), Some("v213"));
644                assert_eq!(
645                    read, 0,
646                    "unsupported servers must receive no StartApi message"
647                );
648            }
649        });
650
651        let mut errors = Vec::new();
652        for _ in 0..2 {
653            errors.push(
654                get_or_connect("127.0.0.1", port, 917, 1, None)
655                    .await
656                    .unwrap_err()
657                    .to_string(),
658            );
659        }
660
661        server.await.unwrap();
662        let expected = format!(
663            "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"
664        );
665        assert_eq!(errors, vec![expected.clone(), expected]);
666        assert_eq!(
667            REGISTRY.ref_count(&ConnectionKey("127.0.0.1".to_string(), port, 917)),
668            None
669        );
670    }
671}
nautilus_interactive_brokers/common/shared_client.rs

nautilus_interactive_brokers/common/
shared_client.rs
