mcp_usage_export/
meter_event.rs

1//! Provider-neutral meter event export with partial-batch retry progress.
2
3use std::collections::{HashSet, VecDeque};
4use std::fmt;
5use std::future::Future;
6use std::pin::Pin;
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::{Mutex, MutexGuard};
9
10use thiserror::Error;
11
12use super::{AggregatedUsage, BatchExporter, ExportError, ExportFuture};
13
14const DEFAULT_DEAD_LETTER_CAPACITY: usize = 1_024;
15
16/// Boxed future returned by a meter event provider.
17pub type MeterEventProviderFuture<'a> = Pin<
18    Box<dyn Future<Output = Result<Vec<MeterEventOutcome>, MeterEventProviderError>> + Send + 'a>,
19>;
20
21/// Submits ordered batches of pre-aggregated usage to a billing provider.
22///
23/// Implementations must return exactly one ordered [`MeterEventOutcome`] for
24/// every submitted aggregate. They must treat [`AggregatedUsage::identifier`]
25/// as an idempotency key because cancellation and batch-wide failures can leave
26/// an accepted event's result ambiguous.
27pub trait MeterEventProvider: Send + Sync {
28    /// Submit a batch and return one outcome for each aggregate in the same order.
29    ///
30    /// A batch-wide error means no per-event outcome is available. Error and
31    /// outcome codes must be static, low-cardinality categories and must not
32    /// contain credentials, customer data, or other identifying values.
33    fn submit<'a>(&'a self, batch: &'a [AggregatedUsage]) -> MeterEventProviderFuture<'a>;
34}
35
36/// The provider's final or retryable disposition for one meter event.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum MeterEventOutcome {
39    /// The provider accepted the event.
40    Accepted,
41    /// The event remains unresolved and must be retried with the same identifier.
42    RetryableFailure {
43        /// Static, sanitized, low-cardinality failure category.
44        code: &'static str,
45    },
46    /// The provider permanently rejected this individual event.
47    PermanentRejection {
48        /// Static, sanitized, low-cardinality rejection category.
49        code: &'static str,
50    },
51}
52
53/// Sanitized batch-wide provider failure without per-event outcomes.
54///
55/// The code is restricted to static data to discourage embedding runtime
56/// credentials or customer values. Provider implementations must use a small,
57/// bounded set of categories such as `unavailable` or `authentication_failed`.
58#[derive(Clone, Copy, PartialEq, Eq, Error)]
59#[error("meter event provider failed: {code}")]
60pub struct MeterEventProviderError {
61    code: &'static str,
62}
63
64impl MeterEventProviderError {
65    /// Construct an error from a static, sanitized, low-cardinality code.
66    #[must_use]
67    pub const fn new(code: &'static str) -> Self {
68        Self { code }
69    }
70
71    /// Return the sanitized failure code.
72    #[must_use]
73    pub const fn code(&self) -> &'static str {
74        self.code
75    }
76}
77
78impl fmt::Debug for MeterEventProviderError {
79    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
80        formatter
81            .debug_struct("MeterEventProviderError")
82            .field("code", &self.code)
83            .finish()
84    }
85}
86
87/// A meter event that needs application-owned reconciliation.
88#[derive(Clone, PartialEq, Eq)]
89pub struct MeterEventDeadLetter {
90    /// The original aggregate, including its stable identifier and timestamp.
91    pub aggregate: AggregatedUsage,
92    /// Why the aggregate was quarantined.
93    pub reason: MeterEventDeadLetterReason,
94}
95
96impl fmt::Debug for MeterEventDeadLetter {
97    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
98        formatter
99            .debug_struct("MeterEventDeadLetter")
100            .field("aggregate", &"[REDACTED]")
101            .field("reason", &self.reason)
102            .finish()
103    }
104}
105
106/// Why a meter event was quarantined for reconciliation.
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub enum MeterEventDeadLetterReason {
109    /// The provider synchronously and permanently rejected the event.
110    SynchronousProviderRejection {
111        /// Static, sanitized, low-cardinality rejection category.
112        code: &'static str,
113    },
114    /// The provider reported a rejection after initially accepting the event.
115    AsynchronousProviderRejection {
116        /// Static, sanitized, low-cardinality rejection category.
117        code: &'static str,
118    },
119}
120
121#[derive(Default)]
122struct ExporterState {
123    dead_letters: VecDeque<MeterEventDeadLetter>,
124    dead_letter_identifiers: HashSet<String>,
125    dropped_dead_letters: u64,
126    retry_progress: Option<RetryProgress>,
127}
128
129struct RetryProgress {
130    batch: Vec<AggregatedUsage>,
131    completed: Vec<bool>,
132}
133
134struct ExportLease<'a> {
135    exporting: &'a AtomicBool,
136}
137
138impl Drop for ExportLease<'_> {
139    fn drop(&mut self) {
140        self.exporting.store(false, Ordering::Release);
141    }
142}
143
144/// Exports pre-aggregated usage through an application-supplied provider.
145///
146/// Confirmed accepted and permanently rejected events are skipped on a later
147/// retry. Retryable and ambiguous events retain their original identifiers.
148/// While partial progress exists, only the identical original batch is accepted.
149pub struct MeterEventExporter<P> {
150    provider: P,
151    dead_letter_capacity: usize,
152    state: Mutex<ExporterState>,
153    exporting: AtomicBool,
154}
155
156impl<P> MeterEventExporter<P> {
157    /// Construct an exporter with a dead letter capacity of 1,024.
158    #[must_use]
159    pub fn new(provider: P) -> Self {
160        Self {
161            provider,
162            dead_letter_capacity: DEFAULT_DEAD_LETTER_CAPACITY,
163            state: Mutex::new(ExporterState::default()),
164            exporting: AtomicBool::new(false),
165        }
166    }
167
168    /// Access the provider for diagnostics and tests.
169    #[must_use]
170    pub const fn provider(&self) -> &P {
171        &self.provider
172    }
173
174    /// Bound how many rejected aggregates are retained for reconciliation.
175    ///
176    /// Reducing the capacity evicts the oldest entries and increments
177    /// [`MeterEventExporter::dropped_dead_letters`]. Zero disables retention
178    /// while continuing to count discarded entries.
179    #[must_use]
180    pub fn with_dead_letter_capacity(mut self, capacity: usize) -> Self {
181        self.dead_letter_capacity = capacity;
182        let state = self
183            .state
184            .get_mut()
185            .unwrap_or_else(std::sync::PoisonError::into_inner);
186        while state.dead_letters.len() > capacity {
187            if let Some(evicted) = state.dead_letters.pop_front() {
188                state
189                    .dead_letter_identifiers
190                    .remove(&evicted.aggregate.identifier);
191                state.dropped_dead_letters = state.dropped_dead_letters.saturating_add(1);
192            }
193        }
194        self
195    }
196
197    /// Number of aggregates retained for reconciliation.
198    #[must_use]
199    pub fn dead_letter_count(&self) -> usize {
200        self.lock_state().dead_letters.len()
201    }
202
203    /// Number of reconciliation aggregates discarded because retention was full.
204    #[must_use]
205    pub fn dropped_dead_letters(&self) -> u64 {
206        self.lock_state().dropped_dead_letters
207    }
208
209    /// Drain retained aggregates for application-owned reconciliation.
210    #[must_use]
211    pub fn take_dead_letters(&self) -> Vec<MeterEventDeadLetter> {
212        let mut state = self.lock_state();
213        state.dead_letter_identifiers.clear();
214        state.dead_letters.drain(..).collect()
215    }
216
217    /// Retain an aggregate that a provider rejected after initial acceptance.
218    ///
219    /// Returns `true` when the aggregate was newly retained. A duplicate
220    /// identifier or zero-capacity queue returns `false`. The code must be a
221    /// static, sanitized, low-cardinality category.
222    #[must_use]
223    pub fn quarantine_async_rejection(
224        &self,
225        aggregate: &AggregatedUsage,
226        code: &'static str,
227    ) -> bool {
228        let mut state = self.lock_state();
229        self.quarantine_locked(
230            &mut state,
231            aggregate,
232            MeterEventDeadLetterReason::AsynchronousProviderRejection { code },
233        )
234    }
235
236    /// Whether an incomplete batch is still owed a retry.
237    ///
238    /// While this is true, [`BatchExporter::export`] accepts only the identical
239    /// original batch and refuses any other.
240    #[must_use]
241    pub fn retry_in_progress(&self) -> bool {
242        self.lock_state().retry_progress.is_some()
243    }
244
245    /// Give up on an incomplete batch so a different one can be exported.
246    ///
247    /// Returns the identifiers that were not confirmed complete. If an export is
248    /// currently running, no progress is abandoned and an empty vector is
249    /// returned. Callers can distinguish that case with [`Self::retry_in_progress`].
250    #[must_use]
251    pub fn abandon_retry_progress(&self) -> Vec<String> {
252        let Ok(_lease) = self.acquire_lease() else {
253            return Vec::new();
254        };
255        let Some(progress) = self.lock_state().retry_progress.take() else {
256            return Vec::new();
257        };
258        progress
259            .batch
260            .into_iter()
261            .zip(progress.completed)
262            .filter_map(|(aggregate, completed)| (!completed).then_some(aggregate.identifier))
263            .collect()
264    }
265
266    fn lock_state(&self) -> MutexGuard<'_, ExporterState> {
267        self.state
268            .lock()
269            .unwrap_or_else(std::sync::PoisonError::into_inner)
270    }
271
272    fn acquire_lease(&self) -> Result<ExportLease<'_>, ExportError> {
273        self.exporting
274            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
275            .map_err(|_| {
276                ExportError::Provider("meter event export already in progress".to_owned())
277            })?;
278        Ok(ExportLease {
279            exporting: &self.exporting,
280        })
281    }
282
283    fn begin_export(&self, batch: &[AggregatedUsage]) -> Result<ExportLease<'_>, ExportError> {
284        let lease = self.acquire_lease()?;
285        let mut state = self.lock_state();
286        match &state.retry_progress {
287            Some(progress) if progress.batch != batch => {
288                return Err(ExportError::Provider(
289                    "meter event exporter requires the incomplete batch to be retried first; \
290                     call abandon_retry_progress to give up on it"
291                        .to_owned(),
292                ));
293            }
294            Some(_) => {}
295            None => {
296                state.retry_progress = Some(RetryProgress {
297                    batch: batch.to_vec(),
298                    completed: vec![false; batch.len()],
299                });
300            }
301        }
302        drop(state);
303        Ok(lease)
304    }
305
306    fn unresolved(&self) -> Vec<(usize, AggregatedUsage)> {
307        let state = self.lock_state();
308        let Some(progress) = &state.retry_progress else {
309            return Vec::new();
310        };
311        progress
312            .batch
313            .iter()
314            .enumerate()
315            .filter(|(index, _)| !progress.completed[*index])
316            .map(|(index, aggregate)| (index, aggregate.clone()))
317            .collect()
318    }
319
320    fn apply_outcomes(
321        &self,
322        unresolved: &[(usize, AggregatedUsage)],
323        outcomes: Vec<MeterEventOutcome>,
324    ) -> bool {
325        let mut state = self.lock_state();
326        let mut has_retryable = false;
327        for ((index, aggregate), outcome) in unresolved.iter().zip(outcomes) {
328            match outcome {
329                MeterEventOutcome::Accepted => {
330                    Self::mark_completed_locked(&mut state, *index);
331                }
332                MeterEventOutcome::RetryableFailure { .. } => {
333                    has_retryable = true;
334                }
335                MeterEventOutcome::PermanentRejection { code } => {
336                    self.quarantine_locked(
337                        &mut state,
338                        aggregate,
339                        MeterEventDeadLetterReason::SynchronousProviderRejection { code },
340                    );
341                    Self::mark_completed_locked(&mut state, *index);
342                }
343            }
344        }
345        has_retryable
346    }
347
348    fn mark_completed_locked(state: &mut ExporterState, index: usize) {
349        if let Some(completed) = state
350            .retry_progress
351            .as_mut()
352            .and_then(|progress| progress.completed.get_mut(index))
353        {
354            *completed = true;
355        }
356    }
357
358    fn quarantine_locked(
359        &self,
360        state: &mut ExporterState,
361        aggregate: &AggregatedUsage,
362        reason: MeterEventDeadLetterReason,
363    ) -> bool {
364        if state
365            .dead_letter_identifiers
366            .contains(&aggregate.identifier)
367        {
368            return false;
369        }
370        if self.dead_letter_capacity == 0 {
371            state.dropped_dead_letters = state.dropped_dead_letters.saturating_add(1);
372            return false;
373        }
374        while state.dead_letters.len() >= self.dead_letter_capacity {
375            if let Some(evicted) = state.dead_letters.pop_front() {
376                state
377                    .dead_letter_identifiers
378                    .remove(&evicted.aggregate.identifier);
379                state.dropped_dead_letters = state.dropped_dead_letters.saturating_add(1);
380            }
381        }
382        state
383            .dead_letter_identifiers
384            .insert(aggregate.identifier.clone());
385        state.dead_letters.push_back(MeterEventDeadLetter {
386            aggregate: aggregate.clone(),
387            reason,
388        });
389        true
390    }
391
392    fn commit_export(&self) {
393        self.lock_state().retry_progress = None;
394    }
395}
396
397impl<P> fmt::Debug for MeterEventExporter<P> {
398    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
399        formatter
400            .debug_struct("MeterEventExporter")
401            .field("provider", &"[REDACTED]")
402            .field("dead_letters", &self.dead_letter_count())
403            .field("dropped_dead_letters", &self.dropped_dead_letters())
404            .field("dead_letter_capacity", &self.dead_letter_capacity)
405            .field(
406                "export_in_progress",
407                &self.exporting.load(Ordering::Relaxed),
408            )
409            .field("retry_in_progress", &self.retry_in_progress())
410            .finish_non_exhaustive()
411    }
412}
413
414impl<P: MeterEventProvider> BatchExporter for MeterEventExporter<P> {
415    fn retry_policy(&self) -> crate::RetryPolicy {
416        crate::RetryPolicy::CompleteBatch
417    }
418
419    fn export<'a>(&'a self, batch: &'a [AggregatedUsage]) -> ExportFuture<'a> {
420        Box::pin(async move {
421            let _lease = self.begin_export(batch)?;
422            let unresolved = self.unresolved();
423            if unresolved.is_empty() {
424                self.commit_export();
425                return Ok(());
426            }
427            let submitted = unresolved
428                .iter()
429                .map(|(_, aggregate)| aggregate.clone())
430                .collect::<Vec<_>>();
431            let outcomes = self
432                .provider
433                .submit(&submitted)
434                .await
435                .map_err(|error| ExportError::Provider(error.to_string()))?;
436            if outcomes.len() != unresolved.len() {
437                return Err(ExportError::Provider(
438                    "meter event provider returned an invalid outcome count".to_owned(),
439                ));
440            }
441            if self.apply_outcomes(&unresolved, outcomes) {
442                return Err(ExportError::Provider(
443                    "meter event provider reported retryable events".to_owned(),
444                ));
445            }
446            self.commit_export();
447            Ok(())
448        })
449    }
450}
mcp_usage_export/meter_event.rs

mcp_usage_export/
meter_event.rs
