mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 17:37:47 +08:00
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:
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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::{
|
||||
|
||||
@@ -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()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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"]);
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
@@ -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(),
|
||||
|
||||
Reference in New Issue
Block a user