Skip to main content

digdigdig3_station/
subscription.rs

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