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

nautilus_polymarket/websocket/
pool.rs
