chore: update gateway pressure observability

crates/aether-usage-runtime/src/settlement.rs

@@ -1,3 +1,7 @@
+use std::collections::hash_map::DefaultHasher;
+use std::hash::{Hash, Hasher};
+use std::sync::OnceLock;
+
 use aether_data_contracts::repository::settlement::{StoredUsageSettlement, UsageSettlementInput};
 use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
 use aether_data_contracts::{DataLayerError, DataLayerError::InvalidInput};
@@ -39,10 +43,50 @@ pub async fn settle_usage_if_needed(
         actual_total_cost_usd: finite_cost(usage.actual_total_cost_usd)?,
         finalized_at_unix_secs,
     };
+    let settlement_key = usage_settlement_lock_key(&input);
+    let _guard = usage_settlement_lock(&settlement_key).lock().await;
     let _ = writer.settle_usage(input).await?;
     Ok(())
 }
 
+const USAGE_SETTLEMENT_LOCK_SHARDS: usize = 4096;
+
+fn usage_settlement_lock(key: &str) -> &'static tokio::sync::Mutex<()> {
+    static LOCKS: OnceLock<Vec<tokio::sync::Mutex<()>>> = OnceLock::new();
+    let locks = LOCKS.get_or_init(|| {
+        (0..USAGE_SETTLEMENT_LOCK_SHARDS)
+            .map(|_| tokio::sync::Mutex::new(()))
+            .collect()
+    });
+    &locks[usage_settlement_lock_shard(key)]
+}
+
+fn usage_settlement_lock_shard(key: &str) -> usize {
+    let mut hasher = DefaultHasher::new();
+    key.hash(&mut hasher);
+    (hasher.finish() as usize) % USAGE_SETTLEMENT_LOCK_SHARDS
+}
+
+fn usage_settlement_lock_key(input: &UsageSettlementInput) -> String {
+    if input.api_key_is_standalone {
+        if let Some(api_key_id) = input.api_key_id.as_deref().and_then(non_empty_trimmed) {
+            return format!("api-key:{api_key_id}");
+        }
+    }
+    if let Some(user_id) = input.user_id.as_deref().and_then(non_empty_trimmed) {
+        return format!("user:{user_id}");
+    }
+    if let Some(api_key_id) = input.api_key_id.as_deref().and_then(non_empty_trimmed) {
+        return format!("api-key:{api_key_id}");
+    }
+    format!("request:{}", input.request_id.trim())
+}
+
+fn non_empty_trimmed(value: &str) -> Option<&str> {
+    let value = value.trim();
+    (!value.is_empty()).then_some(value)
+}
+
 fn usage_api_key_is_standalone(usage: &StoredRequestUsageAudit) -> bool {
     usage
         .request_metadata
@@ -64,7 +108,9 @@ fn finite_cost(value: f64) -> Result<f64, DataLayerError> {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::atomic::{AtomicUsize, Ordering};
     use std::sync::Mutex;
+    use std::time::Duration;
 
     use super::{settle_usage_if_needed, UsageSettlementWriter};
     use aether_data_contracts::repository::settlement::UsageSettlementInput;
@@ -78,6 +124,13 @@ mod tests {
         inputs: Mutex<Vec<UsageSettlementInput>>,
     }
 
+    #[derive(Default)]
+    struct SlowSettlementWriter {
+        active: AtomicUsize,
+        max_active: AtomicUsize,
+        inputs: Mutex<Vec<UsageSettlementInput>>,
+    }
+
     #[async_trait]
     impl UsageSettlementWriter for TestSettlementWriter {
         fn has_usage_settlement_writer(&self) -> bool {
@@ -99,6 +152,31 @@ mod tests {
         }
     }
 
+    #[async_trait]
+    impl UsageSettlementWriter for SlowSettlementWriter {
+        fn has_usage_settlement_writer(&self) -> bool {
+            true
+        }
+
+        async fn settle_usage(
+            &self,
+            input: UsageSettlementInput,
+        ) -> Result<
+            Option<aether_data_contracts::repository::settlement::StoredUsageSettlement>,
+            aether_data_contracts::DataLayerError,
+        > {
+            let active = self.active.fetch_add(1, Ordering::AcqRel) + 1;
+            self.max_active.fetch_max(active, Ordering::AcqRel);
+            tokio::time::sleep(Duration::from_millis(30)).await;
+            self.inputs
+                .lock()
+                .expect("settlement inputs lock")
+                .push(input);
+            self.active.fetch_sub(1, Ordering::AcqRel);
+            Ok(None)
+        }
+    }
+
     fn sample_usage() -> StoredRequestUsageAudit {
         StoredRequestUsageAudit::new(
             "usage-1".to_string(),
@@ -238,4 +316,22 @@ mod tests {
         let inputs = writer.inputs.lock().expect("settlement inputs lock");
         assert!(inputs.is_empty());
     }
+
+    #[tokio::test]
+    async fn serializes_settlements_for_same_billing_subject() {
+        let writer = SlowSettlementWriter::default();
+        let mut first = sample_usage();
+        first.request_id = "req-same-subject-1".to_string();
+        let mut second = sample_usage();
+        second.request_id = "req-same-subject-2".to_string();
+
+        tokio::try_join!(
+            settle_usage_if_needed(&writer, &first),
+            settle_usage_if_needed(&writer, &second)
+        )
+        .expect("settlements should succeed");
+
+        assert_eq!(writer.max_active.load(Ordering::Acquire), 1);
+        assert_eq!(writer.inputs.lock().expect("inputs lock").len(), 2);
+    }
 }