fix: harden concurrency limits and high-RPM runtime paths

Bound request, stream, queue, and shutdown resource lifetimes. Reduce scheduler and Redis hot-path work and isolate database maintenance. Include regression coverage, load probes, and concurrency audit results.
This commit is contained in:
elky
2026-09-10 08:14:58 +08:00
parent 361952ada9
commit ecc16673eb
149 changed files with 27963 additions and 1926 deletions
@@ -623,6 +623,10 @@ fn append_body_capture_metadata_entry(
);
}
pub(crate) fn mark_usage_event_capture_truncated(metadata: &mut Option<Value>, key: &str) {
aether_data_contracts::repository::usage::mark_usage_capture_memory_omitted(metadata, key);
}
fn upsert_body_capture_metadata_value_entry(
metadata: &mut Option<Value>,
key: &str,
+41
View File
@@ -15,6 +15,7 @@ pub struct UsageRuntimeConfig {
pub consumer_group: String,
pub dlq_stream_key: String,
pub stream_maxlen: usize,
pub queue_payload_max_bytes: usize,
pub consumer_batch_size: usize,
pub consumer_block_ms: u64,
pub reclaim_idle_ms: u64,
@@ -47,6 +48,7 @@ impl Default for UsageRuntimeConfig {
consumer_group: "usage_consumers".to_string(),
dlq_stream_key: "usage:events:dlq".to_string(),
stream_maxlen: 200_000,
queue_payload_max_bytes: 1024 * 1024,
consumer_batch_size: 128,
consumer_block_ms: 500,
reclaim_idle_ms: 60_000,
@@ -90,6 +92,11 @@ impl UsageRuntimeConfig {
"usage runtime dlq_stream_key cannot be empty".to_string(),
));
}
if self.stream_key == self.dlq_stream_key {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime stream_key and dlq_stream_key must be different".to_string(),
));
}
if self.worker_count == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime worker_count must be positive".to_string(),
@@ -124,6 +131,11 @@ impl UsageRuntimeConfig {
"usage runtime stream_maxlen must be positive".to_string(),
));
}
if self.queue_payload_max_bytes == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime queue_payload_max_bytes must be positive".to_string(),
));
}
if self.consumer_batch_size == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime consumer_batch_size must be positive".to_string(),
@@ -213,6 +225,19 @@ mod tests {
assert!(config.validate().is_err());
}
#[test]
fn enabled_config_rejects_dead_letter_stream_equal_to_source() {
let mut config = UsageRuntimeConfig::default();
config.dlq_stream_key = config.stream_key.clone();
assert!(config.validate().is_ok());
config.enabled = true;
assert!(matches!(
config.validate(),
Err(aether_data_contracts::DataLayerError::InvalidConfiguration(message))
if message.contains("must be different")
));
}
#[test]
fn enabled_config_rejects_zero_terminal_submission_limit() {
let config = UsageRuntimeConfig {
@@ -222,4 +247,20 @@ mod tests {
};
assert!(config.validate().is_err());
}
#[test]
fn queue_payload_limit_defaults_to_one_mib_and_rejects_zero_when_enabled() {
let mut config = UsageRuntimeConfig::default();
assert_eq!(config.queue_payload_max_bytes, 1024 * 1024);
config.queue_payload_max_bytes = 0;
assert!(config.validate().is_ok());
config.enabled = true;
assert!(matches!(
config.validate(),
Err(aether_data_contracts::DataLayerError::InvalidConfiguration(message))
if message.contains("queue_payload_max_bytes")
));
config.queue_payload_max_bytes = 1;
assert!(config.validate().is_ok());
}
}
@@ -0,0 +1,541 @@
use std::collections::BTreeMap;
use std::io::{self, Write};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, LazyLock};
use aether_data_contracts::DataLayerError;
use aether_runtime_state::RuntimeQueueEntry;
use serde::Serialize;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
const DEFAULT_ENCODING_BUDGET_BYTES: usize = 64 * 1024 * 1024;
const DEFAULT_ENCODING_JOBS: usize = 4;
const MAX_ENCODING_JOBS: usize = 128;
static ENCODING_BUDGET: LazyLock<Arc<DeadLetterEncodingBudget>> = LazyLock::new(|| {
Arc::new(DeadLetterEncodingBudget::new(
configured_limit(
std::env::var("AETHER_USAGE_DLQ_ENCODING_BUDGET_BYTES")
.ok()
.as_deref(),
DEFAULT_ENCODING_BUDGET_BYTES,
maximum_budget_bytes(),
),
configured_limit(
std::env::var("AETHER_USAGE_DLQ_ENCODING_MAX_JOBS")
.ok()
.as_deref(),
DEFAULT_ENCODING_JOBS,
MAX_ENCODING_JOBS,
),
))
});
pub(crate) fn shared_dead_letter_encoding_budget() -> Arc<DeadLetterEncodingBudget> {
Arc::clone(&ENCODING_BUDGET)
}
pub(crate) fn dead_letter_encoding_metrics() -> DeadLetterEncodingSnapshot {
ENCODING_BUDGET.snapshot()
}
fn maximum_budget_bytes() -> usize {
Semaphore::MAX_PERMITS.min(u32::MAX as usize)
}
fn configured_limit(raw: Option<&str>, fallback: usize, maximum: usize) -> usize {
raw.and_then(|raw| raw.trim().parse::<u128>().ok())
.filter(|value| *value > 0)
.map(|value| value.min(maximum as u128) as usize)
.unwrap_or(fallback.min(maximum))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct DeadLetterEncodingSnapshot {
pub(crate) limit_bytes: usize,
pub(crate) job_limit: usize,
pub(crate) reserved_bytes: usize,
pub(crate) active_jobs: usize,
pub(crate) capacity_rejected_total: u64,
pub(crate) oversized_rejected_total: u64,
pub(crate) encoded_total: u64,
}
/// Reserves raw string lengths plus their worst-case JSON encoding before any
/// clone or encoding allocation. This excludes collection/allocation overhead and
/// Redis command, packed-command, and connection buffers; it is not an RSS limit.
pub(crate) struct DeadLetterEncodingBudget {
limit_bytes: usize,
job_limit: usize,
bytes: Arc<Semaphore>,
jobs: Arc<Semaphore>,
reserved_bytes: AtomicUsize,
active_jobs: AtomicUsize,
capacity_rejected_total: AtomicU64,
oversized_rejected_total: AtomicU64,
encoded_total: AtomicU64,
#[cfg(test)]
encode_hook: std::sync::Mutex<Option<Box<dyn FnOnce() + Send>>>,
}
impl DeadLetterEncodingBudget {
pub(crate) fn new(limit_bytes: usize, job_limit: usize) -> Self {
let limit_bytes = limit_bytes.min(maximum_budget_bytes());
let job_limit = job_limit.clamp(1, MAX_ENCODING_JOBS);
Self {
limit_bytes,
job_limit,
bytes: Arc::new(Semaphore::new(limit_bytes)),
jobs: Arc::new(Semaphore::new(job_limit)),
reserved_bytes: AtomicUsize::new(0),
active_jobs: AtomicUsize::new(0),
capacity_rejected_total: AtomicU64::new(0),
oversized_rejected_total: AtomicU64::new(0),
encoded_total: AtomicU64::new(0),
#[cfg(test)]
encode_hook: std::sync::Mutex::new(None),
}
}
pub(crate) fn snapshot(&self) -> DeadLetterEncodingSnapshot {
DeadLetterEncodingSnapshot {
limit_bytes: self.limit_bytes,
job_limit: self.job_limit,
reserved_bytes: self.reserved_bytes.load(Ordering::Relaxed),
active_jobs: self.active_jobs.load(Ordering::Relaxed),
capacity_rejected_total: self.capacity_rejected_total.load(Ordering::Relaxed),
oversized_rejected_total: self.oversized_rejected_total.load(Ordering::Relaxed),
encoded_total: self.encoded_total.load(Ordering::Relaxed),
}
}
pub(crate) fn try_reserve(
self: &Arc<Self>,
entry: &RuntimeQueueEntry,
error: &str,
) -> Result<DeadLetterEncodingReservation, DataLayerError> {
let size = encoding_size(entry, error).filter(|size| size.total <= self.limit_bytes);
let Some(size) = size else {
self.oversized_rejected_total
.fetch_add(1, Ordering::Relaxed);
return Err(DataLayerError::InvalidInput(format!(
"dead-letter raw fields and worst-case JSON exceed the {}-byte encoding budget",
self.limit_bytes
)));
};
let job_permit = Arc::clone(&self.jobs)
.try_acquire_owned()
.map_err(|_| self.capacity_error())?;
let byte_permit = Arc::clone(&self.bytes)
.try_acquire_many_owned(size.total as u32)
.map_err(|_| self.capacity_error())?;
self.reserved_bytes.fetch_add(size.total, Ordering::Relaxed);
self.active_jobs.fetch_add(1, Ordering::Relaxed);
Ok(DeadLetterEncodingReservation {
budget: Arc::clone(self),
size,
_byte_permit: byte_permit,
_job_permit: job_permit,
})
}
fn capacity_error(&self) -> DataLayerError {
self.capacity_rejected_total.fetch_add(1, Ordering::Relaxed);
DataLayerError::TimedOut("dead-letter encoding capacity is exhausted".to_string())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct EncodingSize {
json: usize,
total: usize,
}
fn encoding_size(entry: &RuntimeQueueEntry, error: &str) -> Option<EncodingSize> {
let mut raw = entry.id.len().checked_add(error.len())?;
for (key, value) in &entry.fields {
raw = raw.checked_add(key.len())?.checked_add(value.len())?;
}
checked_encoding_size(raw, entry.fields.len())
}
fn checked_encoding_size(raw: usize, field_count: usize) -> Option<EncodingSize> {
// Every string byte needs at most six bytes (\u00XX). Each map entry adds
// four quotes, a colon and at most one comma to the empty envelope.
const EMPTY_ENVELOPE_BYTES: usize = br#"{"entry_id":"","fields":{},"error":""}"#.len();
let json = raw
.checked_mul(6)?
.checked_add(field_count.checked_mul(6)?)?
.checked_add(EMPTY_ENVELOPE_BYTES)?;
Some(EncodingSize {
json,
total: raw.checked_add(json)?,
})
}
pub(crate) struct DeadLetterEncodingReservation {
budget: Arc<DeadLetterEncodingBudget>,
size: EncodingSize,
_byte_permit: OwnedSemaphorePermit,
_job_permit: OwnedSemaphorePermit,
}
impl DeadLetterEncodingReservation {
pub(crate) fn encode_owned(
self,
entry: RuntimeQueueEntry,
error: String,
) -> impl std::future::Future<Output = Result<EncodedDeadLetter, DataLayerError>> + Send {
let input = EncodingInput {
entry,
error,
reservation: self,
};
async move {
tokio::task::spawn_blocking(move || input.encode())
.await
.map_err(|error| {
DataLayerError::UnexpectedValue(format!(
"dead-letter encoding task failed: {error}"
))
})?
}
}
}
impl Drop for DeadLetterEncodingReservation {
fn drop(&mut self) {
self.budget
.reserved_bytes
.fetch_sub(self.size.total, Ordering::Relaxed);
self.budget.active_jobs.fetch_sub(1, Ordering::Relaxed);
}
}
// Field order also protects cancellation/panic: raw data is dropped before its
// reservation, including when a queued blocking task never starts.
struct EncodingInput {
entry: RuntimeQueueEntry,
error: String,
reservation: DeadLetterEncodingReservation,
}
#[derive(Serialize)]
struct DeadLetterPayload<'a> {
entry_id: &'a str,
fields: &'a BTreeMap<String, String>,
error: &'a str,
}
impl EncodingInput {
fn encode(self) -> Result<EncodedDeadLetter, DataLayerError> {
#[cfg(test)]
{
let hook = self.reservation.budget.encode_hook.lock().unwrap().take();
if let Some(hook) = hook {
hook();
}
}
let mut writer = BoundedJsonWriter::new(self.reservation.size.json);
serde_json::to_writer(
&mut writer,
&DeadLetterPayload {
entry_id: &self.entry.id,
fields: &self.entry.fields,
error: &self.error,
},
)
.map_err(|error| {
DataLayerError::UnexpectedValue(format!(
"failed to encode complete dead-letter fields: {error}"
))
})?;
let payload = String::from_utf8(writer.bytes).map_err(|error| {
DataLayerError::UnexpectedValue(format!("dead-letter JSON was not UTF-8: {error}"))
})?;
self.reservation
.budget
.encoded_total
.fetch_add(1, Ordering::Relaxed);
Ok(EncodedDeadLetter {
entry_id: self.entry.id,
fields: BTreeMap::from([("payload".to_string(), payload)]),
_reservation: self.reservation,
})
}
}
pub(crate) struct EncodedDeadLetter {
pub(crate) entry_id: String,
pub(crate) fields: BTreeMap<String, String>,
// Remains owned by the result until transfer/append finishes.
_reservation: DeadLetterEncodingReservation,
}
struct BoundedJsonWriter {
bytes: Vec<u8>,
max_bytes: usize,
}
impl BoundedJsonWriter {
fn new(max_bytes: usize) -> Self {
Self {
bytes: Vec::new(),
max_bytes,
}
}
}
impl Write for BoundedJsonWriter {
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
if bytes.len() > self.max_bytes.saturating_sub(self.bytes.len()) {
return Err(io::Error::other("dead-letter JSON encoding bound exceeded"));
}
let required = self.bytes.len() + bytes.len();
if required > self.bytes.capacity() {
let capacity = required
.max(self.bytes.capacity().saturating_mul(2))
.min(self.max_bytes);
self.bytes
.try_reserve_exact(capacity - self.bytes.len())
.map_err(|error| {
io::Error::other(format!("dead-letter JSON allocation failed: {error}"))
})?;
}
self.bytes.extend_from_slice(bytes);
Ok(bytes.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
fn test_entry() -> RuntimeQueueEntry {
RuntimeQueueEntry {
id: "10-3".to_string(),
fields: BTreeMap::from([
("payload".to_string(), "historical payload".to_string()),
("extra".to_string(), "original metadata".to_string()),
]),
}
}
#[tokio::test]
async fn dead_letter_encoding_preserves_complete_wire_and_all_escape_forms() {
let control_bytes = (0u8..=31).map(char::from).collect::<String>();
let mut entry = test_entry();
entry.id.push_str("\"\\\n");
entry.fields.insert(
format!("{control_bytes}\"\\"),
format!("{control_bytes}\"\\\u{4e2d}\u{6587}\u{1f600}"),
);
let error = format!("error:{control_bytes}\"\\\u{00e9}");
let expected = serde_json::to_string(&DeadLetterPayload {
entry_id: &entry.id,
fields: &entry.fields,
error: &error,
})
.unwrap();
let size = encoding_size(&entry, &error).unwrap();
assert!(expected.len() <= size.json);
let budget = Arc::new(DeadLetterEncodingBudget::new(size.total, 1));
let encoded = budget
.try_reserve(&entry, &error)
.unwrap()
.encode_owned(entry.clone(), error.clone())
.await
.unwrap();
assert_eq!(encoded.entry_id, entry.id);
assert_eq!(encoded.fields.len(), 1);
assert_eq!(encoded.fields["payload"], expected);
let decoded: serde_json::Value = serde_json::from_str(&encoded.fields["payload"]).unwrap();
assert_eq!(
decoded["fields"],
serde_json::to_value(entry.fields).unwrap()
);
assert_eq!(decoded["error"], error);
assert_eq!(budget.snapshot().encoded_total, 1);
assert_eq!(budget.snapshot().reserved_bytes, size.total);
assert_eq!(budget.snapshot().active_jobs, 1);
drop(encoded);
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert_eq!(budget.snapshot().active_jobs, 0);
}
#[test]
fn dead_letter_encoding_budget_rejects_oversize_and_saturation_without_waiters() {
let entry = test_entry();
let size = encoding_size(&entry, "failure").unwrap();
let too_small = Arc::new(DeadLetterEncodingBudget::new(size.total - 1, 1));
assert!(matches!(
too_small.try_reserve(&entry, "failure"),
Err(DataLayerError::InvalidInput(_))
));
assert_eq!(too_small.snapshot().oversized_rejected_total, 1);
assert_eq!(too_small.snapshot().reserved_bytes, 0);
assert_eq!(too_small.snapshot().active_jobs, 0);
for (limit, jobs) in [(size.total, 2), (size.total * 2, 1)] {
let budget = Arc::new(DeadLetterEncodingBudget::new(limit, jobs));
let first = budget.try_reserve(&entry, "failure").unwrap();
assert!(matches!(
budget.try_reserve(&entry, "failure"),
Err(DataLayerError::TimedOut(_))
));
assert_eq!(budget.snapshot().capacity_rejected_total, 1);
assert_eq!(budget.snapshot().reserved_bytes, size.total);
assert_eq!(budget.snapshot().active_jobs, 1);
assert_eq!(budget.jobs.available_permits(), jobs - 1);
drop(first);
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert_eq!(budget.jobs.available_permits(), jobs);
assert_eq!(budget.bytes.available_permits(), limit);
drop(budget.try_reserve(&entry, "failure").unwrap());
}
}
#[test]
fn dead_letter_encoding_bounds_and_environment_cannot_overflow() {
assert!(checked_encoding_size(usize::MAX, 0).is_none());
assert!(checked_encoding_size(usize::MAX / 6, 0).is_none());
assert!(checked_encoding_size(0, usize::MAX).is_none());
assert!(checked_encoding_size(usize::MAX / 7, 1).is_none());
assert!(checked_encoding_size(0, 0).unwrap().total > 0);
assert_eq!(configured_limit(None, 4, 128), 4);
assert_eq!(configured_limit(Some("0"), 4, 128), 4);
assert_eq!(configured_limit(Some("invalid"), 4, 128), 4);
assert_eq!(configured_limit(Some(" 2 "), 4, 128), 2);
assert_eq!(configured_limit(Some("99999999999"), 4, 128), 128);
assert_eq!(
configured_limit(Some(&u128::MAX.to_string()), 4, maximum_budget_bytes()),
maximum_budget_bytes()
);
let budget = DeadLetterEncodingBudget::new(usize::MAX, usize::MAX);
assert_eq!(budget.snapshot().limit_bytes, maximum_budget_bytes());
assert_eq!(budget.snapshot().job_limit, MAX_ENCODING_JOBS);
}
#[test]
fn dead_letter_encoding_bounded_writer_grows_geometrically_and_stops_at_limit() {
let mut writer = BoundedJsonWriter::new(4096);
let mut allocations = 0;
for _ in 0..4096 {
let previous_capacity = writer.bytes.capacity();
writer.write_all(b"x").unwrap();
allocations += usize::from(previous_capacity != writer.bytes.capacity());
}
assert!(allocations <= 13, "allocations: {allocations}");
assert!(writer.write_all(b"y").is_err());
assert_eq!(writer.bytes.len(), 4096);
assert!(writer.bytes.iter().all(|byte| *byte == b'x'));
}
#[tokio::test]
async fn dead_letter_encoding_unpolled_cancellation_releases_reservation() {
let entry = test_entry();
let budget = Arc::new(DeadLetterEncodingBudget::new(4096, 1));
let reservation = budget.try_reserve(&entry, "failure").unwrap();
let future = reservation.encode_owned(entry, "failure".to_string());
assert_eq!(budget.snapshot().active_jobs, 1);
drop(future);
assert_eq!(budget.snapshot().active_jobs, 0);
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert_eq!(budget.snapshot().encoded_total, 0);
}
#[tokio::test]
async fn dead_letter_encoding_running_cancellation_holds_budget_until_closure_exits() {
let entry = test_entry();
let budget = Arc::new(DeadLetterEncodingBudget::new(4096, 1));
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
let (release_tx, release_rx) = std::sync::mpsc::channel();
*budget.encode_hook.lock().unwrap() = Some(Box::new(move || {
let _ = started_tx.send(());
release_rx.recv_timeout(Duration::from_secs(2)).unwrap();
}));
let reservation = budget.try_reserve(&entry, "failure").unwrap();
let task = tokio::spawn(reservation.encode_owned(entry, "failure".to_string()));
tokio::time::timeout(Duration::from_secs(2), started_rx)
.await
.unwrap()
.unwrap();
task.abort();
assert!(matches!(task.await, Err(error) if error.is_cancelled()));
assert_eq!(budget.snapshot().active_jobs, 1);
assert!(budget.snapshot().reserved_bytes > 0);
assert!(budget.try_reserve(&test_entry(), "failure").is_err());
release_tx.send(()).unwrap();
tokio::time::timeout(Duration::from_secs(2), async {
while budget.snapshot().active_jobs != 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert_eq!(budget.snapshot().encoded_total, 1);
}
#[tokio::test]
async fn dead_letter_encoding_panic_releases_input_and_reservation() {
let entry = test_entry();
let budget = Arc::new(DeadLetterEncodingBudget::new(4096, 1));
*budget.encode_hook.lock().unwrap() = Some(Box::new(|| panic!("encoding test panic")));
let result = budget
.try_reserve(&entry, "failure")
.unwrap()
.encode_owned(entry, "failure".to_string())
.await;
assert!(matches!(result, Err(DataLayerError::UnexpectedValue(_))));
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert_eq!(budget.snapshot().active_jobs, 0);
assert_eq!(budget.snapshot().encoded_total, 0);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn dead_letter_encoding_concurrent_reservations_have_no_hidden_waiting_queue() {
const TASKS: usize = 16;
let entry = Arc::new(test_entry());
let size = encoding_size(&entry, "failure").unwrap();
let budget = Arc::new(DeadLetterEncodingBudget::new(size.total * 2, 2));
let ready = Arc::new(tokio::sync::Barrier::new(TASKS + 1));
let release = Arc::new(tokio::sync::Barrier::new(TASKS + 1));
let mut tasks = tokio::task::JoinSet::new();
for _ in 0..TASKS {
let entry = Arc::clone(&entry);
let budget = Arc::clone(&budget);
let ready = Arc::clone(&ready);
let release = Arc::clone(&release);
tasks.spawn(async move {
let reservation = budget.try_reserve(&entry, "failure");
ready.wait().await;
release.wait().await;
reservation.is_ok()
});
}
tokio::time::timeout(Duration::from_secs(2), ready.wait())
.await
.unwrap();
assert_eq!(budget.snapshot().active_jobs, 2);
assert_eq!(budget.snapshot().reserved_bytes, size.total * 2);
assert_eq!(
budget.snapshot().capacity_rejected_total,
(TASKS - 2) as u64
);
release.wait().await;
let mut admitted = 0;
while let Some(result) = tasks.join_next().await {
admitted += usize::from(result.unwrap());
}
assert_eq!(admitted, 2);
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert_eq!(budget.snapshot().active_jobs, 0);
}
}
+766 -6
View File
@@ -1,4 +1,5 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
@@ -6,8 +7,18 @@ use aether_data_contracts::DataLayerError;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::body_capture::mark_usage_event_capture_truncated;
pub use crate::event_capture_budget::UsageEventCaptureRetention;
use crate::event_capture_budget::{
json_heap_estimate, shared_capture_memory_budget, EventCaptureMemoryBudget,
};
pub const USAGE_EVENT_VERSION: u8 = 1;
#[path = "event_wire.rs"]
mod wire;
pub(crate) use wire::EncodedUsageEvent;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UsageEventType {
@@ -18,7 +29,7 @@ pub enum UsageEventType {
Cancelled,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
#[derive(Debug, PartialEq, Serialize, Deserialize, Default)]
pub struct UsageEventData {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub user_id: Option<String>,
@@ -144,6 +155,171 @@ pub struct UsageEventData {
pub local_execution_runtime_miss_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_metadata: Option<Value>,
#[doc(hidden)]
#[serde(skip)]
pub capture_retention: UsageEventCaptureRetention,
}
impl UsageEventData {
fn capture_heap_estimate(&self) -> usize {
[
&self.request_body,
&self.provider_request_body,
&self.response_body,
&self.client_response_body,
]
.into_iter()
.flatten()
.fold(0usize, |bytes, body| {
bytes
.saturating_add(std::mem::size_of::<Value>())
.saturating_add(json_heap_estimate(body))
})
}
fn captured_fields(&self) -> [bool; 4] {
[
self.request_body.is_some(),
self.provider_request_body.is_some(),
self.response_body.is_some(),
self.client_response_body.is_some(),
]
}
fn mark_capture_omitted(&mut self, captured: [bool; 4]) {
for (present, key, state) in [
(captured[0], "request", &mut self.request_body_state),
(
captured[1],
"provider_request",
&mut self.provider_request_body_state,
),
(captured[2], "response", &mut self.response_body_state),
(
captured[3],
"client_response",
&mut self.client_response_body_state,
),
] {
if present
&& !matches!(
*state,
Some(
UsageBodyCaptureState::None
| UsageBodyCaptureState::Disabled
| UsageBodyCaptureState::Unavailable
)
)
{
*state = Some(UsageBodyCaptureState::Truncated);
mark_usage_event_capture_truncated(&mut self.request_metadata, key);
}
}
}
pub(crate) fn apply_capture_memory_budget(
&mut self,
budget: std::sync::Arc<EventCaptureMemoryBudget>,
) {
let bytes = self.capture_heap_estimate();
if self
.capture_retention
.reserve(std::sync::Arc::clone(&budget), bytes)
{
return;
}
let captured = self.captured_fields();
self.request_body = None;
self.provider_request_body = None;
self.response_body = None;
self.client_response_body = None;
self.mark_capture_omitted(captured);
// The previous lease is released only after the owned JSON bodies are gone.
self.capture_retention.clear(budget);
}
}
impl Clone for UsageEventData {
fn clone(&self) -> Self {
let (capture_retention, retain_bodies) = self
.capture_retention
.clone_for_bodies(|| self.capture_heap_estimate());
// Enumerate every field so additions require an explicit ownership decision.
let mut cloned = Self {
user_id: self.user_id.clone(),
api_key_id: self.api_key_id.clone(),
username: self.username.clone(),
api_key_name: self.api_key_name.clone(),
provider_name: self.provider_name.clone(),
model: self.model.clone(),
target_model: self.target_model.clone(),
model_id: self.model_id.clone(),
global_model_id: self.global_model_id.clone(),
provider_id: self.provider_id.clone(),
provider_endpoint_id: self.provider_endpoint_id.clone(),
provider_api_key_id: self.provider_api_key_id.clone(),
request_type: self.request_type.clone(),
api_format: self.api_format.clone(),
api_family: self.api_family.clone(),
endpoint_kind: self.endpoint_kind.clone(),
endpoint_api_format: self.endpoint_api_format.clone(),
provider_api_family: self.provider_api_family.clone(),
provider_endpoint_kind: self.provider_endpoint_kind.clone(),
has_format_conversion: self.has_format_conversion,
is_stream: self.is_stream,
input_tokens: self.input_tokens,
output_tokens: self.output_tokens,
total_tokens: self.total_tokens,
cache_creation_input_tokens: self.cache_creation_input_tokens,
cache_creation_ephemeral_5m_input_tokens: self.cache_creation_ephemeral_5m_input_tokens,
cache_creation_ephemeral_1h_input_tokens: self.cache_creation_ephemeral_1h_input_tokens,
cache_read_input_tokens: self.cache_read_input_tokens,
cache_creation_cost_usd: self.cache_creation_cost_usd,
cache_read_cost_usd: self.cache_read_cost_usd,
output_price_per_1m: self.output_price_per_1m,
total_cost_usd: self.total_cost_usd,
actual_total_cost_usd: self.actual_total_cost_usd,
status_code: self.status_code,
error_message: self.error_message.clone(),
error_category: self.error_category.clone(),
response_time_ms: self.response_time_ms,
first_byte_time_ms: self.first_byte_time_ms,
request_headers: self.request_headers.clone(),
request_body: retain_bodies.then(|| self.request_body.clone()).flatten(),
request_body_ref: self.request_body_ref.clone(),
request_body_state: self.request_body_state,
provider_request_headers: self.provider_request_headers.clone(),
provider_request_body: retain_bodies
.then(|| self.provider_request_body.clone())
.flatten(),
provider_request_body_ref: self.provider_request_body_ref.clone(),
provider_request_body_state: self.provider_request_body_state,
response_headers: self.response_headers.clone(),
response_body: retain_bodies.then(|| self.response_body.clone()).flatten(),
response_body_ref: self.response_body_ref.clone(),
response_body_state: self.response_body_state,
client_response_headers: self.client_response_headers.clone(),
client_response_body: retain_bodies
.then(|| self.client_response_body.clone())
.flatten(),
client_response_body_ref: self.client_response_body_ref.clone(),
client_response_body_state: self.client_response_body_state,
candidate_id: self.candidate_id.clone(),
candidate_index: self.candidate_index,
key_name: self.key_name.clone(),
planner_kind: self.planner_kind.clone(),
route_family: self.route_family.clone(),
route_kind: self.route_kind.clone(),
execution_path: self.execution_path.clone(),
local_execution_runtime_miss_reason: self.local_execution_runtime_miss_reason.clone(),
request_metadata: self.request_metadata.clone(),
capture_retention,
};
if !retain_bodies {
cloned.mark_capture_omitted(self.captured_fields());
}
cloned
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
@@ -164,6 +340,16 @@ struct UsageEventEnvelope {
data: UsageEventData,
}
#[derive(Serialize)]
struct BorrowedUsageEventEnvelope<'a, T: ?Sized> {
v: u8,
#[serde(rename = "type")]
event_type: UsageEventType,
request_id: &'a str,
timestamp_ms: u64,
data: &'a T,
}
impl UsageEvent {
pub fn new(
event_type: UsageEventType,
@@ -179,12 +365,12 @@ impl UsageEvent {
}
pub fn to_stream_fields(&self) -> Result<BTreeMap<String, String>, DataLayerError> {
let payload = UsageEventEnvelope {
let payload = BorrowedUsageEventEnvelope {
v: USAGE_EVENT_VERSION,
event_type: self.event_type,
request_id: self.request_id.clone(),
request_id: &self.request_id,
timestamp_ms: self.timestamp_ms,
data: self.data.clone(),
data: &self.data,
};
let payload = serde_json::to_string(&payload).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
@@ -194,7 +380,21 @@ impl UsageEvent {
Ok(BTreeMap::from([("payload".to_string(), payload)]))
}
pub(crate) fn to_bounded_stream_fields(
&self,
max_bytes: usize,
) -> Result<EncodedUsageEvent, DataLayerError> {
wire::encode(self, max_bytes)
}
pub fn from_stream_fields(fields: &BTreeMap<String, String>) -> Result<Self, DataLayerError> {
Self::from_stream_fields_with_capture_budget(fields, shared_capture_memory_budget())
}
pub(crate) fn from_stream_fields_with_capture_budget(
fields: &BTreeMap<String, String>,
budget: Arc<EventCaptureMemoryBudget>,
) -> Result<Self, DataLayerError> {
let payload = fields.get("payload").ok_or_else(|| {
DataLayerError::UnexpectedValue(
"usage event stream entry missing payload field".to_string(),
@@ -212,12 +412,16 @@ impl UsageEvent {
)));
}
Ok(Self {
let mut event = Self {
event_type: envelope.event_type,
request_id: envelope.request_id,
timestamp_ms: envelope.timestamp_ms,
data: envelope.data,
})
};
// The wire format has no ownership lease. Preserve billing facts before a decoded
// body can be omitted; the raw Redis response and serde allocation are not budgeted here.
crate::runtime::prepare_decoded_event_capture_memory(&mut event, budget);
Ok(event)
}
}
@@ -230,8 +434,564 @@ pub fn now_ms() -> u64 {
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::sync::Arc;
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
use aether_data_contracts::DataLayerError;
use serde_json::json;
use crate::event_capture_budget::EventCaptureMemoryBudget;
use crate::{
apply_usage_body_capture_policy_to_event, build_upsert_usage_record_from_event,
UsageBodyCapturePolicy,
};
use super::{UsageEvent, UsageEventData, UsageEventType};
fn captured_event() -> UsageEvent {
UsageEvent {
event_type: UsageEventType::Failed,
request_id: "capture-budget-request".to_string(),
timestamp_ms: 123_456,
data: UsageEventData {
provider_name: "provider".to_string(),
model: "model".to_string(),
input_tokens: Some(100),
output_tokens: Some(500),
total_tokens: Some(600),
cache_read_input_tokens: Some(0),
cache_creation_input_tokens: Some(25),
actual_total_cost_usd: Some(1.25),
status_code: Some(502),
error_category: Some("upstream_error".to_string()),
error_message: Some("upstream failed".to_string()),
request_body: Some(json!({"messages": [{"content": "request"}]})),
provider_request_body: Some(json!({"input": "upstream request"})),
response_body: Some(json!({"usage": {"input_tokens": 100, "output_tokens": 500}})),
client_response_body: Some(json!({"error": "client response"})),
request_body_state: Some(UsageBodyCaptureState::Inline),
provider_request_body_state: Some(UsageBodyCaptureState::Inline),
response_body_state: Some(UsageBodyCaptureState::Inline),
client_response_body_state: Some(UsageBodyCaptureState::Inline),
request_metadata: Some(json!({
"requested_reasoning_effort": "high",
"provider_reasoning_effort": "medium",
"provider_service_tier": "priority",
"provider_actual_service_tier": "default",
"provider_cache_ttl_minutes": 60,
"plan_usage_reservation_token": "550e8400-e29b-41d4-a716-446655440000",
"body_capture": {"response": {"state": "inline", "source_bytes": 1000}}
})),
..UsageEventData::default()
},
}
}
#[test]
fn event_capture_budget_zero_preserves_billing_refs_and_database_truncation() {
let budget = Arc::new(EventCaptureMemoryBudget::new(0));
let mut event = captured_event();
event.data.request_body_ref =
Some("usage://capture-budget-request/request_body".to_string());
event.data.response_body_ref =
Some("usage://capture-budget-request/response_body".to_string());
event.data.apply_capture_memory_budget(Arc::clone(&budget));
assert!(event.data.request_body.is_none());
assert!(event.data.provider_request_body.is_none());
assert!(event.data.response_body.is_none());
assert!(event.data.client_response_body.is_none());
let capture_metadata = event
.data
.request_metadata
.as_ref()
.expect("capture metadata");
assert_eq!(
capture_metadata["body_capture"]["response"]["source_bytes"],
1000
);
assert_eq!(
capture_metadata["body_capture"]["response"]["stored_bytes"],
0
);
assert_eq!(
capture_metadata["body_capture"]["response"]["reason"],
"usage_event_memory_budget_exceeded"
);
let record = build_upsert_usage_record_from_event(&event).expect("record mapping");
assert_eq!(record.status, "failed");
assert_eq!(record.input_tokens, Some(100));
assert_eq!(record.output_tokens, Some(500));
assert_eq!(record.cache_read_input_tokens, Some(0));
assert_eq!(record.cache_creation_input_tokens, Some(25));
assert_eq!(record.actual_total_cost_usd, Some(1.25));
assert_eq!(record.error_category.as_deref(), Some("upstream_error"));
assert_eq!(record.request_body_ref, event.data.request_body_ref);
assert_eq!(record.response_body_ref, event.data.response_body_ref);
for state in [
record.request_body_state,
record.provider_request_body_state,
record.response_body_state,
record.client_response_body_state,
] {
assert_eq!(state, Some(UsageBodyCaptureState::Truncated));
}
let metadata = record.request_metadata.expect("preserved metadata");
assert_eq!(metadata["provider_service_tier"], "priority");
assert_eq!(metadata["provider_actual_service_tier"], "default");
assert_eq!(metadata["provider_cache_ttl_minutes"], 60);
assert_eq!(
metadata["plan_usage_reservation_token"],
"550e8400-e29b-41d4-a716-446655440000"
);
// Persistence projects billing metadata; capture state remains in typed columns.
assert!(metadata.get("body_capture").is_none());
assert_eq!(budget.retained_bytes(), 0);
assert_eq!(budget.downgraded_total(), 1);
}
#[test]
fn event_capture_budget_clone_reserves_each_copy_before_cloning_bodies() {
let mut event = captured_event();
let weight = event.data.capture_heap_estimate();
let budget = Arc::new(EventCaptureMemoryBudget::new(weight * 2));
event.data.apply_capture_memory_budget(Arc::clone(&budget));
let copy = event.clone();
assert_eq!(copy, event);
assert_eq!(budget.retained_bytes(), weight * 2);
let downgraded = event.clone();
assert!(event.data.response_body.is_some());
assert!(copy.data.response_body.is_some());
assert!(downgraded.data.response_body.is_none());
assert_eq!(
downgraded.data.response_body_state,
Some(UsageBodyCaptureState::Truncated)
);
assert_eq!(downgraded.data.total_tokens, event.data.total_tokens);
assert_eq!(downgraded.timestamp_ms, event.timestamp_ms);
assert_eq!(downgraded.event_type, event.event_type);
assert_eq!(budget.retained_bytes(), weight * 2);
drop(copy);
assert_eq!(budget.retained_bytes(), weight);
drop((event, downgraded));
assert_eq!(budget.retained_bytes(), 0);
}
#[test]
fn event_capture_budget_serialization_borrows_bodies_without_charging_a_clone() {
let mut event = captured_event();
let weight = event.data.capture_heap_estimate();
let budget = Arc::new(EventCaptureMemoryBudget::new(weight));
crate::runtime::prepare_event_capture_memory(&mut event, Arc::clone(&budget));
let decoded_budget = Arc::new(EventCaptureMemoryBudget::new(usize::MAX));
for _ in 0..3 {
let fields = event.to_stream_fields().expect("wire serialization");
assert!(!fields["payload"].contains("capture_retention"));
let decoded = UsageEvent::from_stream_fields_with_capture_budget(
&fields,
Arc::clone(&decoded_budget),
)
.expect("wire decode");
assert_eq!(decoded, event);
assert_eq!(budget.retained_bytes(), weight);
assert!(decoded_budget.retained_bytes() > 0);
drop(decoded);
assert_eq!(decoded_budget.retained_bytes(), 0);
}
assert_eq!(budget.downgraded_total(), 0);
drop(event);
assert_eq!(budget.retained_bytes(), 0);
}
#[test]
fn from_stream_fields_legacy_body_budget_preserves_billing_and_request_facts() {
let mut event = captured_event();
event.data.model = "gpt-5.6-sol".to_string();
event.data.endpoint_api_format = Some("openai:responses".to_string());
event.data.request_body = Some(json!({"reasoning": {"effort": "high"}}));
event.data.provider_request_body = Some(json!({
"model": "gpt-5.6-sol", "reasoning": {"effort": "medium"},
"service_tier": "priority"
}));
event.data.response_body = Some(json!({"service_tier": "Default"}));
event.data.request_body_state = None;
event.data.provider_request_body_state = None;
event.data.response_body_state = None;
event.data.client_response_body_state = None;
event.data.cache_creation_ephemeral_5m_input_tokens = Some(0);
event.data.cache_creation_ephemeral_1h_input_tokens = Some(25);
event.data.cache_read_cost_usd = Some(0.0);
event.data.request_body_ref = Some("usage://legacy/request".to_string());
event.data.request_metadata = Some(json!({
"plan_usage_reservation_token": "550e8400-e29b-41d4-a716-446655440000"
}));
let fields = event.to_stream_fields().expect("legacy wire serialization");
assert!(!fields["payload"].contains("request_body_state"));
let budget = Arc::new(EventCaptureMemoryBudget::new(0));
let decoded =
UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&budget))
.expect("legacy wire decode");
assert_eq!(decoded.event_type, UsageEventType::Failed);
assert_eq!(decoded.request_id, event.request_id);
assert_eq!(decoded.timestamp_ms, event.timestamp_ms);
assert_eq!(decoded.data.input_tokens, Some(100));
assert_eq!(decoded.data.output_tokens, Some(500));
assert_eq!(decoded.data.total_tokens, Some(600));
assert_eq!(decoded.data.cache_creation_input_tokens, Some(25));
assert_eq!(
decoded.data.cache_creation_ephemeral_5m_input_tokens,
Some(0)
);
assert_eq!(
decoded.data.cache_creation_ephemeral_1h_input_tokens,
Some(25)
);
assert_eq!(decoded.data.cache_read_input_tokens, Some(0));
assert_eq!(decoded.data.cache_read_cost_usd, Some(0.0));
assert_eq!(decoded.data.actual_total_cost_usd, Some(1.25));
assert_eq!(decoded.data.status_code, Some(502));
assert_eq!(
decoded.data.error_category.as_deref(),
Some("upstream_error")
);
assert_eq!(decoded.data.request_body_ref, event.data.request_body_ref);
assert!(decoded.data.request_body.is_none());
assert!(decoded.data.provider_request_body.is_none());
assert!(decoded.data.response_body.is_none());
assert!(decoded.data.client_response_body.is_none());
for state in [
decoded.data.request_body_state,
decoded.data.provider_request_body_state,
decoded.data.response_body_state,
decoded.data.client_response_body_state,
] {
assert_eq!(state, Some(UsageBodyCaptureState::Truncated));
}
let metadata = decoded
.data
.request_metadata
.as_ref()
.expect("preserved facts");
assert_eq!(metadata["requested_reasoning_effort"], "high");
assert_eq!(metadata["provider_reasoning_effort"], "medium");
assert_eq!(metadata["provider_service_tier"], "priority");
assert_eq!(metadata["provider_actual_service_tier"], "default");
assert_eq!(metadata["provider_cache_ttl_minutes"], 30);
assert_eq!(
metadata["plan_usage_reservation_token"],
"550e8400-e29b-41d4-a716-446655440000"
);
assert_eq!(budget.retained_bytes(), 0);
assert_eq!(budget.downgraded_total(), 1);
}
#[test]
fn from_stream_fields_reconstructed_lease_also_bounds_recorder_clones() {
let fields = captured_event()
.to_stream_fields()
.expect("wire serialization");
let probe_budget = Arc::new(EventCaptureMemoryBudget::new(usize::MAX));
let probe =
UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&probe_budget))
.expect("estimate decoded allocation");
let weight = probe_budget.retained_bytes();
assert!(weight > 0);
drop(probe);
assert_eq!(probe_budget.retained_bytes(), 0);
let budget = Arc::new(EventCaptureMemoryBudget::new(weight * 2));
let event =
UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&budget))
.expect("wire decode");
let recorder_copy = event.clone();
assert!(recorder_copy.data.response_body.is_some());
assert_eq!(budget.retained_bytes(), weight * 2);
let omitted_copy = event.clone();
assert!(omitted_copy.data.response_body.is_none());
assert_eq!(omitted_copy.data.total_tokens, Some(600));
assert_eq!(omitted_copy.data.cache_read_input_tokens, Some(0));
assert_eq!(budget.retained_bytes(), weight * 2);
drop(event);
assert_eq!(budget.retained_bytes(), weight);
drop((recorder_copy, omitted_copy));
assert_eq!(budget.retained_bytes(), 0);
}
fn assert_typed_body_clear_is_preserved(event: &UsageEvent, state: UsageBodyCaptureState) {
for (body, reference, actual_state) in [
(
&event.data.request_body,
&event.data.request_body_ref,
event.data.request_body_state,
),
(
&event.data.provider_request_body,
&event.data.provider_request_body_ref,
event.data.provider_request_body_state,
),
(
&event.data.response_body,
&event.data.response_body_ref,
event.data.response_body_state,
),
(
&event.data.client_response_body,
&event.data.client_response_body_ref,
event.data.client_response_body_state,
),
] {
assert!(body.is_none());
assert_eq!(actual_state, Some(state));
assert_eq!(reference.as_deref(), Some("usage://stale/reference"));
}
assert!(event
.data
.request_metadata
.as_ref()
.and_then(|value| value.get("body_capture"))
.is_none());
}
#[test]
fn from_stream_fields_and_clone_budget_preserve_typed_clear_with_residual_bodies() {
for state in [
UsageBodyCaptureState::None,
UsageBodyCaptureState::Disabled,
UsageBodyCaptureState::Unavailable,
] {
let mut source = captured_event();
source.data.request_metadata = None;
source.data.request_body_state = Some(state);
source.data.provider_request_body_state = Some(state);
source.data.response_body_state = Some(state);
source.data.client_response_body_state = Some(state);
source.data.request_body_ref = Some("usage://stale/reference".to_string());
source.data.provider_request_body_ref = Some("usage://stale/reference".to_string());
source.data.response_body_ref = Some("usage://stale/reference".to_string());
source.data.client_response_body_ref = Some("usage://stale/reference".to_string());
let fields = source.to_stream_fields().expect("wire serialization");
let decoded_budget = Arc::new(EventCaptureMemoryBudget::new(0));
let decoded = UsageEvent::from_stream_fields_with_capture_budget(
&fields,
Arc::clone(&decoded_budget),
)
.expect("wire decode");
assert_typed_body_clear_is_preserved(&decoded, state);
assert_eq!(decoded_budget.retained_bytes(), 0);
let clone_budget = Arc::new(EventCaptureMemoryBudget::new(
source.data.capture_heap_estimate(),
));
crate::runtime::prepare_event_capture_memory(&mut source, Arc::clone(&clone_budget));
let cloned = source.clone();
assert_typed_body_clear_is_preserved(&cloned, state);
assert!(source.data.request_body.is_some());
assert_eq!(clone_budget.downgraded_total(), 1);
drop(source);
assert_eq!(clone_budget.retained_bytes(), 0);
}
}
#[test]
fn from_stream_fields_legacy_metadata_only_facts_survive_before_billing() {
for limit in [0, 8192] {
for include_response in [false, true] {
let mut source = legacy_metadata_only_event();
if include_response {
source.data.response_body = Some(json!({"result": "response capture"}));
}
let fields = source
.to_stream_fields()
.expect("legacy wire serialization");
let budget = Arc::new(EventCaptureMemoryBudget::new(limit));
let decoded = UsageEvent::from_stream_fields_with_capture_budget(
&fields,
Arc::clone(&budget),
)
.expect("legacy wire decode");
// The worker enriches this clone before DTO conversion. Missing legacy bodies
// must not erase a previously derived TTL or turn an explicit zero into unknown.
let billing_event = decoded.clone();
let metadata = billing_event
.data
.request_metadata
.as_ref()
.expect("legacy facts");
assert_eq!(metadata["requested_reasoning_effort"], "high");
assert_eq!(metadata["provider_reasoning_effort"], "medium");
assert_eq!(metadata["provider_service_tier"], "priority");
assert_eq!(metadata["provider_actual_service_tier"], "default");
assert_eq!(metadata["provider_cache_ttl_minutes"], 60);
assert_eq!(billing_event.data.input_tokens, Some(0));
assert_eq!(billing_event.data.output_tokens, Some(0));
assert_eq!(billing_event.data.total_tokens, Some(0));
assert_eq!(billing_event.data.cache_read_input_tokens, Some(0));
assert_eq!(billing_event.data.cache_creation_input_tokens, Some(0));
assert_eq!(billing_event.data.actual_total_cost_usd, Some(0.0));
assert_eq!(billing_event.data.request_body_state, None);
assert_eq!(billing_event.data.provider_request_body_state, None);
drop((billing_event, decoded));
assert_eq!(budget.retained_bytes(), 0);
}
}
}
#[test]
fn from_stream_fields_typed_none_still_clears_metadata_only_request_facts() {
for limit in [0, 8192] {
let mut source = legacy_metadata_only_event();
source.data.request_body_state = Some(UsageBodyCaptureState::None);
source.data.provider_request_body_state = Some(UsageBodyCaptureState::None);
let fields = source.to_stream_fields().expect("wire serialization");
let budget = Arc::new(EventCaptureMemoryBudget::new(limit));
let decoded =
UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&budget))
.expect("wire decode");
let metadata = decoded
.data
.request_metadata
.as_ref()
.expect("response facts remain");
for key in [
"requested_reasoning_effort",
"provider_reasoning_effort",
"provider_service_tier",
"provider_cache_ttl_minutes",
] {
assert!(metadata.get(key).is_none(), "typed none must clear {key}");
}
assert_eq!(metadata["provider_actual_service_tier"], "default");
assert_eq!(
decoded.data.request_body_state,
Some(UsageBodyCaptureState::None)
);
assert_eq!(
decoded.data.provider_request_body_state,
Some(UsageBodyCaptureState::None)
);
assert_eq!(decoded.data.cache_read_input_tokens, Some(0));
assert_eq!(budget.retained_bytes(), 0);
assert_eq!(budget.downgraded_total(), 0);
}
}
fn legacy_metadata_only_event() -> UsageEvent {
UsageEvent::new(
UsageEventType::Completed,
"legacy-metadata-only",
UsageEventData {
provider_name: "openai".to_string(),
model: "gpt-5.6-sol".to_string(),
endpoint_api_format: Some("openai:responses".to_string()),
input_tokens: Some(0),
output_tokens: Some(0),
total_tokens: Some(0),
cache_read_input_tokens: Some(0),
cache_creation_input_tokens: Some(0),
actual_total_cost_usd: Some(0.0),
request_metadata: Some(json!({
"requested_reasoning_effort": "high",
"provider_reasoning_effort": "medium",
"provider_service_tier": "priority",
"provider_actual_service_tier": "default",
"provider_cache_ttl_minutes": 60
})),
..UsageEventData::default()
},
)
}
#[test]
fn from_stream_fields_body_omission_keeps_unknown_usage_unknown() {
let budget = Arc::new(EventCaptureMemoryBudget::new(0));
for event_type in [
UsageEventType::Completed,
UsageEventType::Failed,
UsageEventType::Cancelled,
] {
let event = UsageEvent::new(
event_type,
"usage-unavailable",
UsageEventData {
provider_name: "openai".to_string(),
model: "gpt-5".to_string(),
response_body: Some(json!({"error": "usage unavailable"})),
request_metadata: Some(json!({
"usage_available": false,
"usage_pricing_available": false
})),
..UsageEventData::default()
},
);
let fields = event.to_stream_fields().expect("wire serialization");
let decoded =
UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&budget))
.expect("wire decode");
assert_eq!(decoded.event_type, event_type);
assert_eq!(decoded.data.input_tokens, None);
assert_eq!(decoded.data.output_tokens, None);
assert_eq!(decoded.data.total_tokens, None);
assert_eq!(decoded.data.cache_read_input_tokens, None);
assert_eq!(decoded.data.cache_creation_input_tokens, None);
assert_eq!(decoded.data.actual_total_cost_usd, None);
assert_eq!(
decoded.data.request_metadata.as_ref().expect("metadata")["usage_available"],
false
);
assert_eq!(
decoded.data.request_metadata.as_ref().expect("metadata")
["usage_pricing_available"],
false
);
assert_eq!(
decoded.data.response_body_state,
Some(UsageBodyCaptureState::Truncated)
);
}
assert_eq!(budget.retained_bytes(), 0);
assert_eq!(budget.downgraded_total(), 3);
}
#[test]
fn from_stream_fields_invalid_envelopes_do_not_reserve_capture_memory() {
let budget = Arc::new(EventCaptureMemoryBudget::new(1024));
let mut unsupported = captured_event()
.to_stream_fields()
.expect("wire serialization");
let mut payload: serde_json::Value =
serde_json::from_str(&unsupported["payload"]).expect("json");
payload["v"] = json!(99);
unsupported.insert("payload".to_string(), payload.to_string());
for fields in [
BTreeMap::new(),
BTreeMap::from([("payload".to_string(), "not json".to_string())]),
unsupported,
] {
assert!(matches!(
UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&budget)),
Err(DataLayerError::UnexpectedValue(_))
));
assert_eq!(budget.retained_bytes(), 0);
assert_eq!(budget.downgraded_total(), 0);
}
}
#[test]
fn event_capture_budget_basic_policy_needs_no_diagnostic_allocation() {
let mut event = captured_event();
let budget = Arc::new(EventCaptureMemoryBudget::new(0));
apply_usage_body_capture_policy_to_event(UsageBodyCapturePolicy::default(), &mut event);
event.data.apply_capture_memory_budget(Arc::clone(&budget));
assert_eq!(
event.data.response_body_state,
Some(UsageBodyCaptureState::Disabled)
);
assert_eq!(event.data.total_tokens, Some(600));
assert_eq!(budget.retained_bytes(), 0);
assert_eq!(budget.downgraded_total(), 0);
}
#[test]
fn usage_event_round_trips_through_stream_fields() {
let event = UsageEvent::new(
@@ -0,0 +1,102 @@
use std::sync::{Arc, LazyLock};
#[cfg(test)]
use serde_json::Value;
#[doc(hidden)]
pub use aether_data_contracts::repository::usage::UsageCaptureRetention as UsageEventCaptureRetention;
pub(crate) use aether_data_contracts::repository::usage::{
usage_json_heap_estimate as json_heap_estimate,
UsageCaptureMemoryBudget as EventCaptureMemoryBudget,
};
const DEFAULT_CAPTURE_MEMORY_BUDGET_BYTES: usize = 128 * 1024 * 1024;
static CAPTURE_MEMORY_BUDGET: LazyLock<Arc<EventCaptureMemoryBudget>> = LazyLock::new(|| {
let limit = std::env::var("AETHER_USAGE_EVENT_CAPTURE_MEMORY_BUDGET_BYTES")
.ok()
.and_then(|value| value.trim().parse().ok())
.unwrap_or(DEFAULT_CAPTURE_MEMORY_BUDGET_BYTES);
Arc::new(EventCaptureMemoryBudget::new(limit))
});
pub(crate) fn shared_capture_memory_budget() -> Arc<EventCaptureMemoryBudget> {
Arc::clone(&CAPTURE_MEMORY_BUDGET)
}
pub(crate) fn capture_memory_metrics() -> (usize, usize, u64) {
CAPTURE_MEMORY_BUDGET.snapshot()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn event_capture_budget_resize_and_drop_release_estimate() {
let budget = Arc::new(EventCaptureMemoryBudget::new(16));
let mut retention = UsageEventCaptureRetention::default();
assert!(retention.reserve(Arc::clone(&budget), 12));
assert!(!retention.reserve(Arc::clone(&budget), 17));
assert_eq!(budget.retained_bytes(), 12);
assert_eq!(budget.downgraded_total(), 1);
assert!(retention.reserve(Arc::clone(&budget), 4));
assert_eq!(budget.retained_bytes(), 4);
drop(retention);
assert_eq!(budget.retained_bytes(), 0);
}
#[test]
fn event_capture_budget_unmanaged_clone_skips_estimation_and_empty_clone_is_free() {
let unmanaged = UsageEventCaptureRetention::default();
let (_, retained) =
unmanaged.clone_for_bodies(|| panic!("unmanaged JSON must not be scanned"));
assert!(retained);
let budget = Arc::new(EventCaptureMemoryBudget::new(0));
let mut managed = UsageEventCaptureRetention::default();
assert!(managed.reserve(Arc::clone(&budget), 0));
let (cloned, retained) = managed.clone_for_bodies(|| 0);
assert!(retained);
drop((managed, cloned));
assert_eq!(budget.retained_bytes(), 0);
assert_eq!(budget.downgraded_total(), 0);
}
#[test]
fn event_capture_budget_estimate_counts_string_and_array_spare_capacity() {
let mut text = String::with_capacity(1024);
text.push('x');
let mut array = Vec::with_capacity(16);
let expected = text.capacity() + array.capacity() * std::mem::size_of::<Value>();
array.push(Value::String(text));
assert_eq!(json_heap_estimate(&Value::Array(array)), expected);
}
#[test]
fn event_capture_budget_parallel_owners_never_exceed_shared_limit() {
let budget = Arc::new(EventCaptureMemoryBudget::new(1024));
let barrier = Arc::new(std::sync::Barrier::new(8));
std::thread::scope(|scope| {
for _ in 0..8 {
let budget = Arc::clone(&budget);
let barrier = Arc::clone(&barrier);
scope.spawn(move || {
for _ in 0..100 {
let mut retained = UsageEventCaptureRetention::default();
barrier.wait();
let _ = retained.reserve(Arc::clone(&budget), 400);
barrier.wait();
assert_eq!(budget.retained_bytes(), 800);
barrier.wait();
drop(retained);
barrier.wait();
assert_eq!(budget.retained_bytes(), 0);
barrier.wait();
}
});
}
});
assert_eq!(budget.retained_bytes(), 0);
assert_eq!(budget.downgraded_total(), 600);
}
}
@@ -0,0 +1,999 @@
use std::collections::BTreeMap;
use std::io::{self, Write};
use aether_data_contracts::repository::usage::{
resolve_provider_cache_ttl_minutes, UsageBodyCaptureState,
PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY,
PROVIDER_SERVICE_TIER_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY,
};
use aether_data_contracts::DataLayerError;
use serde::ser::{Impossible, SerializeMap, SerializeStruct};
use serde::{Serialize, Serializer};
use serde_json::Value;
use super::{BorrowedUsageEventEnvelope, UsageEvent, UsageEventData, USAGE_EVENT_VERSION};
use crate::body_capture::mark_usage_event_capture_truncated;
use crate::request_metadata::{
attach_client_request_body_metadata, attach_provider_request_body_metadata,
attach_provider_response_body_metadata, clear_client_request_body_metadata,
clear_provider_request_body_metadata, request_body_derived_facts_action,
RequestBodyDerivedFactsAction,
};
const DIAGNOSTIC_FIELDS: [&str; 8] = [
"request_body",
"provider_request_body",
"response_body",
"client_response_body",
"request_headers",
"provider_request_headers",
"response_headers",
"client_response_headers",
];
const BODY_STATE_FIELDS: [&str; 4] = [
"request_body_state",
"provider_request_body_state",
"response_body_state",
"client_response_body_state",
];
const BODY_METADATA_KEYS: [&str; 4] =
["request", "provider_request", "response", "client_response"];
#[derive(Debug)]
pub(crate) struct EncodedUsageEvent {
pub(crate) fields: BTreeMap<String, String>,
pub(crate) diagnostics_omitted: bool,
}
pub(super) fn encode(
event: &UsageEvent,
max_bytes: usize,
) -> Result<EncodedUsageEvent, DataLayerError> {
let mut writer = BoundedJsonWriter::new(max_bytes);
if writer.serialize(&envelope(event, &event.data))? {
return writer.into_event(false);
}
// Conservatively reject oversized original metadata before cloning it, even
// when later fact normalization could make that metadata smaller.
let core = ProjectedData {
data: &event.data,
overrides: None,
};
if !writer.serialize(&envelope(event, &core))? {
return Err(wire_limit_error(max_bytes));
}
let overrides = WireOverrides::new(&event.data)?;
let projected = ProjectedData {
data: &event.data,
overrides: Some(&overrides),
};
if !writer.serialize(&envelope(event, &projected))? {
return Err(wire_limit_error(max_bytes));
}
writer.into_event(true)
}
fn envelope<'a, T: Serialize + ?Sized>(
event: &'a UsageEvent,
data: &'a T,
) -> BorrowedUsageEventEnvelope<'a, T> {
BorrowedUsageEventEnvelope {
v: USAGE_EVENT_VERSION,
event_type: event.event_type,
request_id: &event.request_id,
timestamp_ms: event.timestamp_ms,
data,
}
}
fn wire_limit_error(max_bytes: usize) -> DataLayerError {
DataLayerError::InvalidInput(format!(
"usage event exceeds the {max_bytes}-byte wire limit after omitting diagnostic bodies and headers"
))
}
struct BoundedJsonWriter {
bytes: Vec<u8>,
max_bytes: usize,
exceeded: bool,
}
impl BoundedJsonWriter {
fn new(max_bytes: usize) -> Self {
Self {
bytes: Vec::new(),
max_bytes,
exceeded: false,
}
}
fn serialize<T: Serialize + ?Sized>(&mut self, value: &T) -> Result<bool, DataLayerError> {
self.bytes.clear();
self.exceeded = false;
match serde_json::to_writer(&mut *self, value) {
Ok(()) => Ok(true),
Err(_) if self.exceeded => Ok(false),
Err(error) => Err(DataLayerError::UnexpectedValue(format!(
"failed to serialize usage event payload: {error}"
))),
}
}
fn into_event(self, diagnostics_omitted: bool) -> Result<EncodedUsageEvent, DataLayerError> {
let payload = String::from_utf8(self.bytes).map_err(|error| {
DataLayerError::UnexpectedValue(format!("usage event JSON was not UTF-8: {error}"))
})?;
Ok(EncodedUsageEvent {
fields: BTreeMap::from([("payload".to_string(), payload)]),
diagnostics_omitted,
})
}
}
impl Write for BoundedJsonWriter {
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
if bytes.len() > self.max_bytes.saturating_sub(self.bytes.len()) {
self.exceeded = true;
return Err(io::Error::other("usage event wire limit exceeded"));
}
let required = self.bytes.len() + bytes.len();
if required > self.bytes.capacity() {
let capacity = required
.max(self.bytes.capacity().saturating_mul(2))
.min(self.max_bytes);
self.bytes
.try_reserve_exact(capacity - self.bytes.len())
.map_err(|error| {
io::Error::other(format!("usage event wire allocation failed: {error}"))
})?;
}
self.bytes.extend_from_slice(bytes);
Ok(bytes.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
struct WireOverrides {
truncated: [bool; 4],
metadata: Option<Value>,
}
impl WireOverrides {
fn new(data: &UsageEventData) -> Result<Self, DataLayerError> {
// The full v1 consumer decodes a JSON null Option<Value> as no body.
let request_body = data.request_body.as_ref().filter(|body| !body.is_null());
let provider_request_body = data
.provider_request_body
.as_ref()
.filter(|body| !body.is_null());
let body_cache_ttl = resolve_provider_cache_ttl_minutes(
data.endpoint_api_format
.as_deref()
.or(data.api_format.as_deref()),
data.target_model.as_deref().or(Some(data.model.as_str())),
Some(data.model.as_str()),
provider_request_body,
);
if body_cache_ttl.is_some()
&& data.provider_request_body_state == Some(UsageBodyCaptureState::None)
{
return Err(DataLayerError::InvalidInput(
"usage event cannot omit a provider request body whose cache TTL would be cleared by its explicit none capture state"
.to_string(),
));
}
let mut metadata = data.request_metadata.clone();
match request_body_derived_facts_action(request_body, data.request_body_state) {
RequestBodyDerivedFactsAction::Refresh => {
if request_body.is_some_and(|body| !body.is_object()) {
if let Some(Value::Object(object)) = metadata.as_mut() {
object.remove(REQUESTED_REASONING_EFFORT_METADATA_KEY);
}
} else {
metadata = attach_client_request_body_metadata(metadata, request_body);
}
}
RequestBodyDerivedFactsAction::Clear
if request_body.is_some() || data.request_body_state.is_some() =>
{
metadata = clear_client_request_body_metadata(metadata);
}
RequestBodyDerivedFactsAction::Clear | RequestBodyDerivedFactsAction::Preserve => {}
}
match request_body_derived_facts_action(
provider_request_body,
data.provider_request_body_state,
) {
RequestBodyDerivedFactsAction::Refresh => {
if provider_request_body.is_some_and(|body| !body.is_object()) {
// An authoritative scalar/array has no tier or reasoning,
// but billing still falls back to the metadata's cache TTL.
if let Some(Value::Object(object)) = metadata.as_mut() {
object.remove(PROVIDER_REASONING_EFFORT_METADATA_KEY);
object.remove(PROVIDER_SERVICE_TIER_METADATA_KEY);
}
} else {
metadata = attach_provider_request_body_metadata(
metadata,
data.endpoint_api_format
.as_deref()
.or(data.api_format.as_deref()),
data.target_model.as_deref().or(Some(data.model.as_str())),
Some(data.model.as_str()),
provider_request_body,
);
}
}
RequestBodyDerivedFactsAction::Clear
if provider_request_body.is_some()
|| data.provider_request_body_state.is_some() =>
{
metadata = clear_provider_request_body_metadata(metadata);
}
RequestBodyDerivedFactsAction::Clear | RequestBodyDerivedFactsAction::Preserve => {}
}
metadata = attach_provider_response_body_metadata(metadata, data.response_body.as_ref());
// Billing reads raw-body TTL before metadata regardless of capture state.
// Preserve that precedence independently of reasoning and tier authority.
if let Some(cache_ttl) = body_cache_ttl {
let object = metadata
.get_or_insert_with(|| Value::Object(serde_json::Map::new()))
.as_object_mut()
.ok_or_else(|| {
DataLayerError::InvalidInput(
"usage event cannot preserve provider cache TTL in non-object metadata after omitting diagnostic bodies"
.to_string(),
)
})?;
object.insert(
PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY.to_string(),
Value::Number(cache_ttl.into()),
);
}
let truncated = [
(data.request_body.as_ref(), data.request_body_state),
(
data.provider_request_body.as_ref(),
data.provider_request_body_state,
),
(data.response_body.as_ref(), data.response_body_state),
(
data.client_response_body.as_ref(),
data.client_response_body_state,
),
]
.map(|(body, state)| {
body.is_some_and(|body| !body.is_null())
&& !matches!(
state,
Some(
UsageBodyCaptureState::None
| UsageBodyCaptureState::Disabled
| UsageBodyCaptureState::Unavailable
)
)
});
for (truncated, key) in truncated.into_iter().zip(BODY_METADATA_KEYS) {
if !truncated {
continue;
}
mark_usage_event_capture_truncated(&mut metadata, key);
if let Some(entry) = metadata
.as_mut()
.and_then(|value| value.get_mut("body_capture"))
.and_then(|value| value.get_mut(key))
.and_then(Value::as_object_mut)
{
entry.insert(
"reason".to_string(),
Value::String("wire_limit_exceeded".to_string()),
);
}
}
Ok(Self {
truncated,
metadata,
})
}
}
struct ProjectedData<'a> {
data: &'a UsageEventData,
overrides: Option<&'a WireOverrides>,
}
impl Serialize for ProjectedData<'_> {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
self.data.serialize(FieldProjectionSerializer {
map: serializer.serialize_map(None)?,
overrides: self.overrides,
})
}
}
// Reuse UsageEventData's derived field traversal, including future fields and
// skip_serializing_if rules. Only diagnostic fields and explicit overrides differ.
struct FieldProjectionSerializer<'a, M> {
map: M,
overrides: Option<&'a WireOverrides>,
}
impl<M: SerializeMap> SerializeStruct for FieldProjectionSerializer<'_, M> {
type Ok = M::Ok;
type Error = M::Error;
fn serialize_field<T: Serialize + ?Sized>(
&mut self,
key: &'static str,
value: &T,
) -> Result<(), Self::Error> {
if DIAGNOSTIC_FIELDS.contains(&key) {
return Ok(());
}
if let Some(overrides) = self.overrides {
if key == "request_metadata"
|| BODY_STATE_FIELDS
.iter()
.zip(overrides.truncated)
.any(|(state, truncated)| *state == key && truncated)
{
return Ok(());
}
}
self.map.serialize_entry(key, value)
}
fn end(mut self) -> Result<Self::Ok, Self::Error> {
if let Some(overrides) = self.overrides {
for (key, truncated) in BODY_STATE_FIELDS.into_iter().zip(overrides.truncated) {
if truncated {
self.map
.serialize_entry(key, &UsageBodyCaptureState::Truncated)?;
}
}
if let Some(metadata) = overrides.metadata.as_ref() {
self.map.serialize_entry("request_metadata", metadata)?;
}
}
self.map.end()
}
}
fn expected_struct<E: serde::ser::Error, T>() -> Result<T, E> {
Err(E::custom("usage event data must serialize as a struct"))
}
macro_rules! reject_scalar_serialization {
($($name:ident($value:ident: $ty:ty)),* $(,)?) => {
$(fn $name(self, $value: $ty) -> Result<Self::Ok, Self::Error> {
let _ = $value;
expected_struct()
})*
};
}
impl<M: SerializeMap> Serializer for FieldProjectionSerializer<'_, M> {
type Ok = M::Ok;
type Error = M::Error;
type SerializeSeq = Impossible<Self::Ok, Self::Error>;
type SerializeTuple = Impossible<Self::Ok, Self::Error>;
type SerializeTupleStruct = Impossible<Self::Ok, Self::Error>;
type SerializeTupleVariant = Impossible<Self::Ok, Self::Error>;
type SerializeMap = Impossible<Self::Ok, Self::Error>;
type SerializeStruct = Self;
type SerializeStructVariant = Impossible<Self::Ok, Self::Error>;
reject_scalar_serialization! {
serialize_bool(value: bool), serialize_i8(value: i8), serialize_i16(value: i16),
serialize_i32(value: i32), serialize_i64(value: i64), serialize_i128(value: i128),
serialize_u8(value: u8), serialize_u16(value: u16), serialize_u32(value: u32),
serialize_u64(value: u64), serialize_u128(value: u128), serialize_f32(value: f32),
serialize_f64(value: f64), serialize_char(value: char), serialize_str(value: &str),
serialize_bytes(value: &[u8]),
}
fn serialize_none(self) -> Result<Self::Ok, Self::Error> {
expected_struct()
}
fn serialize_some<T: Serialize + ?Sized>(self, _: &T) -> Result<Self::Ok, Self::Error> {
expected_struct()
}
fn serialize_unit(self) -> Result<Self::Ok, Self::Error> {
expected_struct()
}
fn serialize_unit_struct(self, _: &'static str) -> Result<Self::Ok, Self::Error> {
expected_struct()
}
fn serialize_unit_variant(
self,
_: &'static str,
_: u32,
_: &'static str,
) -> Result<Self::Ok, Self::Error> {
expected_struct()
}
fn serialize_newtype_struct<T: Serialize + ?Sized>(
self,
_: &'static str,
_: &T,
) -> Result<Self::Ok, Self::Error> {
expected_struct()
}
fn serialize_newtype_variant<T: Serialize + ?Sized>(
self,
_: &'static str,
_: u32,
_: &'static str,
_: &T,
) -> Result<Self::Ok, Self::Error> {
expected_struct()
}
fn serialize_seq(self, _: Option<usize>) -> Result<Self::SerializeSeq, Self::Error> {
expected_struct()
}
fn serialize_tuple(self, _: usize) -> Result<Self::SerializeTuple, Self::Error> {
expected_struct()
}
fn serialize_tuple_struct(
self,
_: &'static str,
_: usize,
) -> Result<Self::SerializeTupleStruct, Self::Error> {
expected_struct()
}
fn serialize_tuple_variant(
self,
_: &'static str,
_: u32,
_: &'static str,
_: usize,
) -> Result<Self::SerializeTupleVariant, Self::Error> {
expected_struct()
}
fn serialize_map(self, _: Option<usize>) -> Result<Self::SerializeMap, Self::Error> {
expected_struct()
}
fn serialize_struct(
self,
_: &'static str,
_: usize,
) -> Result<Self::SerializeStruct, Self::Error> {
Ok(self)
}
fn serialize_struct_variant(
self,
_: &'static str,
_: u32,
_: &'static str,
_: usize,
) -> Result<Self::SerializeStructVariant, Self::Error> {
expected_struct()
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use serde_json::json;
use super::*;
use crate::event::UsageEventType;
use crate::event_capture_budget::EventCaptureMemoryBudget;
fn event() -> UsageEvent {
UsageEvent {
event_type: UsageEventType::Completed,
request_id: "wire-request".to_string(),
timestamp_ms: 1_234_567,
data: UsageEventData {
provider_name: "provider".to_string(),
model: "gpt-5.6-sol".to_string(),
endpoint_api_format: Some("openai:responses".to_string()),
user_id: Some("user-id".to_string()),
api_key_id: Some("key-id".to_string()),
provider_id: Some("provider-id".to_string()),
provider_endpoint_id: Some("endpoint-id".to_string()),
provider_api_key_id: Some("provider-key-id".to_string()),
input_tokens: Some(100),
output_tokens: Some(500),
total_tokens: Some(600),
cache_read_input_tokens: Some(0),
cache_creation_input_tokens: Some(25),
cache_creation_ephemeral_5m_input_tokens: Some(0),
cache_creation_ephemeral_1h_input_tokens: Some(25),
cache_read_cost_usd: Some(0.0),
total_cost_usd: Some(1.25),
actual_total_cost_usd: Some(1.25),
status_code: Some(200),
is_stream: Some(false),
candidate_index: Some(0),
first_byte_time_ms: Some(0),
error_message: Some(String::new()),
request_metadata: Some(json!({
"plan_usage_reservation_token": "550e8400-e29b-41d4-a716-446655440000",
"dimensions": {"image_count": 2, "size": "1024x1024", "quality": "high"},
"usage_available": true,
"usage_pricing_available": true
})),
..UsageEventData::default()
},
}
}
fn wire_value(encoded: &EncodedUsageEvent) -> Value {
serde_json::from_str(&encoded.fields["payload"]).expect("complete JSON wire payload")
}
#[test]
fn event_wire_exact_limit_accounts_for_json_escaping_and_utf8() {
let mut event = event();
event.request_id = "escaped\0\n\r\t\"\\\u{03bb}\u{1f600}".to_string();
let original = event.to_stream_fields().expect("original wire payload");
let length = original["payload"].len();
for limit in [length, length + 1] {
let encoded = event.to_bounded_stream_fields(limit).expect("exact fit");
assert_eq!(encoded.fields, original);
assert!(!encoded.diagnostics_omitted);
}
for limit in [0, 1, length - 1] {
assert!(matches!(
event.to_bounded_stream_fields(limit),
Err(DataLayerError::InvalidInput(_))
));
}
}
#[test]
fn event_wire_writer_does_not_append_past_limit_and_reuses_failed_buffer() {
let mut writer = BoundedJsonWriter::new(5);
writer.write_all(b"12345").expect("exact fit");
assert!(writer.write_all(b"6").is_err());
assert_eq!(writer.bytes, b"12345");
assert!(writer.exceeded);
assert!(!writer.serialize(&"\u{0000}").expect("size rejection"));
assert!(writer.bytes.len() <= 5);
assert!(writer.serialize(&"\u{03bb}").expect("valid UTF-8 retry"));
assert_eq!(writer.bytes, "\"\u{03bb}\"".as_bytes());
assert!(!writer.exceeded);
}
#[test]
fn event_wire_projection_transparently_forwards_derived_fields() {
#[derive(Serialize)]
struct FutureFields<'a> {
request_body: &'a Value,
new_billing_field: &'a Value,
zero_count: u64,
enabled: bool,
#[serde(skip_serializing_if = "Option::is_none")]
missing_field: Option<&'a str>,
}
let diagnostic = json!({"large": "x".repeat(8_192)});
let billing = json!({"nested": [0, false, "unchanged"]});
let data = FutureFields {
request_body: &diagnostic,
new_billing_field: &billing,
zero_count: 0,
enabled: false,
missing_field: None,
};
let mut bytes = Vec::new();
let mut serializer = serde_json::Serializer::new(&mut bytes);
data.serialize(FieldProjectionSerializer {
map: (&mut serializer)
.serialize_map(None)
.expect("object serializer"),
overrides: None,
})
.expect("project derived fields");
assert_eq!(
serde_json::from_slice::<Value>(&bytes).expect("projected JSON"),
json!({"new_billing_field": billing, "zero_count": 0, "enabled": false})
);
}
#[test]
fn event_wire_omission_preserves_billing_facts_refs_and_source_ownership() {
let mut event = event();
let padding = "x".repeat(16_384);
event.data.request_body = Some(json!({"reasoning": {"effort": "high"}, "input": padding}));
event.data.provider_request_body = Some(
json!({"model": "gpt-5.6-sol", "reasoning": {"effort": "medium"}, "service_tier": "priority", "input": padding}),
);
event.data.response_body = Some(json!({"service_tier": "Default", "output": padding}));
event.data.client_response_body = Some(json!({"output": padding}));
event.data.request_headers = Some(json!({"x-request": padding}));
event.data.provider_request_headers = Some(json!({"x-provider": padding}));
event.data.response_headers = Some(json!({"x-response": padding}));
event.data.client_response_headers = Some(json!({"x-client": padding}));
event.data.request_body_ref = Some("usage://wire-request/request_body".to_string());
event.data.request_body_state = Some(UsageBodyCaptureState::Inline);
event.data.provider_request_body_state = Some(UsageBodyCaptureState::Inline);
event.data.response_body_state = Some(UsageBodyCaptureState::Inline);
event.data.request_metadata.as_mut().unwrap()["body_capture"] = json!({
"response": {"state": "inline", "source_bytes": 123_456}
});
let weight = event.data.capture_heap_estimate();
let budget = Arc::new(EventCaptureMemoryBudget::new(weight));
event.data.apply_capture_memory_budget(Arc::clone(&budget));
let before = serde_json::to_value(&event).expect("source snapshot");
let original_wire: Value =
serde_json::from_str(&event.to_stream_fields().unwrap()["payload"]).unwrap();
let encoded = event
.to_bounded_stream_fields(8_192)
.expect("diagnostic omission");
assert!(encoded.diagnostics_omitted);
assert!(encoded.fields["payload"].len() <= 8_192);
let value = wire_value(&encoded);
for field in DIAGNOSTIC_FIELDS {
assert!(
value["data"].get(field).is_none(),
"{field} must be omitted"
);
}
for (field, original) in original_wire["data"].as_object().unwrap() {
if !DIAGNOSTIC_FIELDS.contains(&field.as_str())
&& !BODY_STATE_FIELDS.contains(&field.as_str())
&& field != "request_metadata"
{
assert_eq!(&value["data"][field], original, "core field {field}");
}
}
for (field, key) in BODY_STATE_FIELDS.into_iter().zip(BODY_METADATA_KEYS) {
assert_eq!(value["data"][field], "truncated");
assert_eq!(
value["data"]["request_metadata"]["body_capture"][key]["reason"],
"wire_limit_exceeded"
);
assert_eq!(
value["data"]["request_metadata"]["body_capture"][key]["stored_bytes"],
0
);
}
let metadata = &value["data"]["request_metadata"];
assert_eq!(metadata["requested_reasoning_effort"], "high");
assert_eq!(metadata["provider_reasoning_effort"], "medium");
assert_eq!(metadata["provider_service_tier"], "priority");
assert_eq!(metadata["provider_actual_service_tier"], "default");
assert_eq!(metadata["provider_cache_ttl_minutes"], 30);
assert_eq!(
metadata["body_capture"]["response"]["source_bytes"],
123_456
);
assert_eq!(
metadata["dimensions"],
before["data"]["request_metadata"]["dimensions"]
);
assert_eq!(serde_json::to_value(&event).unwrap(), before);
assert_eq!(budget.retained_bytes(), weight);
assert_eq!(budget.downgraded_total(), 0);
let decoded = UsageEvent::from_stream_fields(&encoded.fields).expect("wire decode");
let record = crate::build_upsert_usage_record_from_event(&decoded).expect("record mapping");
assert_eq!(record.input_tokens, Some(100));
assert_eq!(record.output_tokens, Some(500));
assert_eq!(record.cache_read_input_tokens, Some(0));
assert_eq!(record.cache_creation_ephemeral_5m_input_tokens, Some(0));
assert_eq!(record.cache_creation_ephemeral_1h_input_tokens, Some(25));
assert_eq!(record.error_message.as_deref(), Some(""));
assert_eq!(record.request_body_ref, event.data.request_body_ref);
assert_eq!(
record.request_body_state,
Some(UsageBodyCaptureState::Truncated)
);
let metadata = record.request_metadata.expect("record billing metadata");
assert_eq!(metadata["provider_cache_ttl_minutes"], 30);
assert_eq!(metadata["dimensions"]["image_count"], 2);
drop(event);
assert_eq!(budget.retained_bytes(), 0);
}
#[test]
fn event_wire_preserves_explicit_capture_states_and_legacy_metadata() {
for state in [
None,
Some(UsageBodyCaptureState::None),
Some(UsageBodyCaptureState::Disabled),
Some(UsageBodyCaptureState::Unavailable),
Some(UsageBodyCaptureState::Reference),
] {
let mut event = event();
event.data.request_body_state = state;
event.data.provider_request_body_state = state;
event.data.response_body_state = state;
event.data.client_response_body_state = state;
event.data.request_body_ref = Some("usage://wire-request/request_body".to_string());
event.data.response_headers = Some(json!({"large-header": "x".repeat(16_384)}));
event.data.request_metadata = Some(json!({
"requested_reasoning_effort": "high",
"provider_reasoning_effort": "medium",
"provider_service_tier": "priority",
"provider_cache_ttl_minutes": 60,
"provider_actual_service_tier": "flex"
}));
let before = serde_json::to_value(&event).unwrap();
let encoded = event
.to_bounded_stream_fields(4_096)
.expect("headers omitted");
let value = wire_value(&encoded);
for field in BODY_STATE_FIELDS {
assert_eq!(value["data"].get(field), before["data"].get(field));
}
let metadata = &value["data"]["request_metadata"];
if state == Some(UsageBodyCaptureState::None) {
assert!(metadata.get("requested_reasoning_effort").is_none());
assert!(metadata.get("provider_cache_ttl_minutes").is_none());
} else {
assert_eq!(metadata["requested_reasoning_effort"], "high");
assert_eq!(metadata["provider_cache_ttl_minutes"], 60);
}
assert_eq!(metadata["provider_actual_service_tier"], "flex");
assert_eq!(
value["data"]["request_body_ref"],
before["data"]["request_body_ref"]
);
assert_eq!(serde_json::to_value(&event).unwrap(), before);
if matches!(
state,
Some(
UsageBodyCaptureState::None
| UsageBodyCaptureState::Disabled
| UsageBodyCaptureState::Unavailable
)
) {
event.data.endpoint_api_format = Some("claude:messages".to_string());
event.data.request_body = Some(json!({"reasoning": {"effort": "low"}}));
event.data.provider_request_body =
Some(json!({"reasoning": {"effort": "low"}, "service_tier": "default"}));
event.data.response_body = Some(json!({"service_tier": "priority"}));
event.data.client_response_body = Some(json!({"stale": true}));
let baseline = UsageEvent::from_stream_fields_with_capture_budget(
&event.to_stream_fields().expect("full v1 fields"),
Arc::new(EventCaptureMemoryBudget::new(usize::MAX)),
)
.expect("full v1 consumer baseline");
let with_stale_bodies = wire_value(
&event
.to_bounded_stream_fields(4_096)
.expect("explicit capture states override stale bodies"),
);
for field in BODY_STATE_FIELDS {
assert_eq!(
with_stale_bodies["data"].get(field),
value["data"].get(field)
);
}
assert_eq!(
with_stale_bodies["data"]["request_metadata"],
serde_json::to_value(&baseline.data.request_metadata).unwrap(),
"wire omission must preserve the full v1 consumer's billing facts"
);
}
}
}
#[test]
fn event_wire_preserves_raw_body_cache_ttl_across_capture_states() {
use aether_data_contracts::repository::usage::extract_provider_cache_ttl_minutes_from_metadata;
for state in [
None,
Some(UsageBodyCaptureState::Inline),
Some(UsageBodyCaptureState::Reference),
Some(UsageBodyCaptureState::Truncated),
Some(UsageBodyCaptureState::Disabled),
Some(UsageBodyCaptureState::Unavailable),
] {
let mut event = event();
event.data.provider_request_body_state = state;
event.data.provider_request_body = Some(json!({
"prompt_cache_options": {"ttl": "30m"},
"service_tier": "default",
"reasoning": {"effort": "low"}
}));
event.data.response_headers = Some(json!({"large": "x".repeat(16_384)}));
event.data.request_metadata = Some(json!({
"provider_cache_ttl_minutes": 60,
"provider_service_tier": "priority",
"provider_reasoning_effort": "high"
}));
let before = serde_json::to_value(&event).unwrap();
let baseline = UsageEvent::from_stream_fields_with_capture_budget(
&event.to_stream_fields().expect("full v1 fields"),
Arc::new(EventCaptureMemoryBudget::new(usize::MAX)),
)
.expect("full v1 consumer baseline");
let encoded = event.to_bounded_stream_fields(4_096).expect("body omitted");
assert!(encoded.diagnostics_omitted);
let decoded = UsageEvent::from_stream_fields_with_capture_budget(
&encoded.fields,
Arc::new(EventCaptureMemoryBudget::new(usize::MAX)),
)
.expect("projected consumer event");
assert!(decoded.data.provider_request_body.is_none());
assert_eq!(
extract_provider_cache_ttl_minutes_from_metadata(
decoded.data.request_metadata.as_ref()
),
Some(30),
"raw body TTL must win over stale metadata for {state:?}"
);
for field in ["provider_service_tier", "provider_reasoning_effort"] {
assert_eq!(
decoded.data.request_metadata.as_ref().unwrap().get(field),
baseline.data.request_metadata.as_ref().unwrap().get(field),
"TTL preservation must not change {field} authority for {state:?}"
);
}
assert_eq!(serde_json::to_value(&event).unwrap(), before);
}
}
#[test]
fn event_wire_non_object_bodies_match_full_consumer_facts() {
use aether_data_contracts::repository::usage::{
extract_provider_cache_ttl_minutes_from_metadata,
resolve_provider_service_tier_from_request_capture,
};
for body in [
json!(null),
json!("opaque"),
json!([]),
json!(42),
json!(true),
] {
for state in [
None,
Some(UsageBodyCaptureState::None),
Some(UsageBodyCaptureState::Inline),
Some(UsageBodyCaptureState::Reference),
Some(UsageBodyCaptureState::Truncated),
Some(UsageBodyCaptureState::Disabled),
Some(UsageBodyCaptureState::Unavailable),
] {
let mut event = event();
event.data.request_body = Some(body.clone());
event.data.provider_request_body = Some(body.clone());
event.data.request_body_state = state;
event.data.provider_request_body_state = state;
event.data.response_headers = Some(json!({"large": "x".repeat(16_384)}));
event.data.request_metadata = Some(json!({
"requested_reasoning_effort": "medium",
"provider_reasoning_effort": "high",
"provider_service_tier": "priority",
"provider_cache_ttl_minutes": 60,
"unrelated": {"preserved": true}
}));
let original = event.to_stream_fields().expect("full v1 fields");
let baseline = UsageEvent::from_stream_fields_with_capture_budget(
&original,
Arc::new(EventCaptureMemoryBudget::new(usize::MAX)),
)
.expect("full v1 consumer baseline");
let encoded = event
.to_bounded_stream_fields(4_096)
.expect("omit diagnostics");
assert!(encoded.diagnostics_omitted);
let decoded = UsageEvent::from_stream_fields_with_capture_budget(
&encoded.fields,
Arc::new(EventCaptureMemoryBudget::new(usize::MAX)),
)
.expect("projected consumer event");
let tier = |data: &UsageEventData| {
resolve_provider_service_tier_from_request_capture(
data.provider_request_body.as_ref(),
data.provider_request_body_state,
data.request_metadata.as_ref(),
)
};
assert_eq!(
tier(&decoded.data),
tier(&baseline.data),
"provider tier for {body:?}, {state:?}"
);
let metadata = decoded.data.request_metadata.as_ref().unwrap();
let baseline_metadata = baseline.data.request_metadata.as_ref().unwrap();
assert_eq!(
extract_provider_cache_ttl_minutes_from_metadata(Some(metadata)),
extract_provider_cache_ttl_minutes_from_metadata(Some(baseline_metadata)),
"metadata TTL fallback for {body:?}, {state:?}"
);
let authoritative = !body.is_null()
&& matches!(
state,
None | Some(
UsageBodyCaptureState::Inline | UsageBodyCaptureState::Reference
)
);
for field in [
REQUESTED_REASONING_EFFORT_METADATA_KEY,
PROVIDER_REASONING_EFFORT_METADATA_KEY,
] {
assert_eq!(
metadata.get(field),
if authoritative {
None
} else {
baseline_metadata.get(field)
},
"reasoning authority for {body:?}, {state:?}, {field}"
);
}
assert_eq!(metadata["unrelated"], baseline_metadata["unrelated"]);
if body.is_null() {
assert_eq!(
decoded.data.request_body_state,
baseline.data.request_body_state
);
assert_eq!(
decoded.data.provider_request_body_state,
baseline.data.provider_request_body_state
);
}
assert_eq!(event.to_stream_fields().unwrap(), original);
}
}
}
#[test]
fn event_wire_rejects_cache_ttl_loss_from_explicit_none_capture_state() {
let mut event = event();
event.data.provider_request_body_state = Some(UsageBodyCaptureState::None);
event.data.provider_request_body = Some(json!({
"prompt_cache_options": {"ttl": "30m"},
"large": "x".repeat(16_384)
}));
let original = event.to_stream_fields().expect("full v1 fields");
let full = event
.to_bounded_stream_fields(original["payload"].len())
.expect("complete diagnostics remain representable");
assert_eq!(full.fields, original);
assert!(!full.diagnostics_omitted);
assert!(matches!(
event.to_bounded_stream_fields(4_096),
Err(DataLayerError::InvalidInput(_))
));
assert_eq!(event.to_stream_fields().unwrap(), original);
}
#[test]
fn event_wire_rejects_oversized_core_and_post_projection_metadata() {
for field in ["request_metadata", "error_message"] {
let mut event = event();
if field == "request_metadata" {
event.data.request_metadata = Some(json!({"large": "x".repeat(32_768)}));
} else {
event.data.error_message = Some("x".repeat(32_768));
}
let error = event
.to_bounded_stream_fields(1_024)
.expect_err("oversized core must fail");
assert!(matches!(error, DataLayerError::InvalidInput(_)));
assert!(
error.to_string().len() < 200,
"errors must not include payload contents"
);
}
let mut event = event();
event.data.response_body = Some(json!({"large": "x".repeat(8_192)}));
let core = ProjectedData {
data: &event.data,
overrides: None,
};
let core_size = serde_json::to_vec(&envelope(&event, &core)).unwrap().len();
assert!(
matches!(
event.to_bounded_stream_fields(core_size),
Err(DataLayerError::InvalidInput(_))
),
"added capture metadata must also fit the exact wire limit"
);
}
}
+30 -6
View File
@@ -1,5 +1,13 @@
use std::future::Future;
use std::sync::OnceLock;
use std::sync::{Mutex, OnceLock};
use std::time::Duration;
struct UsageBackgroundRuntime {
owner: Mutex<Option<tokio::runtime::Runtime>>,
handle: tokio::runtime::Handle,
}
static RUNTIME: OnceLock<UsageBackgroundRuntime> = OnceLock::new();
const DEFAULT_USAGE_BACKGROUND_RUNTIME_THREADS: usize = 8;
const MAX_USAGE_BACKGROUND_RUNTIME_THREADS: usize = 64;
@@ -18,12 +26,10 @@ where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
usage_background_runtime().handle().spawn(task)
usage_background_runtime().handle.spawn(task)
}
fn usage_background_runtime() -> &'static tokio::runtime::Runtime {
static RUNTIME: OnceLock<&'static tokio::runtime::Runtime> = OnceLock::new();
fn usage_background_runtime() -> &'static UsageBackgroundRuntime {
RUNTIME.get_or_init(|| {
let worker_threads = usage_background_runtime_threads();
let runtime = tokio::runtime::Builder::new_multi_thread()
@@ -36,10 +42,28 @@ fn usage_background_runtime() -> &'static tokio::runtime::Runtime {
.thread_stack_size(USAGE_BACKGROUND_RUNTIME_STACK_BYTES)
.build()
.expect("usage background runtime should build");
Box::leak(Box::new(runtime))
UsageBackgroundRuntime {
handle: runtime.handle().clone(),
owner: Mutex::new(Some(runtime)),
}
})
}
/// Call outside Tokio after every UsageRuntime has drained. Does not start an unused runtime.
pub fn shutdown_usage_background_runtime(timeout: Duration) {
let Some(runtime) = RUNTIME.get() else {
return;
};
let owner = runtime
.owner
.lock()
.unwrap_or_else(|p| p.into_inner())
.take();
if let Some(owner) = owner {
owner.shutdown_timeout(timeout);
}
}
fn usage_background_runtime_threads() -> usize {
parse_usage_background_runtime_threads(
std::env::var(GATEWAY_USAGE_BACKGROUND_RUNTIME_THREADS_ENV)
+6
View File
@@ -1,15 +1,19 @@
mod body_capture;
pub mod config;
mod dead_letter_encoding;
pub mod event;
mod event_capture_budget;
mod executor;
mod keyed_lock;
pub mod queue;
mod queue_read_budget;
pub mod record;
pub mod report;
pub mod report_context;
mod request_metadata;
pub mod runtime;
pub mod settlement;
mod shutdown;
pub mod standardized_usage;
pub mod usage_mapper;
pub mod worker;
@@ -21,6 +25,7 @@ pub use body_capture::{
};
pub use config::UsageRuntimeConfig;
pub use event::{now_ms, UsageEvent, UsageEventData, UsageEventType, USAGE_EVENT_VERSION};
pub use executor::shutdown_usage_background_runtime;
pub use queue::UsageQueue;
pub use record::build_upsert_usage_record_from_event;
pub use report::{
@@ -50,6 +55,7 @@ pub use runtime::{
pub use settlement::{
reconcile_usage_policy_cost_for_event, settle_usage_if_needed, UsageSettlementWriter,
};
pub use shutdown::UsageProducerGuard;
pub use standardized_usage::StandardizedUsage;
pub use usage_mapper::{map_usage, map_usage_from_response, UsageMapper};
pub use worker::{
+556 -21
View File
@@ -1,14 +1,30 @@
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use serde_json::json;
use aether_data_contracts::DataLayerError;
use aether_runtime_state::{
RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStats, RuntimeQueueStore,
RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueReclaimPage, RuntimeQueueStats,
RuntimeQueueStore, RuntimeQueueTransferOutcome,
};
use super::config::UsageRuntimeConfig;
use super::event::UsageEvent;
use super::event::{EncodedUsageEvent, UsageEvent};
use crate::dead_letter_encoding::{shared_dead_letter_encoding_budget, DeadLetterEncodingBudget};
use crate::queue_read_budget::{shared_queue_read_budget, QueueReadBudget, QueueReadReservation};
static PAYLOAD_DOWNGRADED_TOTAL: AtomicU64 = AtomicU64::new(0);
static PAYLOAD_REJECTED_TOTAL: AtomicU64 = AtomicU64::new(0);
pub(crate) fn payload_encoding_totals() -> (u64, u64) {
(
PAYLOAD_DOWNGRADED_TOTAL.load(Ordering::Relaxed),
PAYLOAD_REJECTED_TOTAL.load(Ordering::Relaxed),
)
}
pub(crate) fn is_permanent_enqueue_error(error: &DataLayerError) -> bool {
matches!(error, DataLayerError::InvalidInput(_))
}
#[derive(Clone)]
pub struct UsageQueue {
@@ -17,6 +33,35 @@ pub struct UsageQueue {
stream: String,
group: String,
dlq_stream: String,
read_budget: Arc<QueueReadBudget>,
dead_letter_encoding_budget: Arc<DeadLetterEncodingBudget>,
}
#[derive(Debug)]
pub(crate) enum UsageDeadLetterOutcome {
Transferred {
destination_id: String,
acked: usize,
},
Appended {
destination_id: String,
},
NotPending,
EncodingDeferred {
error: DataLayerError,
},
}
pub(crate) struct ReservedUsageReadBatch {
pub(crate) entries: Vec<RuntimeQueueEntry>,
pub(crate) requested_count: usize,
// Keep last: batch data must be dropped before its reservation is returned.
pub(crate) reservation: QueueReadReservation,
}
pub(crate) struct ReservedUsageReclaimPage {
pub(crate) page: RuntimeQueueReclaimPage,
pub(crate) reservation: QueueReadReservation,
}
impl UsageQueue {
@@ -31,9 +76,26 @@ impl UsageQueue {
group: config.consumer_group.clone(),
dlq_stream: config.dlq_stream_key.clone(),
config,
read_budget: shared_queue_read_budget(),
dead_letter_encoding_budget: shared_dead_letter_encoding_budget(),
})
}
#[cfg(test)]
pub(crate) fn with_read_budget(mut self, read_budget: Arc<QueueReadBudget>) -> Self {
self.read_budget = read_budget;
self
}
#[cfg(test)]
pub(crate) fn with_dead_letter_encoding_budget(
mut self,
budget: Arc<DeadLetterEncodingBudget>,
) -> Self {
self.dead_letter_encoding_budget = budget;
self
}
pub async fn ensure_consumer_group(&self) -> Result<(), DataLayerError> {
self.runner
.ensure_consumer_group(&self.stream, &self.group, "0-0")
@@ -41,12 +103,36 @@ impl UsageQueue {
}
pub async fn enqueue(&self, event: &UsageEvent) -> Result<String, DataLayerError> {
let fields = event.to_stream_fields()?;
let encoded = self.encode_event(event)?;
self.runner
.append_fields_with_maxlen(&self.stream, &fields, Some(self.config.stream_maxlen))
.append_fields_with_maxlen(
&self.stream,
&encoded.fields,
Some(self.config.stream_maxlen),
)
.await
}
pub(crate) fn validate_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
self.encode_event(event).map(|_| ())
}
fn encode_event(&self, event: &UsageEvent) -> Result<EncodedUsageEvent, DataLayerError> {
let encoded = match event.to_bounded_stream_fields(self.config.queue_payload_max_bytes) {
Ok(encoded) => encoded,
Err(error) => {
if is_permanent_enqueue_error(&error) {
PAYLOAD_REJECTED_TOTAL.fetch_add(1, Ordering::Relaxed);
}
return Err(error);
}
};
if encoded.diagnostics_omitted {
PAYLOAD_DOWNGRADED_TOTAL.fetch_add(1, Ordering::Relaxed);
}
Ok(encoded)
}
pub async fn read_group(
&self,
consumer: &str,
@@ -62,13 +148,52 @@ impl UsageQueue {
.await
}
/// Workers retain this lease through processing. The public Vec API remains
/// compatible, but cannot preserve a reservation after returning its entries.
pub(crate) async fn read_group_reserved(
&self,
consumer: &str,
) -> Result<ReservedUsageReadBatch, DataLayerError> {
let (requested_count, mut reservation) = self
.read_budget
.reserve(
self.config.consumer_batch_size,
self.config.queue_payload_max_bytes,
)
.await?;
let entries = self
.runner
.read_group(
&self.stream,
&self.group,
consumer,
requested_count,
Some(self.config.consumer_block_ms.max(1)),
)
.await?;
reservation.observe_entries(&entries, self.config.queue_payload_max_bytes);
Ok(ReservedUsageReadBatch {
entries,
requested_count,
reservation,
})
}
pub async fn claim_stale(
&self,
consumer: &str,
start_id: &str,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
Ok(self.claim_stale_page(consumer, start_id).await?.entries)
}
pub async fn claim_stale_page(
&self,
consumer: &str,
start_id: &str,
) -> Result<RuntimeQueueReclaimPage, DataLayerError> {
self.runner
.claim_stale(
.claim_stale_page(
&self.stream,
&self.group,
consumer,
@@ -81,10 +206,46 @@ impl UsageQueue {
.await
}
pub(crate) async fn claim_stale_page_reserved(
&self,
consumer: &str,
start_id: &str,
) -> Result<ReservedUsageReclaimPage, DataLayerError> {
let (requested_count, mut reservation) = self
.read_budget
.reserve(
self.config.reclaim_count,
self.config.queue_payload_max_bytes,
)
.await?;
let page = self
.runner
.claim_stale_page(
&self.stream,
&self.group,
consumer,
start_id,
RuntimeQueueReclaimConfig {
min_idle_ms: self.config.reclaim_idle_ms,
count: requested_count,
},
)
.await?;
reservation.observe_entries(&page.entries, self.config.queue_payload_max_bytes);
Ok(ReservedUsageReclaimPage { page, reservation })
}
pub async fn ack_and_delete(&self, ids: &[String]) -> Result<(), DataLayerError> {
self.runner.ack(&self.stream, &self.group, ids).await?;
self.ack_and_delete_counted(ids).await.map(|_| ())
}
pub(crate) async fn ack_and_delete_counted(
&self,
ids: &[String],
) -> Result<usize, DataLayerError> {
let acked = self.runner.ack(&self.stream, &self.group, ids).await?;
self.runner.delete(&self.stream, ids).await?;
Ok(())
Ok(acked)
}
pub async fn push_dead_letter(
@@ -92,20 +253,57 @@ impl UsageQueue {
entry: &RuntimeQueueEntry,
error: &str,
) -> Result<String, DataLayerError> {
let fields = std::collections::BTreeMap::from([(
"payload".to_string(),
serde_json::to_string(&json!({
"entry_id": entry.id,
"fields": entry.fields,
"error": error,
}))
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?,
)]);
let reservation = self.dead_letter_encoding_budget.try_reserve(entry, error)?;
let encoded = reservation
.encode_owned(entry.clone(), error.to_string())
.await?;
self.runner
.append_fields_with_maxlen(&self.dlq_stream, &fields, None)
.append_fields_with_maxlen(&self.dlq_stream, &encoded.fields, None)
.await
}
pub(crate) async fn transfer_dead_letter_owned(
&self,
entry: RuntimeQueueEntry,
error: String,
) -> Result<UsageDeadLetterOutcome, DataLayerError> {
let reservation = match self.dead_letter_encoding_budget.try_reserve(&entry, &error) {
Ok(reservation) => reservation,
Err(error) => return Ok(UsageDeadLetterOutcome::EncodingDeferred { error }),
};
let encoded = match reservation.encode_owned(entry, error).await {
Ok(encoded) => encoded,
Err(error) => return Ok(UsageDeadLetterOutcome::EncodingDeferred { error }),
};
match self
.runner
.try_transfer_pending_to_stream(
&self.stream,
&self.group,
&encoded.entry_id,
&self.dlq_stream,
&encoded.fields,
)
.await?
{
Some(RuntimeQueueTransferOutcome::Transferred {
destination_id,
acked,
..
}) => Ok(UsageDeadLetterOutcome::Transferred {
destination_id,
acked,
}),
Some(RuntimeQueueTransferOutcome::NotPending) => Ok(UsageDeadLetterOutcome::NotPending),
None => Ok(UsageDeadLetterOutcome::Appended {
destination_id: self
.runner
.append_fields_with_maxlen(&self.dlq_stream, &encoded.fields, None)
.await?,
}),
}
}
pub async fn stats(&self) -> Result<RuntimeQueueStats, DataLayerError> {
self.runner.stats(&self.stream, Some(&self.group)).await
}
@@ -136,10 +334,216 @@ fn usage_queue_runtime_settings(config: &UsageRuntimeConfig) -> UsageQueueRuntim
#[cfg(test)]
mod tests {
use super::{usage_queue_runtime_settings, UsageQueue, UsageQueueRuntimeSettings};
use super::{
usage_queue_runtime_settings, UsageDeadLetterOutcome, UsageQueue, UsageQueueRuntimeSettings,
};
use crate::dead_letter_encoding::DeadLetterEncodingBudget;
use crate::queue_read_budget::QueueReadBudget;
use crate::UsageRuntimeConfig;
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState};
use aether_data_contracts::DataLayerError;
use aether_runtime_state::{
MemoryRuntimeStateConfig, RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStats,
RuntimeQueueStore, RuntimeQueueTransferOutcome, RuntimeState,
};
use async_trait::async_trait;
use std::collections::BTreeMap;
use std::future::Future;
use std::sync::Arc;
use std::task::Poll;
use std::time::Duration;
struct HeldDeadLetterStore {
started: tokio::sync::Notify,
}
#[async_trait]
impl RuntimeQueueStore for HeldDeadLetterStore {
async fn ensure_consumer_group(
&self,
_stream: &str,
_group: &str,
_start_id: &str,
) -> Result<(), DataLayerError> {
unreachable!("store only exercises dead-letter writes")
}
async fn append_fields_with_maxlen(
&self,
_stream: &str,
fields: &BTreeMap<String, String>,
_maxlen: Option<usize>,
) -> Result<String, DataLayerError> {
assert!(fields.contains_key("payload"));
self.started.notify_one();
std::future::pending().await
}
async fn read_group(
&self,
_stream: &str,
_group: &str,
_consumer: &str,
_count: usize,
_block_ms: Option<u64>,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
unreachable!("store only exercises dead-letter writes")
}
async fn claim_stale(
&self,
_stream: &str,
_group: &str,
_consumer: &str,
_start_id: &str,
_config: RuntimeQueueReclaimConfig,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
unreachable!("store only exercises dead-letter writes")
}
async fn try_transfer_pending_to_stream(
&self,
_source: &str,
_group: &str,
_entry_id: &str,
_destination: &str,
fields: &BTreeMap<String, String>,
) -> Result<Option<RuntimeQueueTransferOutcome>, DataLayerError> {
assert!(fields.contains_key("payload"));
self.started.notify_one();
std::future::pending().await
}
async fn ack(
&self,
_stream: &str,
_group: &str,
_ids: &[String],
) -> Result<usize, DataLayerError> {
unreachable!("store only exercises dead-letter writes")
}
async fn delete(&self, _stream: &str, _ids: &[String]) -> Result<usize, DataLayerError> {
unreachable!("store only exercises dead-letter writes")
}
async fn stats(
&self,
_stream: &str,
_group: Option<&str>,
) -> Result<RuntimeQueueStats, DataLayerError> {
unreachable!("store only exercises dead-letter writes")
}
}
#[tokio::test]
async fn dead_letter_encoding_storage_wait_keeps_budget_until_public_or_owned_call_is_cancelled(
) {
for owned_transfer in [false, true] {
let store = Arc::new(HeldDeadLetterStore {
started: tokio::sync::Notify::new(),
});
let budget = Arc::new(DeadLetterEncodingBudget::new(4096, 1));
let queue = UsageQueue::new(store.clone(), UsageRuntimeConfig::default())
.unwrap()
.with_dead_letter_encoding_budget(Arc::clone(&budget));
let entry = RuntimeQueueEntry {
id: "1-0".to_string(),
fields: BTreeMap::from([("payload".to_string(), "original fields".to_string())]),
};
let task = tokio::spawn(async move {
if owned_transfer {
queue
.transfer_dead_letter_owned(entry, "failure".to_string())
.await
.map(|_| ())
} else {
queue.push_dead_letter(&entry, "failure").await.map(|_| ())
}
});
tokio::time::timeout(Duration::from_secs(2), store.started.notified())
.await
.unwrap();
assert_eq!(budget.snapshot().encoded_total, 1);
assert_eq!(budget.snapshot().active_jobs, 1);
assert!(budget.snapshot().reserved_bytes > 0);
task.abort();
assert!(matches!(task.await, Err(error) if error.is_cancelled()));
assert_eq!(budget.snapshot().active_jobs, 0);
assert_eq!(budget.snapshot().reserved_bytes, 0);
}
}
#[tokio::test]
async fn dead_letter_encoding_owned_defers_oversize_but_preserves_public_and_store_errors() {
let runner = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let queue = UsageQueue::new(runner, UsageRuntimeConfig::default())
.unwrap()
.with_dead_letter_encoding_budget(Arc::new(DeadLetterEncodingBudget::new(1, 1)));
let entry = RuntimeQueueEntry {
id: "1-0".to_string(),
fields: BTreeMap::from([("payload".to_string(), "original fields".to_string())]),
};
assert!(matches!(
queue
.transfer_dead_letter_owned(entry.clone(), "failure".to_string())
.await,
Ok(UsageDeadLetterOutcome::EncodingDeferred {
error: DataLayerError::InvalidInput(_),
})
));
assert!(matches!(
queue.push_dead_letter(&entry, "failure").await,
Err(DataLayerError::InvalidInput(_))
));
let budget = Arc::new(DeadLetterEncodingBudget::new(4096, 1));
let queue = queue.with_dead_letter_encoding_budget(Arc::clone(&budget));
// The absent source causes a native store error after successful encoding.
assert!(matches!(
queue
.transfer_dead_letter_owned(entry, "failure".to_string())
.await,
Err(DataLayerError::InvalidInput(_))
));
assert_eq!(budget.snapshot().encoded_total, 1);
assert_eq!(budget.snapshot().active_jobs, 0);
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert_eq!(queue.dlq_stats().await.unwrap().stream_length, 0);
}
#[tokio::test]
async fn dead_letter_encoding_owned_defers_capacity_without_starting_store_work() {
let store = Arc::new(HeldDeadLetterStore {
started: tokio::sync::Notify::new(),
});
let budget = Arc::new(DeadLetterEncodingBudget::new(4096, 1));
let queue = UsageQueue::new(store, UsageRuntimeConfig::default())
.unwrap()
.with_dead_letter_encoding_budget(Arc::clone(&budget));
let entry = RuntimeQueueEntry {
id: "1-0".to_string(),
fields: BTreeMap::from([("payload".to_string(), "original fields".to_string())]),
};
let held = budget.try_reserve(&entry, "failure").unwrap();
let result = tokio::time::timeout(
Duration::from_secs(2),
queue.transfer_dead_letter_owned(entry, "failure".to_string()),
)
.await
.unwrap();
assert!(matches!(
result,
Ok(UsageDeadLetterOutcome::EncodingDeferred {
error: DataLayerError::TimedOut(_),
})
));
assert_eq!(budget.snapshot().encoded_total, 0);
assert_eq!(budget.snapshot().active_jobs, 1);
assert_eq!(budget.snapshot().capacity_rejected_total, 1);
drop(held);
assert_eq!(budget.snapshot().active_jobs, 0);
assert_eq!(budget.snapshot().reserved_bytes, 0);
}
#[test]
fn usage_queue_applies_runtime_block_and_batch_settings() {
@@ -164,4 +568,135 @@ mod tests {
}
);
}
fn reserved_test_queue(runner: Arc<RuntimeState>, budget: Arc<QueueReadBudget>) -> UsageQueue {
UsageQueue::new(
runner,
UsageRuntimeConfig {
enabled: true,
queue_payload_max_bytes: 8,
consumer_batch_size: 128,
consumer_block_ms: 1,
reclaim_count: 128,
reclaim_idle_ms: 1,
..UsageRuntimeConfig::default()
},
)
.expect("test queue")
.with_read_budget(budget)
}
#[tokio::test]
async fn queue_read_budget_read_and_reclaim_share_a_reservation_across_clones() {
let runtime = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let budget = Arc::new(QueueReadBudget::new(16, 16));
let queue = reserved_test_queue(Arc::clone(&runtime), Arc::clone(&budget));
let other = queue.clone();
queue.ensure_consumer_group().await.unwrap();
for _ in 0..6 {
runtime
.append_fields_with_maxlen(
&queue.stream,
&BTreeMap::from([("payload".to_string(), "12345678".to_string())]),
None,
)
.await
.unwrap();
}
let first = queue.read_group_reserved("reader").await.unwrap();
assert_eq!(first.requested_count, 2);
assert_eq!(first.entries.len(), 2);
let first_ids = first
.entries
.iter()
.map(|entry| entry.id.clone())
.collect::<Vec<_>>();
let extra_pending = runtime
.read_group(&queue.stream, &queue.group, "previous-reader", 2, None)
.await
.unwrap();
assert_eq!(extra_pending.len(), 2);
let next_cursor = extra_pending[0].id.clone();
drop(extra_pending);
assert_eq!(budget.snapshot().reserved_bytes, 16);
let mut reclaim = Box::pin(other.claim_stale_page_reserved("reclaimer", "0-0"));
std::future::poll_fn(|cx| {
assert!(reclaim.as_mut().poll(cx).is_pending());
Poll::Ready(())
})
.await;
assert_eq!(budget.snapshot().waiters, 1);
tokio::time::sleep(Duration::from_millis(5)).await;
drop(first);
let claimed = reclaim.await.unwrap();
assert_eq!(claimed.page.entries.len(), 2);
assert_eq!(claimed.page.next_start_id, next_cursor);
assert_eq!(
claimed
.page
.entries
.iter()
.map(|entry| entry.id.clone())
.collect::<Vec<_>>(),
first_ids
);
assert_eq!(budget.snapshot().reserved_bytes, 16);
assert_eq!(budget.snapshot().waiters, 0);
let mut next_read = Box::pin(queue.read_group_reserved("reader"));
std::future::poll_fn(|cx| {
assert!(next_read.as_mut().poll(cx).is_pending());
Poll::Ready(())
})
.await;
drop(claimed);
let next = next_read.await.unwrap();
assert_eq!(next.entries.len(), 2);
assert!(next
.entries
.iter()
.all(|entry| !first_ids.contains(&entry.id)));
drop(next);
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert_eq!(budget.snapshot().wait_total, 2);
}
#[tokio::test]
async fn queue_read_budget_errors_and_empty_pages_release_reservations() {
let runtime = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let budget = Arc::new(QueueReadBudget::new(16, 16));
let queue = reserved_test_queue(runtime, Arc::clone(&budget));
assert!(queue.read_group_reserved("reader").await.is_err());
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert!(queue
.claim_stale_page_reserved("reader", "0-0")
.await
.is_err());
assert_eq!(budget.snapshot().reserved_bytes, 0);
queue.ensure_consumer_group().await.unwrap();
let empty = queue.read_group_reserved("reader").await.unwrap();
assert!(empty.entries.is_empty());
assert_eq!(budget.snapshot().reserved_bytes, 0);
let page = queue
.claim_stale_page_reserved("reader", "0-0")
.await
.unwrap();
assert!(page.page.entries.is_empty());
assert_eq!(page.page.next_start_id, "0-0");
assert_eq!(budget.snapshot().reserved_bytes, 0);
drop((empty, page));
assert_eq!(budget.snapshot().reserved_bytes, 0);
}
#[test]
fn queue_read_budget_new_queues_and_clones_share_process_budget() {
let runtime = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let first = UsageQueue::new(runtime.clone(), UsageRuntimeConfig::default()).unwrap();
let second = UsageQueue::new(runtime, UsageRuntimeConfig::default()).unwrap();
let cloned = first.clone();
assert!(Arc::ptr_eq(&first.read_budget, &second.read_budget));
assert!(Arc::ptr_eq(&first.read_budget, &cloned.read_budget));
}
}
@@ -0,0 +1,373 @@
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, LazyLock};
use aether_data_contracts::DataLayerError;
use aether_runtime_state::RuntimeQueueEntry;
use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError};
const DEFAULT_READ_PAYLOAD_BUDGET_BYTES: usize = 128 * 1024 * 1024;
const DEFAULT_READ_BATCH_PAYLOAD_BYTES: usize = 8 * 1024 * 1024;
static READ_BUDGET: LazyLock<Arc<QueueReadBudget>> = LazyLock::new(|| {
let limit = std::env::var("AETHER_USAGE_QUEUE_READ_PAYLOAD_BUDGET_BYTES").ok();
let batch = std::env::var("AETHER_USAGE_QUEUE_READ_BATCH_PAYLOAD_BYTES").ok();
Arc::new(QueueReadBudget::new(
configured_bytes(limit.as_deref(), DEFAULT_READ_PAYLOAD_BUDGET_BYTES),
configured_bytes(batch.as_deref(), DEFAULT_READ_BATCH_PAYLOAD_BYTES),
))
});
pub(crate) fn shared_queue_read_budget() -> Arc<QueueReadBudget> {
Arc::clone(&READ_BUDGET)
}
pub(crate) fn queue_read_budget_metrics() -> QueueReadBudgetSnapshot {
READ_BUDGET.snapshot()
}
fn maximum_budget_bytes() -> usize {
Semaphore::MAX_PERMITS.min(u32::MAX as usize)
}
fn configured_bytes(raw: Option<&str>, fallback: usize) -> usize {
raw.and_then(|raw| raw.trim().parse::<u128>().ok())
.filter(|value| *value > 0)
.map(|value| value.min(maximum_budget_bytes() as u128) as usize)
.unwrap_or(fallback.min(maximum_budget_bytes()))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct QueueReadBudgetSnapshot {
pub(crate) limit_bytes: usize,
pub(crate) batch_limit_bytes: usize,
pub(crate) reserved_bytes: usize,
pub(crate) waiters: usize,
pub(crate) wait_total: u64,
/// Observed field names plus values, without cloning the returned strings.
pub(crate) actual_field_bytes_total: u64,
pub(crate) oversized_entries_total: u64,
pub(crate) oversized_batches_total: u64,
}
/// A process-wide reservation based on the current producer payload limit.
/// Historical or externally written messages can exceed the estimate. Field names,
/// allocation capacity, RESP decoding, and decoded JSON are not an RSS bound here.
pub(crate) struct QueueReadBudget {
limit_bytes: usize,
batch_limit_bytes: usize,
permits: Arc<Semaphore>,
reserved_bytes: AtomicUsize,
waiters: AtomicUsize,
wait_total: AtomicU64,
actual_field_bytes_total: AtomicU64,
oversized_entries_total: AtomicU64,
oversized_batches_total: AtomicU64,
}
impl QueueReadBudget {
pub(crate) fn new(limit_bytes: usize, batch_limit_bytes: usize) -> Self {
let limit_bytes = limit_bytes.clamp(1, maximum_budget_bytes());
let batch_limit_bytes = batch_limit_bytes.clamp(1, limit_bytes);
Self {
limit_bytes,
batch_limit_bytes,
permits: Arc::new(Semaphore::new(limit_bytes)),
reserved_bytes: AtomicUsize::new(0),
waiters: AtomicUsize::new(0),
wait_total: AtomicU64::new(0),
actual_field_bytes_total: AtomicU64::new(0),
oversized_entries_total: AtomicU64::new(0),
oversized_batches_total: AtomicU64::new(0),
}
}
pub(crate) fn snapshot(&self) -> QueueReadBudgetSnapshot {
QueueReadBudgetSnapshot {
limit_bytes: self.limit_bytes,
batch_limit_bytes: self.batch_limit_bytes,
reserved_bytes: self.reserved_bytes.load(Ordering::Relaxed),
waiters: self.waiters.load(Ordering::Relaxed),
wait_total: self.wait_total.load(Ordering::Relaxed),
actual_field_bytes_total: self.actual_field_bytes_total.load(Ordering::Relaxed),
oversized_entries_total: self.oversized_entries_total.load(Ordering::Relaxed),
oversized_batches_total: self.oversized_batches_total.load(Ordering::Relaxed),
}
}
pub(crate) async fn reserve(
self: &Arc<Self>,
requested_count: usize,
payload_limit: usize,
) -> Result<(usize, QueueReadReservation), DataLayerError> {
if payload_limit == 0 || payload_limit > self.limit_bytes {
return Err(DataLayerError::InvalidConfiguration(format!(
"usage queue payload limit {payload_limit} must be positive and not exceed the {}-byte read payload budget",
self.limit_bytes
)));
}
// A single valid payload may exceed the preferred batch target, but never
// the total budget. Clamp before multiplying or converting to u32 permits.
let count = requested_count
.max(1)
.min((self.batch_limit_bytes / payload_limit).max(1));
let reserved_bytes = count * payload_limit;
let permits = reserved_bytes as u32;
let permit = match Arc::clone(&self.permits).try_acquire_many_owned(permits) {
Ok(permit) => permit,
Err(TryAcquireError::NoPermits) => {
self.wait_total.fetch_add(1, Ordering::Relaxed);
self.waiters.fetch_add(1, Ordering::Relaxed);
let _waiting = WaitingReservation { budget: self };
Arc::clone(&self.permits)
.acquire_many_owned(permits)
.await
.map_err(|_| closed_budget_error())?
}
Err(TryAcquireError::Closed) => return Err(closed_budget_error()),
};
self.reserved_bytes
.fetch_add(reserved_bytes, Ordering::Relaxed);
Ok((
count,
QueueReadReservation {
budget: Arc::clone(self),
reserved_bytes,
permit: Some(permit),
},
))
}
}
fn closed_budget_error() -> DataLayerError {
DataLayerError::InvalidConfiguration("usage queue read payload budget is closed".to_string())
}
struct WaitingReservation<'a> {
budget: &'a QueueReadBudget,
}
impl Drop for WaitingReservation<'_> {
fn drop(&mut self) {
self.budget.waiters.fetch_sub(1, Ordering::Relaxed);
}
}
// Deliberately not Clone: every concurrently retained batch needs its own lease.
pub(crate) struct QueueReadReservation {
budget: Arc<QueueReadBudget>,
reserved_bytes: usize,
permit: Option<OwnedSemaphorePermit>,
}
impl QueueReadReservation {
pub(crate) fn observe_entries(&mut self, entries: &[RuntimeQueueEntry], payload_limit: usize) {
let mut value_bytes = 0usize;
let mut field_bytes = 0usize;
let mut oversized_entries = 0u64;
for entry in entries {
let mut entry_value_bytes = 0usize;
for (key, value) in &entry.fields {
entry_value_bytes = entry_value_bytes.saturating_add(value.len());
field_bytes = field_bytes
.saturating_add(key.len())
.saturating_add(value.len());
}
value_bytes = value_bytes.saturating_add(entry_value_bytes);
oversized_entries += u64::from(entry_value_bytes > payload_limit);
}
self.budget.actual_field_bytes_total.fetch_add(
u64::try_from(field_bytes).unwrap_or(u64::MAX),
Ordering::Relaxed,
);
self.budget
.oversized_entries_total
.fetch_add(oversized_entries, Ordering::Relaxed);
if value_bytes > self.reserved_bytes {
self.budget
.oversized_batches_total
.fetch_add(1, Ordering::Relaxed);
}
// Shrink unused payload estimates. Never wait for an upgrade after reading
// an oversized historical batch: other batches may hold all remaining bytes.
self.shrink_to(value_bytes.min(self.reserved_bytes));
}
fn shrink_to(&mut self, retained_bytes: usize) {
let released = self.reserved_bytes.saturating_sub(retained_bytes);
if released == 0 {
return;
}
let permit = self
.permit
.as_mut()
.expect("positive reservation must hold a permit")
.split(released)
.expect("released bytes must belong to this reservation");
self.reserved_bytes -= released;
self.budget
.reserved_bytes
.fetch_sub(released, Ordering::Relaxed);
drop(permit);
}
}
impl Drop for QueueReadReservation {
fn drop(&mut self) {
self.budget
.reserved_bytes
.fetch_sub(self.reserved_bytes, Ordering::Relaxed);
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::future::Future;
use std::task::Poll;
use super::*;
#[tokio::test]
async fn queue_read_budget_cancelled_wait_preserves_current_reservations() {
let budget = Arc::new(QueueReadBudget::new(32, 16));
let (_, first) = budget.reserve(8, 8).await.unwrap();
let (_, second) = budget.reserve(8, 8).await.unwrap();
let mut pending = Box::pin(budget.reserve(1, 8));
std::future::poll_fn(|cx| {
assert!(pending.as_mut().poll(cx).is_pending());
Poll::Ready(())
})
.await;
assert_eq!(budget.snapshot().reserved_bytes, 32);
assert_eq!(budget.snapshot().waiters, 1);
drop(pending);
assert_eq!(budget.snapshot().waiters, 0);
assert_eq!(budget.snapshot().wait_total, 1);
drop(first);
let (_, replacement) = budget.reserve(2, 8).await.unwrap();
assert_eq!(budget.snapshot().reserved_bytes, 32);
drop((second, replacement));
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert_eq!(budget.permits.available_permits(), 32);
}
#[tokio::test]
async fn queue_read_budget_counts_payload_values_and_observes_legacy_excess_without_waiting() {
let budget = Arc::new(QueueReadBudget::new(16, 16));
let (count, mut reservation) = budget.reserve(2, 8).await.unwrap();
assert_eq!(count, 2);
let entries = [RuntimeQueueEntry {
id: "1-0".to_string(),
fields: BTreeMap::from([("payload".to_string(), "x".repeat(8))]),
}];
reservation.observe_entries(&entries, 8);
assert_eq!(budget.snapshot().reserved_bytes, 8);
assert_eq!(budget.snapshot().actual_field_bytes_total, 15);
assert_eq!(budget.snapshot().oversized_entries_total, 0);
let (_, mut legacy) = budget.reserve(1, 8).await.unwrap();
let legacy_entries = [RuntimeQueueEntry {
id: "2-0".to_string(),
fields: BTreeMap::from([
("payload".to_string(), "x".repeat(8)),
("extra".to_string(), "y".repeat(24)),
]),
}];
legacy.observe_entries(&legacy_entries, 8);
assert_eq!(budget.snapshot().reserved_bytes, 16);
assert_eq!(budget.snapshot().actual_field_bytes_total, 59);
assert_eq!(budget.snapshot().oversized_entries_total, 1);
assert_eq!(budget.snapshot().oversized_batches_total, 1);
assert_eq!(budget.snapshot().wait_total, 0);
drop((reservation, legacy));
assert_eq!(budget.snapshot().reserved_bytes, 0);
}
#[tokio::test]
async fn queue_read_budget_empty_response_releases_all_reserved_bytes() {
let budget = Arc::new(QueueReadBudget::new(16, 16));
let (_, mut reservation) = budget.reserve(2, 8).await.unwrap();
reservation.observe_entries(&[], 8);
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert_eq!(budget.permits.available_permits(), 16);
drop(reservation);
assert_eq!(budget.snapshot().reserved_bytes, 0);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn queue_read_budget_concurrent_shrink_and_drop_preserve_shared_capacity() {
let budget = Arc::new(QueueReadBudget::new(64, 16));
let barrier = Arc::new(tokio::sync::Barrier::new(16));
let mut tasks = tokio::task::JoinSet::new();
for _ in 0..16 {
let budget = Arc::clone(&budget);
let barrier = Arc::clone(&barrier);
tasks.spawn(async move {
barrier.wait().await;
for _ in 0..16 {
let (count, mut reservation) = budget.reserve(128, 8).await.unwrap();
assert_eq!(count, 2);
assert!(budget.snapshot().reserved_bytes <= 64);
reservation.observe_entries(
&[RuntimeQueueEntry {
id: "1-0".to_string(),
fields: BTreeMap::from([(
"payload".to_string(),
"12345678".to_string(),
)]),
}],
8,
);
tokio::task::yield_now().await;
drop(reservation);
}
});
}
tokio::time::timeout(std::time::Duration::from_secs(5), async {
while let Some(result) = tasks.join_next().await {
result.expect("reservation task");
}
})
.await
.expect("all shared reservations complete");
assert_eq!(budget.snapshot().reserved_bytes, 0);
assert_eq!(budget.snapshot().waiters, 0);
assert_eq!(budget.snapshot().actual_field_bytes_total, 16 * 16 * 15);
assert_eq!(budget.permits.available_permits(), 64);
}
#[tokio::test]
async fn queue_read_budget_large_values_cannot_overflow_or_wait_for_impossible_permits() {
let budget = Arc::new(QueueReadBudget::new(usize::MAX, usize::MAX));
let maximum = maximum_budget_bytes();
assert_eq!(budget.snapshot().limit_bytes, maximum);
let (count, reservation) = budget.reserve(usize::MAX, 1).await.unwrap();
assert_eq!(count, maximum);
assert_eq!(budget.snapshot().reserved_bytes, maximum);
drop(reservation);
assert!(matches!(
budget.reserve(1, maximum + 1).await,
Err(DataLayerError::InvalidConfiguration(_))
));
assert!(matches!(
budget.reserve(1, 0).await,
Err(DataLayerError::InvalidConfiguration(_))
));
let small_batch = Arc::new(QueueReadBudget::new(32, 4));
let (count, reservation) = small_batch.reserve(usize::MAX, 16).await.unwrap();
assert_eq!(count, 1);
assert_eq!(small_batch.snapshot().reserved_bytes, 16);
drop(reservation);
}
#[test]
fn queue_read_budget_env_uses_positive_defaults_and_caps_extreme_values() {
for raw in [None, Some(""), Some("0"), Some("-1"), Some("bad")] {
assert_eq!(configured_bytes(raw, 128), 128);
}
assert_eq!(configured_bytes(Some(" 42 "), 128), 42);
assert_eq!(
configured_bytes(Some(&u128::MAX.to_string()), 128),
maximum_budget_bytes()
);
}
}
+44
View File
@@ -182,6 +182,7 @@ pub fn build_upsert_usage_record_from_event(
finalized_at_unix_secs,
created_at_unix_ms: Some(now_unix_secs),
updated_at_unix_secs: now_unix_secs,
capture_retention: data.capture_retention,
})
}
@@ -262,6 +263,49 @@ mod tests {
use super::build_upsert_usage_record_from_event;
#[test]
fn capture_retention_follows_event_bodies_into_record_and_its_clones() {
use aether_data_contracts::repository::usage::{
usage_json_heap_estimate, UsageCaptureMemoryBudget,
};
use std::sync::Arc;
let body = serde_json::Value::String("retained diagnostic".repeat(8));
let estimate =
4 * (std::mem::size_of::<serde_json::Value>() + usage_json_heap_estimate(&body));
let budget = Arc::new(UsageCaptureMemoryBudget::new(3 * estimate));
let mut event = UsageEvent::new(
UsageEventType::Completed,
"retained-record",
UsageEventData {
provider_name: "provider".to_owned(),
model: "model".to_owned(),
input_tokens: Some(5),
output_tokens: Some(7),
cache_read_input_tokens: Some(0),
request_body: Some(body.clone()),
provider_request_body: Some(body.clone()),
response_body: Some(body.clone()),
client_response_body: Some(body),
..UsageEventData::default()
},
);
event.data.apply_capture_memory_budget(Arc::clone(&budget));
assert_eq!(budget.retained_bytes(), estimate);
let record = build_upsert_usage_record_from_event(&event).unwrap();
assert_eq!(budget.retained_bytes(), 2 * estimate);
drop(event);
assert_eq!(budget.retained_bytes(), estimate);
let cloned = record.clone();
assert_eq!(budget.retained_bytes(), 2 * estimate);
assert_eq!(cloned.cache_read_input_tokens, Some(0));
assert_eq!(cloned.response_body, record.response_body);
drop(record);
assert_eq!(budget.retained_bytes(), estimate);
drop(cloned);
assert_eq!(budget.retained_bytes(), 0);
}
#[test]
fn builds_upsert_record_from_terminal_event() {
let record = build_upsert_usage_record_from_event(&UsageEvent {
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,432 @@
use super::*;
use super::super::{run_usage_enqueue_retry_worker, TerminalPersistenceOutcome};
const PAYLOAD_LIMIT: usize = 4 * 1024;
fn payload_config(name: &str) -> UsageRuntimeConfig {
UsageRuntimeConfig {
enabled: true,
queue_terminal_events: true,
queue_lifecycle_events: true,
stream_key: format!("usage:events:test:payload:{name}"),
consumer_group: format!("usage_consumers_payload_{name}"),
queue_payload_max_bytes: PAYLOAD_LIMIT,
consumer_block_ms: 1,
enqueue_retry_buffer_capacity: 8,
enqueue_retry_workers: 1,
enqueue_retry_initial_backoff_ms: 1,
enqueue_retry_max_backoff_ms: 2,
..UsageRuntimeConfig::default()
}
}
fn payload_event(request_id: &str, oversized: bool) -> UsageEvent {
UsageEvent {
event_type: UsageEventType::Completed,
request_id: request_id.to_string(),
timestamp_ms: 123_000,
data: UsageEventData {
user_id: Some("user-payload".to_string()),
api_key_id: Some("key-payload".to_string()),
provider_name: "openai".to_string(),
provider_id: Some("provider-payload".to_string()),
provider_api_key_id: Some("provider-key-payload".to_string()),
// Model identity is a core field, so omitting diagnostic bodies cannot make this fit.
model: if oversized {
"m".repeat(PAYLOAD_LIMIT * 2)
} else {
"gpt-5".to_string()
},
target_model: Some("gpt-5".to_string()),
api_format: Some("openai:responses".to_string()),
endpoint_api_format: Some("openai:responses".to_string()),
input_tokens: Some(100),
output_tokens: Some(25),
total_tokens: Some(125),
cache_creation_input_tokens: Some(30),
cache_creation_ephemeral_5m_input_tokens: Some(0),
cache_creation_ephemeral_1h_input_tokens: Some(30),
cache_read_input_tokens: Some(0),
total_cost_usd: Some(0.5),
actual_total_cost_usd: Some(0.25),
status_code: Some(200),
error_message: Some("preserve error presence and text".to_string()),
first_byte_time_ms: Some(12),
response_time_ms: Some(34),
request_headers: Some(json!({"x-request": "original"})),
provider_request_headers: Some(json!({"x-provider-request": "original"})),
response_headers: Some(json!({"x-provider-response": "original"})),
client_response_headers: Some(json!({"x-client-response": "original"})),
request_body: Some(json!({"reasoning": {"effort": "high"}, "input": "original"})),
request_body_state: Some(UsageBodyCaptureState::Inline),
provider_request_body: Some(json!({
"model": "gpt-5", "service_tier": "priority",
"prompt_cache_retention": "24h", "input": "original provider input"
})),
provider_request_body_state: Some(UsageBodyCaptureState::Inline),
response_body: Some(json!({"service_tier": "default", "output": "original output"})),
response_body_state: Some(UsageBodyCaptureState::Inline),
client_response_body: Some(json!({"output": "original client output"})),
client_response_body_state: Some(UsageBodyCaptureState::Inline),
request_metadata: Some(json!({
"api_key_is_standalone": true,
"plan_usage_reservation_token": "550e8400-e29b-41d4-a716-446655440000",
"plan_usage_reservation_deferred": true,
"usage_available": true,
"usage_pricing_available": true,
"provider_cache_ttl_minutes": 1440,
"provider_service_tier": "priority",
"provider_actual_service_tier": "default",
"dimensions": {
"image_count": 2, "image_size": "1024x1024", "image_quality": "high",
"image_output_format": "png", "reasoning_tokens": 0
}
})),
..UsageEventData::default()
},
}
}
async fn payload_queue(config: &UsageRuntimeConfig) -> (UsageQueue, Arc<FlakyAppendQueueStore>) {
let inner: Arc<dyn RuntimeQueueStore> =
Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let runner = Arc::new(FlakyAppendQueueStore::new(inner, 0));
let queue = UsageQueue::new(runner.clone(), config.clone()).expect("payload queue");
queue.ensure_consumer_group().await.expect("payload group");
(queue, runner)
}
#[tokio::test]
async fn terminal_oversize_uses_original_event_for_direct_fallback_without_opening_circuit() {
let config = payload_config("terminal_direct");
let (queue, runner) = payload_queue(&config).await;
let store = CloneQueueConfiguredUsageStore {
records: Arc::new(Mutex::new(Vec::new())),
queue: runner.clone(),
};
let runtime = UsageRuntime::new(config).expect("usage runtime");
let event = payload_event("payload-terminal-direct", true);
let expected = crate::build_upsert_usage_record_from_event(&event).expect("original record");
assert!(matches!(
queue.enqueue(&event).await,
Err(DataLayerError::InvalidInput(_))
));
assert_eq!(runner.append_attempts.load(Ordering::Acquire), 0);
assert_eq!(
runtime.enqueue_or_write_terminal(&store, event).await,
TerminalPersistenceOutcome::PersistedDirectly
);
{
let records = store.records.lock().expect("records lock");
assert_eq!(records.as_slice(), &[expected]);
assert_eq!(records[0].cache_read_input_tokens, Some(0));
assert_eq!(records[0].cache_creation_ephemeral_5m_input_tokens, Some(0));
assert_eq!(
records[0].request_body_state,
Some(UsageBodyCaptureState::Inline)
);
assert!(records[0].provider_request_body.is_some());
assert!(records[0].request_headers.is_some());
}
assert_eq!(
runtime
.terminal_enqueue_state
.circuit_open_until_unix_ms
.load(Ordering::Acquire),
0
);
let snapshot = runtime.metrics_snapshot();
assert_eq!(snapshot.terminal_direct_fallback_succeeded_total, 1);
assert_eq!(snapshot.enqueue_retry_scheduled_total, 0);
assert_eq!(snapshot.enqueue_retry_pending, 0);
assert_eq!(
runtime
.enqueue_or_write_terminal(&store, payload_event("payload-after-direct", false))
.await,
TerminalPersistenceOutcome::Queued
);
assert_eq!(runner.successful_appends.load(Ordering::Acquire), 1);
assert_eq!(store.records.lock().expect("records lock").len(), 1);
}
#[tokio::test]
async fn terminal_oversize_direct_failure_preserves_first_byte_and_does_not_retry_or_open_circuit()
{
let config = payload_config("terminal_failed");
let (_, runner) = payload_queue(&config).await;
let store = FailingWriteQueueConfiguredUsageStore {
queue: runner.clone(),
upsert_attempts: Arc::new(AtomicUsize::new(0)),
};
let runtime = UsageRuntime::new(config).expect("usage runtime");
let request_id = "payload-terminal-failed";
let generation = runtime
.lifecycle_coalescer
.mark_first_byte(request_id)
.await
.expect("first byte");
assert_eq!(
runtime
.enqueue_or_write_terminal(&store, payload_event(request_id, true))
.await,
TerminalPersistenceOutcome::Failed
);
assert!(
runtime
.lifecycle_coalescer
.first_byte_is_current(request_id, generation)
.await
);
assert_eq!(runner.append_attempts.load(Ordering::Acquire), 0);
assert_eq!(store.upsert_attempts.load(Ordering::Acquire), 1);
assert_eq!(
runtime
.terminal_enqueue_state
.circuit_open_until_unix_ms
.load(Ordering::Acquire),
0
);
let snapshot = runtime.metrics_snapshot();
assert_eq!(snapshot.terminal_direct_fallback_failed_total, 1);
assert_eq!(snapshot.terminal_enqueue_deferred_dropped_total, 1);
assert_eq!(snapshot.enqueue_retry_scheduled_total, 0);
assert_eq!(snapshot.enqueue_retry_pending, 0);
assert_eq!(
runtime
.enqueue_or_write_terminal(&store, payload_event("payload-after-failed", false))
.await,
TerminalPersistenceOutcome::Queued
);
assert_eq!(runner.successful_appends.load(Ordering::Acquire), 1);
assert_eq!(store.upsert_attempts.load(Ordering::Acquire), 1);
}
#[tokio::test]
async fn terminal_oversize_direct_failure_stays_failed_when_primary_enqueue_is_deferred() {
for (name, circuit_open) in [("circuit_open", true), ("in_flight_limit", false)] {
let mut config = payload_config(name);
config.terminal_enqueue_max_in_flight = 1;
let (_, runner) = payload_queue(&config).await;
let store = FailingWriteQueueConfiguredUsageStore {
queue: runner.clone(),
upsert_attempts: Arc::new(AtomicUsize::new(0)),
};
let runtime = UsageRuntime::new(config).expect("usage runtime");
let request_id = format!("payload-terminal-{name}");
let generation = runtime
.lifecycle_coalescer
.mark_first_byte(&request_id)
.await
.expect("first byte");
let original_deadline = if circuit_open {
let deadline = super::super::now_unix_ms().saturating_add(60_000);
runtime.terminal_enqueue_state.open_circuit(deadline);
deadline
} else {
0
};
let held_guard = if circuit_open {
None
} else {
Some(
runtime
.terminal_enqueue_state
.try_acquire_in_flight(1)
.expect("hold the only enqueue slot"),
)
};
assert_eq!(
timeout(
Duration::from_secs(2),
runtime.enqueue_or_write_terminal(&store, payload_event(&request_id, true)),
)
.await
.expect("bounded terminal fallback"),
TerminalPersistenceOutcome::Failed,
"{name} must not report an oversized event as buffered"
);
assert!(
runtime
.lifecycle_coalescer
.first_byte_is_current(&request_id, generation)
.await
);
assert_eq!(runner.append_attempts.load(Ordering::Acquire), 0);
assert_eq!(store.upsert_attempts.load(Ordering::Acquire), 1);
assert_eq!(
runtime
.terminal_enqueue_state
.circuit_open_until_unix_ms
.load(Ordering::Acquire),
original_deadline
);
let snapshot = runtime.metrics_snapshot();
assert_eq!(snapshot.terminal_direct_fallback_failed_total, 1);
assert_eq!(snapshot.terminal_enqueue_deferred_dropped_total, 1);
assert_eq!(snapshot.terminal_enqueue_deferred_retry_total, 0);
assert_eq!(snapshot.enqueue_retry_permanent_failure_total, 1);
assert_eq!(snapshot.enqueue_retry_scheduled_total, 0);
assert_eq!(snapshot.enqueue_retry_pending, 0);
assert_eq!(
runtime.terminal_enqueue_state.in_flight(),
u64::from(!circuit_open)
);
drop(held_guard);
assert_eq!(runtime.terminal_enqueue_state.in_flight(), 0);
}
}
#[tokio::test]
async fn retry_worker_discards_oversize_and_drains_next_event_on_the_same_shard() {
let config = payload_config("retry_drain");
let (queue, runner) = payload_queue(&config).await;
let (sender, receiver) = mpsc::channel(2);
let dispatcher = UsageEnqueueRetryDispatcher {
senders: vec![sender],
metrics: Arc::new(Default::default()),
};
let metrics_view = UsageEnqueueRetryDispatcher {
senders: Vec::new(),
metrics: Arc::clone(&dispatcher.metrics),
};
for (request_id, oversized) in [
("payload-retry-oversize", true),
("payload-retry-small", false),
] {
// Bypass admission to exercise the worker's defense for an already buffered event.
assert!(dispatcher
.schedule_item(
queue.clone(),
payload_event(request_id, oversized),
"terminal",
Some("prior transient failure"),
)
.is_some());
}
assert_eq!(dispatcher.pending(), 2);
drop(dispatcher);
// Run the real worker as a cancellable future, so a regression cannot leave a detached retry.
timeout(
Duration::from_secs(2),
run_usage_enqueue_retry_worker(0, config, receiver, Arc::clone(&metrics_view.metrics)),
)
.await
.expect("the permanent failure must not block the shard");
assert_eq!(metrics_view.permanent_failure_total(), 1);
assert_eq!(metrics_view.recovered_total(), 1);
assert_eq!(metrics_view.pending(), 0);
assert_eq!(runner.append_attempts.load(Ordering::Acquire), 1);
let entries = queue
.read_group("payload-retry-reader")
.await
.expect("queue read");
assert_eq!(entries.len(), 1);
assert_eq!(
UsageEvent::from_stream_fields(&entries[0].fields)
.expect("queued event")
.request_id,
"payload-retry-small"
);
}
#[tokio::test]
async fn retry_dispatcher_rejects_oversize_for_permanent_and_transient_causes_without_consuming_slots(
) {
let config = payload_config("retry_reject");
let (queue, runner) = payload_queue(&config).await;
let (sender, receiver) = mpsc::channel(1);
let dispatcher = UsageEnqueueRetryDispatcher {
senders: vec![sender],
metrics: Arc::new(Default::default()),
};
let oversized = payload_event("payload-rejected", true);
let error = queue.enqueue(&oversized).await.expect_err("oversize input");
assert!(matches!(error, DataLayerError::InvalidInput(_)));
assert!(!dispatcher.schedule(queue.clone(), oversized, "terminal", error));
for cause in [
DataLayerError::TimedOut("primary enqueue was deferred".to_string()),
DataLayerError::Redis("prior transient failure".to_string()),
] {
assert!(!dispatcher.schedule(
queue.clone(),
payload_event("payload-rejected-transient-cause", true),
"terminal",
cause,
));
}
assert_eq!(dispatcher.permanent_failure_total(), 3);
assert_eq!(dispatcher.pending(), 0);
assert_eq!(dispatcher.scheduled_total(), 0);
assert!(dispatcher.schedule(
queue,
payload_event("payload-after-reject", false),
"terminal",
DataLayerError::Redis("retryable failure".to_string()),
));
let metrics_view = UsageEnqueueRetryDispatcher {
senders: Vec::new(),
metrics: Arc::clone(&dispatcher.metrics),
};
drop(dispatcher);
timeout(
Duration::from_secs(2),
run_usage_enqueue_retry_worker(0, config, receiver, Arc::clone(&metrics_view.metrics)),
)
.await
.expect("retry worker drain");
assert_eq!(metrics_view.permanent_failure_total(), 3);
assert_eq!(metrics_view.recovered_total(), 1);
assert_eq!(metrics_view.pending(), 0);
assert_eq!(runner.append_attempts.load(Ordering::Acquire), 1);
}
#[tokio::test]
async fn lifecycle_oversize_does_not_open_circuit_or_block_the_next_lifecycle_event() {
let config = payload_config("lifecycle");
let (queue, runner) = payload_queue(&config).await;
let store = CloneQueueConfiguredUsageStore {
records: Arc::new(Mutex::new(Vec::new())),
queue: runner.clone(),
};
let runtime = UsageRuntime::new(config).expect("usage runtime");
let mut oversized = payload_event("payload-lifecycle-oversize", true);
oversized.event_type = UsageEventType::Streaming;
assert!(!runtime.enqueue_lifecycle_event(&store, oversized).await);
assert_eq!(
runtime
.lifecycle_enqueue_state
.circuit_open_until_unix_ms
.load(Ordering::Acquire),
0
);
assert_eq!(runner.append_attempts.load(Ordering::Acquire), 0);
let snapshot = runtime.metrics_snapshot();
assert_eq!(snapshot.enqueue_retry_scheduled_total, 0);
assert_eq!(snapshot.enqueue_retry_pending, 0);
assert!(store.records.lock().expect("records lock").is_empty());
let mut small = payload_event("payload-lifecycle-small", false);
small.event_type = UsageEventType::Streaming;
assert!(runtime.enqueue_lifecycle_event(&store, small).await);
assert_eq!(runner.successful_appends.load(Ordering::Acquire), 1);
let entries = queue
.read_group("payload-lifecycle-reader")
.await
.expect("queue read");
assert_eq!(entries.len(), 1);
let queued = UsageEvent::from_stream_fields(&entries[0].fields).expect("queued lifecycle");
assert_eq!(queued.request_id, "payload-lifecycle-small");
assert_eq!(queued.event_type, UsageEventType::Streaming);
assert_eq!(queued.data.first_byte_time_ms, Some(12));
}
@@ -0,0 +1,366 @@
use super::*;
fn config() -> UsageRuntimeConfig {
UsageRuntimeConfig {
enabled: true,
queue_terminal_events: true,
queue_lifecycle_events: true,
worker_count: 2,
consumer_block_ms: 60_000,
enqueue_retry_buffer_capacity: 256,
enqueue_retry_initial_backoff_ms: 60_000,
enqueue_retry_max_backoff_ms: 60_000,
..UsageRuntimeConfig::default()
}
}
fn store() -> CloneQueueConfiguredUsageStore {
CloneQueueConfiguredUsageStore {
records: Arc::new(Mutex::new(Vec::new())),
queue: Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())),
}
}
fn terminal(request_id: &str) -> UsageEvent {
UsageEvent::new(
UsageEventType::Completed,
request_id,
UsageEventData {
provider_name: "openai".to_string(),
model: "test".to_string(),
status_code: Some(200),
input_tokens: Some(3),
output_tokens: Some(7),
total_tokens: Some(10),
..UsageEventData::default()
},
)
}
#[tokio::test]
async fn usage_shutdown_waits_for_a_producer_before_closing_admission() {
let runtime = UsageRuntime::new(config()).unwrap();
let store = store();
let producer = runtime.track_producer();
let copy = runtime.clone();
let task = tokio::spawn(async move { copy.shutdown(Duration::from_secs(3)).await });
sleep(Duration::from_millis(30)).await;
assert!(!task.is_finished());
runtime
.submit_terminal_event(&store, terminal("last-producer"))
.await;
drop(producer);
task.await.unwrap().unwrap();
assert_eq!(runtime.local_work_pending(), 0);
assert_eq!(
store
.queue
.stats(&runtime.config.stream_key, None)
.await
.unwrap()
.stream_length,
1
);
runtime.shutdown(Duration::from_secs(1)).await.unwrap();
runtime
.submit_terminal_event(&store, terminal("closed"))
.await;
runtime
.record_terminal_event(&store, terminal("closed-direct-api"))
.await;
assert_eq!(
store
.queue
.stats(&runtime.config.stream_key, None)
.await
.unwrap()
.stream_length,
1
);
}
#[tokio::test]
async fn usage_shutdown_persists_concurrent_terminal_handoffs() {
let runtime = UsageRuntime::new(config()).unwrap();
let store = store();
let mut tasks = tokio::task::JoinSet::new();
for index in 0..128 {
let runtime = runtime.clone();
let store = store.clone();
let producer = runtime.track_producer();
tasks.spawn(async move {
let _producer = producer;
runtime
.submit_terminal_event(&store, terminal(&format!("drain-{index}")))
.await;
});
}
runtime.shutdown(Duration::from_secs(5)).await.unwrap();
while let Some(result) = tasks.join_next().await {
result.unwrap();
}
assert_eq!(
store
.queue
.stats(&runtime.config.stream_key, None)
.await
.unwrap()
.stream_length,
128
);
assert_eq!(runtime.local_work_pending(), 0);
}
#[tokio::test]
async fn usage_shutdown_wakes_retry_backoff_and_preserves_all_buffered_events() {
let runtime = UsageRuntime::new(config()).unwrap();
let store = store();
let queue = UsageQueue::new(Arc::clone(&store.queue), runtime.config.clone()).unwrap();
for index in 0..32 {
assert!(runtime.enqueue_retry.schedule(
queue.clone(),
terminal(&format!("retry-{index}")),
"terminal",
DataLayerError::Redis("transient".into())
));
}
runtime.shutdown(Duration::from_secs(3)).await.unwrap();
assert_eq!(runtime.enqueue_retry.pending(), 0);
assert_eq!(runtime.enqueue_retry.recovered_total(), 32);
assert_eq!(
store
.queue
.stats(&runtime.config.stream_key, None)
.await
.unwrap()
.stream_length,
32
);
}
#[tokio::test]
async fn usage_shutdown_failure_retains_retry_work_for_a_later_attempt() {
let runtime = UsageRuntime::new(config()).unwrap();
let store = store();
let flaky = Arc::new(FlakyAppendQueueStore::new(
Arc::clone(&store.queue),
usize::MAX,
));
let queue = UsageQueue::new(flaky.clone(), runtime.config.clone()).unwrap();
assert!(runtime.enqueue_retry.schedule(
queue,
terminal("recover-after-deadline"),
"terminal",
DataLayerError::Redis("unavailable".into())
));
let result = runtime.shutdown(Duration::from_millis(150)).await;
assert!(matches!(result, Err(DataLayerError::TimedOut(_))));
assert_eq!(runtime.enqueue_retry.pending(), 1);
assert!(flaky.append_attempts.load(Ordering::Acquire) <= 4);
flaky.remaining_failures.store(0, Ordering::Release);
runtime.shutdown(Duration::from_secs(3)).await.unwrap();
assert_eq!(runtime.enqueue_retry.recovered_total(), 1);
assert_eq!(
store
.queue
.stats(&runtime.config.stream_key, None)
.await
.unwrap()
.stream_length,
1
);
}
#[tokio::test]
async fn usage_shutdown_flushes_delayed_lifecycle_without_waiting_for_its_timer() {
let mut config = config();
config.lifecycle_enqueue_delay_ms = 60_000;
let runtime = UsageRuntime::new(config).unwrap();
let store = store();
let event = UsageEvent::new(
UsageEventType::Pending,
"delayed",
UsageEventData::default(),
);
runtime
.enqueue_lifecycle_event_with_config_delay(&store, event)
.await;
assert!(runtime.local_work_pending() > 0);
runtime.shutdown(Duration::from_secs(3)).await.unwrap();
assert_eq!(runtime.local_work_pending(), 0);
assert_eq!(
store
.queue
.stats(&runtime.config.stream_key, None)
.await
.unwrap()
.stream_length,
1
);
}
#[tokio::test]
async fn usage_shutdown_stops_every_idle_worker_and_supervisor() {
for supervised in [false, true] {
let runtime = UsageRuntime::new(config()).unwrap();
let store = Arc::new(store());
let handles = if supervised {
vec![runtime.spawn_worker_supervisor(store).unwrap()]
} else {
runtime.spawn_workers(store)
};
sleep(Duration::from_millis(30)).await;
runtime.shutdown(Duration::from_secs(3)).await.unwrap();
for handle in handles {
timeout(Duration::from_secs(1), handle)
.await
.unwrap()
.unwrap();
}
assert_eq!(runtime.metrics_snapshot().worker_active_count, 0);
assert_eq!(runtime.shutdown.supervisors.load(Ordering::Acquire), 0);
}
}
#[tokio::test]
async fn usage_shutdown_does_not_cancel_a_write_or_ack_its_unfinished_record() {
let runtime = UsageRuntime::new(config()).unwrap();
let store = BlockingWriteQueueConfiguredUsageStore {
records: Arc::new(Mutex::new(Vec::new())),
queue: store().queue,
write_started: Arc::new(tokio::sync::Notify::new()),
release_writes: Arc::new(tokio::sync::Notify::new()),
writes_completed: Arc::new(AtomicUsize::new(0)),
};
let queue = UsageQueue::new(Arc::clone(&store.queue), runtime.config.clone()).unwrap();
queue.ensure_consumer_group().await.unwrap();
queue.enqueue(&terminal("in-flight-worker")).await.unwrap();
let worker = runtime.spawn_worker(Arc::new(store.clone())).unwrap();
timeout(Duration::from_secs(3), store.write_started.notified())
.await
.unwrap();
assert!(runtime.shutdown(Duration::from_millis(50)).await.is_err());
assert!(!worker.is_finished());
assert_eq!(
store
.queue
.stats(
&runtime.config.stream_key,
Some(&runtime.config.consumer_group)
)
.await
.unwrap()
.group_pending,
1
);
store.release_writes.notify_one();
runtime.shutdown(Duration::from_secs(3)).await.unwrap();
worker.await.unwrap();
assert_eq!(store.writes_completed.load(Ordering::Acquire), 1);
assert_eq!(
store
.queue
.stats(
&runtime.config.stream_key,
Some(&runtime.config.consumer_group)
)
.await
.unwrap()
.group_pending,
0
);
}
#[tokio::test]
async fn usage_shutdown_disabled_runtime_is_immediate() {
UsageRuntime::disabled()
.shutdown(Duration::from_secs(1))
.await
.unwrap();
}
#[tokio::test]
async fn usage_shutdown_consumes_process_local_queue_before_stopping_workers() {
let runtime = UsageRuntime::new(config()).unwrap();
let store = store();
let worker = runtime
.spawn_worker_supervisor(Arc::new(store.clone()))
.unwrap();
for index in 0..64 {
runtime
.submit_terminal_event(&store, terminal(&format!("local-{index}")))
.await;
}
runtime
.shutdown_with_local_queue(Duration::from_secs(5), Some(Arc::clone(&store.queue)))
.await
.unwrap();
worker.await.unwrap();
assert_eq!(store.records.lock().unwrap().len(), 64);
assert_eq!(
store
.queue
.stats(
&runtime.config.stream_key,
Some(&runtime.config.consumer_group)
)
.await
.unwrap()
.group_pending,
0
);
}
#[tokio::test]
async fn usage_shutdown_rejects_unconsumed_memory_queue_as_success() {
let runtime = UsageRuntime::new(config()).unwrap();
let store = store();
runtime
.submit_terminal_event(&store, terminal("unconsumed"))
.await;
let result = runtime
.shutdown_with_local_queue(Duration::from_millis(50), Some(Arc::clone(&store.queue)))
.await;
assert!(result.is_err());
let worker = runtime.spawn_worker(Arc::new(store.clone())).unwrap();
runtime
.shutdown_with_local_queue(Duration::from_secs(3), Some(Arc::clone(&store.queue)))
.await
.unwrap();
worker.await.unwrap();
assert_eq!(store.records.lock().unwrap().len(), 1);
}
#[tokio::test]
async fn usage_shutdown_does_not_miss_accepted_pending_to_terminal_handoffs() {
let mut config = config();
config.queue_terminal_events = false;
let runtime = UsageRuntime::new(config).unwrap();
let store = store();
for index in 0..32 {
let id = format!("ordered-{index}");
let plan = terminal_test_plan(&id);
runtime.record_pending(&store, build_lifecycle_usage_seed(&plan, None));
runtime.record_stream_started(
&store,
&build_lifecycle_usage_seed(&plan, None),
200,
Some(&ExecutionTelemetry {
ttfb_ms: Some(5),
elapsed_ms: None,
upstream_bytes: None,
}),
);
runtime.submit_terminal_event(&store, terminal(&id)).await;
}
runtime.shutdown(Duration::from_secs(5)).await.unwrap();
let records = store.records.lock().unwrap();
for index in 0..32 {
let statuses: Vec<_> = records
.iter()
.filter(|r| r.request_id == format!("ordered-{index}"))
.map(|r| r.status.as_str())
.collect();
assert_eq!(statuses, ["pending", "streaming", "completed"]);
}
}
+63 -27
View File
@@ -36,8 +36,19 @@ pub async fn reconcile_usage_policy_cost_for_event(
writer: &dyn UsageSettlementWriter,
event: &UsageEvent,
) -> Result<(), DataLayerError> {
reconcile_usage_policy_cost_for_event_with_result(writer, event)
.await
.map(|_| ())
}
pub(crate) struct ReconciledUsagePolicyCost(ReconcileUsagePolicyCostInput);
pub(crate) async fn reconcile_usage_policy_cost_for_event_with_result(
writer: &dyn UsageSettlementWriter,
event: &UsageEvent,
) -> Result<Option<ReconciledUsagePolicyCost>, DataLayerError> {
if !writer.has_usage_settlement_writer() {
return Ok(());
return Ok(None);
}
let terminal_state = match event.event_type {
UsageEventType::Completed => UsagePolicyCostReservationState::Finalized,
@@ -49,16 +60,16 @@ pub async fn reconcile_usage_policy_cost_for_event(
UsageEventType::Failed | UsageEventType::Cancelled => {
UsagePolicyCostReservationState::Released
}
UsageEventType::Pending | UsageEventType::Streaming => return Ok(()),
UsageEventType::Pending | UsageEventType::Streaming => return Ok(None),
};
if plan_usage_reservation_reconciliation_is_deferred(event.data.request_metadata.as_ref()) {
return Ok(());
return Ok(None);
}
let Some(subject_id) = event.data.user_id.as_deref().and_then(non_empty_trimmed) else {
return Ok(());
return Ok(None);
};
let Some(reservation_token) = event_usage_policy_reservation_token(event) else {
return Ok(());
return Ok(None);
};
let actual_cost_units = if terminal_state == UsagePolicyCostReservationState::Finalized {
let actual_cost_usd = event.data.actual_total_cost_usd.ok_or_else(|| {
@@ -75,22 +86,39 @@ pub async fn reconcile_usage_policy_cost_for_event(
0
};
let _ = writer
.reconcile_usage_policy_cost(ReconcileUsagePolicyCostInput {
request_id: event.request_id.clone(),
subject_id: subject_id.to_string(),
reservation_token: reservation_token.to_string(),
actual_cost_units,
terminal_state,
finalized_at_unix_secs: event.timestamp_ms / 1_000,
})
.await?;
Ok(())
let input = ReconcileUsagePolicyCostInput {
request_id: event.request_id.clone(),
subject_id: subject_id.to_string(),
reservation_token: reservation_token.to_string(),
actual_cost_units,
terminal_state,
finalized_at_unix_secs: event.timestamp_ms / 1_000,
};
let stored = writer.reconcile_usage_policy_cost(input.clone()).await?;
// A successful call alone is insufficient: None or a different existing terminal
// reservation must not suppress reconciliation of the subsequent stored usage row.
let matches = stored.is_some_and(|stored| {
stored.request_id == input.request_id
&& stored.subject_id == input.subject_id
&& stored.reservation_token == input.reservation_token
&& stored.actual_cost_units == Some(input.actual_cost_units)
&& stored.state == input.terminal_state
&& stored.finalized_at_unix_secs == Some(input.finalized_at_unix_secs)
});
Ok(matches.then_some(ReconciledUsagePolicyCost(input)))
}
pub async fn settle_usage_if_needed(
writer: &dyn UsageSettlementWriter,
usage: &StoredRequestUsageAudit,
) -> Result<(), DataLayerError> {
settle_usage_with_reconciled_cost(writer, usage, None).await
}
pub(crate) async fn settle_usage_with_reconciled_cost(
writer: &dyn UsageSettlementWriter,
usage: &StoredRequestUsageAudit,
reconciled: Option<ReconciledUsagePolicyCost>,
) -> Result<(), DataLayerError> {
if !writer.has_usage_settlement_writer() {
return Ok(());
@@ -132,17 +160,21 @@ pub async fn settle_usage_if_needed(
} else {
(UsagePolicyCostReservationState::Released, 0)
};
let _ = writer
.reconcile_usage_policy_cost(ReconcileUsagePolicyCostInput {
request_id: usage.request_id.clone(),
subject_id: subject_id.to_string(),
reservation_token: reservation_token.to_string(),
actual_cost_units,
terminal_state,
finalized_at_unix_secs: finalized_at_unix_secs
.unwrap_or(usage.updated_at_unix_secs),
})
.await?;
let input = ReconcileUsagePolicyCostInput {
request_id: usage.request_id.clone(),
subject_id: subject_id.to_string(),
reservation_token: reservation_token.to_string(),
actual_cost_units,
terminal_state,
finalized_at_unix_secs: finalized_at_unix_secs
.unwrap_or(usage.updated_at_unix_secs),
};
if !reconciled
.as_ref()
.is_some_and(|previous| previous.0 == input)
{
let _ = writer.reconcile_usage_policy_cost(input).await?;
}
}
}
@@ -243,6 +275,10 @@ fn finite_cost(value: f64) -> Result<f64, DataLayerError> {
#[cfg(test)]
mod tests {
mod reconciliation_reuse {
include!("settlement_reuse_tests.rs");
}
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;
use std::time::Duration;
@@ -0,0 +1,517 @@
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use aether_data::repository::settlement::InMemorySettlementRepository;
use aether_data_contracts::repository::settlement::{
ReconcileUsagePolicyCostInput, StoredUsagePolicyCostReservation, StoredUsageSettlement,
UsagePolicyCostReservationState,
};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::DataLayerError;
use aether_runtime_state::RuntimeQueueStore;
use async_trait::async_trait;
use serde_json::json;
use super::sample_usage as base_usage;
use crate::settlement::{UsageSettlementInput, UsageSettlementWriter};
use crate::worker::write_event_record;
use crate::{
UsageBillingEventEnricher, UsageEvent, UsageEventData, UsageEventType, UsageRecordWriter,
UsageRuntime, UsageRuntimeAccess, UsageRuntimeConfig,
};
const RESERVATION_TOKEN: &str = "550e8400-e29b-41d4-a716-446655440000";
fn sample_usage() -> StoredRequestUsageAudit {
let mut usage = base_usage();
usage.request_metadata = Some(json!({"plan_usage_reservation_token": RESERVATION_TOKEN}));
usage
}
fn runtime() -> UsageRuntime {
UsageRuntime::new(UsageRuntimeConfig {
enabled: true,
..Default::default()
})
.unwrap()
}
#[derive(Default)]
enum ReconcileResponse {
#[default]
Exact,
Missing,
Changed(fn(&mut StoredUsagePolicyCostReservation)),
Error,
}
#[derive(Default)]
struct ReuseStore {
repository: Option<InMemorySettlementRepository>,
response: ReconcileResponse,
stored_override: Option<StoredRequestUsageAudit>,
fail_next_upsert: AtomicBool,
upserts: AtomicUsize,
reconciliations: Mutex<Vec<ReconcileUsagePolicyCostInput>>,
settlements: Mutex<Vec<UsageSettlementInput>>,
}
#[async_trait]
impl UsageSettlementWriter for ReuseStore {
fn has_usage_settlement_writer(&self) -> bool {
true
}
async fn reconcile_usage_policy_cost(
&self,
input: ReconcileUsagePolicyCostInput,
) -> Result<Option<StoredUsagePolicyCostReservation>, DataLayerError> {
input.validate()?;
self.reconciliations.lock().unwrap().push(input.clone());
tokio::task::yield_now().await;
if let Some(repository) = self.repository.as_ref() {
return aether_data_contracts::repository::settlement::SettlementWriteRepository::reconcile_usage_policy_cost(repository, input).await;
}
let mut stored = StoredUsagePolicyCostReservation {
request_id: input.request_id,
subject_id: input.subject_id,
reservation_token: input.reservation_token,
admitted_at_unix_secs: 100,
reserved_cost_units: 100_000_000,
actual_cost_units: Some(input.actual_cost_units),
state: input.terminal_state,
reservation_expires_at_unix_secs: 500,
retain_until_unix_secs: 1_000,
finalized_at_unix_secs: Some(input.finalized_at_unix_secs),
};
match self.response {
ReconcileResponse::Exact => {}
ReconcileResponse::Missing => return Ok(None),
ReconcileResponse::Changed(change) => change(&mut stored),
ReconcileResponse::Error => {
return Err(DataLayerError::TimedOut("reconciliation".to_string()));
}
}
Ok(Some(stored))
}
async fn settle_usage(
&self,
input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
self.settlements.lock().unwrap().push(input.clone());
if let Some(repository) = self.repository.as_ref() {
return aether_data_contracts::repository::settlement::SettlementWriteRepository::settle_usage(repository, input).await;
}
Ok(None)
}
}
#[async_trait]
impl UsageRecordWriter for ReuseStore {
async fn upsert_usage_record(
&self,
record: UpsertUsageRecord,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
self.upserts.fetch_add(1, Ordering::Relaxed);
if self.fail_next_upsert.swap(false, Ordering::Relaxed) {
return Err(DataLayerError::TimedOut("upsert".to_string()));
}
if let Some(stored) = self.stored_override.as_ref() {
return Ok(Some(stored.clone()));
}
let mut stored = sample_usage();
stored.request_id = record.request_id;
stored.user_id = record.user_id;
stored.api_key_id = record.api_key_id;
stored.provider_id = record.provider_id;
stored.status = record.status;
stored.billing_status = record.billing_status;
stored.total_cost_usd = record.total_cost_usd.unwrap_or_default();
stored.actual_total_cost_usd = record.actual_total_cost_usd.unwrap_or_default();
stored.request_metadata = record.request_metadata;
stored.updated_at_unix_secs = record.updated_at_unix_secs;
stored.finalized_at_unix_secs = record.finalized_at_unix_secs;
Ok(Some(stored))
}
}
#[async_trait]
impl UsageBillingEventEnricher for ReuseStore {
async fn enrich_usage_event(&self, _event: &mut UsageEvent) -> Result<(), DataLayerError> {
Ok(())
}
}
impl UsageRuntimeAccess for ReuseStore {
fn has_usage_writer(&self) -> bool {
true
}
fn has_usage_worker_queue(&self) -> bool {
false
}
fn usage_worker_queue(&self) -> Option<Arc<dyn RuntimeQueueStore>> {
None
}
}
fn event() -> UsageEvent {
let mut event = UsageEvent::new(
UsageEventType::Completed,
"req-1",
UsageEventData {
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
provider_name: "openai".to_string(),
model: "gpt-5".to_string(),
total_cost_usd: Some(1.25),
actual_total_cost_usd: Some(0.75),
request_metadata: Some(json!({"plan_usage_reservation_token": RESERVATION_TOKEN})),
..Default::default()
},
);
event.timestamp_ms = 200_999;
event
}
async fn write(store: &ReuseStore, event: UsageEvent, direct: bool) {
if direct {
runtime().record_terminal_event_direct(store, event).await;
} else {
write_event_record(store, &event).await.unwrap();
}
}
#[tokio::test]
async fn worker_and_direct_writes_reuse_confirmed_reservation_and_still_settle_wallet() {
for direct in [false, true] {
let store = ReuseStore::default();
write(&store, event(), direct).await;
assert_eq!(store.upserts.load(Ordering::Relaxed), 1);
let reconciliations = store.reconciliations.lock().unwrap();
assert_eq!(reconciliations.len(), 1, "direct={direct}");
assert_eq!(reconciliations[0].actual_cost_units, 75_000_000);
assert_eq!(reconciliations[0].finalized_at_unix_secs, 200);
let settlements = store.settlements.lock().unwrap();
assert_eq!(settlements.len(), 1);
assert_eq!(settlements[0].request_id, "req-1");
assert_eq!(settlements[0].actual_total_cost_usd, 0.75);
}
}
#[tokio::test]
async fn missing_or_different_reconciliation_results_keep_stored_usage_reconciliation() {
let changes: [fn(&mut StoredUsagePolicyCostReservation); 9] = [
|row| row.request_id = "other-request".to_string(),
|row| row.subject_id = "other-user".to_string(),
|row| row.reservation_token = "other-token".to_string(),
|row| row.actual_cost_units = Some(1),
|row| row.actual_cost_units = None,
|row| row.state = UsagePolicyCostReservationState::Reserved,
|row| row.state = UsagePolicyCostReservationState::Released,
|row| row.finalized_at_unix_secs = Some(199),
|row| row.finalized_at_unix_secs = None,
];
for direct in [false, true] {
for response in std::iter::once(ReconcileResponse::Missing)
.chain(changes.into_iter().map(ReconcileResponse::Changed))
{
let store = ReuseStore {
response,
..Default::default()
};
write(&store, event(), direct).await;
assert_eq!(store.reconciliations.lock().unwrap().len(), 2);
assert_eq!(store.settlements.lock().unwrap().len(), 1);
}
}
}
#[tokio::test]
async fn changed_stored_usage_is_reconciled_using_its_own_identity_cost_and_terminal_state() {
let changes: [fn(&mut StoredRequestUsageAudit); 6] = [
|row| row.request_id = "other-request".to_string(),
|row| row.user_id = Some("other-user".to_string()),
|row| {
row.request_metadata.as_mut().unwrap()["plan_usage_reservation_token"] =
json!("other-token")
},
|row| row.actual_total_cost_usd = 0.25,
|row| row.status = "failed".to_string(),
|row| row.finalized_at_unix_secs = Some(199),
];
for direct in [false, true] {
for change in changes {
let mut stored = sample_usage();
change(&mut stored);
let store = ReuseStore {
stored_override: Some(stored.clone()),
..Default::default()
};
write(&store, event(), direct).await;
let reconciliations = store.reconciliations.lock().unwrap();
assert_eq!(reconciliations.len(), 2);
assert_eq!(reconciliations[1].request_id, stored.request_id);
assert_eq!(
reconciliations[1].subject_id,
stored.user_id.as_ref().unwrap().as_str()
);
assert_eq!(
reconciliations[1].reservation_token,
stored.request_metadata.as_ref().unwrap()["plan_usage_reservation_token"]
.as_str()
.unwrap()
);
assert_ne!(reconciliations[0], reconciliations[1]);
let settlements = store.settlements.lock().unwrap();
assert_eq!(settlements.len(), 1);
assert_eq!(
settlements[0].actual_total_cost_usd,
stored.actual_total_cost_usd
);
assert_eq!(settlements[0].status, stored.status);
}
}
}
#[tokio::test]
async fn cancellation_release_billable_cancellation_and_zero_cost_preserve_settlement_rules() {
for direct in [false, true] {
for (event_type, billable_cancel, cost, terminal_state, wallets) in [
(
UsageEventType::Failed,
false,
0.0,
UsagePolicyCostReservationState::Released,
0,
),
(
UsageEventType::Cancelled,
false,
0.75,
UsagePolicyCostReservationState::Released,
0,
),
(
UsageEventType::Cancelled,
true,
0.75,
UsagePolicyCostReservationState::Finalized,
1,
),
(
UsageEventType::Completed,
false,
0.0,
UsagePolicyCostReservationState::Finalized,
1,
),
] {
let store = ReuseStore::default();
let mut event = event();
event.event_type = event_type;
event.data.actual_total_cost_usd = Some(cost);
event.data.request_metadata.as_mut().unwrap()["cancelled_request_fee"] =
json!(billable_cancel);
write(&store, event, direct).await;
let reconciliations = store.reconciliations.lock().unwrap();
assert_eq!(reconciliations.len(), 1);
assert_eq!(reconciliations[0].terminal_state, terminal_state);
assert_eq!(
reconciliations[0].actual_cost_units,
if billable_cancel { 75_000_000 } else { 0 }
);
assert_eq!(store.settlements.lock().unwrap().len(), wallets);
}
}
}
#[tokio::test]
async fn reconciliation_failure_stops_both_writes_before_upsert_and_wallet_settlement() {
for direct in [false, true] {
let store = ReuseStore {
response: ReconcileResponse::Error,
..Default::default()
};
if direct {
write(&store, event(), true).await;
} else {
assert!(write_event_record(&store, &event()).await.is_err());
}
assert_eq!(store.reconciliations.lock().unwrap().len(), 1);
assert_eq!(store.upserts.load(Ordering::Relaxed), 0);
assert!(store.settlements.lock().unwrap().is_empty());
}
}
#[tokio::test]
async fn retry_after_upsert_failure_reconciles_again_before_settling() {
for direct in [false, true] {
let store = ReuseStore {
fail_next_upsert: AtomicBool::new(true),
..Default::default()
};
if direct {
write(&store, event(), true).await;
} else {
assert!(write_event_record(&store, &event()).await.is_err());
}
assert!(store.settlements.lock().unwrap().is_empty());
write(&store, event(), direct).await;
assert_eq!(store.reconciliations.lock().unwrap().len(), 2);
assert_eq!(store.upserts.load(Ordering::Relaxed), 2);
assert_eq!(store.settlements.lock().unwrap().len(), 1);
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_same_user_worker_and_direct_writes_each_reconcile_once() {
const REQUESTS: usize = 256;
let store = Arc::new(ReuseStore::default());
let runtime = runtime();
let barrier = Arc::new(tokio::sync::Barrier::new(REQUESTS));
let mut tasks = tokio::task::JoinSet::new();
for index in 0..REQUESTS {
let store = store.clone();
let runtime = runtime.clone();
let barrier = barrier.clone();
tasks.spawn(async move {
let mut event = event();
event.request_id = format!("reuse-concurrent-{index}");
event.data.request_metadata.as_mut().unwrap()["plan_usage_reservation_token"] =
json!(format!("550e8400-e29b-41d4-a716-{index:012x}"));
barrier.wait().await;
if index % 2 == 0 {
runtime
.record_terminal_event_direct(store.as_ref(), event)
.await;
} else {
write_event_record(store.as_ref(), &event).await.unwrap();
}
});
}
tokio::time::timeout(std::time::Duration::from_secs(10), async {
while let Some(result) = tasks.join_next().await {
result.unwrap();
}
})
.await
.unwrap();
let reconciliations = store.reconciliations.lock().unwrap();
assert_eq!(reconciliations.len(), REQUESTS);
let unique_tokens: std::collections::HashSet<_> = reconciliations
.iter()
.map(|input| &input.reservation_token)
.collect();
assert_eq!(unique_tokens.len(), REQUESTS);
assert_eq!(store.upserts.load(Ordering::Relaxed), REQUESTS);
let settlements = store.settlements.lock().unwrap();
assert_eq!(settlements.len(), REQUESTS);
let unique_requests: std::collections::HashSet<_> =
settlements.iter().map(|input| &input.request_id).collect();
assert_eq!(unique_requests.len(), REQUESTS);
assert_eq!(
settlements
.iter()
.map(|input| input.actual_total_cost_usd)
.sum::<f64>(),
REQUESTS as f64 * 0.75
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_duplicate_delivery_debits_real_memory_wallet_only_once() {
use aether_data::repository::wallet::{
InMemoryWalletRepository, StoredWalletSnapshot, WalletLookupKey, WalletReadRepository,
};
use aether_data_contracts::repository::settlement::{
ReserveUsagePolicyCostInput, ReserveUsagePolicyCostOutcome, SettlementWriteRepository,
UsagePolicyCostWindow,
};
let wallet = StoredWalletSnapshot::new(
"wallet-1".to_string(),
Some("user-1".to_string()),
None,
10.0,
2.0,
"finite".to_string(),
"USD".to_string(),
"active".to_string(),
0.0,
0.0,
0.0,
0.0,
100,
)
.unwrap();
let wallets = Arc::new(InMemoryWalletRepository::seed([wallet]));
let repository = InMemorySettlementRepository::from_wallet_repository(wallets.clone());
let reservation = ReserveUsagePolicyCostInput {
request_id: "req-1".to_string(),
subject_id: "user-1".to_string(),
reservation_token: RESERVATION_TOKEN.to_string(),
admitted_at_unix_secs: 100,
reserved_cost_units: 100_000_000,
reservation_expires_at_unix_secs: 500,
retain_until_unix_secs: 1_000,
windows: vec![UsagePolicyCostWindow {
window_id: "window-1".to_string(),
starts_at_unix_secs: 0,
ends_at_unix_secs: 1_000,
limit_cost_units: 1_000_000_000,
}],
};
assert!(matches!(
repository
.reserve_usage_policy_cost(reservation.clone())
.await
.unwrap(),
ReserveUsagePolicyCostOutcome::Allowed { .. }
));
let store = Arc::new(ReuseStore {
repository: Some(repository),
..Default::default()
});
let runtime = runtime();
let mut tasks = tokio::task::JoinSet::new();
for index in 0..32 {
let store = store.clone();
let runtime = runtime.clone();
tasks.spawn(async move {
if index % 2 == 0 {
write_event_record(store.as_ref(), &event()).await.unwrap();
} else {
runtime
.record_terminal_event_direct(store.as_ref(), event())
.await;
}
});
}
while let Some(result) = tasks.join_next().await {
result.unwrap();
}
assert_eq!(store.reconciliations.lock().unwrap().len(), 32);
assert_eq!(store.settlements.lock().unwrap().len(), 32);
let wallet = wallets
.find(WalletLookupKey::UserId("user-1"))
.await
.unwrap()
.unwrap();
assert_eq!(wallet.balance + wallet.gift_balance, 11.25);
assert_eq!(wallet.total_consumed, 0.75);
assert!(matches!(
store
.repository
.as_ref()
.unwrap()
.reserve_usage_policy_cost(reservation)
.await
.unwrap(),
ReserveUsagePolicyCostOutcome::AlreadyTerminal {
state: UsagePolicyCostReservationState::Finalized
}
));
}
@@ -0,0 +1,99 @@
use std::future::Future;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::watch;
use tokio::task::JoinHandle;
#[derive(Debug, Default)]
pub(crate) struct UsageBackgroundTasks {
handles: Mutex<Vec<JoinHandle<()>>>,
}
impl UsageBackgroundTasks {
pub(crate) fn spawn(&self, task: impl Future<Output = ()> + Send + 'static) {
self.handles
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.push(crate::executor::spawn_on_usage_background_runtime(task));
}
pub(crate) async fn stop_idle(&self) {
let handles = std::mem::take(
&mut *self
.handles
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner()),
);
for handle in &handles {
handle.abort();
}
for handle in handles {
let _ = handle.await;
}
}
}
impl Drop for UsageBackgroundTasks {
fn drop(&mut self) {
for handle in self.handles.get_mut().unwrap_or_else(|p| p.into_inner()) {
handle.abort();
}
}
}
#[derive(Debug)]
pub(crate) struct UsageShutdownState {
pub(crate) producers: Arc<AtomicUsize>,
pub(crate) drain: watch::Sender<bool>,
pub(crate) tasks: UsageBackgroundTasks,
pub(crate) worker_control: crate::worker::UsageWorkerControl,
pub(crate) supervisors: Arc<AtomicUsize>,
pub(crate) lock: tokio::sync::Mutex<()>,
}
impl Default for UsageShutdownState {
fn default() -> Self {
Self {
producers: Arc::new(AtomicUsize::new(0)),
drain: watch::channel(false).0,
tasks: UsageBackgroundTasks::default(),
worker_control: crate::worker::UsageWorkerControl::default(),
supervisors: Arc::new(AtomicUsize::new(0)),
lock: tokio::sync::Mutex::new(()),
}
}
}
/// Retain across an owned request finalizer, including any detached handoff.
#[derive(Debug)]
pub struct UsageProducerGuard(pub(crate) Arc<AtomicUsize>);
impl Drop for UsageProducerGuard {
fn drop(&mut self) {
self.0.fetch_sub(1, Ordering::AcqRel);
}
}
pub(crate) async fn wait_for_drain(signal: &mut watch::Receiver<bool>) {
loop {
if *signal.borrow_and_update() {
return;
}
if signal.changed().await.is_err() {
std::future::pending::<()>().await;
}
}
}
pub(crate) async fn retry_delay(delay: Duration, signal: &mut watch::Receiver<bool>) {
if *signal.borrow_and_update() {
tokio::time::sleep(delay.min(Duration::from_millis(100))).await;
} else {
tokio::select! {
_ = tokio::time::sleep(delay) => {},
_ = wait_for_drain(signal) => {},
}
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,449 @@
use super::*;
use crate::dead_letter_encoding::DeadLetterEncodingBudget;
use crate::worker::UsageWorkerObservation;
use aether_runtime_state::RuntimeQueueTransferOutcome;
struct TransferProbe {
inner: RuntimeState,
native: bool,
fail_write: AtomicBool,
lose_reply: AtomicBool,
ack_calls: Mutex<Vec<Vec<String>>>,
append_calls: AtomicUsize,
}
impl TransferProbe {
fn new(native: bool) -> Self {
Self {
inner: RuntimeState::memory(MemoryRuntimeStateConfig::default()),
native,
fail_write: AtomicBool::new(false),
lose_reply: AtomicBool::new(false),
ack_calls: Mutex::new(Vec::new()),
append_calls: AtomicUsize::new(0),
}
}
}
#[async_trait]
impl RuntimeQueueStore for TransferProbe {
async fn ensure_consumer_group(
&self,
stream: &str,
group: &str,
start_id: &str,
) -> Result<(), DataLayerError> {
self.inner
.ensure_consumer_group(stream, group, start_id)
.await
}
async fn append_fields_with_maxlen(
&self,
stream: &str,
fields: &BTreeMap<String, String>,
maxlen: Option<usize>,
) -> Result<String, DataLayerError> {
self.append_calls.fetch_add(1, Ordering::Relaxed);
if self.fail_write.load(Ordering::Acquire) {
return Err(DataLayerError::TimedOut("test append failure".to_string()));
}
self.inner
.append_fields_with_maxlen(stream, fields, maxlen)
.await
}
async fn try_transfer_pending_to_stream(
&self,
source: &str,
group: &str,
entry_id: &str,
destination: &str,
fields: &BTreeMap<String, String>,
) -> Result<Option<RuntimeQueueTransferOutcome>, DataLayerError> {
if !self.native {
return Ok(None);
}
if self.fail_write.load(Ordering::Acquire) {
return Err(DataLayerError::TimedOut(
"test atomic transfer failure".to_string(),
));
}
let result = self
.inner
.try_transfer_pending_to_stream(source, group, entry_id, destination, fields)
.await?;
if self.lose_reply.swap(false, Ordering::AcqRel) {
return Err(DataLayerError::TimedOut(
"test committed transfer reply lost".to_string(),
));
}
Ok(result)
}
async fn read_group(
&self,
stream: &str,
group: &str,
consumer: &str,
count: usize,
block_ms: Option<u64>,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
self.inner
.read_group(stream, group, consumer, count, block_ms)
.await
}
async fn claim_stale(
&self,
stream: &str,
group: &str,
consumer: &str,
start_id: &str,
config: RuntimeQueueReclaimConfig,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
self.inner
.claim_stale(stream, group, consumer, start_id, config)
.await
}
async fn ack(
&self,
stream: &str,
group: &str,
ids: &[String],
) -> Result<usize, DataLayerError> {
self.ack_calls.lock().expect("ack calls").push(ids.to_vec());
self.inner.ack(stream, group, ids).await
}
async fn delete(&self, stream: &str, ids: &[String]) -> Result<usize, DataLayerError> {
self.inner.delete(stream, ids).await
}
async fn stats(
&self,
stream: &str,
group: Option<&str>,
) -> Result<RuntimeQueueStats, DataLayerError> {
self.inner.stats(stream, group).await
}
}
async fn transfer_worker(
native: bool,
) -> (
Arc<TransferProbe>,
UsageQueueWorker,
Arc<SelectiveFailingRecorder>,
tokio::sync::mpsc::Receiver<UsageWorkerObservation>,
) {
let runner = Arc::new(TransferProbe::new(native));
let recorder = Arc::new(SelectiveFailingRecorder::default());
let (telemetry, observations) = tokio::sync::mpsc::channel(32);
let mut worker = UsageQueueWorker::new(
runner.clone(),
recorder.clone(),
UsageRuntimeConfig {
enabled: true,
consumer_batch_size: 10,
consumer_block_ms: 1,
..UsageRuntimeConfig::default()
},
None,
)
.expect("worker")
.with_supervisor(UsageWorkerControl::default(), telemetry);
worker.queue =
worker
.queue
.with_dead_letter_encoding_budget(Arc::new(DeadLetterEncodingBudget::new(
64 * 1024 * 1024,
4,
)));
worker.queue.ensure_consumer_group().await.expect("group");
(runner, worker, recorder, observations)
}
async fn malformed_entries(
runner: &TransferProbe,
worker: &UsageQueueWorker,
) -> Vec<RuntimeQueueEntry> {
runner
.inner
.append_fields_with_maxlen(
&worker.config.stream_key,
&BTreeMap::from([
("payload".to_string(), "malformed\u{0000}\n\"\\".to_string()),
(
"legacy".to_string(),
"preserve all original fields".to_string(),
),
]),
None,
)
.await
.expect("raw append");
worker
.queue
.read_group(&worker.consumer)
.await
.expect("read")
}
fn observed_totals(
observations: &mut tokio::sync::mpsc::Receiver<UsageWorkerObservation>,
) -> (usize, usize) {
let mut totals = (0, 0);
while let Ok(observation) = observations.try_recv() {
totals.0 += observation.acked_entries;
totals.1 += observation.dead_lettered_entries;
}
totals
}
#[tokio::test]
async fn native_transfer_replay_does_not_append_or_ack_twice() {
let (runner, worker, recorder, mut observations) = transfer_worker(true).await;
let entries = malformed_entries(&runner, &worker).await;
worker
.process_entries(entries.clone())
.await
.expect("transfer");
worker.process_entries(entries).await.expect("stale replay");
assert_eq!(worker.queue.stats().await.expect("stats").group_pending, 0);
assert_eq!(
worker.queue.dlq_stats().await.expect("dlq").stream_length,
1
);
assert_eq!(observed_totals(&mut observations), (1, 1));
assert!(runner.ack_calls.lock().expect("acks").is_empty());
assert_eq!(runner.append_calls.load(Ordering::Relaxed), 0);
assert!(recorder.calls.lock().expect("record calls").is_empty());
}
#[tokio::test]
async fn committed_transfer_lost_reply_retries_without_duplicate_or_false_metrics() {
let (runner, worker, _, mut observations) = transfer_worker(true).await;
let entries = malformed_entries(&runner, &worker).await;
runner.lose_reply.store(true, Ordering::Release);
assert!(matches!(
worker.process_entries(entries.clone()).await,
Err(DataLayerError::TimedOut(_))
));
worker.process_entries(entries).await.expect("retry");
assert_eq!(
worker.queue.dlq_stats().await.expect("dlq").stream_length,
1
);
assert_eq!(worker.queue.stats().await.expect("stats").stream_length, 0);
assert_eq!(observed_totals(&mut observations), (0, 0));
assert!(runner.ack_calls.lock().expect("acks").is_empty());
assert_eq!(runner.append_calls.load(Ordering::Relaxed), 0);
}
#[tokio::test]
async fn source_no_longer_pending_is_not_reported_as_archived_or_deleted() {
let (runner, worker, _, mut observations) = transfer_worker(true).await;
let entries = malformed_entries(&runner, &worker).await;
runner
.inner
.ack(
&worker.config.stream_key,
&worker.config.consumer_group,
&[entries[0].id.clone()],
)
.await
.expect("external ack");
worker
.process_entries(entries)
.await
.expect("no longer pending");
assert_eq!(worker.queue.stats().await.expect("stats").stream_length, 1);
assert_eq!(
worker.queue.dlq_stats().await.expect("dlq").stream_length,
0
);
assert_eq!(observed_totals(&mut observations), (0, 0));
assert!(runner.ack_calls.lock().expect("acks").is_empty());
}
#[tokio::test]
async fn failed_transfer_acknowledges_successful_prefix_and_preserves_suffix_for_retry() {
let (runner, worker, recorder, mut observations) = transfer_worker(true).await;
for request_id in ["prefix", "req-worker-poison", "suffix"] {
let mut event = sample_event();
event.request_id = request_id.to_string();
worker.queue.enqueue(&event).await.expect("enqueue");
}
let entries = worker
.queue
.read_group(&worker.consumer)
.await
.expect("read batch");
let retry = entries[1..].to_vec();
let prefix_id = entries[0].id.clone();
runner.fail_write.store(true, Ordering::Release);
assert!(worker.process_entries(entries).await.is_err());
assert_eq!(
recorder.calls.lock().expect("calls").as_slice(),
["prefix", "req-worker-poison"]
);
assert_eq!(
*runner.ack_calls.lock().expect("acks"),
vec![vec![prefix_id]]
);
assert_eq!(worker.queue.stats().await.expect("stats").group_pending, 2);
assert_eq!(
worker.queue.dlq_stats().await.expect("dlq").stream_length,
0
);
assert_eq!(observed_totals(&mut observations), (1, 0));
runner.fail_write.store(false, Ordering::Release);
worker.process_entries(retry).await.expect("retry suffix");
assert_eq!(worker.queue.stats().await.expect("stats").stream_length, 0);
assert_eq!(
worker.queue.dlq_stats().await.expect("dlq").stream_length,
1
);
assert_eq!(observed_totals(&mut observations), (2, 1));
assert_eq!(runner.append_calls.load(Ordering::Relaxed), 3);
}
#[tokio::test]
async fn legacy_transfer_fallback_only_acknowledges_after_append_succeeds() {
let (runner, worker, _, mut observations) = transfer_worker(false).await;
let entries = malformed_entries(&runner, &worker).await;
runner.fail_write.store(true, Ordering::Release);
assert!(worker.process_entries(entries.clone()).await.is_err());
assert!(runner.ack_calls.lock().expect("acks").is_empty());
assert_eq!(worker.queue.stats().await.expect("stats").group_pending, 1);
assert_eq!(observed_totals(&mut observations), (0, 0));
runner.fail_write.store(false, Ordering::Release);
worker.process_entries(entries).await.expect("legacy retry");
assert_eq!(
worker.queue.dlq_stats().await.expect("dlq").stream_length,
1
);
assert_eq!(worker.queue.stats().await.expect("stats").stream_length, 0);
assert_eq!(runner.ack_calls.lock().expect("acks").len(), 1);
assert_eq!(observed_totals(&mut observations), (1, 1));
}
#[tokio::test]
async fn encoding_rejection_preserves_pending_original_until_budget_allows_retry() {
let (runner, mut worker, _, mut observations) = transfer_worker(true).await;
let small = Arc::new(DeadLetterEncodingBudget::new(1, 1));
worker.queue = worker.queue.with_dead_letter_encoding_budget(small.clone());
let entries = malformed_entries(&runner, &worker).await;
let original = entries[0].fields.clone();
assert!(worker.process_entries(entries.clone()).await.is_err());
assert_eq!(small.snapshot().reserved_bytes, 0);
assert_eq!(small.snapshot().active_jobs, 0);
assert_eq!(worker.queue.stats().await.expect("stats").group_pending, 1);
assert_eq!(
worker.queue.dlq_stats().await.expect("dlq").stream_length,
0
);
assert_eq!(observed_totals(&mut observations), (0, 0));
worker.queue = worker
.queue
.with_dead_letter_encoding_budget(Arc::new(DeadLetterEncodingBudget::new(64 * 1024, 1)));
worker
.process_entries(entries)
.await
.expect("retry with capacity");
runner
.inner
.ensure_consumer_group(&worker.config.dlq_stream_key, "inspect", "0-0")
.await
.expect("dlq group");
let dlq = runner
.inner
.read_group(
&worker.config.dlq_stream_key,
"inspect",
"inspector",
1,
Some(1),
)
.await
.expect("dlq read");
let payload: serde_json::Value =
serde_json::from_str(&dlq[0].fields["payload"]).expect("wire JSON");
assert_eq!(
payload["fields"],
serde_json::to_value(original).expect("original JSON")
);
assert_eq!(observed_totals(&mut observations), (1, 1));
}
#[tokio::test]
async fn normal_record_replay_reports_actual_ack_count() {
let (_, worker, _, mut observations) = transfer_worker(true).await;
worker
.queue
.enqueue(&sample_event())
.await
.expect("enqueue");
let entries = worker
.queue
.read_group(&worker.consumer)
.await
.expect("read");
worker
.process_entries(entries.clone())
.await
.expect("first record");
worker.process_entries(entries).await.expect("replay");
assert_eq!(observed_totals(&mut observations), (1, 0));
}
#[tokio::test]
async fn oversized_dead_letter_does_not_block_healthy_entries_in_the_same_batch() {
let (runner, mut worker, recorder, mut observations) = transfer_worker(true).await;
let bad = malformed_entries(&runner, &worker).await;
for request_id in ["healthy-first", "healthy-second"] {
let mut event = sample_event();
event.request_id = request_id.to_string();
worker.queue.enqueue(&event).await.expect("healthy enqueue");
}
let mut entries = bad.clone();
entries.extend(
worker
.queue
.read_group(&worker.consumer)
.await
.expect("healthy read"),
);
worker.queue = worker
.queue
.with_dead_letter_encoding_budget(Arc::new(DeadLetterEncodingBudget::new(1, 1)));
assert!(worker.process_entries(entries).await.is_err());
assert_eq!(
recorder.calls.lock().expect("calls").as_slice(),
["healthy-first", "healthy-second"]
);
let stats = worker.queue.stats().await.expect("stats");
assert_eq!((stats.group_pending, stats.stream_length), (1, 1));
assert_eq!(
worker.queue.dlq_stats().await.expect("dlq").stream_length,
0
);
assert_eq!(observed_totals(&mut observations), (2, 0));
assert!(worker.process_entries(bad.clone()).await.is_err());
assert_eq!(observed_totals(&mut observations), (0, 0));
worker.queue = worker
.queue
.with_dead_letter_encoding_budget(Arc::new(DeadLetterEncodingBudget::new(64 * 1024, 1)));
worker
.process_entries(bad)
.await
.expect("archive after raising budget");
assert_eq!(worker.queue.stats().await.expect("stats").group_pending, 0);
assert_eq!(
worker.queue.dlq_stats().await.expect("dlq").stream_length,
1
);
assert_eq!(observed_totals(&mut observations), (1, 1));
}
+3
View File
@@ -581,6 +581,7 @@ fn build_lifecycle_usage_event_from_record(
execution_path: record.execution_path,
local_execution_runtime_miss_reason: record.local_execution_runtime_miss_reason,
request_metadata: record.request_metadata,
capture_retention: record.capture_retention,
..UsageEventData::default()
},
}
@@ -1534,6 +1535,7 @@ fn build_lifecycle_usage_record_owned(
};
Ok(UpsertUsageRecord {
capture_retention: Default::default(),
request_id,
user_id,
api_key_id,
@@ -1631,6 +1633,7 @@ fn build_lifecycle_usage_record_impl(
};
Ok(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: seed.request_id.clone(),
user_id: seed.user_id.clone(),
api_key_id: seed.api_key_id.clone(),