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

nautilus_polymarket/websocket/
pool.rs
