digdigdig3 0.3.10

//! # WebSocketPool — Thread-safe per-(exchange, account_type) WS pool
//!
//! Many exchanges (Binance, Bybit) have separate WS endpoints for spot vs
//! futures. The pool keys on `(ExchangeId, AccountType)` for that reason.

use dashmap::DashMap;
use std::sync::Arc;

use crate::core::traits::WebSocketConnector;
use crate::core::types::{AccountType, ExchangeId};

/// Thread-safe WebSocket connection pool keyed on `(ExchangeId, AccountType)`.
///
/// Mirrors `ConnectorPool` but for `Arc<dyn WebSocketConnector>` — REST and WS
/// pools are kept separate because they have different lifecycles and
/// per-account-type WS endpoints (e.g. Binance spot vs futures WS URLs differ).
///
/// # Examples
///
/// ```ignore
/// let pool = WebSocketPool::new();
/// let ws = ConnectorFactory::create_websocket(ExchangeId::Binance, AccountType::Spot, false).await?;
/// pool.insert(ExchangeId::Binance, AccountType::Spot, ws);
/// if let Some(ws) = pool.get(ExchangeId::Binance, AccountType::Spot) {
///     ws.connect(AccountType::Spot).await?;
/// }
/// ```
#[derive(Clone)]
pub(crate) struct WebSocketPool {
    sockets: Arc<DashMap<(ExchangeId, AccountType), Arc<dyn WebSocketConnector>>>,
}

impl WebSocketPool {
    /// Create a new empty pool.
    pub(crate) fn new() -> Self {
        Self {
            sockets: Arc::new(DashMap::new()),
        }
    }

    /// Insert a WebSocket connector into the pool.
    ///
    /// If an entry already exists for `(id, account_type)`, it is replaced and
    /// the previous value is returned.
    pub(crate) fn insert(
        &self,
        id: ExchangeId,
        account_type: AccountType,
        socket: Arc<dyn WebSocketConnector>,
    ) -> Option<Arc<dyn WebSocketConnector>> {
        self.sockets.insert((id, account_type), socket)
    }

    /// Get a WebSocket connector by `(exchange, account_type)` (lock-free read).
    pub(crate) fn get(
        &self,
        id: ExchangeId,
        account_type: AccountType,
    ) -> Option<Arc<dyn WebSocketConnector>> {
        self.sockets.get(&(id, account_type)).map(|e| e.value().clone())
    }

    /// Remove a WebSocket connector from the pool.
    pub(crate) fn remove(
        &self,
        id: ExchangeId,
        account_type: AccountType,
    ) -> Option<Arc<dyn WebSocketConnector>> {
        self.sockets.remove(&(id, account_type)).map(|(_, v)| v)
    }

    /// Number of entries in the pool.
    pub(crate) fn len(&self) -> usize {
        self.sockets.len()
    }

    /// Returns `true` if the pool contains no entries.
    pub(crate) fn is_empty(&self) -> bool {
        self.sockets.is_empty()
    }

    /// Remove all entries from the pool.
    pub(crate) fn clear(&self) {
        self.sockets.clear();
    }
}

impl Default for WebSocketPool {
    fn default() -> Self {
        Self::new()
    }
}