nautilus_polymarket/websocket/
pool.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//! Market-channel WebSocket connection pool for the Polymarket CLOB API.
17//!
18//! [`PolymarketWebSocketClient`] stays a single-channel, single-connection
19//! primitive. This pool owns a set of market-channel connections (shards) and
20//! spreads unique asset subscriptions across them so no single connection carries
21//! more than `ws_max_subscriptions` assets. See [`WS_DEFAULT_SUBSCRIPTIONS`] for
22//! why that bound exists.
23//!
24//! The pool grows lazily: it starts with one shard and opens another only when the
25//! current shards are full at subscribe time. A secondary shard closes once it owns
26//! no assets; the primary shard (which carries new-market discovery) always
27//! persists. Each shard replays only its own subscriptions on reconnect because
28//! that state lives inside its own [`PolymarketWebSocketClient`]. When custom
29//! features are enabled, every shard requests asset-scoped best-bid/ask events,
30//! while secondary shards discard global discovery and resolution events.
31
32use std::sync::{
33    Arc,
34    atomic::{AtomicBool, Ordering},
35};
36
37use ahash::AHashMap;
38use nautilus_live::{
39    SocketControlFactory,
40    task::{TaskJoinOutcome, TaskSlot, TaskSpawnError, finish_task},
41};
42use nautilus_network::websocket::{TransportBackend, proxy::ProxyUrl};
43use parking_lot::Mutex;
44use ustr::Ustr;
45
46use super::{
47    MARKET_STREAMS_ENDPOINT,
48    client::{PolymarketWebSocketClient, WsSubscriptionHandle},
49    messages::{MarketWsMessage, PolymarketWsMessage},
50};
51use crate::common::consts::WS_DEFAULT_SUBSCRIPTIONS;
52
53// Primary shard carries new-market discovery and never auto-closes.
54const PRIMARY_SHARD_ID: usize = 0;
55
56/// A pool of market-channel WebSocket connections that shards asset subscriptions.
57#[derive(Debug)]
58pub struct PolymarketMarketConnectionPool {
59    inner: Arc<PoolInner>,
60}
61
62/// Cloneable routing handle used from spawned subscription tasks.
63///
64/// Routes each asset to its owning shard and grows the pool on demand.
65#[derive(Clone, Debug)]
66pub struct PolymarketMarketPoolHandle {
67    inner: Arc<PoolInner>,
68}
69
70#[derive(Debug)]
71struct PoolInner {
72    base_url: Option<String>,
73    proxy_url: Option<ProxyUrl>,
74    transport_backend: TransportBackend,
75    subscribe_new_markets: bool,
76    max_subscriptions: usize,
77    // Serializes routing and shard growth; held across the async wire sends.
78    wire_mutex: tokio::sync::Mutex<()>,
79    // Never locked across an await, so routing futures stay `Send`.
80    state: Mutex<PoolState>,
81    out_tx: Mutex<Option<tokio::sync::mpsc::UnboundedSender<PolymarketWsMessage>>>,
82    out_rx: Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<PolymarketWsMessage>>>,
83    socket_factory: Mutex<Option<SocketControlFactory>>,
84    closed: AtomicBool,
85}
86
87#[derive(Debug)]
88struct PoolState {
89    shards: AHashMap<usize, ShardEntry>,
90    assignments: AHashMap<Ustr, usize>,
91    shutdown_errors: Vec<String>,
92}
93
94struct PoolDrain<'a> {
95    owner: &'a Mutex<PoolState>,
96    state: PoolState,
97}
98
99impl<'a> PoolDrain<'a> {
100    fn take(owner: &'a Mutex<PoolState>) -> Self {
101        let state = std::mem::replace(&mut *owner.lock(), PoolState::new());
102        Self { owner, state }
103    }
104}
105
106impl Drop for PoolDrain<'_> {
107    fn drop(&mut self) {
108        *self.owner.lock() = std::mem::replace(&mut self.state, PoolState::new());
109    }
110}
111
112impl PoolState {
113    fn new() -> Self {
114        Self {
115            shards: AHashMap::new(),
116            assignments: AHashMap::new(),
117            shutdown_errors: Vec::new(),
118        }
119    }
120}
121
122#[derive(Debug)]
123struct ShardEntry {
124    client: PolymarketWebSocketClient,
125    handle: WsSubscriptionHandle,
126    forwarder: TaskSlot<()>,
127    owned: usize,
128    closing: bool,
129}
130
131enum ReleaseOutcome {
132    NotOwned,
133    Unsubscribe(WsSubscriptionHandle),
134    CloseShard(usize, Box<ShardEntry>),
135}
136
137struct ShardClose<'a> {
138    owner: &'a Mutex<PoolState>,
139    id: usize,
140    shard: Option<Box<ShardEntry>>,
141}
142
143impl<'a> ShardClose<'a> {
144    fn new(owner: &'a Mutex<PoolState>, id: usize, mut shard: Box<ShardEntry>) -> Self {
145        shard.closing = true;
146        Self {
147            owner,
148            id,
149            shard: Some(shard),
150        }
151    }
152
153    fn shard_mut(&mut self) -> &mut ShardEntry {
154        self.shard.as_deref_mut().expect("closing shard present")
155    }
156
157    fn complete(mut self) {
158        self.shard.take();
159    }
160}
161
162impl Drop for ShardClose<'_> {
163    fn drop(&mut self) {
164        if let Some(shard) = self.shard.take() {
165            let replaced = self.owner.lock().shards.insert(self.id, *shard);
166            assert!(replaced.is_none(), "closing shard ID is already present");
167        }
168    }
169}
170
171#[allow(
172    clippy::missing_panics_doc,
173    reason = "internal mutex locks and shard-state invariants are not expected to panic"
174)]
175impl PolymarketMarketConnectionPool {
176    /// Creates a new market connection pool (unconnected).
177    ///
178    /// A `max_subscriptions` of `0` is invalid and clamps to
179    /// [`WS_DEFAULT_SUBSCRIPTIONS`] with a warning.
180    #[must_use]
181    pub fn new(
182        base_url: Option<String>,
183        subscribe_new_markets: bool,
184        transport_backend: TransportBackend,
185        max_subscriptions: usize,
186    ) -> Self {
187        Self::new_with_proxy(
188            base_url,
189            subscribe_new_markets,
190            transport_backend,
191            max_subscriptions,
192            None,
193        )
194    }
195
196    /// Creates a new market connection pool with an optional validated proxy URL.
197    #[must_use]
198    pub fn new_with_proxy(
199        base_url: Option<String>,
200        subscribe_new_markets: bool,
201        transport_backend: TransportBackend,
202        max_subscriptions: usize,
203        proxy_url: Option<ProxyUrl>,
204    ) -> Self {
205        Self {
206            inner: Arc::new(PoolInner::new_with_proxy(
207                base_url,
208                transport_backend,
209                subscribe_new_markets,
210                max_subscriptions,
211                proxy_url,
212            )),
213        }
214    }
215
216    /// Configures socket state reporting and reconnect control for every connection in the pool.
217    #[must_use]
218    pub(crate) fn with_socket_factory(self, factory: SocketControlFactory) -> Self {
219        *self.inner.socket_factory.lock() = Some(factory);
220        self
221    }
222
223    #[cfg(test)]
224    pub(crate) fn proxy_url(&self) -> Option<&ProxyUrl> {
225        self.inner.proxy_url.as_ref()
226    }
227
228    /// Returns a cloneable routing handle for use in spawned subscription tasks.
229    #[must_use]
230    pub fn handle(&self) -> PolymarketMarketPoolHandle {
231        PolymarketMarketPoolHandle {
232            inner: Arc::clone(&self.inner),
233        }
234    }
235
236    /// Opens the primary shard and prepares the merged message stream.
237    ///
238    /// # Errors
239    ///
240    /// Returns an error if the primary connection cannot be established.
241    pub async fn connect(&self) -> anyhow::Result<()> {
242        if self.inner.closed.load(Ordering::Acquire) {
243            self.disconnect().await?;
244        }
245
246        let _wire = self.inner.wire_mutex.lock().await;
247
248        if !self.inner.closed.load(Ordering::Acquire) && !self.inner.state.lock().shards.is_empty()
249        {
250            log::warn!("Polymarket market pool already connected");
251            return Ok(());
252        }
253
254        {
255            let _state = self.inner.state.lock();
256            self.inner.closed.store(false, Ordering::Release);
257        }
258
259        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel();
260        *self.inner.out_tx.lock() = Some(out_tx);
261        *self.inner.out_rx.lock() = Some(out_rx);
262
263        self.inner.connect_new_shard(true).await?;
264        Ok(())
265    }
266
267    /// Sends the new-market discovery subscribe on the primary shard.
268    ///
269    /// # Errors
270    ///
271    /// Returns an error if no primary shard is available.
272    pub async fn subscribe_new_markets_feed(&self) -> anyhow::Result<()> {
273        let _wire = self.inner.wire_mutex.lock().await;
274
275        let handle = {
276            let state = self.inner.state.lock();
277            if self.inner.closed.load(Ordering::Acquire) {
278                anyhow::bail!("Market connection pool is closed");
279            }
280            state
281                .shards
282                .get(&PRIMARY_SHARD_ID)
283                .map(|shard| shard.handle.clone())
284        };
285
286        match handle {
287            Some(handle) => handle.subscribe_market(vec![]).await,
288            None => anyhow::bail!("No primary market shard available for new-market discovery"),
289        }
290    }
291
292    /// Takes the merged message receiver, leaving `None` in its place.
293    #[must_use]
294    pub fn take_message_receiver(
295        &self,
296    ) -> Option<tokio::sync::mpsc::UnboundedReceiver<PolymarketWsMessage>> {
297        self.inner.out_rx.lock().take()
298    }
299
300    /// Disconnects every shard and clears routing state.
301    ///
302    /// # Errors
303    ///
304    /// Returns an error after attempting every shard when a task or connection does not stop.
305    pub async fn disconnect(&self) -> anyhow::Result<()> {
306        self.inner.begin_shutdown();
307        let _wire = self.inner.wire_mutex.lock().await;
308
309        let mut drain = PoolDrain::take(&self.inner.state);
310        let shard_ids = drain.state.shards.keys().copied().collect::<Vec<_>>();
311        for shard_id in shard_ids {
312            let shard = drain
313                .state
314                .shards
315                .get_mut(&shard_id)
316                .expect("market shard ID collected from pool state");
317            let mut shard_failed = false;
318
319            shard.forwarder.abort();
320            if let Some(outcome) = finish_task(
321                &mut shard.forwarder,
322                std::time::Duration::ZERO,
323                std::time::Duration::from_secs(2),
324            )
325            .await
326            {
327                match outcome {
328                    TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
329                    TaskJoinOutcome::Failed(error) => {
330                        shard_failed = true;
331                        drain
332                            .state
333                            .shutdown_errors
334                            .push(format!("market shard {shard_id} forwarder failed: {error}"));
335                    }
336                    TaskJoinOutcome::Incomplete => {
337                        shard_failed = true;
338                        drain.state.shutdown_errors.push(format!(
339                            "market shard {shard_id} forwarder did not stop after abort"
340                        ));
341                    }
342                }
343            }
344
345            if let Err(e) = shard.client.disconnect().await {
346                shard_failed = true;
347                drain
348                    .state
349                    .shutdown_errors
350                    .push(format!("market shard {shard_id} disconnect failed: {e}"));
351            }
352
353            if !shard_failed {
354                drain.state.shards.remove(&shard_id);
355                drain
356                    .state
357                    .assignments
358                    .retain(|_, assigned_id| *assigned_id != shard_id);
359            }
360        }
361
362        if !drain.state.shutdown_errors.is_empty() {
363            let errors = std::mem::take(&mut drain.state.shutdown_errors);
364            anyhow::bail!(
365                "Polymarket market pool shutdown failed: {}",
366                errors.join("; ")
367            );
368        }
369
370        *self.inner.out_tx.lock() = None;
371        *self.inner.out_rx.lock() = None;
372        Ok(())
373    }
374
375    pub(crate) fn begin_shutdown(&self) {
376        self.inner.begin_shutdown();
377    }
378
379    /// Clears retained reconnect-replay state on any remaining shards.
380    pub(crate) fn clear_reconnect_state(&self) {
381        let state = self.inner.state.lock();
382        for shard in state.shards.values() {
383            shard.client.clear_reconnect_state();
384        }
385    }
386
387    /// Returns the number of open shard connections.
388    #[must_use]
389    pub fn connection_count(&self) -> usize {
390        self.inner.state.lock().shards.len()
391    }
392
393    /// Returns the number of unique assets assigned across all shards.
394    #[must_use]
395    pub fn subscription_count(&self) -> usize {
396        self.inner.state.lock().assignments.len()
397    }
398}
399
400impl Drop for PoolInner {
401    fn drop(&mut self) {
402        self.closed.store(true, Ordering::Release);
403
404        for shard in self.state.get_mut().shards.values_mut() {
405            shard.forwarder.abort();
406            shard.client.abort();
407        }
408    }
409}
410
411#[allow(
412    clippy::missing_panics_doc,
413    reason = "internal mutex locks and shard-state invariants are not expected to panic"
414)]
415impl PolymarketMarketPoolHandle {
416    pub(crate) fn begin_shutdown(&self) {
417        self.inner.begin_shutdown();
418    }
419
420    /// Subscribes to market data for the given asset IDs, sharding across connections.
421    ///
422    /// # Errors
423    ///
424    /// Returns an error if a shard cannot be opened or a subscribe send fails.
425    pub async fn subscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
426        let _wire = self.inner.wire_mutex.lock().await;
427        self.inner.ensure_open()?;
428        for asset_id in asset_ids {
429            self.inner.subscribe_one(asset_id).await?;
430        }
431        Ok(())
432    }
433
434    /// Removes asset IDs from their owning shards, closing emptied secondary shards.
435    ///
436    /// # Errors
437    ///
438    /// Returns an error if an unsubscribe send fails.
439    pub async fn unsubscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
440        let _wire = self.inner.wire_mutex.lock().await;
441        self.inner.ensure_open()?;
442        for asset_id in asset_ids {
443            self.inner.unsubscribe_one(asset_id).await?;
444        }
445        Ok(())
446    }
447}
448
449impl PoolInner {
450    fn begin_shutdown(&self) {
451        let state = self.state.lock();
452        self.closed.store(true, Ordering::Release);
453
454        for shard in state.shards.values() {
455            shard.client.begin_shutdown();
456        }
457    }
458
459    fn ensure_open(&self) -> anyhow::Result<()> {
460        let _state = self.state.lock();
461
462        if self.closed.load(Ordering::Acquire) {
463            anyhow::bail!("Market connection pool is closed");
464        }
465        Ok(())
466    }
467}
468
469impl PoolInner {
470    #[cfg(test)]
471    fn new(
472        base_url: Option<String>,
473        transport_backend: TransportBackend,
474        subscribe_new_markets: bool,
475        max_subscriptions: usize,
476    ) -> Self {
477        Self::new_with_proxy(
478            base_url,
479            transport_backend,
480            subscribe_new_markets,
481            max_subscriptions,
482            None,
483        )
484    }
485
486    fn new_with_proxy(
487        base_url: Option<String>,
488        transport_backend: TransportBackend,
489        subscribe_new_markets: bool,
490        max_subscriptions: usize,
491        proxy_url: Option<ProxyUrl>,
492    ) -> Self {
493        let max_subscriptions = if max_subscriptions == 0 {
494            log::warn!(
495                "PolymarketDataClientConfig.ws_max_subscriptions=0 is invalid, using {WS_DEFAULT_SUBSCRIPTIONS}"
496            );
497            WS_DEFAULT_SUBSCRIPTIONS
498        } else {
499            max_subscriptions
500        };
501
502        Self {
503            base_url,
504            proxy_url,
505            transport_backend,
506            subscribe_new_markets,
507            max_subscriptions,
508            wire_mutex: tokio::sync::Mutex::new(()),
509            state: Mutex::new(PoolState::new()),
510            out_tx: Mutex::new(None),
511            out_rx: Mutex::new(None),
512            socket_factory: Mutex::new(None),
513            closed: AtomicBool::new(false),
514        }
515    }
516
517    // Callers hold `wire_mutex`.
518    async fn subscribe_one(&self, asset_id: String) -> anyhow::Result<()> {
519        let token = Ustr::from(asset_id.as_str());
520
521        let Some(handle) = self.assign(token).await? else {
522            return Ok(());
523        };
524
525        if let Err(e) = self.ensure_open() {
526            if let ReleaseOutcome::CloseShard(id, shard) = self.release(token)
527                && let Err(close_error) = self.close_shard(id, shard).await
528            {
529                anyhow::bail!("{e}; subscription rollback failed: {close_error}");
530            }
531            return Err(e);
532        }
533
534        if let Err(e) = handle.subscribe_market(vec![asset_id]).await {
535            // Roll back so a failed send leaves no stale assignment or empty shard.
536            if let ReleaseOutcome::CloseShard(id, shard) = self.release(token)
537                && let Err(close_error) = self.close_shard(id, shard).await
538            {
539                anyhow::bail!("{e}; subscription rollback failed: {close_error}");
540            }
541            return Err(e);
542        }
543        Ok(())
544    }
545
546    // Callers hold `wire_mutex`.
547    async fn unsubscribe_one(&self, asset_id: String) -> anyhow::Result<()> {
548        self.ensure_open()?;
549        let token = Ustr::from(asset_id.as_str());
550
551        match self.release(token) {
552            ReleaseOutcome::NotOwned => Ok(()),
553            ReleaseOutcome::Unsubscribe(handle) => handle.unsubscribe_market(vec![asset_id]).await,
554            ReleaseOutcome::CloseShard(id, shard) => {
555                // Disconnect drops the shard's subscriptions; no unsubscribe send needed.
556                self.close_shard(id, shard).await
557            }
558        }
559    }
560
561    // Returns `None` when the token is already owned by a shard.
562    async fn assign(&self, token: Ustr) -> anyhow::Result<Option<WsSubscriptionHandle>> {
563        {
564            let mut state = self.state.lock();
565
566            if self.closed.load(Ordering::Acquire) {
567                anyhow::bail!("Market connection pool is closed");
568            }
569
570            if state.assignments.contains_key(&token) {
571                return Ok(None);
572            }
573
574            if let Some(id) = smallest_shard_with_capacity(&state, self.max_subscriptions) {
575                let handle = {
576                    let shard = state.shards.get_mut(&id).expect("shard present");
577                    shard.owned += 1;
578                    shard.handle.clone()
579                };
580                state.assignments.insert(token, id);
581                return Ok(Some(handle));
582            }
583        }
584
585        let id = self.connect_new_shard(false).await?;
586
587        let rejected_shard = {
588            let mut state = self.state.lock();
589
590            if self.closed.load(Ordering::Acquire) {
591                Some(
592                    state
593                        .shards
594                        .remove(&id)
595                        .expect("new shard retained for shutdown"),
596                )
597            } else {
598                let handle = {
599                    let shard = state.shards.get_mut(&id).expect("new shard present");
600                    shard.owned += 1;
601                    shard.handle.clone()
602                };
603                state.assignments.insert(token, id);
604                return Ok(Some(handle));
605            }
606        };
607
608        if let Some(shard) = rejected_shard {
609            if let Err(e) = self.close_shard(id, Box::new(shard)).await {
610                anyhow::bail!("Market connection pool is closed; shard rollback failed: {e}");
611            }
612            anyhow::bail!("Market connection pool is closed");
613        }
614        unreachable!("open pool returned from assignment")
615    }
616
617    fn release(&self, token: Ustr) -> ReleaseOutcome {
618        let mut state = self.state.lock();
619
620        let Some(id) = state.assignments.remove(&token) else {
621            return ReleaseOutcome::NotOwned;
622        };
623
624        let owned = {
625            let Some(shard) = state.shards.get_mut(&id) else {
626                return ReleaseOutcome::NotOwned;
627            };
628            shard.owned = shard.owned.saturating_sub(1);
629            shard.owned
630        };
631
632        if id != PRIMARY_SHARD_ID && owned == 0 {
633            let shard = state.shards.remove(&id).expect("shard present");
634            ReleaseOutcome::CloseShard(id, Box::new(shard))
635        } else {
636            let handle = state.shards.get(&id).expect("shard present").handle.clone();
637            ReleaseOutcome::Unsubscribe(handle)
638        }
639    }
640
641    async fn connect_new_shard(&self, is_primary: bool) -> anyhow::Result<usize> {
642        if self.closed.load(Ordering::Acquire) {
643            anyhow::bail!("Market connection pool is closed");
644        }
645
646        let id = if is_primary {
647            PRIMARY_SHARD_ID
648        } else {
649            let state = self.state.lock();
650            available_shard_id(&state)
651        };
652
653        let mut client = self.market_client(self.subscribe_new_markets, id);
654        client.connect().await?;
655
656        let handle = client.clone_subscription_handle();
657        let rx = client
658            .take_message_receiver()
659            .ok_or_else(|| anyhow::anyhow!("Market shard receiver unavailable after connect"))?;
660        let forwarder = match self.spawn_forwarder(rx, is_primary) {
661            Ok(forwarder) => forwarder,
662            Err((e, forwarder)) => {
663                let shard = Box::new(ShardEntry {
664                    client,
665                    handle,
666                    forwarder,
667                    owned: 0,
668                    closing: true,
669                });
670
671                if let Err(close_error) = self.close_shard(id, shard).await {
672                    anyhow::bail!(
673                        "Failed to start market shard forwarder: {e}; startup rollback failed: \
674                         {close_error}"
675                    );
676                }
677                anyhow::bail!("Failed to start market shard forwarder: {e}");
678            }
679        };
680
681        let shard = ShardEntry {
682            client,
683            handle,
684            forwarder,
685            owned: 0,
686            closing: false,
687        };
688        let rejected_shard = {
689            let mut state = self.state.lock();
690
691            if self.closed.load(Ordering::Acquire) {
692                Some(shard)
693            } else {
694                state.shards.insert(id, shard);
695                None
696            }
697        };
698
699        if let Some(shard) = rejected_shard {
700            if let Err(e) = self.close_shard(id, Box::new(shard)).await {
701                anyhow::bail!("Market connection pool is closed; shard rollback failed: {e}");
702            }
703            anyhow::bail!("Market connection pool is closed");
704        }
705
706        log::debug!("Opened Polymarket market shard {id}");
707        Ok(id)
708    }
709
710    fn market_client(
711        &self,
712        subscribe_new_markets: bool,
713        shard_id: usize,
714    ) -> PolymarketWebSocketClient {
715        let client = PolymarketWebSocketClient::new_market_with_proxy(
716            self.base_url.clone(),
717            subscribe_new_markets,
718            self.transport_backend,
719            self.proxy_url.clone(),
720        );
721        let factory = self.socket_factory.lock().clone();
722
723        if let Some(factory) = factory {
724            let endpoint = if shard_id == PRIMARY_SHARD_ID {
725                MARKET_STREAMS_ENDPOINT.to_string()
726            } else {
727                format!("{MARKET_STREAMS_ENDPOINT}-{shard_id}")
728            };
729            client.with_socket_control(factory.control(endpoint))
730        } else {
731            client
732        }
733    }
734
735    fn spawn_forwarder(
736        &self,
737        mut rx: tokio::sync::mpsc::UnboundedReceiver<PolymarketWsMessage>,
738        is_primary: bool,
739    ) -> Result<TaskSlot<()>, (TaskSpawnError, TaskSlot<()>)> {
740        let out_tx = self.out_tx.lock().clone();
741
742        let mut forwarder = TaskSlot::new();
743        if let Err(e) = forwarder.spawn(async move {
744            let Some(out_tx) = out_tx else {
745                return;
746            };
747
748            while let Some(msg) = rx.recv().await {
749                if !should_forward_from_shard(&msg, is_primary) {
750                    continue;
751                }
752
753                if out_tx.send(msg).is_err() {
754                    break;
755                }
756            }
757        }) {
758            return Err((e, forwarder));
759        }
760        Ok(forwarder)
761    }
762
763    async fn close_shard(&self, id: usize, shard: Box<ShardEntry>) -> anyhow::Result<()> {
764        let mut close = ShardClose::new(&self.state, id, shard);
765        close.shard_mut().forwarder.abort();
766        let forwarder_stopped = match finish_task(
767            &mut close.shard_mut().forwarder,
768            std::time::Duration::ZERO,
769            std::time::Duration::from_secs(2),
770        )
771        .await
772        {
773            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => true,
774            Some(TaskJoinOutcome::Failed(error)) => {
775                let error = format!("market shard {id} forwarder failed: {error}");
776                self.state.lock().shutdown_errors.push(error);
777                true
778            }
779            Some(TaskJoinOutcome::Incomplete) => {
780                let error = format!("market shard {id} forwarder did not stop after abort");
781                self.state.lock().shutdown_errors.push(error);
782                false
783            }
784        };
785
786        if let Err(e) = close.shard_mut().client.disconnect().await {
787            let error = format!("market shard {id} disconnect failed: {e}");
788            self.state.lock().shutdown_errors.push(error);
789        }
790
791        if forwarder_stopped && !close.shard_mut().client.has_task() {
792            close.complete();
793        }
794
795        let errors = std::mem::take(&mut self.state.lock().shutdown_errors);
796
797        if errors.is_empty() {
798            Ok(())
799        } else {
800            anyhow::bail!(errors.join("; "))
801        }
802    }
803
804    #[cfg(test)]
805    fn subscription_count_for_test(&self) -> usize {
806        self.state.lock().assignments.len()
807    }
808}
809
810fn should_forward_from_shard(message: &PolymarketWsMessage, is_primary: bool) -> bool {
811    is_primary
812        || !matches!(
813            message,
814            PolymarketWsMessage::Market(
815                MarketWsMessage::NewMarket(_) | MarketWsMessage::MarketResolved(_)
816            )
817        )
818}
819
820fn smallest_shard_with_capacity(state: &PoolState, max_subscriptions: usize) -> Option<usize> {
821    state
822        .shards
823        .iter()
824        .filter(|(_, shard)| !shard.closing && shard.owned < max_subscriptions)
825        .map(|(id, _)| *id)
826        .min()
827}
828
829fn available_shard_id(state: &PoolState) -> usize {
830    let mut id = PRIMARY_SHARD_ID + 1;
831    while state.shards.contains_key(&id) {
832        id = id.checked_add(1).expect("market shard ID space exhausted");
833    }
834    id
835}
836
837#[cfg(test)]
838impl PolymarketMarketPoolHandle {
839    /// In-memory single-shard handle backed by `sender`, `assigned` tokens
840    /// pre-owned. Never connected, so growth is never triggered.
841    pub(crate) fn test_single_shard(
842        sender: tokio::sync::mpsc::UnboundedSender<super::handler::HandlerCommand>,
843        assigned: &[&str],
844    ) -> Self {
845        let inner = PoolInner::new(
846            None,
847            TransportBackend::default(),
848            false,
849            WS_DEFAULT_SUBSCRIPTIONS,
850        );
851        {
852            let mut state = inner.state.lock();
853            state.shards.insert(
854                PRIMARY_SHARD_ID,
855                ShardEntry {
856                    client: PolymarketWebSocketClient::new_market(
857                        None,
858                        false,
859                        TransportBackend::default(),
860                    ),
861                    handle: WsSubscriptionHandle::from_sender(sender),
862                    forwarder: TaskSlot::new(),
863                    owned: assigned.len(),
864                    closing: false,
865                },
866            );
867
868            for token in assigned {
869                state
870                    .assignments
871                    .insert(Ustr::from(token), PRIMARY_SHARD_ID);
872            }
873        }
874        Self {
875            inner: Arc::new(inner),
876        }
877    }
878}
879
880#[cfg(test)]
881mod tests {
882    use std::{net::SocketAddr, sync::Arc as StdArc, time::Duration};
883
884    use PolymarketMarketPoolHandle as Handle;
885    use axum::{
886        Router,
887        extract::ws::{WebSocket, WebSocketUpgrade},
888        response::Response,
889        routing::get,
890    };
891    use nautilus_common::{
892        live::runner::replace_system_event_sender,
893        messages::{SystemEvent, system::SocketState},
894    };
895    use nautilus_live::{SocketReconnectRegistry, SocketReconnectRequestOutcome};
896    use nautilus_model::identifiers::ClientId;
897    use parking_lot::{Condvar, Mutex as TestMutex};
898    use rstest::rstest;
899
900    use super::*;
901    use crate::websocket::handler::HandlerCommand;
902
903    struct BlockingDrop(StdArc<(TestMutex<(bool, bool)>, Condvar)>);
904
905    impl Drop for BlockingDrop {
906        fn drop(&mut self) {
907            let (state, wake) = &*self.0;
908            let mut state = state.lock();
909            state.0 = true;
910            wake.notify_all();
911            while !state.1 {
912                wake.wait(&mut state);
913            }
914        }
915    }
916
917    async fn handle_socket_upgrade(ws: WebSocketUpgrade) -> Response {
918        ws.on_upgrade(handle_socket)
919    }
920
921    async fn handle_socket(mut socket: WebSocket) {
922        while socket.recv().await.is_some() {}
923    }
924
925    async fn start_socket_server() -> SocketAddr {
926        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
927            .await
928            .expect("bind test websocket server");
929        let addr = listener.local_addr().expect("test websocket address");
930        let router = Router::new().route("/ws/market", get(handle_socket_upgrade));
931
932        tokio::spawn(async move {
933            axum::serve(listener, router)
934                .await
935                .expect("test websocket server failed");
936        });
937
938        addr
939    }
940
941    // Bare state with unconnected shards for pure capacity-accounting tests.
942    fn state_with_shards(owned: &[usize]) -> PoolState {
943        let mut state = PoolState::new();
944
945        for (id, owned) in owned.iter().enumerate() {
946            let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
947            state.shards.insert(
948                id,
949                ShardEntry {
950                    client: PolymarketWebSocketClient::new_market(
951                        None,
952                        false,
953                        TransportBackend::default(),
954                    ),
955                    handle: WsSubscriptionHandle::from_sender(tx),
956                    forwarder: TaskSlot::new(),
957                    owned: *owned,
958                    closing: false,
959                },
960            );
961        }
962        state
963    }
964
965    fn market_message(filename: &str) -> PolymarketWsMessage {
966        let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
967            .join("test_data")
968            .join(filename);
969        let json = std::fs::read_to_string(path).unwrap();
970        PolymarketWsMessage::Market(serde_json::from_str(&json).unwrap())
971    }
972
973    #[rstest]
974    #[case::primary_new_market("ws_market_new_market_msg.json", true, true)]
975    #[case::secondary_new_market("ws_market_new_market_msg.json", false, false)]
976    #[case::primary_resolution("ws_market_resolved_msg.json", true, true)]
977    #[case::secondary_resolution("ws_market_resolved_msg.json", false, false)]
978    #[case::secondary_best_bid_ask("ws_market_best_bid_ask_msg.json", false, true)]
979    fn shard_forwarding_keeps_global_events_on_primary(
980        #[case] filename: &str,
981        #[case] is_primary: bool,
982        #[case] expected: bool,
983    ) {
984        let message = market_message(filename);
985        assert_eq!(should_forward_from_shard(&message, is_primary), expected);
986    }
987
988    #[rstest]
989    #[tokio::test]
990    async fn secondary_forwarder_drops_global_events_and_keeps_best_bid_ask() {
991        let inner = PoolInner::new(None, TransportBackend::default(), true, 1);
992        let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
993        *inner.out_tx.lock() = Some(out_tx);
994        let (shard_tx, shard_rx) = tokio::sync::mpsc::unbounded_channel();
995        let mut forwarder = inner
996            .spawn_forwarder(shard_rx, false)
997            .expect("spawn forwarder");
998
999        shard_tx
1000            .send(market_message("ws_market_new_market_msg.json"))
1001            .unwrap();
1002        shard_tx
1003            .send(market_message("ws_market_resolved_msg.json"))
1004            .unwrap();
1005        shard_tx
1006            .send(market_message("ws_market_best_bid_ask_msg.json"))
1007            .unwrap();
1008        drop(shard_tx);
1009        let outcome = finish_task(
1010            &mut forwarder,
1011            Duration::from_secs(1),
1012            Duration::from_secs(1),
1013        )
1014        .await
1015        .expect("forwarder task");
1016        assert!(matches!(outcome, TaskJoinOutcome::Completed(())));
1017
1018        let forwarded = out_rx.try_recv().unwrap();
1019        let PolymarketWsMessage::Market(MarketWsMessage::BestBidAsk(message)) = forwarded else {
1020            panic!("unexpected forwarded message: {forwarded:?}");
1021        };
1022        assert_eq!(
1023            message.asset_id,
1024            Ustr::from(
1025                "85354956062430465315924116860125388538595433819574542752031640332592237464430"
1026            ),
1027        );
1028        assert!(out_rx.try_recv().is_err());
1029    }
1030
1031    #[rstest]
1032    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1033    async fn canceled_shard_close_restores_unfinished_ownership() {
1034        let inner = Arc::new(PoolInner::new(None, TransportBackend::default(), false, 1));
1035        let blocking = StdArc::new((TestMutex::new((false, false)), Condvar::new()));
1036        let blocking_task = StdArc::clone(&blocking);
1037        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1038        let mut forwarder = TaskSlot::new();
1039        forwarder
1040            .spawn(async move {
1041                let _blocking = BlockingDrop(blocking_task);
1042                let _ = started_tx.send(());
1043                std::future::pending::<()>().await;
1044            })
1045            .expect("spawn forwarder");
1046        started_rx.await.expect("forwarder should start");
1047
1048        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1049        let shard = Box::new(ShardEntry {
1050            client: PolymarketWebSocketClient::new_market(None, false, TransportBackend::default()),
1051            handle: WsSubscriptionHandle::from_sender(cmd_tx),
1052            forwarder,
1053            owned: 0,
1054            closing: false,
1055        });
1056        let close_inner = Arc::clone(&inner);
1057        let close = tokio::spawn(async move {
1058            let _result = close_inner.close_shard(1, shard).await;
1059        });
1060
1061        loop {
1062            if blocking.0.lock().0 {
1063                break;
1064            }
1065            tokio::task::yield_now().await;
1066        }
1067        close.abort();
1068        let _ = close.await;
1069
1070        {
1071            let restored = inner.state.lock();
1072            let shard = restored.shards.get(&1).expect("closing shard restored");
1073            assert!(shard.closing);
1074            assert!(shard.forwarder.is_some());
1075            assert_eq!(smallest_shard_with_capacity(&restored, 1), None);
1076        }
1077
1078        {
1079            let (state, wake) = &*blocking;
1080            state.lock().1 = true;
1081            wake.notify_all();
1082        }
1083        let shard = inner
1084            .state
1085            .lock()
1086            .shards
1087            .remove(&1)
1088            .expect("restored shard");
1089        inner
1090            .close_shard(1, Box::new(shard))
1091            .await
1092            .expect("restored shard should close");
1093
1094        assert!(inner.state.lock().shards.is_empty());
1095    }
1096
1097    #[rstest]
1098    fn zero_max_subscriptions_clamps_to_default() {
1099        let inner = PoolInner::new(None, TransportBackend::default(), false, 0);
1100        assert_eq!(inner.max_subscriptions, WS_DEFAULT_SUBSCRIPTIONS);
1101    }
1102
1103    #[rstest]
1104    fn pool_retains_proxy_for_lazily_created_shards() {
1105        const PROXY_URL: &str = "http://pool-user:pool-proxy-secret@127.0.0.1:18088";
1106        let pool = PolymarketMarketConnectionPool::new_with_proxy(
1107            Some("ws://market.example/ws".to_string()),
1108            true,
1109            TransportBackend::Tungstenite,
1110            17,
1111            Some(ProxyUrl::parse(PROXY_URL).unwrap()),
1112        );
1113        let primary = pool.inner.market_client(true, PRIMARY_SHARD_ID);
1114        let secondary = pool.inner.market_client(false, PRIMARY_SHARD_ID + 1);
1115        let debug = format!("{pool:?}");
1116
1117        assert_eq!(pool.inner.proxy_url.as_ref().unwrap().expose(), PROXY_URL);
1118        assert_eq!(primary.proxy_url().unwrap().expose(), PROXY_URL);
1119        assert_eq!(secondary.proxy_url().unwrap().expose(), PROXY_URL);
1120        assert_eq!(pool.inner.max_subscriptions, 17);
1121        assert!(!debug.contains("pool-proxy-secret"));
1122    }
1123
1124    #[rstest]
1125    #[tokio::test]
1126    async fn pool_assigns_distinct_endpoint_sinks_and_handles_before_connect() {
1127        let addr = start_socket_server().await;
1128        let (system_tx, mut system_rx) = tokio::sync::mpsc::unbounded_channel();
1129        replace_system_event_sender(system_tx);
1130        let registry = SocketReconnectRegistry::default();
1131        let factory = SocketControlFactory::with_registry(
1132            ClientId::from("POLYMARKET"),
1133            Some(*crate::common::consts::POLYMARKET_VENUE),
1134            &registry,
1135        );
1136        let pool = PolymarketMarketConnectionPool::new(
1137            Some(format!("ws://{addr}/ws/market")),
1138            false,
1139            TransportBackend::Tungstenite,
1140            1,
1141        )
1142        .with_socket_factory(factory);
1143
1144        pool.connect().await.expect("connect primary shard");
1145        pool.handle()
1146            .subscribe_market(vec!["asset-0".to_string(), "asset-1".to_string()])
1147            .await
1148            .expect("open secondary shard");
1149
1150        let mut connected = Vec::new();
1151        while connected.len() < 2 {
1152            let event = tokio::time::timeout(Duration::from_secs(2), system_rx.recv())
1153                .await
1154                .expect("wait for socket state event")
1155                .expect("system event channel closed");
1156            let SystemEvent::SocketState(change) = event;
1157            if change.state == SocketState::Connected {
1158                connected.push(change.endpoint);
1159            }
1160        }
1161        connected.sort_unstable();
1162
1163        assert_eq!(
1164            connected,
1165            vec![
1166                Ustr::from(MARKET_STREAMS_ENDPOINT),
1167                Ustr::from("polymarket-market-streams-1"),
1168            ],
1169        );
1170        let client_id = ClientId::from("POLYMARKET");
1171        let primary = registry
1172            .handle(client_id, Ustr::from(MARKET_STREAMS_ENDPOINT))
1173            .expect("primary reconnect handle should be registered");
1174        let secondary = registry
1175            .handle(client_id, Ustr::from("polymarket-market-streams-1"))
1176            .expect("secondary reconnect handle should be registered");
1177        assert_eq!(
1178            primary.request_reconnect(),
1179            SocketReconnectRequestOutcome::Accepted,
1180        );
1181        let event = system_rx
1182            .try_recv()
1183            .expect("selected shard should report reconnect state");
1184        let SystemEvent::SocketState(change) = event;
1185        assert_eq!(change.client_id, client_id);
1186        assert_eq!(change.endpoint, Ustr::from(MARKET_STREAMS_ENDPOINT));
1187        assert_eq!(change.state, SocketState::Disconnected);
1188        assert_eq!(
1189            secondary.request_reconnect(),
1190            SocketReconnectRequestOutcome::Accepted,
1191        );
1192
1193        pool.disconnect().await.expect("disconnect pool");
1194        assert!(
1195            registry
1196                .handle(client_id, Ustr::from(MARKET_STREAMS_ENDPOINT))
1197                .is_none()
1198        );
1199        assert!(
1200            registry
1201                .handle(client_id, Ustr::from("polymarket-market-streams-1"))
1202                .is_none()
1203        );
1204    }
1205
1206    #[rstest]
1207    #[case::first_has_room(&[0, 200], 200, Some(0))]
1208    #[case::prefers_lowest_id(&[200, 5, 5], 200, Some(1))]
1209    #[case::all_full(&[200, 200], 200, None)]
1210    #[case::exact_boundary_is_full(&[1], 1, None)]
1211    fn smallest_shard_with_capacity_picks_lowest_open_id(
1212        #[case] owned: &[usize],
1213        #[case] max: usize,
1214        #[case] expected: Option<usize>,
1215    ) {
1216        let state = state_with_shards(owned);
1217        assert_eq!(smallest_shard_with_capacity(&state, max), expected);
1218    }
1219
1220    #[rstest]
1221    fn available_shard_id_reuses_lowest_closed_shard() {
1222        let mut state = state_with_shards(&[1, 1, 1]);
1223        state.shards.remove(&(PRIMARY_SHARD_ID + 1));
1224
1225        assert_eq!(available_shard_id(&state), PRIMARY_SHARD_ID + 1);
1226    }
1227
1228    #[rstest]
1229    #[tokio::test]
1230    async fn subscribe_routes_command_and_tracks_assignment() {
1231        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1232        let handle = Handle::test_single_shard(tx, &[]);
1233
1234        handle
1235            .subscribe_market(vec!["token-a".to_string()])
1236            .await
1237            .expect("subscribe");
1238
1239        match rx.try_recv().expect("expected SubscribeMarket") {
1240            HandlerCommand::SubscribeMarket(ids) => assert_eq!(ids, vec!["token-a".to_string()]),
1241            other => panic!("unexpected command: {other:?}"),
1242        }
1243        assert_eq!(handle.inner.subscription_count_for_test(), 1);
1244    }
1245
1246    #[rstest]
1247    #[tokio::test]
1248    async fn duplicate_subscribe_does_not_consume_capacity_or_resend() {
1249        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1250        let handle = Handle::test_single_shard(tx, &[]);
1251
1252        handle
1253            .subscribe_market(vec!["token-a".to_string()])
1254            .await
1255            .expect("first subscribe");
1256        handle
1257            .subscribe_market(vec!["token-a".to_string()])
1258            .await
1259            .expect("duplicate subscribe");
1260
1261        assert!(matches!(
1262            rx.try_recv(),
1263            Ok(HandlerCommand::SubscribeMarket(_))
1264        ));
1265        assert!(rx.try_recv().is_err(), "duplicate must not resend");
1266        assert_eq!(handle.inner.subscription_count_for_test(), 1);
1267    }
1268
1269    #[rstest]
1270    #[tokio::test]
1271    async fn unsubscribe_routes_command_and_releases_assignment() {
1272        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1273        let handle = Handle::test_single_shard(tx, &["token-a"]);
1274
1275        handle
1276            .unsubscribe_market(vec!["token-a".to_string()])
1277            .await
1278            .expect("unsubscribe");
1279
1280        match rx.try_recv().expect("expected UnsubscribeMarket") {
1281            HandlerCommand::UnsubscribeMarket(ids) => assert_eq!(ids, vec!["token-a".to_string()]),
1282            other => panic!("unexpected command: {other:?}"),
1283        }
1284        assert_eq!(handle.inner.subscription_count_for_test(), 0);
1285    }
1286
1287    #[rstest]
1288    #[tokio::test]
1289    async fn unsubscribe_unknown_token_is_noop() {
1290        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1291        let handle = Handle::test_single_shard(tx, &[]);
1292
1293        handle
1294            .unsubscribe_market(vec!["token-a".to_string()])
1295            .await
1296            .expect("unsubscribe");
1297
1298        assert!(rx.try_recv().is_err());
1299    }
1300
1301    #[rstest]
1302    #[tokio::test]
1303    async fn subscribe_send_failure_rolls_back_assignment() {
1304        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1305        drop(rx);
1306        let handle = Handle::test_single_shard(tx, &[]);
1307
1308        let result = handle.subscribe_market(vec!["token-a".to_string()]).await;
1309
1310        assert!(result.is_err());
1311        assert_eq!(handle.inner.subscription_count_for_test(), 0);
1312    }
1313}
nautilus_polymarket/websocket/pool.rs

nautilus_polymarket/websocket/
pool.rs
