Skip to main content

brokapi_bybit/core/
ud.rs

1//! The user data core: frames of the private stream (own orders, fills, positions, wallet) to events.
2
3use std::collections::{HashMap, HashSet, VecDeque};
4use std::time::Instant;
5
6use brokapi_bybit_protocol::http::models::{
7    ExecutionRecord, OrderRecord, PositionRecord, WalletRecord,
8};
9use brokapi_bybit_protocol::ws::models::WsIncoming;
10use brokapi_core::{ClientOrderId, Event, SubscriptionId};
11use brokapi_sansio::{
12    BrokerCore, CoreEvent, DesiredSubscription, EventSink, Outbound, SubscriptionKind,
13    SubscriptionSet,
14};
15use tokio::sync::Notify;
16use tracing::{debug, warn};
17
18use super::registry::{Registry, SubState};
19use crate::cid::{cid_from_link_id, foreign_cid};
20use crate::map::{execution, order, portfolio};
21use crate::subs::UdSubKind;
22
23const SEEN_EXECUTIONS_CAP: usize = 20_000;
24const ORDER_IDS_CAP: usize = 20_000;
25
26/// The topic of the own orders of a category.
27pub fn order_topic(category: &str) -> String {
28    format!("order.{category}")
29}
30
31/// The topic of the own fills of a category.
32pub fn execution_topic(category: &str) -> String {
33    format!("execution.{category}")
34}
35
36/// The topic of the positions of a category.
37pub fn position_topic(category: &str) -> String {
38    format!("position.{category}")
39}
40
41/// The topic of the wallet.
42pub const WALLET_TOPIC: &str = "wallet";
43
44/// The sans-io core of the user data connector.
45#[derive(Debug)]
46pub struct UdCore {
47    reg: Registry<UdSubKind>,
48    category: String,
49    seen_executions: HashSet<String>,
50    seen_order: VecDeque<String>,
51    order_ids: HashMap<String, ClientOrderId>,
52    order_ids_order: VecDeque<String>,
53    connects: u32,
54    desired: SubscriptionSet,
55}
56
57impl UdCore {
58    /// An empty core for the category of the account.
59    pub fn new(category: impl Into<String>) -> Self {
60        Self {
61            reg: Registry::default(),
62            category: category.into(),
63            seen_executions: HashSet::new(),
64            seen_order: VecDeque::new(),
65            order_ids: HashMap::new(),
66            order_ids_order: VecDeque::new(),
67            connects: 0,
68            desired: SubscriptionSet::new(),
69        }
70    }
71
72    /// Woken whenever a subscription changes state.
73    pub fn notifier(&self) -> std::sync::Arc<Notify> {
74        self.reg.notifier()
75    }
76
77    /// The state of a subscription.
78    pub fn sub_state(&self, id: SubscriptionId) -> (SubState, Option<String>) {
79        self.reg.state(id).unwrap_or((SubState::Pending, None))
80    }
81
82    /// Registers a subscription for `topics`; `reqs` are (request id, topic) pairs of the requests that are about to
83    /// be sent.
84    pub fn register(
85        &mut self,
86        id: SubscriptionId,
87        kind: UdSubKind,
88        topics: &[String],
89        reqs: &[(String, String)],
90    ) {
91        self.reg.register(id, kind, topics.to_vec());
92        for (r, t) in reqs {
93            self.reg.add_request(r, t);
94        }
95        self.refresh_desired();
96    }
97
98    /// Forgets a subscription. Returns the topics nobody listens to any more.
99    pub fn forget(&mut self, id: SubscriptionId) -> Vec<String> {
100        let orphans = self.reg.forget(id);
101        self.refresh_desired();
102        orphans
103    }
104
105    /// Whether any subscription listens to the topic.
106    pub fn has_topic(&self, topic: &str) -> bool {
107        self.reg.has_topic(topic)
108    }
109
110    fn refresh_desired(&mut self) {
111        let mut subs: Vec<DesiredSubscription> = self
112            .reg
113            .iter()
114            .map(|(id, kind)| DesiredSubscription {
115                id,
116                kind: match kind {
117                    UdSubKind::OrderState => SubscriptionKind::OrderState,
118                    UdSubKind::Portfolio => SubscriptionKind::Portfolio,
119                    UdSubKind::Executions => SubscriptionKind::Executions,
120                },
121            })
122            .collect();
123        subs.sort_by_key(|s| s.id.value());
124        self.desired = SubscriptionSet::with_subscriptions(subs);
125    }
126
127    /// Remembers that the order with this number is the order of `cid`.
128    pub fn learn_order_id(&mut self, order_id: &str, cid: ClientOrderId) {
129        if order_id.is_empty() || cid.is_nil() {
130            return;
131        }
132        if self.order_ids.insert(order_id.to_string(), cid).is_none() {
133            self.order_ids_order.push_back(order_id.to_string());
134            if self.order_ids_order.len() > ORDER_IDS_CAP
135                && let Some(old) = self.order_ids_order.pop_front()
136            {
137                self.order_ids.remove(&old);
138            }
139        }
140    }
141
142    /// The cid of the order with this number, if the core has seen it.
143    pub fn cid_of_order_id(&self, order_id: &str) -> Option<ClientOrderId> {
144        self.order_ids.get(order_id).copied()
145    }
146
147    /// The cid of an order of the exchange: the one in its `orderLinkId`, else the one learned for its number, else a
148    /// stable made-up one (the order is not ours).
149    pub fn resolve_cid(&mut self, order_id: &str, link_id: &str) -> ClientOrderId {
150        if let Some(cid) = cid_from_link_id(Some(link_id)) {
151            self.learn_order_id(order_id, cid);
152            return cid;
153        }
154        self.cid_of_order_id(order_id)
155            .unwrap_or_else(|| foreign_cid(order_id))
156    }
157
158    fn on_connected(&mut self, out: &mut EventSink) {
159        self.connects += 1;
160        if self.connects < 2 {
161            return;
162        }
163        // A new connection: what happened while it was down is not in the stream. The subscriptions are asked for
164        // again, and the readers of orders and fills are told where the gap starts.
165        self.reg.reset();
166        let gap_from = crate::map::ms_to_nanos(chrono::Utc::now().timestamp_millis());
167        for (id, kind) in self.reg.iter().map(|(i, k)| (i, *k)).collect::<Vec<_>>() {
168            if matches!(kind, UdSubKind::OrderState | UdSubKind::Executions) {
169                out.push(CoreEvent {
170                    sub: Some(id),
171                    payload: Event::Reconnected { sub: id, gap_from },
172                });
173            }
174        }
175    }
176
177    fn on_data(
178        &mut self,
179        topic: &str,
180        data: &serde_json::Value,
181        now: Instant,
182        out: &mut EventSink,
183    ) {
184        self.reg.on_data(topic);
185        let subs = self.reg.subs_of(topic);
186        if topic == order_topic(&self.category) {
187            let Some(records) = read::<Vec<OrderRecord>>(data) else {
188                return;
189            };
190            for r in &records {
191                // A conditional order that waits is the business of the stop order list, not of the order stream.
192                // It joins the stream when it fires.
193                if matches!(
194                    r.order_status.as_str(),
195                    "Untriggered" | "Triggered" | "Deactivated"
196                ) {
197                    self.resolve_cid(&r.order_id, &r.order_link_id);
198                    continue;
199                }
200                let cid = self.resolve_cid(&r.order_id, &r.order_link_id);
201                if let Some(snap) = order::order_snapshot(r, cid, &self.category, now) {
202                    for id in &subs {
203                        out.push(CoreEvent {
204                            sub: Some(*id),
205                            payload: Event::Order(snap.clone()),
206                        });
207                    }
208                }
209            }
210        } else if topic == execution_topic(&self.category) {
211            let Some(records) = read::<Vec<ExecutionRecord>>(data) else {
212                return;
213            };
214            for r in &records {
215                // A fill the stream repeats is told once.
216                if !self.seen_executions.insert(r.exec_id.clone()) {
217                    continue;
218                }
219                self.seen_order.push_back(r.exec_id.clone());
220                if self.seen_order.len() > SEEN_EXECUTIONS_CAP
221                    && let Some(old) = self.seen_order.pop_front()
222                {
223                    self.seen_executions.remove(&old);
224                }
225                let cid = cid_from_link_id(Some(&r.order_link_id))
226                    .inspect(|c| self.learn_order_id(&r.order_id, *c))
227                    .or_else(|| self.cid_of_order_id(&r.order_id));
228                if let Some(e) = execution::execution(r, cid, &self.category, now) {
229                    for id in &subs {
230                        out.push(CoreEvent {
231                            sub: Some(*id),
232                            payload: Event::Execution(e.clone()),
233                        });
234                    }
235                }
236            }
237        } else if topic == position_topic(&self.category) {
238            let Some(records) = read::<Vec<PositionRecord>>(data) else {
239                return;
240            };
241            let positions: Vec<_> = records
242                .iter()
243                .map(|r| portfolio::position(r, &self.category, now))
244                .collect();
245            if !positions.is_empty() {
246                for id in &subs {
247                    out.push(CoreEvent {
248                        sub: Some(*id),
249                        payload: Event::Positions(positions.clone()),
250                    });
251                }
252            }
253        } else if topic == WALLET_TOPIC {
254            let Some(records) = read::<Vec<WalletRecord>>(data) else {
255                return;
256            };
257            let balances = portfolio::balances(&records, now);
258            if !balances.is_empty() {
259                for id in &subs {
260                    out.push(CoreEvent {
261                        sub: Some(*id),
262                        payload: Event::Balances(balances.clone()),
263                    });
264                }
265            }
266        }
267    }
268}
269
270fn read<T: serde::de::DeserializeOwned>(data: &serde_json::Value) -> Option<T> {
271    match serde_json::from_value(data.clone()) {
272        Ok(v) => Some(v),
273        Err(e) => {
274            warn!("a data frame cannot be read: {e}");
275            None
276        }
277    }
278}
279
280impl BrokerCore for UdCore {
281    type Frame = WsIncoming;
282
283    fn on_frame(&mut self, frame: Self::Frame, now: Instant, out: &mut EventSink) {
284        match frame {
285            WsIncoming::Connected => self.on_connected(out),
286            WsIncoming::Data { topic, data, .. } => {
287                if !self.reg.has_topic(&topic) {
288                    debug!("a frame of an unknown topic {topic} is dropped");
289                    return;
290                }
291                self.on_data(&topic, &data, now, out);
292            }
293            WsIncoming::Ack {
294                op,
295                req_id: Some(req_id),
296                success,
297                message,
298            } if op == "subscribe" => {
299                self.reg.on_ack(&req_id, success, &message);
300            }
301            WsIncoming::Ack {
302                op,
303                success: false,
304                message,
305                ..
306            } if op == "auth" => {
307                self.reg
308                    .fail_pending(&format!("the private stream refused the login: {message}"));
309            }
310            WsIncoming::Ack { .. } | WsIncoming::Pong | WsIncoming::Unknown(_) => {}
311        }
312    }
313
314    fn poll_transmit(&mut self) -> Option<Outbound> {
315        None
316    }
317
318    fn poll_timeout(&self) -> Option<Instant> {
319        None
320    }
321
322    fn on_timeout(&mut self, _now: Instant, _out: &mut EventSink) {}
323
324    fn apply(&mut self, desired: &SubscriptionSet) {
325        self.desired = desired.clone();
326    }
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332    use brokapi_bybit_protocol::ws::models::parse_incoming;
333    use serde_json::json;
334    use uuid::Uuid;
335
336    fn frame(v: serde_json::Value) -> WsIncoming {
337        parse_incoming(&v.to_string())
338    }
339
340    fn core_with(kind: UdSubKind, id: u64, topics: &[&str]) -> (UdCore, SubscriptionId) {
341        let mut c = UdCore::new("linear");
342        let sid = SubscriptionId::new(id);
343        let topics: Vec<String> = topics.iter().map(|t| t.to_string()).collect();
344        let reqs: Vec<(String, String)> = topics
345            .iter()
346            .enumerate()
347            .map(|(i, t)| (format!("r{i}"), t.clone()))
348            .collect();
349        c.register(sid, kind, &topics, &reqs);
350        (c, sid)
351    }
352
353    fn order_frame(status: &str, link: &str) -> WsIncoming {
354        frame(
355            json!({"id":"1","topic":"order.linear","creationTime":5,"data":[
356            {"category":"linear","orderId":"77","orderLinkId":link,"symbol":"BTCUSDT","side":"Buy","orderType":"Limit",
357             "orderStatus":status,"qty":"0.2","price":"100","cumExecQty":"0","updatedTime":"1700000000000"}]}),
358        )
359    }
360
361    #[test]
362    fn the_cid_comes_from_the_link_id_and_is_learned_for_the_number() {
363        let (mut c, id) = core_with(UdSubKind::OrderState, 1001, &["order.linear"]);
364        let cid = Uuid::new_v4();
365        let mut out = EventSink::new();
366        let now = Instant::now();
367        c.on_frame(order_frame("New", &cid.to_string()), now, &mut out);
368        // Later frames of the same order may come without the link id: the number is known.
369        c.on_frame(order_frame("Cancelled", ""), now, &mut out);
370        let events = out.drain();
371        assert_eq!(events.len(), 2);
372        for e in &events {
373            let Event::Order(s) = &e.payload else {
374                panic!("an order")
375            };
376            assert_eq!(s.cid, cid);
377            assert_eq!(e.sub, Some(id));
378        }
379    }
380
381    #[test]
382    fn an_order_that_is_not_ours_gets_a_stable_made_up_cid() {
383        let (mut c, _) = core_with(UdSubKind::OrderState, 1, &["order.linear"]);
384        let a = c.resolve_cid("9", "");
385        assert_eq!(a, c.resolve_cid("9", "from the app"));
386        assert_ne!(a, c.resolve_cid("10", ""));
387    }
388
389    #[test]
390    fn a_conditional_order_that_waits_is_not_in_the_order_stream_but_one_that_fired_is() {
391        let (mut c, _) = core_with(UdSubKind::OrderState, 1, &["order.linear"]);
392        let cid = Uuid::new_v4();
393        let mut out = EventSink::new();
394        let now = Instant::now();
395        c.on_frame(order_frame("Untriggered", &cid.to_string()), now, &mut out);
396        assert!(out.drain().is_empty());
397        c.on_frame(order_frame("New", ""), now, &mut out);
398        let events = out.drain();
399        assert_eq!(events.len(), 1);
400        let Event::Order(s) = &events[0].payload else {
401            panic!("an order")
402        };
403        assert_eq!(s.cid, cid, "the number was learned while it waited");
404    }
405
406    #[test]
407    fn a_fill_is_told_once_with_the_cid_of_its_order() {
408        let (mut c, id) = core_with(UdSubKind::Executions, 2, &["execution.linear"]);
409        let cid = Uuid::new_v4();
410        let mut out = EventSink::new();
411        let now = Instant::now();
412        let fill = |exec_id: &str, kind: &str| {
413            frame(json!({"topic":"execution.linear","creationTime":1,"data":[
414                {"category":"linear","symbol":"BTCUSDT","orderId":"77","orderLinkId":cid.to_string(),"side":"Buy",
415                 "execId":exec_id,"execPrice":"100","execQty":"0.1","execFee":"0.01","execType":kind,"execTime":"1700000000000"}]}))
416        };
417        c.on_frame(fill("e1", "Trade"), now, &mut out);
418        c.on_frame(fill("e1", "Trade"), now, &mut out);
419        c.on_frame(fill("e2", "Funding"), now, &mut out);
420        let events = out.drain();
421        assert_eq!(
422            events.len(),
423            1,
424            "the repeat and the funding payment are not told"
425        );
426        let Event::Execution(e) = &events[0].payload else {
427            panic!("a fill")
428        };
429        assert_eq!((e.cid, events[0].sub), (Some(cid), Some(id)));
430    }
431
432    #[test]
433    fn positions_and_the_wallet_go_to_the_portfolio_subscription() {
434        let (mut c, id) = core_with(UdSubKind::Portfolio, 3, &["position.linear", "wallet"]);
435        let mut out = EventSink::new();
436        let now = Instant::now();
437        c.on_frame(
438            frame(json!({"topic":"position.linear","creationTime":1,"data":[
439                {"category":"linear","symbol":"BTCUSDT","side":"Sell","size":"0.5","entryPrice":"81000","markPrice":"80000"}]})),
440            now,
441            &mut out,
442        );
443        c.on_frame(
444            frame(json!({"topic":"wallet","creationTime":1,"data":[{"accountType":"UNIFIED","coin":[
445                {"coin":"USDT","equity":"100","walletBalance":"100","totalOrderIM":"","totalPositionIM":"20"}]}]})),
446            now,
447            &mut out,
448        );
449        let events = out.drain();
450        assert_eq!(events.len(), 2);
451        let Event::Positions(p) = &events[0].payload else {
452            panic!("positions")
453        };
454        assert_eq!(
455            (p[0].qty.to_string(), p[0].avg_price.unwrap().to_string()),
456            ("-0.5".into(), "81000".into())
457        );
458        let Event::Balances(b) = &events[1].payload else {
459            panic!("balances")
460        };
461        assert_eq!(
462            (b[0].blocked.to_string(), b[0].available.to_string()),
463            ("20".into(), "80".into())
464        );
465        assert!(events.iter().all(|e| e.sub == Some(id)));
466    }
467
468    #[test]
469    fn a_reconnect_tells_the_order_and_fill_readers_where_the_gap_starts() {
470        let mut c = UdCore::new("linear");
471        let (o, p) = (SubscriptionId::new(1), SubscriptionId::new(2));
472        c.register(o, UdSubKind::OrderState, &["order.linear".into()], &[]);
473        c.register(p, UdSubKind::Portfolio, &["wallet".into()], &[]);
474        let mut out = EventSink::new();
475        c.on_frame(WsIncoming::Connected, Instant::now(), &mut out);
476        assert!(out.drain().is_empty());
477        c.on_frame(WsIncoming::Connected, Instant::now(), &mut out);
478        let events = out.drain();
479        assert_eq!(events.len(), 1);
480        assert!(matches!(events[0].payload, Event::Reconnected { sub, .. } if sub == o));
481        assert_eq!(c.sub_state(o).0, SubState::Pending);
482    }
483
484    #[test]
485    fn a_refused_login_fails_the_subscriptions_that_wait() {
486        let (mut c, id) = core_with(UdSubKind::OrderState, 1, &["order.linear"]);
487        let mut out = EventSink::new();
488        c.on_frame(
489            WsIncoming::Ack {
490                op: "auth".into(),
491                req_id: None,
492                success: false,
493                message: "Params Error".into(),
494            },
495            Instant::now(),
496            &mut out,
497        );
498        let (state, reason) = c.sub_state(id);
499        assert_eq!(state, SubState::Failed);
500        assert!(reason.unwrap().contains("Params Error"));
501        assert!(c.poll_transmit().is_none() && c.poll_timeout().is_none());
502    }
503}