Skip to main content

brokapi_alor/
pump.rs

1//! WS event pumps: parse in the receiver task, never block on a slow consumer.
2//!
3//! The frames of the WebSocket session come in through a channel. The pump gives each one to its core and hands
4//! the events to the [`Emitter`]: per-subscription buffers in front of the bot's channel. A bot that reads
5//! slowly fills the buffer of a subscription, and the policy of that subscription applies: `Conflate` keeps the
6//! latest, `DropNewest` counts drops, `Lossless` overflows (`Event::Overflow`, connector `Degraded`).
7
8use std::sync::Arc;
9use std::time::{Duration, Instant};
10
11use brokapi_alor_protocol::ws::models::WsIncoming;
12use brokapi_runtime::Emitter;
13use brokapi_sansio::{BrokerCore, EventSink};
14use kanal::AsyncReceiver;
15use tokio::sync::Mutex;
16use tokio_util::sync::CancellationToken;
17
18/// How often the events that wait in the buffers are offered to the bot's channel again.
19const FLUSH_EVERY: Duration = Duration::from_millis(20);
20
21/// Runs an event pump until `cancel` fires or the frame channel closes.
22///
23/// Every frame from `rx` goes through `core` and the resulting events to `emitter`. `before` runs first for each
24/// frame (the user data connector uses it to look up what a frame needs). Events that a slow bot left in the
25/// subscription buffers are offered again every 20 ms.
26pub async fn run_pump<C, F, Fut>(
27    core: Arc<Mutex<C>>,
28    emitter: Arc<Emitter>,
29    rx: AsyncReceiver<WsIncoming>,
30    cancel: CancellationToken,
31    before: F,
32) where
33    C: BrokerCore<Frame = WsIncoming>,
34    F: Fn(WsIncoming) -> Fut,
35    Fut: std::future::Future<Output = WsIncoming>,
36{
37    loop {
38        let backlog = emitter.has_backlog();
39        tokio::select! {
40            _ = cancel.cancelled() => break,
41            msg = rx.recv() => {
42                let Ok(msg) = msg else { break };
43                let msg = before(msg).await;
44                let now: Instant = brokapi_runtime::now();
45                let mut sink = EventSink::new();
46                core.lock().await.on_frame(msg, now, &mut sink);
47                emitter.emit(sink.drain());
48            }
49            _ = tokio::time::sleep(FLUSH_EVERY), if backlog => emitter.flush(),
50        }
51    }
52}
53
54/// Gives the bot's channel time to take what the buffers still hold (shutdown).
55pub async fn drain_backlog(emitter: &Emitter, deadline: Duration) {
56    let end = brokapi_runtime::now() + deadline;
57    while emitter.has_backlog() && brokapi_runtime::now() < end {
58        emitter.flush();
59        tokio::time::sleep(Duration::from_millis(5)).await;
60    }
61}