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 futures_util::{StreamExt, stream::FuturesUnordered};
39use nautilus_live::{
40    SocketControlFactory,
41    book::snapshot::SnapshotGate,
42    task::{TaskJoinOutcome, TaskSlot, TaskSpawnError, finish_task},
43};
44use nautilus_network::websocket::{TransportBackend, proxy::ProxyUrl};
45use parking_lot::Mutex;
46use ustr::Ustr;
47
48use super::{
49    MARKET_STREAMS_ENDPOINT,
50    client::{PolymarketWebSocketClient, WsSubscriptionHandle},
51    handler::CycleMarketOutcome,
52    messages::{MarketWsMessage, PolymarketWsMessage},
53};
54use crate::common::consts::WS_DEFAULT_SUBSCRIPTIONS;
55
56// Primary shard carries new-market discovery and never auto-closes.
57const PRIMARY_SHARD_ID: usize = 0;
58
59/// A pool of market-channel WebSocket connections that shards asset subscriptions.
60#[derive(Debug)]
61pub struct PolymarketMarketConnectionPool {
62    inner: Arc<PoolInner>,
63}
64
65/// Cloneable routing handle used from spawned subscription tasks.
66///
67/// Routes each asset to its owning shard and grows the pool on demand.
68#[derive(Clone, Debug)]
69pub struct PolymarketMarketPoolHandle {
70    inner: Arc<PoolInner>,
71}
72
73#[derive(Debug)]
74struct PoolInner {
75    base_url: Option<String>,
76    proxy_url: Option<ProxyUrl>,
77    transport_backend: TransportBackend,
78    subscribe_new_markets: bool,
79    max_subscriptions: usize,
80    connect_timeout_secs: Mutex<Option<u64>>,
81    // Serializes routing and shard growth; held across the async wire sends.
82    wire_mutex: tokio::sync::Mutex<()>,
83    // Never locked across an await, so routing futures stay `Send`.
84    state: Mutex<PoolState>,
85    out_tx: Mutex<Option<tokio::sync::mpsc::UnboundedSender<PolymarketWsMessage>>>,
86    out_rx: Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<PolymarketWsMessage>>>,
87    socket_factory: Mutex<Option<SocketControlFactory>>,
88    closed: AtomicBool,
89}
90
91#[derive(Debug, Default)]
92struct PoolState {
93    shards: AHashMap<usize, ShardEntry>,
94    assignments: AHashMap<Ustr, usize>,
95    shutdown_errors: Vec<String>,
96}
97
98struct PoolDrain<'a> {
99    owner: &'a Mutex<PoolState>,
100    state: PoolState,
101    completed: Vec<usize>,
102}
103
104impl<'a> PoolDrain<'a> {
105    fn take(owner: &'a Mutex<PoolState>) -> Self {
106        let state = std::mem::take(&mut *owner.lock());
107        Self {
108            owner,
109            state,
110            completed: Vec::new(),
111        }
112    }
113}
114
115impl Drop for PoolDrain<'_> {
116    fn drop(&mut self) {
117        for id in &self.completed {
118            self.state.shards.remove(id);
119        }
120
121        // Prune once, including when cancellation interrupts the drain
122        self.state
123            .assignments
124            .retain(|_, id| self.state.shards.contains_key(id));
125        *self.owner.lock() = std::mem::take(&mut self.state);
126    }
127}
128
129#[derive(Debug)]
130struct ShardEntry {
131    client: PolymarketWebSocketClient,
132    handle: WsSubscriptionHandle,
133    forwarder: TaskSlot<()>,
134    owned: usize,
135    closing: bool,
136}
137
138enum ReleaseOutcome {
139    NotOwned,
140    Unsubscribe(WsSubscriptionHandle),
141    CloseShard(usize, Box<ShardEntry>),
142}
143
144struct ShardClose<'a> {
145    owner: &'a Mutex<PoolState>,
146    id: usize,
147    shard: Option<Box<ShardEntry>>,
148}
149
150impl<'a> ShardClose<'a> {
151    fn new(owner: &'a Mutex<PoolState>, id: usize, mut shard: Box<ShardEntry>) -> Self {
152        shard.closing = true;
153        Self {
154            owner,
155            id,
156            shard: Some(shard),
157        }
158    }
159
160    fn shard_mut(&mut self) -> &mut ShardEntry {
161        self.shard.as_deref_mut().expect("closing shard present")
162    }
163
164    fn complete(mut self) {
165        self.shard.take();
166    }
167}
168
169impl Drop for ShardClose<'_> {
170    fn drop(&mut self) {
171        if let Some(shard) = self.shard.take() {
172            let replaced = self.owner.lock().shards.insert(self.id, *shard);
173            assert!(replaced.is_none(), "closing shard ID is already present");
174        }
175    }
176}
177
178#[allow(
179    clippy::missing_panics_doc,
180    reason = "internal mutex locks and shard-state invariants are not expected to panic"
181)]
182impl PolymarketMarketConnectionPool {
183    /// Creates a new market connection pool (unconnected).
184    ///
185    /// A `max_subscriptions` of `0` is invalid and clamps to
186    /// [`WS_DEFAULT_SUBSCRIPTIONS`] with a warning.
187    #[must_use]
188    pub fn new(
189        base_url: Option<String>,
190        subscribe_new_markets: bool,
191        transport_backend: TransportBackend,
192        max_subscriptions: usize,
193    ) -> Self {
194        Self::new_with_proxy(
195            base_url,
196            subscribe_new_markets,
197            transport_backend,
198            max_subscriptions,
199            None,
200        )
201    }
202
203    /// Creates a new market connection pool with an optional validated proxy URL.
204    #[must_use]
205    pub fn new_with_proxy(
206        base_url: Option<String>,
207        subscribe_new_markets: bool,
208        transport_backend: TransportBackend,
209        max_subscriptions: usize,
210        proxy_url: Option<ProxyUrl>,
211    ) -> Self {
212        Self {
213            inner: Arc::new(PoolInner::new_with_proxy(
214                base_url,
215                transport_backend,
216                subscribe_new_markets,
217                max_subscriptions,
218                proxy_url,
219            )),
220        }
221    }
222
223    pub(crate) fn with_connect_timeout(self, timeout_secs: u64) -> Self {
224        *self.inner.connect_timeout_secs.lock() = Some(timeout_secs);
225        self
226    }
227
228    /// Configures socket state reporting and reconnect control for every connection in the pool.
229    #[must_use]
230    pub(crate) fn with_socket_factory(self, factory: SocketControlFactory) -> Self {
231        *self.inner.socket_factory.lock() = Some(factory);
232        self
233    }
234
235    #[cfg(test)]
236    pub(crate) fn proxy_url(&self) -> Option<&ProxyUrl> {
237        self.inner.proxy_url.as_ref()
238    }
239
240    /// Returns a cloneable routing handle for use in spawned subscription tasks.
241    #[must_use]
242    pub fn handle(&self) -> PolymarketMarketPoolHandle {
243        PolymarketMarketPoolHandle {
244            inner: Arc::clone(&self.inner),
245        }
246    }
247
248    pub(crate) async fn prepare(&self) -> anyhow::Result<()> {
249        if self.inner.closed.load(Ordering::Acquire) {
250            self.disconnect().await?;
251        }
252
253        let _wire = self.inner.wire_mutex.lock().await;
254        self.inner.prepare_channel();
255        Ok(())
256    }
257
258    /// Opens the primary shard and prepares the merged message stream.
259    ///
260    /// # Errors
261    ///
262    /// Returns an error if the primary connection cannot be established.
263    pub async fn connect(&self) -> anyhow::Result<()> {
264        if self.inner.closed.load(Ordering::Acquire) {
265            self.disconnect().await?;
266        }
267
268        let _wire = self.inner.wire_mutex.lock().await;
269
270        if !self.inner.closed.load(Ordering::Acquire) && !self.inner.state.lock().shards.is_empty()
271        {
272            log::warn!("Polymarket market pool already connected");
273            return Ok(());
274        }
275
276        self.inner.prepare_channel();
277        self.inner.connect_new_shard(true).await?;
278        Ok(())
279    }
280
281    /// Sends the new-market discovery subscribe on the primary shard.
282    ///
283    /// # Errors
284    ///
285    /// Returns an error if no primary shard is available.
286    pub async fn subscribe_new_markets_feed(&self) -> anyhow::Result<()> {
287        let _wire = self.inner.wire_mutex.lock().await;
288
289        let handle = {
290            let state = self.inner.state.lock();
291            if self.inner.closed.load(Ordering::Acquire) {
292                anyhow::bail!("Market connection pool is closed");
293            }
294            state
295                .shards
296                .get(&PRIMARY_SHARD_ID)
297                .map(|shard| shard.handle.clone())
298        };
299
300        match handle {
301            Some(handle) => handle.subscribe_market(vec![]).await,
302            None => anyhow::bail!("No primary market shard available for new-market discovery"),
303        }
304    }
305
306    /// Takes the merged message receiver, leaving `None` in its place.
307    #[must_use]
308    pub fn take_message_receiver(
309        &self,
310    ) -> Option<tokio::sync::mpsc::UnboundedReceiver<PolymarketWsMessage>> {
311        self.inner.out_rx.lock().take()
312    }
313
314    /// Disconnects every shard and clears routing state.
315    ///
316    /// # Errors
317    ///
318    /// Returns an error if a task or connection does not stop. Attempts every shard before
319    /// returning the error.
320    pub async fn disconnect(&self) -> anyhow::Result<()> {
321        self.inner.begin_shutdown();
322        let _wire = self.inner.wire_mutex.lock().await;
323
324        let mut drain = PoolDrain::take(&self.inner.state);
325        {
326            let errors = Mutex::new(&mut drain.state.shutdown_errors);
327            let mut disconnects = drain
328                .state
329                .shards
330                .iter_mut()
331                .map(|(&shard_id, shard)| {
332                    let errors = &errors;
333                    async move {
334                        let mut failed = false;
335                        shard.forwarder.abort();
336                        if let Some(outcome) = finish_task(
337                            &mut shard.forwarder,
338                            std::time::Duration::ZERO,
339                            std::time::Duration::from_secs(2),
340                        )
341                        .await
342                        {
343                            match outcome {
344                                TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
345                                TaskJoinOutcome::Failed(error) => {
346                                    failed = true;
347                                    errors.lock().push(format!(
348                                        "market shard {shard_id} forwarder failed: {error}"
349                                    ));
350                                }
351                                TaskJoinOutcome::Incomplete => {
352                                    failed = true;
353                                    errors.lock().push(format!(
354                                        "market shard {shard_id} forwarder did not stop after abort"
355                                    ));
356                                }
357                            }
358                        }
359
360                        if let Err(e) = shard.client.disconnect().await {
361                            failed = true;
362                            errors
363                                .lock()
364                                .push(format!("market shard {shard_id} disconnect failed: {e}"));
365                        }
366                        (shard_id, failed)
367                    }
368                })
369                .collect::<FuturesUnordered<_>>();
370
371            while let Some((shard_id, failed)) = disconnects.next().await {
372                if !failed {
373                    drain.completed.push(shard_id);
374                }
375            }
376        }
377
378        if !drain.state.shutdown_errors.is_empty() {
379            let errors = std::mem::take(&mut drain.state.shutdown_errors);
380            anyhow::bail!(
381                "Polymarket market pool shutdown failed: {}",
382                errors.join("; ")
383            );
384        }
385
386        *self.inner.out_tx.lock() = None;
387        *self.inner.out_rx.lock() = None;
388        Ok(())
389    }
390
391    pub(crate) fn begin_shutdown(&self) {
392        self.inner.begin_shutdown();
393    }
394
395    /// Clears retained reconnect-replay state on any remaining shards.
396    pub(crate) fn clear_reconnect_state(&self) {
397        let state = self.inner.state.lock();
398        for shard in state.shards.values() {
399            shard.client.clear_reconnect_state();
400        }
401    }
402
403    /// Returns the number of open shard connections.
404    #[must_use]
405    pub fn connection_count(&self) -> usize {
406        self.inner.state.lock().shards.len()
407    }
408
409    /// Returns the number of unique assets assigned across all shards.
410    #[must_use]
411    pub fn subscription_count(&self) -> usize {
412        self.inner.state.lock().assignments.len()
413    }
414}
415
416impl Drop for PoolInner {
417    fn drop(&mut self) {
418        self.closed.store(true, Ordering::Release);
419
420        for shard in self.state.get_mut().shards.values_mut() {
421            shard.forwarder.abort();
422            shard.client.abort();
423        }
424    }
425}
426
427#[allow(
428    clippy::missing_panics_doc,
429    reason = "internal mutex locks and shard-state invariants are not expected to panic"
430)]
431impl PolymarketMarketPoolHandle {
432    pub(crate) fn begin_shutdown(&self) {
433        self.inner.begin_shutdown();
434    }
435
436    /// Subscribes to market data for the given asset IDs, sharding across connections.
437    ///
438    /// # Errors
439    ///
440    /// Returns an error if a shard cannot be opened or a subscribe send fails.
441    pub async fn subscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
442        let _wire = self.inner.wire_mutex.lock().await;
443        self.inner.ensure_open()?;
444        for asset_id in asset_ids {
445            self.inner.subscribe_one(asset_id).await?;
446        }
447        Ok(())
448    }
449
450    /// Removes asset IDs from their owning shards, closing emptied secondary shards.
451    ///
452    /// # Errors
453    ///
454    /// Returns an error if an unsubscribe send fails.
455    pub async fn unsubscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
456        let _wire = self.inner.wire_mutex.lock().await;
457        self.inner.ensure_open()?;
458        for asset_id in asset_ids {
459            self.inner.unsubscribe_one(asset_id).await?;
460        }
461        Ok(())
462    }
463
464    /// Cycles owned assets (unsubscribe then subscribe) to trigger fresh snapshots.
465    ///
466    /// Unlike [`Self::subscribe_market`], this never changes shard assignment: it routes
467    /// each asset to its owning shard, which re-subscribes it on the wire without
468    /// touching desired subscription state. The venue ignores a duplicate subscribe,
469    /// so only a real cycle produces a fresh `book` snapshot. Used by book recovery.
470    ///
471    /// # Errors
472    ///
473    /// Returns an error if the pool is closed, an asset is not owned by any shard,
474    /// or a cycle command send fails.
475    pub async fn resubscribe_market(
476        &self,
477        asset_ids: Vec<String>,
478        cancel: &tokio_util::sync::CancellationToken,
479        gate: &SnapshotGate,
480    ) -> anyhow::Result<CycleMarketOutcome> {
481        let mut outcome = CycleMarketOutcome::Completed;
482
483        for asset_id in asset_ids {
484            outcome = self.inner.cycle_one(&asset_id, cancel, gate).await?;
485
486            if !matches!(outcome, CycleMarketOutcome::Completed) {
487                break;
488            }
489        }
490
491        Ok(outcome)
492    }
493
494    /// Returns the tokens currently assigned to `shard_id`.
495    pub(crate) fn tokens_for_shard(&self, shard_id: usize) -> Vec<Ustr> {
496        let state = self.inner.state.lock();
497        state
498            .assignments
499            .iter()
500            .filter_map(|(token, assigned)| (*assigned == shard_id).then_some(*token))
501            .collect()
502    }
503}
504
505impl PoolInner {
506    fn begin_shutdown(&self) {
507        let state = self.state.lock();
508        self.closed.store(true, Ordering::Release);
509
510        for shard in state.shards.values() {
511            shard.client.begin_shutdown();
512        }
513    }
514
515    fn prepare_channel(&self) {
516        {
517            let _state = self.state.lock();
518            self.closed.store(false, Ordering::Release);
519        }
520
521        if self.out_tx.lock().is_some() {
522            return;
523        }
524
525        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel();
526        *self.out_tx.lock() = Some(out_tx);
527        *self.out_rx.lock() = Some(out_rx);
528    }
529
530    fn ensure_open(&self) -> anyhow::Result<()> {
531        let _state = self.state.lock();
532
533        if self.closed.load(Ordering::Acquire) {
534            anyhow::bail!("Market connection pool is closed");
535        }
536        Ok(())
537    }
538}
539
540impl PoolInner {
541    #[cfg(test)]
542    fn new(
543        base_url: Option<String>,
544        transport_backend: TransportBackend,
545        subscribe_new_markets: bool,
546        max_subscriptions: usize,
547    ) -> Self {
548        Self::new_with_proxy(
549            base_url,
550            transport_backend,
551            subscribe_new_markets,
552            max_subscriptions,
553            None,
554        )
555    }
556
557    fn new_with_proxy(
558        base_url: Option<String>,
559        transport_backend: TransportBackend,
560        subscribe_new_markets: bool,
561        max_subscriptions: usize,
562        proxy_url: Option<ProxyUrl>,
563    ) -> Self {
564        let max_subscriptions = if max_subscriptions == 0 {
565            log::warn!(
566                "PolymarketDataClientConfig.ws_max_subscriptions=0 is invalid, using {WS_DEFAULT_SUBSCRIPTIONS}"
567            );
568            WS_DEFAULT_SUBSCRIPTIONS
569        } else {
570            max_subscriptions
571        };
572
573        Self {
574            base_url,
575            proxy_url,
576            transport_backend,
577            subscribe_new_markets,
578            max_subscriptions,
579            connect_timeout_secs: Mutex::new(None),
580            wire_mutex: tokio::sync::Mutex::new(()),
581            state: Mutex::new(PoolState::default()),
582            out_tx: Mutex::new(None),
583            out_rx: Mutex::new(None),
584            socket_factory: Mutex::new(None),
585            closed: AtomicBool::new(false),
586        }
587    }
588
589    // Callers hold `wire_mutex`.
590    async fn subscribe_one(&self, asset_id: String) -> anyhow::Result<()> {
591        let token = Ustr::from(asset_id.as_str());
592
593        let Some(handle) = self.assign(token).await? else {
594            return Ok(());
595        };
596
597        if let Err(e) = self.ensure_open() {
598            if let ReleaseOutcome::CloseShard(id, shard) = self.release(token)
599                && let Err(close_error) = self.close_shard(id, shard).await
600            {
601                anyhow::bail!("{e}; subscription rollback failed: {close_error}");
602            }
603            return Err(e);
604        }
605
606        if let Err(e) = handle.subscribe_market(vec![asset_id]).await {
607            // Roll back so a failed send leaves no stale assignment or empty shard.
608            if let ReleaseOutcome::CloseShard(id, shard) = self.release(token)
609                && let Err(close_error) = self.close_shard(id, shard).await
610            {
611                anyhow::bail!("{e}; subscription rollback failed: {close_error}");
612            }
613            return Err(e);
614        }
615        Ok(())
616    }
617
618    // Callers hold `wire_mutex`.
619    async fn unsubscribe_one(&self, asset_id: String) -> anyhow::Result<()> {
620        self.ensure_open()?;
621        let token = Ustr::from(asset_id.as_str());
622
623        match self.release(token) {
624            ReleaseOutcome::NotOwned => Ok(()),
625            ReleaseOutcome::Unsubscribe(handle) => handle.unsubscribe_market(vec![asset_id]).await,
626            ReleaseOutcome::CloseShard(id, shard) => {
627                // Disconnect drops the shard's subscriptions; no unsubscribe send needed.
628                self.close_shard(id, shard).await
629            }
630        }
631    }
632
633    async fn cycle_one(
634        &self,
635        asset_id: &str,
636        cancel: &tokio_util::sync::CancellationToken,
637        gate: &SnapshotGate,
638    ) -> anyhow::Result<CycleMarketOutcome> {
639        let responder_rx = {
640            let _wire = self.wire_mutex.lock().await;
641            self.ensure_open()?;
642
643            let handle = {
644                let state = self.state.lock();
645                let token = Ustr::from(asset_id);
646
647                let Some(id) = state.assignments.get(&token).copied() else {
648                    anyhow::bail!("Market asset {asset_id} is not owned by any shard");
649                };
650
651                state
652                    .shards
653                    .get(&id)
654                    .map(|shard| shard.handle.clone())
655                    .ok_or_else(|| anyhow::anyhow!("Market shard {id} is not available"))?
656            };
657
658            // The wire mutex is released before awaiting completion: the queued
659            // command is ordered, and awaiting must not block pool operations.
660            let (responder_tx, responder_rx) = tokio::sync::oneshot::channel();
661            handle
662                .cycle_market_subscription(
663                    vec![asset_id.to_string()],
664                    cancel.clone(),
665                    responder_tx,
666                    gate.clone(),
667                )
668                .await?;
669            responder_rx
670        };
671
672        responder_rx
673            .await
674            .map_err(|_| anyhow::anyhow!("Market subscription cycle response lost"))
675    }
676
677    // Returns `None` when the token is already owned by a shard.
678    async fn assign(&self, token: Ustr) -> anyhow::Result<Option<WsSubscriptionHandle>> {
679        {
680            let mut state = self.state.lock();
681
682            if self.closed.load(Ordering::Acquire) {
683                anyhow::bail!("Market connection pool is closed");
684            }
685
686            if state.assignments.contains_key(&token) {
687                return Ok(None);
688            }
689
690            if let Some(id) = smallest_shard_with_capacity(&state, self.max_subscriptions) {
691                let handle = {
692                    let shard = state.shards.get_mut(&id).expect("shard present");
693                    shard.owned += 1;
694                    shard.handle.clone()
695                };
696                state.assignments.insert(token, id);
697                return Ok(Some(handle));
698            }
699        }
700
701        let id = self.connect_new_shard(false).await?;
702
703        let rejected_shard = {
704            let mut state = self.state.lock();
705
706            if self.closed.load(Ordering::Acquire) {
707                Some(
708                    state
709                        .shards
710                        .remove(&id)
711                        .expect("new shard retained for shutdown"),
712                )
713            } else {
714                let handle = {
715                    let shard = state.shards.get_mut(&id).expect("new shard present");
716                    shard.owned += 1;
717                    shard.handle.clone()
718                };
719                state.assignments.insert(token, id);
720                return Ok(Some(handle));
721            }
722        };
723
724        if let Some(shard) = rejected_shard {
725            if let Err(e) = self.close_shard(id, Box::new(shard)).await {
726                anyhow::bail!("Market connection pool is closed; shard rollback failed: {e}");
727            }
728            anyhow::bail!("Market connection pool is closed");
729        }
730        unreachable!("open pool returned from assignment")
731    }
732
733    fn release(&self, token: Ustr) -> ReleaseOutcome {
734        let mut state = self.state.lock();
735
736        let Some(id) = state.assignments.remove(&token) else {
737            return ReleaseOutcome::NotOwned;
738        };
739
740        let owned = {
741            let Some(shard) = state.shards.get_mut(&id) else {
742                return ReleaseOutcome::NotOwned;
743            };
744            shard.owned = shard.owned.saturating_sub(1);
745            shard.owned
746        };
747
748        if id != PRIMARY_SHARD_ID && owned == 0 {
749            let shard = state.shards.remove(&id).expect("shard present");
750            ReleaseOutcome::CloseShard(id, Box::new(shard))
751        } else {
752            let handle = state.shards.get(&id).expect("shard present").handle.clone();
753            ReleaseOutcome::Unsubscribe(handle)
754        }
755    }
756
757    async fn connect_new_shard(&self, is_primary: bool) -> anyhow::Result<usize> {
758        if self.closed.load(Ordering::Acquire) {
759            anyhow::bail!("Market connection pool is closed");
760        }
761
762        let id = if is_primary {
763            PRIMARY_SHARD_ID
764        } else {
765            let state = self.state.lock();
766            available_shard_id(&state)
767        };
768
769        let mut client = self.market_client(self.subscribe_new_markets, id);
770        client.connect().await?;
771
772        let handle = client.clone_subscription_handle();
773        let rx = client
774            .take_message_receiver()
775            .ok_or_else(|| anyhow::anyhow!("Market shard receiver unavailable after connect"))?;
776
777        let forwarder = match self.spawn_forwarder(rx, id) {
778            Ok(forwarder) => forwarder,
779            Err((e, forwarder)) => {
780                let shard = Box::new(ShardEntry {
781                    client,
782                    handle,
783                    forwarder,
784                    owned: 0,
785                    closing: true,
786                });
787
788                if let Err(close_error) = self.close_shard(id, shard).await {
789                    anyhow::bail!(
790                        "Failed to start market shard forwarder: {e}; startup rollback failed: \
791                         {close_error}"
792                    );
793                }
794                anyhow::bail!("Failed to start market shard forwarder: {e}");
795            }
796        };
797
798        let shard = ShardEntry {
799            client,
800            handle,
801            forwarder,
802            owned: 0,
803            closing: false,
804        };
805        let rejected_shard = {
806            let mut state = self.state.lock();
807
808            if self.closed.load(Ordering::Acquire) {
809                Some(shard)
810            } else {
811                state.shards.insert(id, shard);
812                None
813            }
814        };
815
816        if let Some(shard) = rejected_shard {
817            if let Err(e) = self.close_shard(id, Box::new(shard)).await {
818                anyhow::bail!("Market connection pool is closed; shard rollback failed: {e}");
819            }
820            anyhow::bail!("Market connection pool is closed");
821        }
822
823        log::debug!("Opened Polymarket market shard {id}");
824        Ok(id)
825    }
826
827    fn market_client(
828        &self,
829        subscribe_new_markets: bool,
830        shard_id: usize,
831    ) -> PolymarketWebSocketClient {
832        let client = PolymarketWebSocketClient::new_market_with_proxy(
833            self.base_url.clone(),
834            subscribe_new_markets,
835            self.transport_backend,
836            self.proxy_url.clone(),
837        );
838
839        let client = match *self.connect_timeout_secs.lock() {
840            Some(timeout_secs) => client.with_connect_timeout(timeout_secs),
841            None => client,
842        };
843
844        let factory = self.socket_factory.lock().clone();
845
846        if let Some(factory) = factory {
847            let endpoint = if shard_id == PRIMARY_SHARD_ID {
848                MARKET_STREAMS_ENDPOINT.to_string()
849            } else {
850                format!("{MARKET_STREAMS_ENDPOINT}-{shard_id}")
851            };
852            client.with_socket_control(factory.control(endpoint))
853        } else {
854            client
855        }
856    }
857
858    fn spawn_forwarder(
859        &self,
860        mut rx: tokio::sync::mpsc::UnboundedReceiver<PolymarketWsMessage>,
861        shard_id: usize,
862    ) -> Result<TaskSlot<()>, (TaskSpawnError, TaskSlot<()>)> {
863        let out_tx = self.out_tx.lock().clone();
864        let is_primary = shard_id == PRIMARY_SHARD_ID;
865
866        let mut forwarder = TaskSlot::new();
867        if let Err(e) = forwarder.spawn(async move {
868            let Some(out_tx) = out_tx else {
869                return;
870            };
871
872            while let Some(msg) = rx.recv().await {
873                if !should_forward_from_shard(&msg, is_primary) {
874                    continue;
875                }
876
877                let msg = match msg {
878                    PolymarketWsMessage::Reconnected { .. } => PolymarketWsMessage::Reconnected {
879                        shard_id: Some(shard_id),
880                    },
881                    other => other,
882                };
883
884                if out_tx.send(msg).is_err() {
885                    break;
886                }
887            }
888        }) {
889            return Err((e, forwarder));
890        }
891        Ok(forwarder)
892    }
893
894    async fn close_shard(&self, id: usize, shard: Box<ShardEntry>) -> anyhow::Result<()> {
895        let mut close = ShardClose::new(&self.state, id, shard);
896        close.shard_mut().forwarder.abort();
897        let forwarder_stopped = match finish_task(
898            &mut close.shard_mut().forwarder,
899            std::time::Duration::ZERO,
900            std::time::Duration::from_secs(2),
901        )
902        .await
903        {
904            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => true,
905            Some(TaskJoinOutcome::Failed(error)) => {
906                let error = format!("market shard {id} forwarder failed: {error}");
907                self.state.lock().shutdown_errors.push(error);
908                true
909            }
910            Some(TaskJoinOutcome::Incomplete) => {
911                let error = format!("market shard {id} forwarder did not stop after abort");
912                self.state.lock().shutdown_errors.push(error);
913                false
914            }
915        };
916
917        if let Err(e) = close.shard_mut().client.disconnect().await {
918            let error = format!("market shard {id} disconnect failed: {e}");
919            self.state.lock().shutdown_errors.push(error);
920        }
921
922        if forwarder_stopped && !close.shard_mut().client.has_task() {
923            close.complete();
924        }
925
926        let errors = std::mem::take(&mut self.state.lock().shutdown_errors);
927
928        if errors.is_empty() {
929            Ok(())
930        } else {
931            anyhow::bail!(errors.join("; "))
932        }
933    }
934
935    #[cfg(test)]
936    fn subscription_count_for_test(&self) -> usize {
937        self.state.lock().assignments.len()
938    }
939}
940
941fn should_forward_from_shard(message: &PolymarketWsMessage, is_primary: bool) -> bool {
942    is_primary
943        || !matches!(
944            message,
945            PolymarketWsMessage::Market(
946                MarketWsMessage::NewMarket(_) | MarketWsMessage::MarketResolved(_)
947            )
948        )
949}
950
951fn smallest_shard_with_capacity(state: &PoolState, max_subscriptions: usize) -> Option<usize> {
952    state
953        .shards
954        .iter()
955        .filter(|(_, shard)| !shard.closing && shard.owned < max_subscriptions)
956        .map(|(id, _)| *id)
957        .min()
958}
959
960fn available_shard_id(state: &PoolState) -> usize {
961    let mut id = PRIMARY_SHARD_ID + 1;
962    while state.shards.contains_key(&id) {
963        id = id.checked_add(1).expect("market shard ID space exhausted");
964    }
965    id
966}
967
968#[cfg(test)]
969impl PolymarketMarketPoolHandle {
970    /// In-memory single-shard handle backed by `sender`, `assigned` tokens
971    /// pre-owned. Never connected, so growth is never triggered.
972    pub(crate) fn test_single_shard(
973        sender: tokio::sync::mpsc::UnboundedSender<super::handler::HandlerCommand>,
974        assigned: &[&str],
975    ) -> Self {
976        let inner = PoolInner::new(
977            None,
978            TransportBackend::default(),
979            false,
980            WS_DEFAULT_SUBSCRIPTIONS,
981        );
982        {
983            let mut state = inner.state.lock();
984            state.shards.insert(
985                PRIMARY_SHARD_ID,
986                ShardEntry {
987                    client: PolymarketWebSocketClient::new_market(
988                        None,
989                        false,
990                        TransportBackend::default(),
991                    ),
992                    handle: WsSubscriptionHandle::from_sender(sender),
993                    forwarder: TaskSlot::new(),
994                    owned: assigned.len(),
995                    closing: false,
996                },
997            );
998
999            for token in assigned {
1000                state
1001                    .assignments
1002                    .insert(Ustr::from(token), PRIMARY_SHARD_ID);
1003            }
1004        }
1005        Self {
1006            inner: Arc::new(inner),
1007        }
1008    }
1009}
1010
1011#[cfg(test)]
1012mod tests {
1013    use std::{net::SocketAddr, sync::Arc as StdArc, time::Duration};
1014
1015    use PolymarketMarketPoolHandle as Handle;
1016    use axum::{
1017        Router,
1018        extract::ws::{WebSocket, WebSocketUpgrade},
1019        response::Response,
1020        routing::get,
1021    };
1022    use nautilus_common::{
1023        live::runner::replace_system_event_sender,
1024        messages::{SystemEvent, system::SocketState},
1025    };
1026    use nautilus_live::{SocketReconnectRegistry, SocketReconnectRequestOutcome};
1027    use nautilus_model::identifiers::ClientId;
1028    use nautilus_network::transport::TransportError;
1029    use parking_lot::{Condvar, Mutex as TestMutex};
1030    use rstest::rstest;
1031    use tokio_util::sync::CancellationToken;
1032
1033    use super::*;
1034    use crate::websocket::handler::HandlerCommand;
1035
1036    struct BlockingDrop(StdArc<(TestMutex<(bool, bool)>, Condvar)>);
1037
1038    impl Drop for BlockingDrop {
1039        fn drop(&mut self) {
1040            let (state, wake) = &*self.0;
1041            let mut state = state.lock();
1042            state.0 = true;
1043            wake.notify_all();
1044            while !state.1 {
1045                wake.wait(&mut state);
1046            }
1047        }
1048    }
1049
1050    async fn handle_socket_upgrade(ws: WebSocketUpgrade) -> Response {
1051        ws.on_upgrade(handle_socket)
1052    }
1053
1054    async fn handle_socket(mut socket: WebSocket) {
1055        while socket.recv().await.is_some() {}
1056    }
1057
1058    async fn start_socket_server() -> SocketAddr {
1059        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
1060            .await
1061            .expect("bind test websocket server");
1062        let addr = listener.local_addr().expect("test websocket address");
1063        let router = Router::new().route("/ws/market", get(handle_socket_upgrade));
1064
1065        tokio::spawn(async move {
1066            axum::serve(listener, router)
1067                .await
1068                .expect("test websocket server failed");
1069        });
1070
1071        addr
1072    }
1073
1074    // Bare state with unconnected shards for pure capacity-accounting tests.
1075    fn state_with_shards(owned: &[usize]) -> PoolState {
1076        let mut state = PoolState::default();
1077
1078        for (id, owned) in owned.iter().enumerate() {
1079            let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
1080            state.shards.insert(
1081                id,
1082                ShardEntry {
1083                    client: PolymarketWebSocketClient::new_market(
1084                        None,
1085                        false,
1086                        TransportBackend::default(),
1087                    ),
1088                    handle: WsSubscriptionHandle::from_sender(tx),
1089                    forwarder: TaskSlot::new(),
1090                    owned: *owned,
1091                    closing: false,
1092                },
1093            );
1094        }
1095        state
1096    }
1097
1098    fn market_message(filename: &str) -> PolymarketWsMessage {
1099        let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
1100            .join("test_data")
1101            .join(filename);
1102        let json = std::fs::read_to_string(path).unwrap();
1103        PolymarketWsMessage::Market(serde_json::from_str(&json).unwrap())
1104    }
1105
1106    #[rstest]
1107    #[case::primary_new_market("ws_market_new_market_msg.json", true, true)]
1108    #[case::secondary_new_market("ws_market_new_market_msg.json", false, false)]
1109    #[case::primary_resolution("ws_market_resolved_msg.json", true, true)]
1110    #[case::secondary_resolution("ws_market_resolved_msg.json", false, false)]
1111    #[case::secondary_best_bid_ask("ws_market_best_bid_ask_msg.json", false, true)]
1112    fn shard_forwarding_keeps_global_events_on_primary(
1113        #[case] filename: &str,
1114        #[case] is_primary: bool,
1115        #[case] expected: bool,
1116    ) {
1117        let message = market_message(filename);
1118        assert_eq!(should_forward_from_shard(&message, is_primary), expected);
1119    }
1120
1121    #[rstest]
1122    #[tokio::test]
1123    async fn secondary_forwarder_drops_global_events_and_keeps_best_bid_ask() {
1124        let inner = PoolInner::new(None, TransportBackend::default(), true, 1);
1125        let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
1126        *inner.out_tx.lock() = Some(out_tx);
1127        let (shard_tx, shard_rx) = tokio::sync::mpsc::unbounded_channel();
1128        let mut forwarder = inner
1129            .spawn_forwarder(shard_rx, PRIMARY_SHARD_ID + 1)
1130            .expect("spawn forwarder");
1131
1132        shard_tx
1133            .send(market_message("ws_market_new_market_msg.json"))
1134            .unwrap();
1135        shard_tx
1136            .send(market_message("ws_market_resolved_msg.json"))
1137            .unwrap();
1138        shard_tx
1139            .send(market_message("ws_market_best_bid_ask_msg.json"))
1140            .unwrap();
1141        drop(shard_tx);
1142        let outcome = finish_task(
1143            &mut forwarder,
1144            Duration::from_secs(1),
1145            Duration::from_secs(1),
1146        )
1147        .await
1148        .expect("forwarder task");
1149        assert!(matches!(outcome, TaskJoinOutcome::Completed(())));
1150
1151        let forwarded = out_rx.try_recv().unwrap();
1152        let PolymarketWsMessage::Market(MarketWsMessage::BestBidAsk(message)) = forwarded else {
1153            panic!("unexpected forwarded message: {forwarded:?}");
1154        };
1155        assert_eq!(
1156            message.asset_id,
1157            Ustr::from(
1158                "85354956062430465315924116860125388538595433819574542752031640332592237464430"
1159            ),
1160        );
1161        assert!(out_rx.try_recv().is_err());
1162    }
1163
1164    #[rstest]
1165    #[tokio::test]
1166    async fn pool_drain_restores_remaining_state_on_drop() {
1167        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1168        let handle = Handle::test_single_shard(cmd_tx, &["token-a"]);
1169        handle
1170            .inner
1171            .state
1172            .lock()
1173            .shutdown_errors
1174            .push("before drain".to_string());
1175
1176        {
1177            let mut drain = PoolDrain::take(&handle.inner.state);
1178            let owner = handle.inner.state.lock();
1179            assert!(owner.shards.is_empty());
1180            assert!(owner.assignments.is_empty());
1181            assert!(owner.shutdown_errors.is_empty());
1182            drop(owner);
1183
1184            assert_eq!(drain.state.shards.len(), 1);
1185            assert_eq!(
1186                drain.state.assignments,
1187                AHashMap::from([(Ustr::from("token-a"), PRIMARY_SHARD_ID)]),
1188            );
1189            assert_eq!(drain.state.shutdown_errors, vec!["before drain"]);
1190            drain.state.shutdown_errors.push("during drain".to_string());
1191        }
1192
1193        let restored = handle.inner.state.lock();
1194        assert_eq!(restored.shards.len(), 1);
1195        let shard = restored
1196            .shards
1197            .get(&PRIMARY_SHARD_ID)
1198            .expect("restored shard");
1199        assert_eq!(shard.owned, 1);
1200        assert!(!shard.closing);
1201        assert_eq!(
1202            restored.assignments,
1203            AHashMap::from([(Ustr::from("token-a"), PRIMARY_SHARD_ID)]),
1204        );
1205        assert_eq!(
1206            restored.shutdown_errors,
1207            vec!["before drain", "during drain"]
1208        );
1209    }
1210
1211    async fn blocked_forwarder() -> (TaskSlot<()>, std::sync::mpsc::Sender<()>) {
1212        let (release, blocked) = std::sync::mpsc::channel();
1213        let (started, ready) = tokio::sync::oneshot::channel();
1214        let handle = tokio::task::spawn_blocking(move || {
1215            let _ = started.send(());
1216            let _ = blocked.recv();
1217        });
1218        ready.await.unwrap();
1219        (TaskSlot::from_handle(handle), release)
1220    }
1221
1222    #[rstest]
1223    #[case::single(1)]
1224    #[case::many(64)]
1225    #[tokio::test(start_paused = true)]
1226    async fn pool_disconnects_concurrently_and_preserves_failures(#[case] count: usize) {
1227        let pool = PolymarketMarketConnectionPool::new(None, false, TransportBackend::default(), 1);
1228        let mut state = state_with_shards(&vec![1; count]);
1229        let mut releases = Vec::new();
1230
1231        for (&id, shard) in &mut state.shards {
1232            let (forwarder, release) = blocked_forwarder().await;
1233            shard.forwarder = forwarder;
1234            releases.push(release);
1235            state.assignments.insert(Ustr::from(&id.to_string()), id);
1236        }
1237        *pool.inner.state.lock() = state;
1238
1239        let started = tokio::time::Instant::now();
1240        let expected = Duration::from_secs(2);
1241        let mut disconnect = std::pin::pin!(pool.disconnect());
1242        let result = loop {
1243            if let std::task::Poll::Ready(result) = futures_util::poll!(disconnect.as_mut()) {
1244                break result;
1245            }
1246            assert!(started.elapsed() < expected + Duration::from_secs(1));
1247            tokio::time::advance(Duration::from_millis(100)).await;
1248        };
1249        assert!(started.elapsed() >= expected);
1250        let error = result.unwrap_err().to_string();
1251        let failures = error.matches("forwarder did not stop after abort");
1252        assert_eq!(failures.count(), count);
1253        assert_eq!(pool.connection_count(), count);
1254        assert_eq!(pool.subscription_count(), count);
1255
1256        drop(releases);
1257        pool.disconnect().await.expect("cleanup retry");
1258        assert_eq!(pool.connection_count(), 0);
1259        assert_eq!(pool.subscription_count(), 0);
1260    }
1261
1262    #[rstest]
1263    #[tokio::test]
1264    async fn canceled_pool_disconnect_restores_only_unfinished_ownership() {
1265        let pool = PolymarketMarketConnectionPool::new(None, false, TransportBackend::default(), 1);
1266        let mut state = state_with_shards(&[1, 1]);
1267        let (forwarder, release) = blocked_forwarder().await;
1268        state.shards.get_mut(&1).unwrap().forwarder = forwarder;
1269        state.assignments = AHashMap::from([(Ustr::from("done"), 0), (Ustr::from("pending"), 1)]);
1270        *pool.inner.state.lock() = state;
1271
1272        {
1273            let mut disconnect = std::pin::pin!(pool.disconnect());
1274            assert!(futures_util::poll!(disconnect.as_mut()).is_pending());
1275        }
1276        assert_eq!(pool.connection_count(), 1);
1277        assert_eq!(
1278            pool.inner.state.lock().assignments,
1279            AHashMap::from([(Ustr::from("pending"), 1)])
1280        );
1281
1282        drop(release);
1283        pool.disconnect().await.expect("cleanup retry");
1284        assert_eq!(pool.connection_count(), 0);
1285        assert_eq!(pool.subscription_count(), 0);
1286    }
1287
1288    #[rstest]
1289    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1290    async fn canceled_shard_close_restores_unfinished_ownership() {
1291        let inner = Arc::new(PoolInner::new(None, TransportBackend::default(), false, 1));
1292        let blocking = StdArc::new((TestMutex::new((false, false)), Condvar::new()));
1293        let blocking_task = StdArc::clone(&blocking);
1294        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1295        let mut forwarder = TaskSlot::new();
1296        forwarder
1297            .spawn(async move {
1298                let _blocking = BlockingDrop(blocking_task);
1299                let _ = started_tx.send(());
1300                std::future::pending::<()>().await;
1301            })
1302            .expect("spawn forwarder");
1303        started_rx.await.expect("forwarder should start");
1304
1305        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1306        let shard = Box::new(ShardEntry {
1307            client: PolymarketWebSocketClient::new_market(None, false, TransportBackend::default()),
1308            handle: WsSubscriptionHandle::from_sender(cmd_tx),
1309            forwarder,
1310            owned: 0,
1311            closing: false,
1312        });
1313        let close_inner = Arc::clone(&inner);
1314        let close = tokio::spawn(async move {
1315            let _result = close_inner.close_shard(1, shard).await;
1316        });
1317
1318        loop {
1319            if blocking.0.lock().0 {
1320                break;
1321            }
1322            tokio::task::yield_now().await;
1323        }
1324        close.abort();
1325        let _ = close.await;
1326
1327        {
1328            let restored = inner.state.lock();
1329            let shard = restored.shards.get(&1).expect("closing shard restored");
1330            assert!(shard.closing);
1331            assert!(shard.forwarder.is_some());
1332            assert_eq!(smallest_shard_with_capacity(&restored, 1), None);
1333        }
1334
1335        {
1336            let (state, wake) = &*blocking;
1337            state.lock().1 = true;
1338            wake.notify_all();
1339        }
1340        let shard = inner
1341            .state
1342            .lock()
1343            .shards
1344            .remove(&1)
1345            .expect("restored shard");
1346        inner
1347            .close_shard(1, Box::new(shard))
1348            .await
1349            .expect("restored shard should close");
1350
1351        assert!(inner.state.lock().shards.is_empty());
1352    }
1353
1354    #[rstest]
1355    fn zero_max_subscriptions_clamps_to_default() {
1356        let inner = PoolInner::new(None, TransportBackend::default(), false, 0);
1357        assert_eq!(inner.max_subscriptions, WS_DEFAULT_SUBSCRIPTIONS);
1358    }
1359
1360    #[rstest]
1361    #[tokio::test]
1362    async fn secondary_market_connection_timeout() {
1363        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1364        let url = format!("ws://{}", listener.local_addr().unwrap());
1365
1366        let server = tokio::spawn(async move {
1367            let mut connections = Vec::new();
1368
1369            loop {
1370                let (stream, _) = listener.accept().await.unwrap();
1371                connections.push(stream);
1372            }
1373        });
1374
1375        let pool =
1376            PolymarketMarketConnectionPool::new(Some(url), false, TransportBackend::default(), 1)
1377                .with_connect_timeout(1);
1378        let mut ws_client = pool.inner.market_client(false, PRIMARY_SHARD_ID + 1);
1379
1380        let error = tokio::time::timeout(Duration::from_secs(20), ws_client.connect())
1381            .await
1382            .expect("secondary shard must use the configured connection timeout")
1383            .unwrap_err();
1384        server.abort();
1385
1386        assert!(
1387            matches!(error.downcast_ref::<TransportError>(),
1388                Some(TransportError::Io(e)) if e.kind() == std::io::ErrorKind::TimedOut),
1389            "{error:#}",
1390        );
1391    }
1392
1393    #[rstest]
1394    fn pool_retains_proxy_for_lazily_created_shards() {
1395        const PROXY_URL: &str = "http://pool-user:pool-proxy-secret@127.0.0.1:18088";
1396        let pool = PolymarketMarketConnectionPool::new_with_proxy(
1397            Some("ws://market.example/ws".to_string()),
1398            true,
1399            TransportBackend::Tungstenite,
1400            17,
1401            Some(ProxyUrl::parse(PROXY_URL).unwrap()),
1402        );
1403        let primary = pool.inner.market_client(true, PRIMARY_SHARD_ID);
1404        let secondary = pool.inner.market_client(false, PRIMARY_SHARD_ID + 1);
1405        let debug = format!("{pool:?}");
1406
1407        assert_eq!(pool.inner.proxy_url.as_ref().unwrap().expose(), PROXY_URL);
1408        assert_eq!(primary.proxy_url().unwrap().expose(), PROXY_URL);
1409        assert_eq!(secondary.proxy_url().unwrap().expose(), PROXY_URL);
1410        assert_eq!(pool.inner.max_subscriptions, 17);
1411        assert!(!debug.contains("pool-proxy-secret"));
1412    }
1413
1414    #[rstest]
1415    #[tokio::test]
1416    async fn pool_assigns_distinct_endpoint_sinks_and_handles_before_connect() {
1417        let addr = start_socket_server().await;
1418        let (system_tx, mut system_rx) = tokio::sync::mpsc::unbounded_channel();
1419        replace_system_event_sender(system_tx);
1420        let registry = SocketReconnectRegistry::default();
1421        let factory = SocketControlFactory::with_registry(
1422            ClientId::from("POLYMARKET"),
1423            Some(*crate::common::consts::POLYMARKET_VENUE),
1424            &registry,
1425        );
1426        let pool = PolymarketMarketConnectionPool::new(
1427            Some(format!("ws://{addr}/ws/market")),
1428            false,
1429            TransportBackend::Tungstenite,
1430            1,
1431        )
1432        .with_socket_factory(factory);
1433
1434        pool.connect().await.expect("connect primary shard");
1435        pool.handle()
1436            .subscribe_market(vec!["asset-0".to_string(), "asset-1".to_string()])
1437            .await
1438            .expect("open secondary shard");
1439
1440        let mut connected = Vec::new();
1441        while connected.len() < 2 {
1442            let event = tokio::time::timeout(Duration::from_secs(2), system_rx.recv())
1443                .await
1444                .expect("wait for socket state event")
1445                .expect("system event channel closed");
1446            let SystemEvent::SocketState(change) = event;
1447            if change.state == SocketState::Connected {
1448                connected.push(change.endpoint);
1449            }
1450        }
1451        connected.sort_unstable();
1452
1453        assert_eq!(
1454            connected,
1455            vec![
1456                Ustr::from(MARKET_STREAMS_ENDPOINT),
1457                Ustr::from("polymarket-market-streams-1"),
1458            ],
1459        );
1460        let client_id = ClientId::from("POLYMARKET");
1461        let primary = registry
1462            .handle(client_id, Ustr::from(MARKET_STREAMS_ENDPOINT))
1463            .expect("primary reconnect handle should be registered");
1464        let secondary = registry
1465            .handle(client_id, Ustr::from("polymarket-market-streams-1"))
1466            .expect("secondary reconnect handle should be registered");
1467        assert_eq!(
1468            primary.request_reconnect(),
1469            SocketReconnectRequestOutcome::Accepted,
1470        );
1471        let event = system_rx
1472            .try_recv()
1473            .expect("selected shard should report reconnect state");
1474        let SystemEvent::SocketState(change) = event;
1475        assert_eq!(change.client_id, client_id);
1476        assert_eq!(change.endpoint, Ustr::from(MARKET_STREAMS_ENDPOINT));
1477        assert_eq!(change.state, SocketState::Disconnected);
1478        assert_eq!(
1479            secondary.request_reconnect(),
1480            SocketReconnectRequestOutcome::Accepted,
1481        );
1482
1483        pool.disconnect().await.expect("disconnect pool");
1484        assert!(
1485            registry
1486                .handle(client_id, Ustr::from(MARKET_STREAMS_ENDPOINT))
1487                .is_none()
1488        );
1489        assert!(
1490            registry
1491                .handle(client_id, Ustr::from("polymarket-market-streams-1"))
1492                .is_none()
1493        );
1494    }
1495
1496    #[rstest]
1497    #[case::first_has_room(&[0, 200], 200, Some(0))]
1498    #[case::prefers_lowest_id(&[200, 5, 5], 200, Some(1))]
1499    #[case::all_full(&[200, 200], 200, None)]
1500    #[case::exact_boundary_is_full(&[1], 1, None)]
1501    fn smallest_shard_with_capacity_picks_lowest_open_id(
1502        #[case] owned: &[usize],
1503        #[case] max: usize,
1504        #[case] expected: Option<usize>,
1505    ) {
1506        let state = state_with_shards(owned);
1507        assert_eq!(smallest_shard_with_capacity(&state, max), expected);
1508    }
1509
1510    #[rstest]
1511    fn available_shard_id_reuses_lowest_closed_shard() {
1512        let mut state = state_with_shards(&[1, 1, 1]);
1513        state.shards.remove(&(PRIMARY_SHARD_ID + 1));
1514
1515        assert_eq!(available_shard_id(&state), PRIMARY_SHARD_ID + 1);
1516    }
1517
1518    #[rstest]
1519    #[tokio::test]
1520    async fn subscribe_routes_command_and_tracks_assignment() {
1521        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1522        let handle = Handle::test_single_shard(tx, &[]);
1523
1524        handle
1525            .subscribe_market(vec!["token-a".to_string()])
1526            .await
1527            .expect("subscribe");
1528
1529        match rx.try_recv().expect("expected SubscribeMarket") {
1530            HandlerCommand::SubscribeMarket(ids) => assert_eq!(ids, vec!["token-a".to_string()]),
1531            other => panic!("unexpected command: {other:?}"),
1532        }
1533        assert_eq!(handle.inner.subscription_count_for_test(), 1);
1534    }
1535
1536    #[rstest]
1537    #[tokio::test]
1538    async fn duplicate_subscribe_does_not_consume_capacity_or_resend() {
1539        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1540        let handle = Handle::test_single_shard(tx, &[]);
1541
1542        handle
1543            .subscribe_market(vec!["token-a".to_string()])
1544            .await
1545            .expect("first subscribe");
1546        handle
1547            .subscribe_market(vec!["token-a".to_string()])
1548            .await
1549            .expect("duplicate subscribe");
1550
1551        assert!(matches!(
1552            rx.try_recv(),
1553            Ok(HandlerCommand::SubscribeMarket(_))
1554        ));
1555        assert!(rx.try_recv().is_err(), "duplicate must not resend");
1556        assert_eq!(handle.inner.subscription_count_for_test(), 1);
1557    }
1558
1559    #[rstest]
1560    #[tokio::test]
1561    async fn unsubscribe_routes_command_and_releases_assignment() {
1562        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1563        let handle = Handle::test_single_shard(tx, &["token-a"]);
1564
1565        handle
1566            .unsubscribe_market(vec!["token-a".to_string()])
1567            .await
1568            .expect("unsubscribe");
1569
1570        match rx.try_recv().expect("expected UnsubscribeMarket") {
1571            HandlerCommand::UnsubscribeMarket(ids) => assert_eq!(ids, vec!["token-a".to_string()]),
1572            other => panic!("unexpected command: {other:?}"),
1573        }
1574        assert_eq!(handle.inner.subscription_count_for_test(), 0);
1575    }
1576
1577    #[rstest]
1578    #[tokio::test]
1579    async fn unsubscribe_unknown_token_is_noop() {
1580        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1581        let handle = Handle::test_single_shard(tx, &[]);
1582
1583        handle
1584            .unsubscribe_market(vec!["token-a".to_string()])
1585            .await
1586            .expect("unsubscribe");
1587
1588        assert!(rx.try_recv().is_err());
1589    }
1590
1591    #[rstest]
1592    #[tokio::test]
1593    async fn subscribe_send_failure_rolls_back_assignment() {
1594        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1595        drop(rx);
1596        let handle = Handle::test_single_shard(tx, &[]);
1597
1598        let result = handle.subscribe_market(vec!["token-a".to_string()]).await;
1599
1600        assert!(result.is_err());
1601        assert_eq!(handle.inner.subscription_count_for_test(), 0);
1602    }
1603
1604    #[rstest]
1605    #[tokio::test]
1606    async fn resubscribe_market_cycles_without_changing_assignment() {
1607        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1608        let handle = Handle::test_single_shard(tx, &["token-a"]);
1609        let cancel = CancellationToken::new();
1610
1611        let task = {
1612            let handle = handle.clone();
1613            let cancel = cancel.clone();
1614
1615            tokio::spawn(async move {
1616                handle
1617                    .resubscribe_market(
1618                        vec!["token-a".to_string()],
1619                        &cancel,
1620                        &SnapshotGate::default(),
1621                    )
1622                    .await
1623            })
1624        };
1625
1626        match rx.recv().await.expect("expected cycle command") {
1627            HandlerCommand::CycleMarketSubscription {
1628                asset_ids,
1629                responder,
1630                ..
1631            } => {
1632                assert_eq!(asset_ids, vec!["token-a".to_string()]);
1633                responder
1634                    .send(CycleMarketOutcome::Completed)
1635                    .expect("answer cycle");
1636            }
1637            other => panic!("unexpected command: {other:?}"),
1638        }
1639
1640        assert!(matches!(
1641            task.await.expect("cycle task").expect("cycle owned asset"),
1642            CycleMarketOutcome::Completed
1643        ));
1644        assert_eq!(
1645            handle.tokens_for_shard(PRIMARY_SHARD_ID),
1646            vec![Ustr::from("token-a")]
1647        );
1648        assert!(handle.tokens_for_shard(PRIMARY_SHARD_ID + 1).is_empty());
1649        assert_eq!(handle.inner.subscription_count_for_test(), 1);
1650    }
1651
1652    #[rstest]
1653    #[tokio::test]
1654    async fn resubscribe_market_propagates_cycle_outcome() {
1655        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1656        let handle = Handle::test_single_shard(tx, &["token-a"]);
1657        let cancel = CancellationToken::new();
1658
1659        let task = {
1660            let handle = handle.clone();
1661            let cancel = cancel.clone();
1662
1663            tokio::spawn(async move {
1664                handle
1665                    .resubscribe_market(
1666                        vec!["token-a".to_string()],
1667                        &cancel,
1668                        &SnapshotGate::default(),
1669                    )
1670                    .await
1671            })
1672        };
1673
1674        match rx.recv().await.expect("expected cycle command") {
1675            HandlerCommand::CycleMarketSubscription { responder, .. } => {
1676                responder
1677                    .send(CycleMarketOutcome::ConnectionChanged)
1678                    .expect("answer cycle");
1679            }
1680            other => panic!("unexpected command: {other:?}"),
1681        }
1682
1683        assert!(matches!(
1684            task.await.expect("cycle task").expect("cycle reports"),
1685            CycleMarketOutcome::ConnectionChanged
1686        ));
1687        // Ownership is preserved even when the cycle does not complete.
1688        assert_eq!(
1689            handle.tokens_for_shard(PRIMARY_SHARD_ID),
1690            vec![Ustr::from("token-a")]
1691        );
1692    }
1693
1694    #[rstest]
1695    #[tokio::test]
1696    async fn resubscribe_market_fails_when_cycle_response_lost() {
1697        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1698        let handle = Handle::test_single_shard(tx, &["token-a"]);
1699        let cancel = CancellationToken::new();
1700
1701        let task = {
1702            let handle = handle.clone();
1703            let cancel = cancel.clone();
1704
1705            tokio::spawn(async move {
1706                handle
1707                    .resubscribe_market(
1708                        vec!["token-a".to_string()],
1709                        &cancel,
1710                        &SnapshotGate::default(),
1711                    )
1712                    .await
1713            })
1714        };
1715
1716        // Simulate handler death: receive the command, drop the responder.
1717        assert!(matches!(
1718            rx.recv().await.expect("expected cycle command"),
1719            HandlerCommand::CycleMarketSubscription { .. }
1720        ));
1721
1722        let err = task.await.expect("cycle task").expect_err("lost response");
1723
1724        assert!(
1725            err.to_string().contains("response lost"),
1726            "unexpected error: {err}"
1727        );
1728    }
1729
1730    #[rstest]
1731    #[tokio::test]
1732    async fn resubscribe_market_rejects_unowned_asset() {
1733        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1734        let handle = Handle::test_single_shard(tx, &["token-a"]);
1735
1736        let err = handle
1737            .resubscribe_market(
1738                vec!["token-unknown".to_string()],
1739                &CancellationToken::new(),
1740                &SnapshotGate::default(),
1741            )
1742            .await
1743            .expect_err("unowned asset must fail");
1744
1745        assert!(
1746            err.to_string().contains("not owned by any shard"),
1747            "unexpected error: {err}"
1748        );
1749        assert!(rx.try_recv().is_err());
1750    }
1751
1752    #[rstest]
1753    #[tokio::test]
1754    async fn resubscribe_market_rejects_missing_shard() {
1755        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1756        let handle = Handle::test_single_shard(tx, &["token-a"]);
1757        handle.inner.state.lock().shards.remove(&PRIMARY_SHARD_ID);
1758
1759        let err = handle
1760            .resubscribe_market(
1761                vec!["token-a".to_string()],
1762                &CancellationToken::new(),
1763                &SnapshotGate::default(),
1764            )
1765            .await
1766            .expect_err("missing shard must fail");
1767
1768        assert!(
1769            err.to_string().contains("is not available"),
1770            "unexpected error: {err}"
1771        );
1772    }
1773
1774    #[rstest]
1775    #[tokio::test]
1776    async fn resubscribe_market_rejects_closed_pool() {
1777        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1778        let handle = Handle::test_single_shard(tx, &["token-a"]);
1779        handle.inner.closed.store(true, Ordering::SeqCst);
1780
1781        let err = handle
1782            .resubscribe_market(
1783                vec!["token-a".to_string()],
1784                &CancellationToken::new(),
1785                &SnapshotGate::default(),
1786            )
1787            .await
1788            .expect_err("closed pool must fail");
1789
1790        assert!(
1791            err.to_string().contains("closed"),
1792            "unexpected error: {err}"
1793        );
1794    }
1795
1796    #[rstest]
1797    #[tokio::test]
1798    async fn forwarder_stamps_reconnected_with_shard_id() {
1799        let inner = PoolInner::new(None, TransportBackend::default(), true, 1);
1800        let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
1801        *inner.out_tx.lock() = Some(out_tx);
1802        let (shard_tx, shard_rx) = tokio::sync::mpsc::unbounded_channel();
1803        let mut forwarder = inner
1804            .spawn_forwarder(shard_rx, PRIMARY_SHARD_ID + 1)
1805            .expect("spawn forwarder");
1806
1807        shard_tx
1808            .send(PolymarketWsMessage::Reconnected { shard_id: None })
1809            .unwrap();
1810        shard_tx
1811            .send(market_message("ws_market_best_bid_ask_msg.json"))
1812            .unwrap();
1813        drop(shard_tx);
1814        let outcome = finish_task(
1815            &mut forwarder,
1816            Duration::from_secs(1),
1817            Duration::from_secs(1),
1818        )
1819        .await
1820        .expect("forwarder task");
1821        assert!(matches!(outcome, TaskJoinOutcome::Completed(())));
1822
1823        assert!(
1824            matches!(
1825                out_rx.try_recv().unwrap(),
1826                PolymarketWsMessage::Reconnected { shard_id: Some(id) }
1827                if id == PRIMARY_SHARD_ID + 1
1828            ),
1829            "forwarder must stamp the reconnecting shard"
1830        );
1831        assert!(matches!(
1832            out_rx.try_recv().unwrap(),
1833            PolymarketWsMessage::Market(MarketWsMessage::BestBidAsk(_))
1834        ));
1835        assert!(out_rx.try_recv().is_err());
1836    }
1837}
nautilus_polymarket/websocket/pool.rs

nautilus_polymarket/websocket/
pool.rs
