diff --git a/crates/engines/src/mock_transport.rs b/crates/engines/src/mock_transport.rs index dc4721b..9b5154f 100644 --- a/crates/engines/src/mock_transport.rs +++ b/crates/engines/src/mock_transport.rs @@ -55,11 +55,9 @@ impl MockTransport { } fn ok_json(v: Value) -> GResult { - let body = serde_json::to_vec(&v) - .map_err(|e| GatewayError::internal("mock: encode response").with_source(e))?; Ok(UpstreamResponse { status: 200, - body: UpstreamBody::Json(body.into()), + body: UpstreamBody::Json(bytes::Bytes::from(v.to_string())), headers: HeaderMap::new(), }) } diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 147a1d5..a99eef6 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -2655,9 +2655,10 @@ mod tests { assert_eq!(pending.status, gw_state::BatchStatus::Pending); let drainer = OfflineHandler::new(online); + let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); let drain = tokio::spawn(async move { drainer - .drain_forever(120, std::time::Duration::from_millis(50)) + .drain_until(120, std::time::Duration::from_millis(50), shutdown_rx) .await }); let mut completed = None; @@ -2670,7 +2671,11 @@ mod tests { } tokio::time::sleep(std::time::Duration::from_millis(20)).await; } - drain.abort(); + shutdown_tx.send_replace(true); + tokio::time::timeout(std::time::Duration::from_secs(2), drain) + .await + .expect("drain stops once shutdown is signalled") + .unwrap(); let j = completed.expect("drain completed the batch"); assert_eq!(j.results.len(), 2, "both items executed exactly once"); assert!(j.results.iter().all(|r| r.ok && r.total_tokens > 0)); diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index c735c1a..10c774f 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -198,11 +198,29 @@ impl OfflineHandler { } } - /// Fleet drain loop for distributed stores: claim and execute pending batches, requeuing stale ones. - pub async fn drain_forever(&self, stale_secs: i64, poll: std::time::Duration) { + /// Fleet drain loop for distributed stores: stop claiming on shutdown, finish the claimed batch. + pub async fn drain_until( + &self, + stale_secs: i64, + poll: std::time::Duration, + mut shutdown: tokio::sync::watch::Receiver, + ) { let store = self.online.state().store.clone(); loop { - match store.batch_claim_pending(stale_secs).await { + if *shutdown.borrow() { + return; + } + let claimed = tokio::select! { + biased; + changed = shutdown.changed() => { + if stopping(changed, &shutdown) { + return; + } + continue; + } + claimed = store.batch_claim_pending(stale_secs) => claimed, + }; + match claimed { Ok(Some((job, claim))) => { // a key revoked/banned/expired since submit stops its queued work let ak = match self.online.state().auth.authenticate(&job.ak).await { @@ -242,16 +260,40 @@ impl OfflineHandler { ) .await; } - Ok(None) => tokio::time::sleep(poll).await, + Ok(None) => { + if pause_or_stop(&mut shutdown, poll).await { + return; + } + } Err(e) => { tracing::warn!(error = %e, "batch claim failed; backing off"); - tokio::time::sleep(poll).await; + if pause_or_stop(&mut shutdown, poll).await { + return; + } } } } } } +fn stopping( + changed: Result<(), tokio::sync::watch::error::RecvError>, + shutdown: &tokio::sync::watch::Receiver, +) -> bool { + changed.is_err() || *shutdown.borrow() +} + +async fn pause_or_stop( + shutdown: &mut tokio::sync::watch::Receiver, + poll: std::time::Duration, +) -> bool { + tokio::select! { + biased; + changed = shutdown.changed() => stopping(changed, shutdown), + _ = tokio::time::sleep(poll) => false, + } +} + fn failed_item(index: usize, message: String, user: String) -> BatchItemResult { BatchItemResult { index, diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 503776f..0827fc8 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -125,11 +125,14 @@ async fn main() -> anyhow::Result<()> { } // fleet batch drain: on a distributed store any instance claims submitted batches + let (batch_shutdown_tx, batch_shutdown_rx) = tokio::sync::watch::channel(false); let batch_task = if distributed_batches { let offline = app_state.offline.clone(); tracing::info!("batch drain loop started (distributed store)"); Some(tokio::spawn(async move { - offline.drain_forever(BATCH_STALE_SECS, BATCH_POLL).await + offline + .drain_until(BATCH_STALE_SECS, BATCH_POLL, batch_shutdown_rx) + .await })) } else { None @@ -200,9 +203,17 @@ async fn main() -> anyhow::Result<()> { listener, router.into_make_service_with_connect_info::(), ) - .with_graceful_shutdown(shutdown_signal()) + .with_graceful_shutdown(async move { + shutdown_signal().await; + batch_shutdown_tx.send_replace(true); + }) .await?; + if let Some(task) = batch_task + && let Err(e) = task.await + { + tracing::error!(error = %e, "batch drain task failed during shutdown"); + } gw_state::admission::flush_billing(&shared.load().state).await; quota_task.abort(); purge_task.abort(); @@ -210,9 +221,6 @@ async fn main() -> anyhow::Result<()> { avail_task.abort(); alert_task.abort(); avail_alert_task.abort(); - if let Some(task) = batch_task { - task.abort(); - } tracing::info!("gw drained and exiting"); Ok(()) } diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 0f119d4..787d116 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -66,19 +66,26 @@ impl BillingLedger { // The bounded worker owns accepted rows through caller cancellation. tokio::spawn(async move { let mut batch = Vec::with_capacity(LEDGER_BATCH_MAX); + let mut row_acks = Vec::with_capacity(LEDGER_BATCH_MAX); let mut ack = None; let mut dropped = 0u64; while let Some(msg) = pending.recv().await { - msg.take(&mut batch, &mut ack); + msg.take(&mut batch, &mut row_acks, &mut ack); while ack.is_none() && batch.len() < LEDGER_BATCH_MAX { let Ok(next) = pending.try_recv() else { break; }; - next.take(&mut batch, &mut ack); + next.take(&mut batch, &mut row_acks, &mut ack); } - let rows = batch.len() as u64; - if !Self::commit(&worker_store, &mut batch).await { - dropped += rows; + let unacked = (batch.len() - row_acks.len()) as u64; + let committed = Self::commit(&worker_store, &mut batch).await; + if !committed { + dropped += unacked; + } + for tx in row_acks.drain(..) { + if tx.send(committed).is_err() && !committed { + dropped += 1; + } } if let Some(tx) = ack.take() { let _ = tx.send(std::mem::take(&mut dropped)); @@ -104,13 +111,20 @@ impl BillingLedger { rx.await.unwrap_or(0) } - // deferred on SQL backends: a SIGKILL between the response and the next flush loses the queue plus the batch in flight (4352 rows at the defaults) async fn write(&self, record: &BillingRecord) { if self.deferred && let Some(queue) = &self.queue - && queue.try_send(LedgerWrite::Row(record.clone())).is_ok() { - return; + let queued = if record.user_id.is_empty() { + queue + .try_send(LedgerWrite::Row(record.clone(), None)) + .is_ok() + } else { + Self::queue_attributed(queue, record).await + }; + if queued { + return; + } } let Err(e) = self.store.ledger_add(record).await else { return; @@ -121,11 +135,27 @@ impl BillingLedger { return; }; tracing::error!(error = %e, "billing ledger write failed; queued for repair"); - if queue.send(LedgerWrite::Row(record.clone())).await.is_err() { + if queue + .send(LedgerWrite::Row(record.clone(), None)) + .await + .is_err() + { tracing::error!("billing ledger repair worker stopped"); } } + async fn queue_attributed(queue: &mpsc::Sender, record: &BillingRecord) -> bool { + let (tx, rx) = oneshot::channel(); + if queue + .send(LedgerWrite::Row(record.clone(), Some(tx))) + .await + .is_err() + { + return false; + } + rx.await.unwrap_or(false) + } + async fn commit(store: &Arc, batch: &mut Vec) -> bool { if batch.is_empty() { return true; @@ -155,14 +185,24 @@ impl BillingLedger { #[allow(clippy::large_enum_variant)] #[derive(Debug)] enum LedgerWrite { - Row(BillingRecord), + Row(BillingRecord, Option>), Flush(oneshot::Sender), } impl LedgerWrite { - fn take(self, batch: &mut Vec, ack: &mut Option>) { + fn take( + self, + batch: &mut Vec, + row_acks: &mut Vec>, + ack: &mut Option>, + ) { match self { - LedgerWrite::Row(r) => batch.push(r), + LedgerWrite::Row(r, tx) => { + batch.push(r); + if let Some(tx) = tx { + row_acks.push(tx); + } + } LedgerWrite::Flush(tx) => *ack = Some(tx), } } @@ -407,7 +447,7 @@ mod tests { ak: "ak".into(), product: "p".into(), tenant: "default".into(), - user_id: "u".into(), + user_id: String::new(), request_id: request_id.into(), created_at_epoch_secs: 1, model: "m".into(), @@ -447,6 +487,60 @@ mod tests { assert_eq!(rows[0].request_id, "req-repair"); } + #[tokio::test] + async fn attributed_write_returns_after_the_batch_commits() { + let store = Arc::new(crate::MemoryStore::default()); + let ledger = BillingLedger { + deferred: true, + ..BillingLedger::repairing(store.clone()) + }; + let mut row = record("req-attributed"); + row.user_id = "user-42".into(); + + ledger.write(&row).await; + + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!(count, 1); + assert_eq!(rows[0].request_id, "req-attributed"); + } + + #[tokio::test(start_paused = true)] + async fn attributed_write_falls_back_to_a_direct_write_when_the_batch_is_dropped() { + let store = Arc::new(crate::MemoryStore::default()); + store.fail_next_ledger_writes(LEDGER_RETRY_ATTEMPTS); + let ledger = BillingLedger { + deferred: true, + ..BillingLedger::repairing(store.clone()) + }; + let mut row = record("req-fallback"); + row.user_id = "user-42".into(); + + ledger.write(&row).await; + + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!(count, 1); + assert_eq!(rows[0].request_id, "req-fallback"); + } + + #[tokio::test(start_paused = true)] + async fn canceled_attributed_write_is_reported_as_dropped() { + let store = Arc::new(crate::MemoryStore::default()); + store.fail_next_ledger_writes(LEDGER_RETRY_ATTEMPTS); + let ledger = BillingLedger::repairing(store); + let (tx, rx) = oneshot::channel(); + drop(rx); + + ledger + .queue + .as_ref() + .unwrap() + .send(LedgerWrite::Row(record("req-canceled"), Some(tx))) + .await + .unwrap(); + + assert_eq!(ledger.flush().await, 1); + } + #[tokio::test] async fn billing_ledger_backpressures_at_capacity_then_repairs_every_row() { let store = Arc::new(crate::MemoryStore::default());