brokapi_bybit/core/
registry.rs

1//! Which subscription listens to which topic, and which topics the exchange has confirmed.
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use brokapi_core::SubscriptionId;
7use tokio::sync::Notify;
8
9/// Where a subscription is with the exchange.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum SubState {
12    /// Asked for, no answer yet.
13    Pending,
14    /// The exchange confirmed it, or data came.
15    Active,
16    /// The exchange refused it.
17    Failed,
18}
19
20#[derive(Debug)]
21struct TopicEntry {
22    state: SubState,
23    reason: Option<String>,
24    subs: Vec<SubscriptionId>,
25}
26
27/// The subscriptions of a core, their topics and the state of each topic.
28///
29/// Several subscriptions may listen to one topic (two bots, or a quote stream and a book of depth 1): the exchange is
30/// asked once, the frames go to all of them.
31#[derive(Debug)]
32pub struct Registry<K> {
33    subs: HashMap<SubscriptionId, (K, Vec<String>)>,
34    topics: HashMap<String, TopicEntry>,
35    /// `req_id` -> topic, for the answers of the exchange.
36    reqs: HashMap<String, String>,
37    ack: Arc<Notify>,
38}
39
40impl<K> Default for Registry<K> {
41    fn default() -> Self {
42        Self {
43            subs: HashMap::new(),
44            topics: HashMap::new(),
45            reqs: HashMap::new(),
46            ack: Arc::new(Notify::new()),
47        }
48    }
49}
50
51impl<K> Registry<K> {
52    /// Woken whenever a topic changes state. The connector waits on it for the confirmation of a new subscription.
53    pub fn notifier(&self) -> Arc<Notify> {
54        Arc::clone(&self.ack)
55    }
56
57    /// Registers a subscription for `topics`. A topic that already has subscriptions keeps its state.
58    pub fn register(&mut self, id: SubscriptionId, kind: K, topics: Vec<String>) {
59        for t in &topics {
60            self.topics
61                .entry(t.clone())
62                .or_insert(TopicEntry {
63                    state: SubState::Pending,
64                    reason: None,
65                    subs: Vec::new(),
66                })
67                .subs
68                .push(id);
69        }
70        self.subs.insert(id, (kind, topics));
71    }
72
73    /// Remembers that the answer carrying `req_id` is about `topic`.
74    pub fn add_request(&mut self, req_id: &str, topic: &str) {
75        self.reqs.insert(req_id.to_string(), topic.to_string());
76    }
77
78    /// Forgets a subscription. Returns the topics nobody listens to any more.
79    pub fn forget(&mut self, id: SubscriptionId) -> Vec<String> {
80        let mut orphans = Vec::new();
81        if let Some((_, topics)) = self.subs.remove(&id) {
82            for t in topics {
83                if let Some(e) = self.topics.get_mut(&t) {
84                    e.subs.retain(|s| *s != id);
85                    if e.subs.is_empty() {
86                        self.topics.remove(&t);
87                        self.reqs.retain(|_, v| *v != t);
88                        orphans.push(t);
89                    }
90                }
91            }
92        }
93        orphans
94    }
95
96    /// The kind of a subscription.
97    pub fn kind(&self, id: SubscriptionId) -> Option<&K> {
98        self.subs.get(&id).map(|(k, _)| k)
99    }
100
101    /// The topics of a subscription.
102    pub fn topics_of(&self, id: SubscriptionId) -> &[String] {
103        self.subs.get(&id).map(|(_, t)| t.as_slice()).unwrap_or(&[])
104    }
105
106    /// The subscriptions that listen to a topic.
107    pub fn subs_of(&self, topic: &str) -> Vec<SubscriptionId> {
108        self.topics
109            .get(topic)
110            .map(|e| e.subs.clone())
111            .unwrap_or_default()
112    }
113
114    /// Whether anyone listens to the topic.
115    pub fn has_topic(&self, topic: &str) -> bool {
116        self.topics.contains_key(topic)
117    }
118
119    /// Whether the exchange has confirmed the topic.
120    pub fn topic_active(&self, topic: &str) -> bool {
121        self.topics
122            .get(topic)
123            .is_some_and(|e| e.state == SubState::Active)
124    }
125
126    /// The state of a subscription: failed when any topic was refused, active when all are, else pending. With the
127    /// reason of the refusal.
128    pub fn state(&self, id: SubscriptionId) -> Option<(SubState, Option<String>)> {
129        let (_, topics) = self.subs.get(&id)?;
130        let mut all_active = true;
131        for t in topics {
132            match self.topics.get(t) {
133                Some(e) if e.state == SubState::Failed => {
134                    return Some((SubState::Failed, e.reason.clone()));
135                }
136                Some(e) if e.state == SubState::Active => {}
137                _ => all_active = false,
138            }
139        }
140        Some((
141            if all_active {
142                SubState::Active
143            } else {
144                SubState::Pending
145            },
146            None,
147        ))
148    }
149
150    /// An answer of the exchange to a subscription request.
151    pub fn on_ack(&mut self, req_id: &str, success: bool, message: &str) {
152        let Some(topic) = self.reqs.get(req_id).cloned() else {
153            return;
154        };
155        if let Some(e) = self.topics.get_mut(&topic) {
156            if success {
157                e.state = SubState::Active;
158                e.reason = None;
159            } else {
160                e.state = SubState::Failed;
161                e.reason = Some(if message.is_empty() {
162                    "refused by the exchange".to_string()
163                } else {
164                    message.to_string()
165                });
166            }
167            self.ack.notify_waiters();
168        }
169    }
170
171    /// Data came on a topic: it is open, whatever the answer said.
172    pub fn on_data(&mut self, topic: &str) {
173        if let Some(e) = self.topics.get_mut(topic)
174            && e.state != SubState::Active
175        {
176            e.state = SubState::Active;
177            e.reason = None;
178            self.ack.notify_waiters();
179        }
180    }
181
182    /// The login was refused: every topic still waiting fails with the reason.
183    pub fn fail_pending(&mut self, reason: &str) {
184        for e in self.topics.values_mut() {
185            if e.state == SubState::Pending {
186                e.state = SubState::Failed;
187                e.reason = Some(reason.to_string());
188            }
189        }
190        self.ack.notify_waiters();
191    }
192
193    /// A connection was lost and made again: every topic is asked for again and waits for its answer.
194    pub fn reset(&mut self) {
195        for e in self.topics.values_mut() {
196            e.state = SubState::Pending;
197            e.reason = None;
198        }
199    }
200
201    /// Every subscription with its kind.
202    pub fn iter(&self) -> impl Iterator<Item = (SubscriptionId, &K)> {
203        self.subs.iter().map(|(id, (k, _))| (*id, k))
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210
211    fn id(n: u64) -> SubscriptionId {
212        SubscriptionId::new(n)
213    }
214
215    #[test]
216    fn a_subscription_with_two_topics_is_active_when_both_are() {
217        let mut r: Registry<()> = Registry::default();
218        r.register(id(1), (), vec!["position.linear".into(), "wallet".into()]);
219        r.add_request("a", "position.linear");
220        r.add_request("b", "wallet");
221        assert_eq!(r.state(id(1)).unwrap().0, SubState::Pending);
222        r.on_ack("a", true, "");
223        assert_eq!(r.state(id(1)).unwrap().0, SubState::Pending);
224        r.on_ack("b", true, "");
225        assert_eq!(r.state(id(1)).unwrap().0, SubState::Active);
226        r.on_ack("b", false, "denied");
227        assert_eq!(
228            r.state(id(1)),
229            Some((SubState::Failed, Some("denied".into())))
230        );
231        r.reset();
232        assert_eq!(r.state(id(1)).unwrap().0, SubState::Pending);
233    }
234
235    #[test]
236    fn a_topic_is_dropped_when_its_last_listener_goes() {
237        let mut r: Registry<()> = Registry::default();
238        r.register(id(1), (), vec!["t".into()]);
239        r.register(id(2), (), vec!["t".into()]);
240        r.add_request("x", "t");
241        assert_eq!(r.subs_of("t"), vec![id(1), id(2)]);
242        assert!(r.forget(id(1)).is_empty());
243        assert!(r.has_topic("t"));
244        assert_eq!(r.forget(id(2)), vec!["t".to_string()]);
245        assert!(!r.has_topic("t"));
246        // An answer to a request nobody asked for any more is ignored.
247        r.on_ack("x", true, "");
248    }
249
250    #[test]
251    fn data_is_the_proof_that_a_topic_is_open_and_a_failed_login_fails_the_waiting_ones() {
252        let mut r: Registry<()> = Registry::default();
253        r.register(id(1), (), vec!["t".into()]);
254        r.register(id(2), (), vec!["u".into()]);
255        r.on_data("t");
256        assert_eq!(r.state(id(1)).unwrap().0, SubState::Active);
257        r.fail_pending("bad key");
258        assert_eq!(r.state(id(1)).unwrap().0, SubState::Active);
259        assert_eq!(
260            r.state(id(2)),
261            Some((SubState::Failed, Some("bad key".into())))
262        );
263    }
264}
brokapi_bybit/core/registry.rs

brokapi_bybit/core/
registry.rs
