Skip to main content

brokapi_bybit/
pump.rs

1//! The task that feeds frames of the stream to a core and the events of the core to the emitter.
2
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6use brokapi_bybit_protocol::ws::models::WsIncoming;
7use brokapi_runtime::Emitter;
8use brokapi_sansio::{BrokerCore, EventSink};
9use kanal::AsyncReceiver;
10use tokio::sync::Mutex;
11use tokio_util::sync::CancellationToken;
12
13const FLUSH_EVERY: Duration = Duration::from_millis(20);
14
15/// Reads frames from `rx` until `cancel`, runs them through `core`, and gives the events to `emitter`. `after` runs
16/// after every frame with the core locked out, so a connector can act on what the core noted (a book to ask for
17/// again).
18pub async fn run_pump<C, F, Fut>(
19    core: Arc<Mutex<C>>,
20    emitter: Arc<Emitter>,
21    rx: AsyncReceiver<WsIncoming>,
22    cancel: CancellationToken,
23    after: F,
24) where
25    C: BrokerCore<Frame = WsIncoming>,
26    F: Fn() -> Fut,
27    Fut: std::future::Future<Output = ()>,
28{
29    loop {
30        let backlog = emitter.has_backlog();
31        tokio::select! {
32            _ = cancel.cancelled() => break,
33            msg = rx.recv() => {
34                let Ok(msg) = msg else { break };
35                let now: Instant = brokapi_runtime::now();
36                let mut sink = EventSink::new();
37                core.lock().await.on_frame(msg, now, &mut sink);
38                emitter.emit(sink.drain());
39                after().await;
40            }
41            _ = tokio::time::sleep(FLUSH_EVERY), if backlog => emitter.flush(),
42        }
43    }
44}
45
46/// Waits until the emitter has nothing buffered, or `deadline` has passed.
47pub async fn drain_backlog(emitter: &Emitter, deadline: Duration) {
48    let end = brokapi_runtime::now() + deadline;
49    while emitter.has_backlog() && brokapi_runtime::now() < end {
50        emitter.flush();
51        tokio::time::sleep(Duration::from_millis(5)).await;
52    }
53}