1use 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
26pub fn order_topic(category: &str) -> String {
28 format!("order.{category}")
29}
30
31pub fn execution_topic(category: &str) -> String {
33 format!("execution.{category}")
34}
35
36pub fn position_topic(category: &str) -> String {
38 format!("position.{category}")
39}
40
41pub const WALLET_TOPIC: &str = "wallet";
43
44#[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 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 pub fn notifier(&self) -> std::sync::Arc<Notify> {
74 self.reg.notifier()
75 }
76
77 pub fn sub_state(&self, id: SubscriptionId) -> (SubState, Option<String>) {
79 self.reg.state(id).unwrap_or((SubState::Pending, None))
80 }
81
82 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 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 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 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 pub fn cid_of_order_id(&self, order_id: &str) -> Option<ClientOrderId> {
144 self.order_ids.get(order_id).copied()
145 }
146
147 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 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 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 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 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}