Skip to main content

bsv_overlay_engine/
meter.rs

1//! Poll-scoped counters: concurrent futures cannot charge each other's runs.
2use serde::{Deserialize, Serialize};
3use std::{
4    cell::{Cell, RefCell},
5    future::Future,
6    pin::Pin,
7    rc::Rc,
8    task::{Context, Poll},
9};
10
11/// Storage figures count attempted trait calls, including failures. They are
12/// operations, not SQL rows; discovery lookups count as peer requests.
13#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(default)]
15pub struct SyncCost {
16    pub calls_out: u64,
17    pub peer_requests: u64,
18    pub courier_rungs: u64,
19    pub storage_reads: u64,
20    pub storage_writes: u64,
21    pub d1_queries: u64,
22    pub courier_not_found_memo_hits: u64,
23    pub ms: u64,
24    pub outputs_walked: u64,
25    pub ms_per_1000_outputs: Option<u64>,
26}
27
28impl SyncCost {
29    pub fn since(self, before: Self) -> Self {
30        let mut cost = Self {
31            calls_out: self.calls_out.saturating_sub(before.calls_out),
32            peer_requests: self.peer_requests.saturating_sub(before.peer_requests),
33            courier_rungs: self.courier_rungs.saturating_sub(before.courier_rungs),
34            storage_reads: self.storage_reads.saturating_sub(before.storage_reads),
35            storage_writes: self.storage_writes.saturating_sub(before.storage_writes),
36            d1_queries: self.d1_queries.saturating_sub(before.d1_queries),
37            courier_not_found_memo_hits: self
38                .courier_not_found_memo_hits
39                .saturating_sub(before.courier_not_found_memo_hits),
40            ms: self.ms.saturating_sub(before.ms),
41            outputs_walked: self.outputs_walked.saturating_sub(before.outputs_walked),
42            ms_per_1000_outputs: None,
43        };
44        cost.refresh_speed();
45        cost
46    }
47
48    pub fn refresh_speed(&mut self) {
49        self.ms_per_1000_outputs = (self.outputs_walked != 0).then(|| {
50            ((u128::from(self.ms) * 1000) / u128::from(self.outputs_walked))
51                .min(u128::from(u64::MAX)) as u64
52        });
53    }
54
55    pub fn add(&mut self, other: Self) {
56        self.calls_out = self.calls_out.saturating_add(other.calls_out);
57        self.peer_requests = self.peer_requests.saturating_add(other.peer_requests);
58        self.courier_rungs = self.courier_rungs.saturating_add(other.courier_rungs);
59        self.storage_reads = self.storage_reads.saturating_add(other.storage_reads);
60        self.storage_writes = self.storage_writes.saturating_add(other.storage_writes);
61        self.d1_queries = self.d1_queries.saturating_add(other.d1_queries);
62        self.courier_not_found_memo_hits = self
63            .courier_not_found_memo_hits
64            .saturating_add(other.courier_not_found_memo_hits);
65        self.ms = self.ms.saturating_add(other.ms);
66        self.outputs_walked = self.outputs_walked.saturating_add(other.outputs_walked);
67        self.refresh_speed();
68    }
69}
70
71struct Data {
72    cost: Cell<SyncCost>,
73    started: f64,
74    clock: Rc<dyn Fn() -> f64>,
75}
76#[derive(Clone)]
77pub struct LoadMeter(Rc<Data>);
78impl Default for LoadMeter {
79    fn default() -> Self {
80        Self::with_clock(Rc::new(crate::engine::current_timestamp_ms))
81    }
82}
83impl LoadMeter {
84    /// Inject a monotonic millisecond clock for deterministic host tests.
85    pub fn with_clock(clock: Rc<dyn Fn() -> f64>) -> Self {
86        Self(Rc::new(Data {
87            cost: Cell::default(),
88            started: clock(),
89            clock,
90        }))
91    }
92    pub fn snapshot(&self) -> SyncCost {
93        let mut cost = self.0.cost.get();
94        cost.ms = ((self.0.clock)() - self.0.started).max(0.0) as u64;
95        cost.refresh_speed();
96        cost
97    }
98    pub fn scope<F: Future>(&self, future: F) -> Measured<F> {
99        Measured {
100            meter: self.clone(),
101            future: Box::pin(future),
102        }
103    }
104}
105thread_local! {
106    static CURRENT: RefCell<Option<LoadMeter>> = const { RefCell::new(None) };
107    static ROOT: RefCell<Option<Rc<RootLimit>>> = const { RefCell::new(None) };
108}
109struct Restore(Option<LoadMeter>);
110impl Drop for Restore {
111    fn drop(&mut self) {
112        CURRENT.with(|c| *c.borrow_mut() = self.0.take());
113    }
114}
115pub struct Measured<F> {
116    meter: LoadMeter,
117    future: Pin<Box<F>>,
118}
119impl<F: Future> Future for Measured<F> {
120    type Output = F::Output;
121    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
122        let _restore = Restore(CURRENT.with(|c| c.borrow_mut().replace(self.meter.clone())));
123        self.future.as_mut().poll(cx)
124    }
125}
126pub fn snapshot() -> SyncCost {
127    CURRENT.with(|c| {
128        c.borrow()
129            .as_ref()
130            .map_or_else(SyncCost::default, LoadMeter::snapshot)
131    })
132}
133fn note(f: impl FnOnce(&mut SyncCost)) {
134    CURRENT.with(|c| {
135        if let Some(m) = c.borrow().as_ref() {
136            let mut v = m.0.cost.get();
137            f(&mut v);
138            m.0.cost.set(v);
139        }
140    });
141}
142pub fn walk_output() {
143    note(|v| v.outputs_walked += 1);
144}
145pub fn peer_request() {
146    note(|v| {
147        v.calls_out += 1;
148        v.peer_requests += 1;
149    });
150}
151pub fn storage_read() {
152    note(|v| v.storage_reads += 1);
153}
154pub fn storage_write() {
155    note(|v| v.storage_writes += 1);
156}
157/// Count attempted SQL statements at the worker adapter, including batch members.
158pub fn d1_queries(count: u64) {
159    note(|v| v.d1_queries = v.d1_queries.saturating_add(count));
160}
161pub fn not_found_memo_hit() {
162    note(|v| v.courier_not_found_memo_hits += 1);
163}
164
165/// Twelve courier HTTP attempts at at most 1,000 ms each fit in a 15,000 ms
166/// graph slice, leaving 3,000 ms for peer/storage work. With the 30,000 ms
167/// peer slice and 9,900 ms resume share, the page retains 20,100 ms. These
168/// are stop bounds, not promises about finalize/storage latency.
169pub const ROOT_MAX_COURIER_RUNGS: u32 = 12;
170pub const COURIER_RUNG_TIMEOUT_MS: u64 = 1_000;
171pub const ROOT_MAX_MS: u64 = 15_000;
172#[derive(Default)]
173struct RootLimit {
174    calls: Cell<u32>,
175    exhausted: Cell<bool>,
176}
177pub struct RootScope<F> {
178    limit: Rc<RootLimit>,
179    future: Pin<Box<F>>,
180}
181pub fn root_scope<F: Future>(future: F) -> RootScope<F> {
182    RootScope {
183        limit: Rc::default(),
184        future: Box::pin(future),
185    }
186}
187struct RestoreRoot(Option<Rc<RootLimit>>);
188impl Drop for RestoreRoot {
189    fn drop(&mut self) {
190        ROOT.with(|c| *c.borrow_mut() = self.0.take());
191    }
192}
193impl<F: Future> Future for RootScope<F> {
194    type Output = F::Output;
195    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
196        let _restore = RestoreRoot(ROOT.with(|c| c.borrow_mut().replace(self.limit.clone())));
197        self.future.as_mut().poll(cx)
198    }
199}
200pub fn root_scoped() -> bool {
201    ROOT.with(|c| c.borrow().is_some())
202}
203pub fn root_exhausted() -> bool {
204    ROOT.with(|c| c.borrow().as_ref().is_some_and(|r| r.exhausted.get()))
205}
206/// Charge before awaiting HTTP. A refused charge sends no request.
207pub fn courier_request() -> Result<(), &'static str> {
208    let allowed = ROOT.with(|c| {
209        c.borrow().as_ref().is_none_or(|r| {
210            if r.calls.get() >= ROOT_MAX_COURIER_RUNGS {
211                r.exhausted.set(true);
212                false
213            } else {
214                r.calls.set(r.calls.get() + 1);
215                true
216            }
217        })
218    });
219    if !allowed {
220        return Err("root courier rung budget exhausted");
221    }
222    note(|v| {
223        v.calls_out += 1;
224        v.courier_rungs += 1;
225    });
226    Ok(())
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    #[tokio::test]
233    async fn load_0112_cancelled_run_keeps_its_attempts() {
234        let meter = LoadMeter::default();
235        let mut task = Box::pin(meter.scope(async {
236            peer_request();
237            storage_write();
238            std::future::pending::<()>().await;
239        }));
240        std::future::poll_fn(|cx| {
241            assert!(task.as_mut().poll(cx).is_pending());
242            Poll::Ready(())
243        })
244        .await;
245        drop(task);
246        assert_eq!(
247            (
248                meter.snapshot().peer_requests,
249                meter.snapshot().storage_writes
250            ),
251            (1, 1)
252        );
253        assert_eq!(snapshot(), SyncCost::default());
254    }
255
256    #[tokio::test]
257    async fn load_0112_poll_scopes_keep_interleaved_runs_separate() {
258        let a = LoadMeter::default();
259        let b = LoadMeter::default();
260        let ((), ()) = tokio::join!(
261            a.scope(async {
262                peer_request();
263                tokio::task::yield_now().await;
264                storage_read();
265            }),
266            b.scope(async {
267                storage_write();
268                tokio::task::yield_now().await;
269                courier_request().unwrap();
270            })
271        );
272        assert_eq!(
273            (
274                a.snapshot().peer_requests,
275                a.snapshot().storage_reads,
276                a.snapshot().storage_writes
277            ),
278            (1, 1, 0)
279        );
280        assert_eq!(
281            (
282                b.snapshot().courier_rungs,
283                b.snapshot().storage_writes,
284                b.snapshot().storage_reads
285            ),
286            (1, 1, 0)
287        );
288        assert_eq!(snapshot(), SyncCost::default());
289    }
290}