lean_ctx_proxy/
usage_sink.rs

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

lean_ctx_proxy/
usage_sink.rs
