Skip to main content

mcp_usage_export/
lib.rs

1//! Buffered usage accounting for MCP calls.
2//!
3//! Recording is synchronous, memory-only, and never performs network I/O. A
4//! separate [`BillingPipeline::flush`] drains aggregates into an exporter. If an
5//! export fails, the exact same event identifiers and quantities are retained
6//! for retry, which makes ambiguous provider failures safe.
7
8#![forbid(unsafe_code)]
9#![warn(missing_docs, clippy::pedantic)]
10#![allow(clippy::module_name_repetitions)]
11
12use std::collections::{HashMap, HashSet, VecDeque};
13use std::fmt;
14use std::future::Future;
15use std::pin::Pin;
16use std::sync::atomic::{AtomicBool, Ordering};
17use std::sync::{Arc, Mutex, MutexGuard};
18use std::time::{SystemTime, UNIX_EPOCH};
19
20use serde::Serialize;
21use thiserror::Error;
22
23mod meter_event;
24
25pub use meter_event::{
26    MeterEventDeadLetter, MeterEventDeadLetterReason, MeterEventExporter, MeterEventOutcome,
27    MeterEventProvider, MeterEventProviderError, MeterEventProviderFuture,
28};
29
30/// Boxed future returned by object-safe billing exporters.
31pub type ExportFuture<'a> = Pin<Box<dyn Future<Output = Result<(), ExportError>> + Send + 'a>>;
32
33/// One delivered unit event emitted by the request path.
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub struct UsageEvent {
36    /// Internal tenant identifier used for deduplication and reconciliation.
37    pub tenant_id: String,
38    /// Billing-provider customer identifier.
39    pub customer_id: String,
40    /// Provider meter event name.
41    pub meter: String,
42    /// Integer quantity to add to the meter.
43    pub units: u64,
44    /// Stable key for once-only work such as a completed durable task.
45    pub idempotency_key: Option<String>,
46    /// Event time in Unix seconds.
47    pub timestamp: u64,
48}
49
50impl UsageEvent {
51    /// Construct a usage event stamped with the current wall-clock time.
52    #[must_use]
53    pub fn now(
54        tenant_id: impl Into<String>,
55        customer_id: impl Into<String>,
56        meter: impl Into<String>,
57        units: u64,
58        idempotency_key: Option<String>,
59    ) -> Self {
60        let timestamp = SystemTime::now()
61            .duration_since(UNIX_EPOCH)
62            .unwrap_or_default()
63            .as_secs();
64        Self {
65            tenant_id: tenant_id.into(),
66            customer_id: customer_id.into(),
67            meter: meter.into(),
68            units,
69            idempotency_key,
70            timestamp,
71        }
72    }
73}
74
75/// One pre-aggregated provider event ready for export.
76#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
77pub struct AggregatedUsage {
78    /// Globally unique provider-side idempotency identifier.
79    pub identifier: String,
80    /// Provider customer identifier.
81    pub customer_id: String,
82    /// Provider meter event name.
83    pub meter: String,
84    /// Summed integer quantity.
85    pub units: u64,
86    /// Timestamp of the oldest event represented by the aggregate.
87    pub timestamp: u64,
88}
89
90/// Result of attempting to record an event.
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92pub enum RecordOutcome {
93    /// The event was added to its aggregate.
94    Recorded,
95    /// The event's idempotency key was already observed.
96    Duplicate,
97    /// Zero units need no provider event.
98    ZeroUnits,
99}
100
101/// A synchronous usage sink suitable for the HTTP hot path.
102pub trait UsageRecorder: Send + Sync {
103    /// Add one event without performing I/O.
104    ///
105    /// # Errors
106    ///
107    /// Returns [`RecordError`] only when the local synchronization primitive is
108    /// unusable; provider outages cannot surface on this path.
109    fn record(&self, event: UsageEvent) -> Result<RecordOutcome, RecordError>;
110}
111
112/// Exports a prepared batch away from the request path.
113pub trait BatchExporter: Send + Sync {
114    /// Export every aggregate in `batch`.
115    ///
116    /// Retrying the same batch must be safe because identifiers are stable.
117    fn export<'a>(&'a self, batch: &'a [AggregatedUsage]) -> ExportFuture<'a>;
118}
119
120/// Convenient trait-object alias for composing exporters at runtime.
121pub type SharedExporter = Arc<dyn BatchExporter>;
122
123/// Adapter for an application-owned async export function.
124pub struct FnExporter<F> {
125    export: F,
126}
127
128impl<F> FnExporter<F> {
129    /// Wrap an async export function without defining a new type.
130    #[must_use]
131    pub const fn new(export: F) -> Self {
132        Self { export }
133    }
134}
135
136impl<F> fmt::Debug for FnExporter<F> {
137    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
138        formatter.debug_struct("FnExporter").finish_non_exhaustive()
139    }
140}
141
142impl<F> BatchExporter for FnExporter<F>
143where
144    F: for<'a> Fn(&'a [AggregatedUsage]) -> ExportFuture<'a> + Send + Sync,
145{
146    fn export<'a>(&'a self, batch: &'a [AggregatedUsage]) -> ExportFuture<'a> {
147        (self.export)(batch)
148    }
149}
150
151/// Fan-out exporter for logs, billing, analytics, and audit sinks.
152///
153/// Exporters run in registration order. If one fails, a later retry starts from
154/// the first exporter. Provider implementations must therefore honor the stable
155/// aggregate identifiers, as required by [`BatchExporter`].
156#[derive(Default)]
157pub struct CompositeExporter {
158    exporters: Vec<SharedExporter>,
159}
160
161impl CompositeExporter {
162    /// Construct an empty fan-out exporter.
163    #[must_use]
164    pub const fn new() -> Self {
165        Self {
166            exporters: Vec::new(),
167        }
168    }
169
170    /// Append an exporter to the fan-out sequence.
171    #[must_use]
172    pub fn with_exporter(mut self, exporter: SharedExporter) -> Self {
173        self.exporters.push(exporter);
174        self
175    }
176
177    /// Number of configured exporters.
178    #[must_use]
179    pub fn len(&self) -> usize {
180        self.exporters.len()
181    }
182
183    /// Whether no exporters are configured.
184    #[must_use]
185    pub fn is_empty(&self) -> bool {
186        self.exporters.is_empty()
187    }
188}
189
190impl fmt::Debug for CompositeExporter {
191    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
192        formatter
193            .debug_struct("CompositeExporter")
194            .field("exporters", &self.exporters.len())
195            .finish()
196    }
197}
198
199impl BatchExporter for CompositeExporter {
200    fn export<'a>(&'a self, batch: &'a [AggregatedUsage]) -> ExportFuture<'a> {
201        Box::pin(async move {
202            for exporter in &self.exporters {
203                exporter.export(batch).await?;
204            }
205            Ok(())
206        })
207    }
208}
209
210/// Recording failures are internal synchronization failures only.
211#[derive(Debug, Error)]
212pub enum RecordError {
213    /// A synchronization primitive was poisoned by a panicking thread.
214    #[error("usage buffer lock was poisoned")]
215    Poisoned,
216    /// Adding an event would exceed the representable integer quantity.
217    #[error("usage aggregate exceeds u64::MAX units")]
218    UnitsOverflow,
219}
220
221/// Provider export failure.
222#[derive(Debug, Error, Clone, PartialEq, Eq)]
223pub enum ExportError {
224    /// The provider or transport rejected an event.
225    #[error("billing export failed: {0}")]
226    Provider(String),
227    /// Another flush already owns the current retry batch.
228    #[error("a billing flush is already in progress")]
229    FlushInProgress,
230}
231
232#[derive(Debug, Clone, PartialEq, Eq, Hash)]
233struct AggregateKey {
234    customer_id: String,
235    meter: String,
236}
237
238#[derive(Debug, Clone)]
239struct PendingAggregate {
240    units: u64,
241    oldest_timestamp: u64,
242}
243
244#[derive(Debug, Default)]
245struct BufferState {
246    pending: HashMap<AggregateKey, PendingAggregate>,
247    retry: Vec<AggregatedUsage>,
248    seen: HashSet<(String, String)>,
249    seen_order: VecDeque<(String, String)>,
250}
251
252/// Thread-safe in-memory aggregator with bounded idempotency memory.
253pub struct UsageBuffer {
254    state: Mutex<BufferState>,
255    max_idempotency_keys: usize,
256}
257
258impl fmt::Debug for UsageBuffer {
259    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
260        let state = self
261            .state
262            .lock()
263            .unwrap_or_else(std::sync::PoisonError::into_inner);
264        formatter
265            .debug_struct("UsageBuffer")
266            .field(
267                "pending_buckets",
268                &(state.pending.len() + state.retry.len()),
269            )
270            .field("retained_idempotency_keys", &state.seen.len())
271            .field("max_idempotency_keys", &self.max_idempotency_keys)
272            .finish_non_exhaustive()
273    }
274}
275
276impl Default for UsageBuffer {
277    fn default() -> Self {
278        Self::new(100_000)
279    }
280}
281
282impl UsageBuffer {
283    /// Construct a buffer retaining at most `max_idempotency_keys` once-only
284    /// keys. Zero disables in-process deduplication.
285    #[must_use]
286    pub fn new(max_idempotency_keys: usize) -> Self {
287        Self {
288            state: Mutex::new(BufferState::default()),
289            max_idempotency_keys,
290        }
291    }
292
293    fn lock(&self) -> Result<MutexGuard<'_, BufferState>, RecordError> {
294        self.state.lock().map_err(|_| RecordError::Poisoned)
295    }
296
297    /// Record one event into a `(customer, meter)` aggregate.
298    ///
299    /// # Errors
300    ///
301    /// Returns [`RecordError::Poisoned`] if another thread panicked while
302    /// mutating the buffer, or [`RecordError::UnitsOverflow`] rather than
303    /// silently saturating an invoice quantity.
304    pub fn record(&self, event: UsageEvent) -> Result<RecordOutcome, RecordError> {
305        if event.units == 0 {
306            return Ok(RecordOutcome::ZeroUnits);
307        }
308
309        let mut state = self.lock()?;
310        let scoped_idempotency = event
311            .idempotency_key
312            .map(|idempotency_key| (event.tenant_id, idempotency_key));
313        if scoped_idempotency
314            .as_ref()
315            .is_some_and(|scoped_key| state.seen.contains(scoped_key))
316        {
317            return Ok(RecordOutcome::Duplicate);
318        }
319        let key = AggregateKey {
320            customer_id: event.customer_id,
321            meter: event.meter,
322        };
323        let next_units = state
324            .pending
325            .get(&key)
326            .map_or(Some(event.units), |aggregate| {
327                aggregate.units.checked_add(event.units)
328            })
329            .ok_or(RecordError::UnitsOverflow)?;
330
331        if let Some(scoped_key) = scoped_idempotency
332            && self.max_idempotency_keys > 0
333        {
334            state.seen.insert(scoped_key.clone());
335            state.seen_order.push_back(scoped_key);
336            while state.seen_order.len() > self.max_idempotency_keys {
337                if let Some(expired) = state.seen_order.pop_front() {
338                    state.seen.remove(&expired);
339                }
340            }
341        }
342
343        let aggregate = state.pending.entry(key).or_insert(PendingAggregate {
344            units: 0,
345            oldest_timestamp: event.timestamp,
346        });
347        aggregate.units = next_units;
348        aggregate.oldest_timestamp = aggregate.oldest_timestamp.min(event.timestamp);
349        Ok(RecordOutcome::Recorded)
350    }
351
352    fn prepare_flush(&self) -> Result<Vec<AggregatedUsage>, RecordError> {
353        let mut state = self.lock()?;
354        if !state.retry.is_empty() {
355            return Ok(std::mem::take(&mut state.retry));
356        }
357        let mut batch: Vec<_> = std::mem::take(&mut state.pending)
358            .into_iter()
359            .map(|(key, aggregate)| AggregatedUsage {
360                identifier: format!("mcp_usage_{}", uuid::Uuid::new_v4().simple()),
361                customer_id: key.customer_id,
362                meter: key.meter,
363                units: aggregate.units,
364                timestamp: aggregate.oldest_timestamp,
365            })
366            .collect();
367        batch.sort_by(|a, b| {
368            a.customer_id
369                .cmp(&b.customer_id)
370                .then_with(|| a.meter.cmp(&b.meter))
371        });
372        Ok(batch)
373    }
374
375    fn restore_after_interruption(&self, batch: Vec<AggregatedUsage>) {
376        let mut state = self
377            .state
378            .lock()
379            .unwrap_or_else(std::sync::PoisonError::into_inner);
380        if state.retry.is_empty() {
381            state.retry = batch;
382        } else {
383            state.retry.extend(batch);
384        }
385    }
386
387    /// Number of live aggregate buckets plus retry events.
388    #[must_use]
389    pub fn pending_buckets(&self) -> usize {
390        self.state
391            .lock()
392            .map(|state| state.pending.len() + state.retry.len())
393            .unwrap_or_default()
394    }
395}
396
397/// Coordinates a [`UsageBuffer`] with an off-path exporter.
398pub struct BillingPipeline<E> {
399    buffer: UsageBuffer,
400    exporter: E,
401    flush_in_progress: AtomicBool,
402}
403
404impl<E> fmt::Debug for BillingPipeline<E> {
405    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
406        formatter
407            .debug_struct("BillingPipeline")
408            .field("pending_buckets", &self.pending_buckets())
409            .field(
410                "flush_in_progress",
411                &self.flush_in_progress.load(Ordering::Relaxed),
412            )
413            .finish_non_exhaustive()
414    }
415}
416
417impl<E> BillingPipeline<E> {
418    /// Construct a pipeline with the default bounded deduplication store.
419    #[must_use]
420    pub fn new(exporter: E) -> Self {
421        Self {
422            buffer: UsageBuffer::default(),
423            exporter,
424            flush_in_progress: AtomicBool::new(false),
425        }
426    }
427
428    /// Construct a pipeline with an explicit idempotency-key bound.
429    #[must_use]
430    pub fn with_idempotency_capacity(exporter: E, capacity: usize) -> Self {
431        Self {
432            buffer: UsageBuffer::new(capacity),
433            exporter,
434            flush_in_progress: AtomicBool::new(false),
435        }
436    }
437
438    /// Number of buffered aggregate buckets.
439    #[must_use]
440    pub fn pending_buckets(&self) -> usize {
441        self.buffer.pending_buckets()
442    }
443
444    /// Access the exporter for diagnostics and tests.
445    #[must_use]
446    pub const fn exporter(&self) -> &E {
447        &self.exporter
448    }
449}
450
451impl<E: BatchExporter> BillingPipeline<E> {
452    /// Drain and export one batch. Failed batches retain their exact identifiers
453    /// for an idempotent retry.
454    ///
455    /// # Errors
456    ///
457    /// Returns [`ExportError::FlushInProgress`] for overlapping flushes, or the
458    /// provider error after restoring the batch for retry.
459    pub async fn flush(&self) -> Result<usize, ExportError> {
460        if self
461            .flush_in_progress
462            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
463            .is_err()
464        {
465            return Err(ExportError::FlushInProgress);
466        }
467        let _guard = FlushGuard(&self.flush_in_progress);
468        let batch = self
469            .buffer
470            .prepare_flush()
471            .map_err(|error| ExportError::Provider(error.to_string()))?;
472        if batch.is_empty() {
473            return Ok(0);
474        }
475        let retry = RetryGuard::new(&self.buffer, batch);
476        self.exporter.export(retry.batch()).await?;
477        Ok(retry.commit())
478    }
479}
480
481struct FlushGuard<'a>(&'a AtomicBool);
482
483impl Drop for FlushGuard<'_> {
484    fn drop(&mut self) {
485        self.0.store(false, Ordering::Release);
486    }
487}
488
489/// Owns a prepared batch until the exporter commits successfully.
490///
491/// Dropping a flush future through cancellation or unwinding drops this guard,
492/// which restores the exact batch before another flush can acquire ownership.
493struct RetryGuard<'a> {
494    buffer: &'a UsageBuffer,
495    batch: Option<Vec<AggregatedUsage>>,
496}
497
498impl<'a> RetryGuard<'a> {
499    fn new(buffer: &'a UsageBuffer, batch: Vec<AggregatedUsage>) -> Self {
500        Self {
501            buffer,
502            batch: Some(batch),
503        }
504    }
505
506    fn batch(&self) -> &[AggregatedUsage] {
507        self.batch.as_deref().unwrap_or_default()
508    }
509
510    fn commit(mut self) -> usize {
511        self.batch.take().map_or(0, |batch| batch.len())
512    }
513}
514
515impl Drop for RetryGuard<'_> {
516    fn drop(&mut self) {
517        if let Some(batch) = self.batch.take() {
518            self.buffer.restore_after_interruption(batch);
519        }
520    }
521}
522
523impl<E: BatchExporter> UsageRecorder for BillingPipeline<E> {
524    fn record(&self, event: UsageEvent) -> Result<RecordOutcome, RecordError> {
525        self.buffer.record(event)
526    }
527}
528
529/// Stateless exporter that emits structured log events.
530#[derive(Debug, Default, Clone, Copy)]
531pub struct LogExporter;
532
533impl LogExporter {
534    /// Construct a logging exporter.
535    #[must_use]
536    pub const fn new() -> Self {
537        Self
538    }
539}
540
541impl BatchExporter for LogExporter {
542    fn export<'a>(&'a self, batch: &'a [AggregatedUsage]) -> ExportFuture<'a> {
543        Box::pin(async move {
544            for usage in batch {
545                tracing::info!(
546                    identifier = %usage.identifier,
547                    meter = %usage.meter,
548                    units = usage.units,
549                    timestamp = usage.timestamp,
550                    "MCP billing usage"
551                );
552            }
553            Ok(())
554        })
555    }
556}
557
558/// Shared, no-op recorder for deployments that want attribution metrics only.
559#[derive(Debug, Default)]
560pub struct NoopRecorder;
561
562impl UsageRecorder for NoopRecorder {
563    fn record(&self, event: UsageEvent) -> Result<RecordOutcome, RecordError> {
564        Ok(if event.units == 0 {
565            RecordOutcome::ZeroUnits
566        } else {
567            RecordOutcome::Recorded
568        })
569    }
570}
571
572/// Convenient trait-object alias used by the edge layer.
573pub type SharedRecorder = Arc<dyn UsageRecorder>;
574
575impl fmt::Display for RecordOutcome {
576    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
577        f.write_str(match self {
578            Self::Recorded => "recorded",
579            Self::Duplicate => "duplicate",
580            Self::ZeroUnits => "zero_units",
581        })
582    }
583}
584
585#[cfg(test)]
586mod tests {
587    use super::*;
588    use std::sync::atomic::AtomicUsize;
589    use std::thread;
590
591    #[derive(Debug, Default)]
592    struct CaptureExporter {
593        exported: Mutex<Vec<AggregatedUsage>>,
594    }
595
596    impl CaptureExporter {
597        fn exported(&self) -> Vec<AggregatedUsage> {
598            self.exported
599                .lock()
600                .unwrap_or_else(std::sync::PoisonError::into_inner)
601                .clone()
602        }
603    }
604
605    impl BatchExporter for CaptureExporter {
606        fn export<'a>(&'a self, batch: &'a [AggregatedUsage]) -> ExportFuture<'a> {
607            Box::pin(async move {
608                self.exported
609                    .lock()
610                    .unwrap_or_else(std::sync::PoisonError::into_inner)
611                    .extend_from_slice(batch);
612                Ok(())
613            })
614        }
615    }
616
617    fn event(tenant: &str, units: u64) -> UsageEvent {
618        UsageEvent {
619            tenant_id: tenant.to_owned(),
620            customer_id: format!("cus_{tenant}"),
621            meter: "mcp_units".to_owned(),
622            units,
623            idempotency_key: None,
624            timestamp: 1_800_000_000,
625        }
626    }
627
628    #[tokio::test]
629    async fn aggregates_by_customer_and_meter() {
630        let pipeline = BillingPipeline::new(CaptureExporter::default());
631        pipeline.record(event("acme", 2)).unwrap();
632        pipeline.record(event("acme", 3)).unwrap();
633        pipeline.record(event("other", 7)).unwrap();
634        assert_eq!(pipeline.pending_buckets(), 2);
635        assert_eq!(pipeline.flush().await.unwrap(), 2);
636
637        let exported = pipeline.exporter().exported();
638        assert_eq!(exported.len(), 2);
639        assert_eq!(exported.iter().map(|e| e.units).sum::<u64>(), 12);
640    }
641
642    #[test]
643    fn duplicate_task_completions_are_suppressed_per_tenant() {
644        let pipeline = BillingPipeline::new(CaptureExporter::default());
645        let mut first = event("acme", 10);
646        first.idempotency_key = Some("task-1".to_owned());
647        assert_eq!(
648            pipeline.record(first.clone()).unwrap(),
649            RecordOutcome::Recorded
650        );
651        assert_eq!(pipeline.record(first).unwrap(), RecordOutcome::Duplicate);
652
653        let mut other_tenant = event("other", 10);
654        other_tenant.idempotency_key = Some("task-1".to_owned());
655        assert_eq!(
656            pipeline.record(other_tenant).unwrap(),
657            RecordOutcome::Recorded
658        );
659    }
660
661    #[test]
662    fn aggregate_overflow_is_reported_without_consuming_the_idempotency_key() {
663        let pipeline = BillingPipeline::new(CaptureExporter::default());
664        pipeline.record(event("acme", u64::MAX)).unwrap();
665        let mut overflowing = event("acme", 1);
666        overflowing.idempotency_key = Some("task-overflow".to_owned());
667        assert!(matches!(
668            pipeline.record(overflowing.clone()),
669            Err(RecordError::UnitsOverflow)
670        ));
671        assert!(matches!(
672            pipeline.record(overflowing),
673            Err(RecordError::UnitsOverflow)
674        ));
675    }
676
677    #[derive(Debug, Default)]
678    struct FailOnceExporter {
679        failed: AtomicBool,
680        attempts: Mutex<Vec<Vec<AggregatedUsage>>>,
681    }
682
683    impl BatchExporter for FailOnceExporter {
684        fn export<'a>(&'a self, batch: &'a [AggregatedUsage]) -> ExportFuture<'a> {
685            Box::pin(async move {
686                self.attempts.lock().unwrap().push(batch.to_vec());
687                if self.failed.swap(true, Ordering::SeqCst) {
688                    Ok(())
689                } else {
690                    Err(ExportError::Provider("temporary outage".to_owned()))
691                }
692            })
693        }
694    }
695
696    #[tokio::test]
697    async fn failed_flush_retries_identical_provider_events() {
698        let pipeline = BillingPipeline::new(FailOnceExporter::default());
699        pipeline.record(event("acme", 5)).unwrap();
700        assert!(pipeline.flush().await.is_err());
701        assert_eq!(pipeline.pending_buckets(), 1);
702        assert_eq!(pipeline.flush().await.unwrap(), 1);
703
704        let attempts = pipeline.exporter().attempts.lock().unwrap();
705        assert_eq!(attempts.len(), 2);
706        assert_eq!(attempts[0], attempts[1]);
707    }
708
709    #[derive(Debug, Default)]
710    struct CancelOnceExporter {
711        attempts: Mutex<Vec<Vec<AggregatedUsage>>>,
712        successful: AtomicUsize,
713    }
714
715    impl BatchExporter for CancelOnceExporter {
716        fn export<'a>(&'a self, batch: &'a [AggregatedUsage]) -> ExportFuture<'a> {
717            Box::pin(async move {
718                let attempt = {
719                    let mut attempts = self.attempts.lock().unwrap();
720                    attempts.push(batch.to_vec());
721                    attempts.len()
722                };
723                if attempt == 1 {
724                    std::future::pending::<()>().await;
725                }
726                self.successful.fetch_add(1, Ordering::SeqCst);
727                Ok(())
728            })
729        }
730    }
731
732    #[tokio::test]
733    async fn cancelled_flush_restores_the_identical_batch() {
734        let pipeline = Arc::new(BillingPipeline::new(CancelOnceExporter::default()));
735        pipeline.record(event("acme", 9)).unwrap();
736        let flush = tokio::spawn({
737            let pipeline = Arc::clone(&pipeline);
738            async move { pipeline.flush().await }
739        });
740        while pipeline.exporter().attempts.lock().unwrap().is_empty() {
741            tokio::task::yield_now().await;
742        }
743        flush.abort();
744        assert!(flush.await.unwrap_err().is_cancelled());
745
746        assert_eq!(pipeline.pending_buckets(), 1);
747        assert_eq!(pipeline.flush().await.unwrap(), 1);
748        let attempts = pipeline.exporter().attempts.lock().unwrap();
749        assert_eq!(attempts.len(), 2);
750        assert_eq!(attempts[0], attempts[1]);
751        assert_eq!(pipeline.exporter().successful.load(Ordering::SeqCst), 1);
752    }
753
754    #[derive(Debug, Default)]
755    struct PanicOnceExporter {
756        attempts: Mutex<Vec<Vec<AggregatedUsage>>>,
757        panicked: AtomicBool,
758        successful: AtomicUsize,
759    }
760
761    impl BatchExporter for PanicOnceExporter {
762        fn export<'a>(&'a self, batch: &'a [AggregatedUsage]) -> ExportFuture<'a> {
763            Box::pin(async move {
764                self.attempts.lock().unwrap().push(batch.to_vec());
765                assert!(
766                    self.panicked.swap(true, Ordering::SeqCst),
767                    "intentional exporter panic"
768                );
769                self.successful.fetch_add(1, Ordering::SeqCst);
770                Ok(())
771            })
772        }
773    }
774
775    #[tokio::test]
776    async fn panicking_flush_restores_the_identical_batch() {
777        let pipeline = Arc::new(BillingPipeline::new(PanicOnceExporter::default()));
778        pipeline.record(event("acme", 11)).unwrap();
779        let flush = tokio::spawn({
780            let pipeline = Arc::clone(&pipeline);
781            async move { pipeline.flush().await }
782        });
783        assert!(flush.await.unwrap_err().is_panic());
784
785        assert_eq!(pipeline.pending_buckets(), 1);
786        assert_eq!(pipeline.flush().await.unwrap(), 1);
787        let attempts = pipeline.exporter().attempts.lock().unwrap();
788        assert_eq!(attempts.len(), 2);
789        assert_eq!(attempts[0], attempts[1]);
790        assert_eq!(pipeline.exporter().successful.load(Ordering::SeqCst), 1);
791    }
792
793    #[test]
794    fn debug_output_contains_counts_but_not_buffered_identifiers() {
795        #[derive(Debug)]
796        struct SensitiveExporter(&'static str);
797
798        impl BatchExporter for SensitiveExporter {
799            fn export<'a>(&'a self, _batch: &'a [AggregatedUsage]) -> ExportFuture<'a> {
800                Box::pin(async { Ok(()) })
801            }
802        }
803
804        let pipeline = BillingPipeline::new(SensitiveExporter("https://secret.example"));
805        let mut sensitive = event("tenant-private", 3);
806        sensitive.customer_id = "customer-private".to_owned();
807        sensitive.meter = "meter-private".to_owned();
808        sensitive.idempotency_key = Some("task-private".to_owned());
809        pipeline.record(sensitive).unwrap();
810
811        let debug = format!("{pipeline:?}");
812        for secret in [
813            "tenant-private",
814            "customer-private",
815            "meter-private",
816            "task-private",
817            "https://secret.example",
818        ] {
819            assert!(!debug.contains(secret));
820        }
821        assert!(debug.contains("pending_buckets"));
822        assert_eq!(pipeline.exporter().0, "https://secret.example");
823    }
824
825    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
826    async fn concurrent_recording_reconciles_exactly() {
827        let pipeline = Arc::new(BillingPipeline::new(CaptureExporter::default()));
828        let threads = 8;
829        let per_thread = 2_000;
830        let mut handles = Vec::new();
831        for _ in 0..threads {
832            let pipeline = pipeline.clone();
833            handles.push(thread::spawn(move || {
834                for _ in 0..per_thread {
835                    pipeline.record(event("acme", 1)).unwrap();
836                }
837            }));
838        }
839        for handle in handles {
840            handle.join().unwrap();
841        }
842        pipeline.flush().await.unwrap();
843        let total: u64 = pipeline.exporter().exported().iter().map(|e| e.units).sum();
844        assert_eq!(total, threads * per_thread);
845    }
846}