lean_ctx/proxy/
usage_sink.rs

1//! Process-wide usage event sink (enterprise#17).
2//!
3//! `usage_meter::record` is the single choke-point every measured turn flows
4//! through (streaming and non-streaming alike). This module lets a run-mode
5//! (the self-hosted gateway) subscribe to that stream without the proxy knowing
6//! anything about Postgres: the gateway installs an `mpsc` sender at startup,
7//! and `push` forwards each finalized [`RealUsage`].
8//!
9//! Fail-open by construction (enterprise#12): `push` never blocks and never
10//! errors the request path — when no sink is installed it is a no-op, and when
11//! the writer falls behind the event is dropped (counted, visible in logs)
12//! rather than back-pressuring live LLM traffic.
13
14use std::sync::OnceLock;
15use std::sync::atomic::{AtomicU64, Ordering};
16
17use super::usage::RealUsage;
18
19static SINK: OnceLock<tokio::sync::mpsc::Sender<RealUsage>> = OnceLock::new();
20static DROPPED: AtomicU64 = AtomicU64::new(0);
21
22/// Installs the process-wide sink. First caller wins (one gateway run-mode per
23/// process); later calls return `false` and change nothing.
24pub fn install(sender: tokio::sync::mpsc::Sender<RealUsage>) -> bool {
25    SINK.set(sender).is_ok()
26}
27
28/// True once a sink is installed (the gateway run-mode is active).
29#[must_use]
30pub fn installed() -> bool {
31    SINK.get().is_some()
32}
33
34/// Forwards one finalized usage record to the sink, if any. Never blocks: on a
35/// full or closed channel the event is dropped and counted.
36pub fn push(usage: &RealUsage) {
37    let Some(tx) = SINK.get() else { return };
38    if tx.try_send(usage.clone()).is_err() {
39        let n = DROPPED.fetch_add(1, Ordering::Relaxed) + 1;
40        // Log sparsely (powers of two) so a stalled writer can't flood stderr.
41        if n.is_power_of_two() {
42            tracing::warn!("usage sink backlogged: {n} event(s) dropped so far");
43        }
44    }
45}
46
47/// Events dropped because the sink was full/closed (observability, #34).
48#[must_use]
49pub fn dropped_count() -> u64 {
50    DROPPED.load(Ordering::Relaxed)
51}
52
53/// Events currently queued (sent but not yet consumed by the writer). Used by
54/// the gateway's graceful shutdown to drain metering before exit (#51).
55#[must_use]
56pub fn pending_count() -> usize {
57    SINK.get().map_or(0, |tx| tx.max_capacity() - tx.capacity())
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63
64    #[test]
65    fn push_without_sink_is_noop() {
66        // No sink installed in this test binary at this point: must not panic.
67        push(&RealUsage {
68            model: "m".into(),
69            input_tokens: 1,
70            ..Default::default()
71        });
72    }
73}
lean_ctx/proxy/usage_sink.rs

lean_ctx/proxy/
usage_sink.rs
