digdigdig3_station/
subscription.rs

1use digdigdig3::core::types::{AccountType, ExchangeId};
2use digdigdig3::core::websocket::KlineInterval;
3
4use crate::data::{
5    AggTradePoint, BarPoint, FundingRatePoint, LiquidationPoint, MarkPricePoint, ObSnapshotPoint,
6    OpenInterestPoint, TickerPoint, TradePoint,
7};
8use crate::series::{Kind, SeriesKey};
9
10/// User-facing stream class to request in a `SubscriptionSet`.
11///
12/// `Kline` carries a typed `KlineInterval` (e.g. `KlineInterval::new("1m")`).
13/// All other variants are parameterless.
14#[derive(Debug, Clone, PartialEq, Eq, Hash)]
15pub enum Stream {
16    Ticker,
17    Trade,
18    Orderbook,
19    Kline(KlineInterval),
20    MarkPrice,
21    FundingRate,
22    OpenInterest,
23    Liquidation,
24    AggTrade,
25}
26
27impl Stream {
28    pub(crate) fn to_kind(&self) -> Kind {
29        match self {
30            Stream::Trade => Kind::Trade,
31            Stream::AggTrade => Kind::AggTrade,
32            Stream::Kline(iv) => Kind::Kline(iv.clone()),
33            Stream::Ticker => Kind::Ticker,
34            Stream::Orderbook => Kind::Orderbook,
35            Stream::MarkPrice => Kind::MarkPrice,
36            Stream::FundingRate => Kind::FundingRate,
37            Stream::OpenInterest => Kind::OpenInterest,
38            Stream::Liquidation => Kind::Liquidation,
39        }
40    }
41}
42
43#[derive(Debug, Clone)]
44pub(crate) struct Entry {
45    pub(crate) exchange: ExchangeId,
46    pub(crate) symbol: String,
47    pub(crate) account_type: AccountType,
48    pub(crate) streams: Vec<Stream>,
49}
50
51/// Declarative subscription request — built up fluently, consumed by
52/// [`crate::Station::subscribe`].
53#[derive(Debug, Default, Clone)]
54pub struct SubscriptionSet {
55    pub(crate) entries: Vec<Entry>,
56}
57
58impl SubscriptionSet {
59    pub fn new() -> Self { Self::default() }
60
61    pub fn add(
62        mut self,
63        exchange: ExchangeId,
64        symbol: impl Into<String>,
65        account_type: AccountType,
66        streams: impl IntoIterator<Item = Stream>,
67    ) -> Self {
68        self.entries.push(Entry {
69            exchange,
70            symbol: symbol.into(),
71            account_type,
72            streams: streams.into_iter().collect(),
73        });
74        self
75    }
76
77    pub fn len(&self) -> usize { self.entries.len() }
78    pub fn is_empty(&self) -> bool { self.entries.is_empty() }
79}
80
81/// Events forwarded to consumers. One variant per market-data class.
82#[derive(Debug, Clone)]
83pub enum Event {
84    Trade {
85        exchange: ExchangeId,
86        symbol: String,
87        point: TradePoint,
88    },
89    AggTrade {
90        exchange: ExchangeId,
91        symbol: String,
92        point: AggTradePoint,
93    },
94    Bar {
95        exchange: ExchangeId,
96        symbol: String,
97        timeframe: KlineInterval,
98        point: BarPoint,
99    },
100    Ticker {
101        exchange: ExchangeId,
102        symbol: String,
103        point: TickerPoint,
104    },
105    OrderbookSnapshot {
106        exchange: ExchangeId,
107        symbol: String,
108        point: ObSnapshotPoint,
109    },
110    MarkPrice {
111        exchange: ExchangeId,
112        symbol: String,
113        point: MarkPricePoint,
114    },
115    FundingRate {
116        exchange: ExchangeId,
117        symbol: String,
118        point: FundingRatePoint,
119    },
120    OpenInterest {
121        exchange: ExchangeId,
122        symbol: String,
123        point: OpenInterestPoint,
124    },
125    Liquidation {
126        exchange: ExchangeId,
127        symbol: String,
128        point: LiquidationPoint,
129    },
130}
131
132impl Event {
133    pub fn exchange(&self) -> ExchangeId {
134        match self {
135            Event::Trade { exchange, .. } | Event::AggTrade { exchange, .. } |
136            Event::Bar { exchange, .. } | Event::Ticker { exchange, .. } |
137            Event::OrderbookSnapshot { exchange, .. } | Event::MarkPrice { exchange, .. } |
138            Event::FundingRate { exchange, .. } | Event::OpenInterest { exchange, .. } |
139            Event::Liquidation { exchange, .. } => *exchange,
140        }
141    }
142    pub fn symbol(&self) -> &str {
143        match self {
144            Event::Trade { symbol, .. } | Event::AggTrade { symbol, .. } |
145            Event::Bar { symbol, .. } | Event::Ticker { symbol, .. } |
146            Event::OrderbookSnapshot { symbol, .. } | Event::MarkPrice { symbol, .. } |
147            Event::FundingRate { symbol, .. } | Event::OpenInterest { symbol, .. } |
148            Event::Liquidation { symbol, .. } => symbol,
149        }
150    }
151
152    /// Replace the symbol label on this event in-place.
153    ///
154    /// Used by `Station::subscribe` so each `SubscriptionHandle` sees the
155    /// user-input symbol it passed in `SubscriptionSet::add(...)`, regardless
156    /// of which other consumer first established the underlying multiplex.
157    /// The routing key (raw exchange-native) is unaffected; this only changes
158    /// the cosmetic label that `Event.symbol()` returns to the consumer.
159    pub(crate) fn set_symbol(&mut self, new_symbol: String) {
160        match self {
161            Event::Trade { symbol, .. }
162            | Event::AggTrade { symbol, .. }
163            | Event::Bar { symbol, .. }
164            | Event::Ticker { symbol, .. }
165            | Event::OrderbookSnapshot { symbol, .. }
166            | Event::MarkPrice { symbol, .. }
167            | Event::FundingRate { symbol, .. }
168            | Event::OpenInterest { symbol, .. }
169            | Event::Liquidation { symbol, .. } => *symbol = new_symbol,
170        }
171    }
172    pub fn timestamp_ms(&self) -> i64 {
173        use crate::series::DataPoint;
174        match self {
175            Event::Trade { point, .. } => point.timestamp_ms(),
176            Event::AggTrade { point, .. } => point.timestamp_ms(),
177            Event::Bar { point, .. } => point.timestamp_ms(),
178            Event::Ticker { point, .. } => point.timestamp_ms(),
179            Event::OrderbookSnapshot { point, .. } => point.timestamp_ms(),
180            Event::MarkPrice { point, .. } => point.timestamp_ms(),
181            Event::FundingRate { point, .. } => point.timestamp_ms(),
182            Event::OpenInterest { point, .. } => point.timestamp_ms(),
183            Event::Liquidation { point, .. } => point.timestamp_ms(),
184        }
185    }
186}
187
188/// RAII handle returned from `Station::subscribe`. Dropping releases the
189/// per-StreamKey consumer ref count; when count hits zero the multiplexer
190/// shuts down.
191pub struct SubscriptionHandle {
192    pub(crate) rx: tokio::sync::mpsc::UnboundedReceiver<Event>,
193    pub(crate) _refs: Vec<MultiplexRef>,
194}
195
196impl std::fmt::Debug for SubscriptionHandle {
197    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
198        f.debug_struct("SubscriptionHandle").finish()
199    }
200}
201
202impl SubscriptionHandle {
203    pub async fn recv(&mut self) -> Option<Event> {
204        self.rx.recv().await
205    }
206}
207
208pub(crate) struct MultiplexRef {
209    pub(crate) station: std::sync::Weak<crate::station::StationInner>,
210    pub(crate) key: SeriesKey,
211}
212
213impl Drop for MultiplexRef {
214    fn drop(&mut self) {
215        if let Some(inner) = self.station.upgrade() {
216            inner.release_consumer(&self.key);
217        }
218    }
219}
digdigdig3_station/subscription.rs

digdigdig3_station/
subscription.rs
