df_hyperliquid_core/live/store/
interest.rs1use std::{
6 collections::{BTreeMap, HashMap, HashSet},
7 sync::mpsc::{self, Receiver, Sender},
8};
9
10use crate::live::{Limits, Refused, Subscription};
11use crate::markets::Interval;
12use crate::utils::Address;
13
14#[derive(Debug, Clone, PartialEq, Eq, Hash)]
17pub enum Interest {
18 Book {
20 coin: String,
21 },
22 Trades {
24 coin: String,
25 },
26 Candles {
28 coin: String,
29 interval: Interval,
30 },
31 Mids {
33 dex: String,
34 },
35 Context {
37 coin: String,
38 },
39 PerpContexts,
41 SpotContexts,
43 State {
46 user: Address,
47 },
48 Orders {
50 user: Address,
51 dex: String,
52 },
53 Fills {
54 user: Address,
55 },
56 Fundings {
57 user: Address,
58 },
59 Ledger {
61 user: Address,
62 },
63}
64
65impl Interest {
66 #[must_use]
68 pub fn subscriptions(&self) -> Vec<Subscription> {
69 match self {
70 Interest::Book { coin } => {
71 vec![Subscription::L2Book { coin: coin.clone(), n_sig_figs: None, mantissa: None, fast: false }]
72 }
73 Interest::Trades { coin } => vec![Subscription::Trades { coin: coin.clone() }],
74 Interest::Candles { coin, interval } => {
75 vec![Subscription::Candle { coin: coin.clone(), interval: *interval }]
76 }
77 Interest::Mids { dex } => vec![Subscription::AllMids { dex: dex.clone() }],
78 Interest::Context { coin } => vec![Subscription::ActiveAssetCtx { coin: coin.clone() }],
79 Interest::PerpContexts => vec![Subscription::AllDexsAssetCtxs],
80 Interest::SpotContexts => vec![Subscription::SpotAssetCtxs],
81 Interest::State { user } => vec![
82 Subscription::AllDexsClearinghouseState { user: *user },
83 Subscription::SpotState { user: *user, ignore_portfolio_margin: false },
84 Subscription::WebData3 { user: *user },
85 ],
86 Interest::Orders { user, dex } => vec![Subscription::OpenOrders { user: *user, dex: dex.clone() }],
87 Interest::Fills { user } => vec![Subscription::UserFills { user: *user, aggregate_by_time: false }],
88 Interest::Fundings { user } => vec![Subscription::UserFundings { user: *user }],
89 Interest::Ledger { user } => vec![Subscription::UserNonFundingLedgerUpdates { user: *user }],
90 }
91 }
92
93 #[must_use]
95 pub const fn user(&self) -> Option<Address> {
96 match self {
97 Interest::State { user }
98 | Interest::Orders { user, .. }
99 | Interest::Fills { user }
100 | Interest::Fundings { user }
101 | Interest::Ledger { user } => Some(*user),
102 _ => None,
103 }
104 }
105}
106
107#[must_use = "dropping a handle ends its interest"]
110#[derive(Debug)]
111pub struct Handle {
112 id: u64,
113 interest: Interest,
114 released: Sender<u64>,
115}
116
117impl Handle {
118 #[must_use]
119 pub const fn interest(&self) -> &Interest {
120 &self.interest
121 }
122}
123
124impl Drop for Handle {
125 fn drop(&mut self) {
126 let _ = self.released.send(self.id);
128 }
129}
130
131#[derive(Debug, Clone, Default, PartialEq, Eq)]
134pub struct Changes {
135 pub unsubscribe: Vec<Subscription>,
136 pub subscribe: Vec<Subscription>,
137}
138
139impl Changes {
140 #[must_use]
141 pub fn is_empty(&self) -> bool {
142 self.unsubscribe.is_empty() && self.subscribe.is_empty()
143 }
144}
145
146#[derive(Debug)]
148pub(crate) struct Interests {
149 limits: Limits,
150 next: u64,
151 held: BTreeMap<u64, Interest>,
152 counts: HashMap<Subscription, usize>,
154 order: Vec<Subscription>,
156 active: Vec<Subscription>,
158 users: Vec<Address>,
160 released: Receiver<u64>,
161 sender: Sender<u64>,
162}
163
164impl Interests {
165 pub(crate) fn new(limits: Limits) -> Interests {
166 let (sender, released) = mpsc::channel();
167 Interests {
168 limits,
169 next: 1,
170 held: BTreeMap::new(),
171 counts: HashMap::new(),
172 order: Vec::new(),
173 active: Vec::new(),
174 users: Vec::new(),
175 released,
176 sender,
177 }
178 }
179
180 pub(crate) fn set_limits(&mut self, limits: Limits) {
181 self.limits = limits;
182 }
183
184 pub(crate) fn register(&mut self, interest: Interest) -> Result<Handle, Refused> {
187 let needed = interest.subscriptions();
188 let mut new: Vec<&Subscription> = Vec::new();
189 for subscription in &needed {
190 subscription.check()?;
191 if !self.counts.contains_key(subscription) && !new.contains(&subscription) {
192 new.push(subscription);
193 }
194 }
195 if self.order.len().saturating_add(new.len()) > self.limits.subscriptions {
196 return Err(Refused::Subscriptions { limit: self.limits.subscriptions });
197 }
198 if let Some(user) = interest.user()
199 && !self.users.contains(&user)
200 {
201 if self.users.len() >= self.limits.users {
202 return Err(Refused::Users { limit: self.limits.users });
203 }
204 self.users.push(user);
205 }
206 for subscription in needed {
207 let count = self.counts.entry(subscription.clone()).or_insert(0);
208 if *count == 0 {
209 self.order.push(subscription);
210 }
211 *count = count.saturating_add(1);
212 }
213 let id = self.next;
214 self.next = self.next.saturating_add(1);
215 self.held.insert(id, interest.clone());
216 Ok(Handle { id, interest, released: self.sender.clone() })
217 }
218
219 pub(crate) fn drain(&mut self) -> bool {
221 let mut any = false;
222 while let Ok(id) = self.released.try_recv() {
223 let Some(interest) = self.held.remove(&id) else { continue };
224 any = true;
225 for subscription in interest.subscriptions() {
226 let Some(count) = self.counts.get_mut(&subscription) else { continue };
227 *count = count.saturating_sub(1);
228 if *count == 0 {
229 self.counts.remove(&subscription);
230 self.order.retain(|wanted| *wanted != subscription);
231 }
232 }
233 }
234 let order = &self.order;
235 self.users.retain(|user| order.iter().any(|wanted| wanted.user() == Some(*user)));
236 any
237 }
238
239 pub(crate) fn wants(&self, subscription: &Subscription) -> bool {
241 self.counts.contains_key(subscription)
242 }
243
244 pub(crate) fn wanted(&self) -> &[Subscription] {
246 &self.order
247 }
248
249 pub(crate) fn changes(&mut self) -> Changes {
251 let held: HashSet<&Subscription> = self.active.iter().collect();
252 let unsubscribe = self.active.iter().filter(|held| !self.counts.contains_key(*held)).cloned().collect();
253 let subscribe = self.order.iter().filter(|wanted| !held.contains(wanted)).cloned().collect();
254 self.active.clone_from(&self.order);
255 Changes { unsubscribe, subscribe }
256 }
257
258 pub(crate) fn refused(&mut self, subscription: &Subscription) {
260 self.active.retain(|held| held != subscription);
261 }
262
263 pub(crate) fn users(&self) -> &[Address] {
265 &self.users
266 }
267}
268
269#[cfg(test)]
270#[path = "tests/interest.rs"]
271mod tests;