From 30846761dfa1f302a3a3f21ddda618877021e980 Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 3 Sep 2026 16:27:33 +0800 Subject: [PATCH 1/6] fix: close billing and batch shutdown races --- crates/handler/src/lib.rs | 3 +- crates/handler/src/offline.rs | 46 +++++++++++++++++--- crates/server/src/main.rs | 18 +++++--- crates/state/src/admission.rs | 79 ++++++++++++++++++++++++++++++----- 4 files changed, 124 insertions(+), 22 deletions(-) diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 147a1d5..36b7e49 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; diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index c735c1a..e3bdba0 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 changed.is_err() || *shutdown.borrow() { + 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,10 +260,28 @@ impl OfflineHandler { ) .await; } - Ok(None) => tokio::time::sleep(poll).await, + Ok(None) => { + tokio::select! { + biased; + changed = shutdown.changed() => { + if changed.is_err() || *shutdown.borrow() { + return; + } + } + _ = tokio::time::sleep(poll) => {} + } + } Err(e) => { tracing::warn!(error = %e, "batch claim failed; backing off"); - tokio::time::sleep(poll).await; + tokio::select! { + biased; + changed = shutdown.changed() => { + if changed.is_err() || *shutdown.borrow() { + return; + } + } + _ = tokio::time::sleep(poll) => {} + } } } } 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..9ba5eee 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -66,20 +66,25 @@ 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 { + let committed = Self::commit(&worker_store, &mut batch).await; + if !committed { dropped += rows; } + for tx in row_acks.drain(..) { + let _ = tx.send(()); + } if let Some(tx) = ack.take() { let _ = tx.send(std::mem::take(&mut dropped)); } @@ -104,13 +109,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; + if record.user_id.is_empty() { + if queue + .try_send(LedgerWrite::Row(record.clone(), None)) + .is_ok() + { + return; + } + } else if Self::queue_attributed(queue, record).await { + return; + } } let Err(e) = self.store.ledger_add(record).await else { return; @@ -121,11 +133,31 @@ impl BillingLedger { return; }; tracing::error!(error = %e, "billing ledger write failed; queued for repair"); - if queue.send(LedgerWrite::Row(record.clone())).await.is_err() { + let queued = if record.user_id.is_empty() { + queue + .send(LedgerWrite::Row(record.clone(), None)) + .await + .is_ok() + } else { + Self::queue_attributed(queue, record).await + }; + if !queued { 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.is_ok() + } + async fn commit(store: &Arc, batch: &mut Vec) -> bool { if batch.is_empty() { return true; @@ -155,14 +187,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 +449,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 +489,21 @@ mod tests { assert_eq!(rows[0].request_id, "req-repair"); } + #[tokio::test] + async fn attributed_write_returns_after_repair_commits() { + let store = Arc::new(crate::MemoryStore::default()); + store.fail_next_ledger_writes(1); + let ledger = 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] async fn billing_ledger_backpressures_at_capacity_then_repairs_every_row() { let store = Arc::new(crate::MemoryStore::default()); From 5d44be7b8125c8a9ddcf0bc20a6bf91ea7953382 Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 3 Sep 2026 16:49:35 +0800 Subject: [PATCH 2/6] fix: an attributed ledger ack carries the commit result The row ack fired after every batch attempt, so an attributed request returned confirmed even when the batch had been dropped after its retries. The ack now carries committed; on false the writer makes one direct ledger_add and only then hands the row to the repair queue, which stays best-effort for every row. The three shutdown selects in drain_until fold into stopping and pause_or_stop. Hot path: an attributed request now waits for one batched ledger commit (shared with concurrent rows) before it completes; unattributed rows are unchanged. With the store down the wait is bounded by the worker's retry schedule. --- crates/handler/src/offline.rs | 40 +++++++++++++++++++-------------- crates/state/src/admission.rs | 42 +++++++++++++++++------------------ 2 files changed, 43 insertions(+), 39 deletions(-) diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index e3bdba0..10c774f 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -213,7 +213,7 @@ impl OfflineHandler { let claimed = tokio::select! { biased; changed = shutdown.changed() => { - if changed.is_err() || *shutdown.borrow() { + if stopping(changed, &shutdown) { return; } continue; @@ -261,26 +261,14 @@ impl OfflineHandler { .await; } Ok(None) => { - tokio::select! { - biased; - changed = shutdown.changed() => { - if changed.is_err() || *shutdown.borrow() { - return; - } - } - _ = tokio::time::sleep(poll) => {} + if pause_or_stop(&mut shutdown, poll).await { + return; } } Err(e) => { tracing::warn!(error = %e, "batch claim failed; backing off"); - tokio::select! { - biased; - changed = shutdown.changed() => { - if changed.is_err() || *shutdown.borrow() { - return; - } - } - _ = tokio::time::sleep(poll) => {} + if pause_or_stop(&mut shutdown, poll).await { + return; } } } @@ -288,6 +276,24 @@ impl OfflineHandler { } } +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/state/src/admission.rs b/crates/state/src/admission.rs index 9ba5eee..a8cfdcd 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -83,7 +83,7 @@ impl BillingLedger { dropped += rows; } for tx in row_acks.drain(..) { - let _ = tx.send(()); + let _ = tx.send(committed); } if let Some(tx) = ack.take() { let _ = tx.send(std::mem::take(&mut dropped)); @@ -113,14 +113,14 @@ impl BillingLedger { if self.deferred && let Some(queue) = &self.queue { - if record.user_id.is_empty() { - if queue + let queued = if record.user_id.is_empty() { + queue .try_send(LedgerWrite::Row(record.clone(), None)) .is_ok() - { - return; - } - } else if Self::queue_attributed(queue, record).await { + } else { + Self::queue_attributed(queue, record).await + }; + if queued { return; } } @@ -133,15 +133,11 @@ impl BillingLedger { return; }; tracing::error!(error = %e, "billing ledger write failed; queued for repair"); - let queued = if record.user_id.is_empty() { - queue - .send(LedgerWrite::Row(record.clone(), None)) - .await - .is_ok() - } else { - Self::queue_attributed(queue, record).await - }; - if !queued { + if queue + .send(LedgerWrite::Row(record.clone(), None)) + .await + .is_err() + { tracing::error!("billing ledger repair worker stopped"); } } @@ -155,7 +151,7 @@ impl BillingLedger { { return false; } - rx.await.is_ok() + rx.await.unwrap_or(false) } async fn commit(store: &Arc, batch: &mut Vec) -> bool { @@ -187,7 +183,7 @@ impl BillingLedger { #[allow(clippy::large_enum_variant)] #[derive(Debug)] enum LedgerWrite { - Row(BillingRecord, Option>), + Row(BillingRecord, Option>), Flush(oneshot::Sender), } @@ -195,7 +191,7 @@ impl LedgerWrite { fn take( self, batch: &mut Vec, - row_acks: &mut Vec>, + row_acks: &mut Vec>, ack: &mut Option>, ) { match self { @@ -490,10 +486,12 @@ mod tests { } #[tokio::test] - async fn attributed_write_returns_after_repair_commits() { + async fn attributed_write_returns_after_the_batch_commits() { let store = Arc::new(crate::MemoryStore::default()); - store.fail_next_ledger_writes(1); - let ledger = BillingLedger::repairing(store.clone()); + let ledger = BillingLedger { + deferred: true, + ..BillingLedger::repairing(store.clone()) + }; let mut row = record("req-attributed"); row.user_id = "user-42".into(); From d87a3947437f706fc64bd7ea657531b48a45b65d Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 3 Sep 2026 16:56:11 +0800 Subject: [PATCH 3/6] test: cover the attributed fallback and the drain shutdown path The fallback test runs under a paused tokio clock so the worker's retry backoff completes instantly; the drain test now flips the shutdown watch and requires the loop to exit within two seconds instead of aborting it. --- crates/handler/src/lib.rs | 8 ++++++-- crates/state/src/admission.rs | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 36b7e49..a99eef6 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -2655,7 +2655,7 @@ 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 (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); let drain = tokio::spawn(async move { drainer .drain_until(120, std::time::Duration::from_millis(50), shutdown_rx) @@ -2671,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/state/src/admission.rs b/crates/state/src/admission.rs index a8cfdcd..f7e223b 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -502,6 +502,24 @@ mod tests { 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] async fn billing_ledger_backpressures_at_capacity_then_repairs_every_row() { let store = Arc::new(crate::MemoryStore::default()); From 5fe45b43c7501829333e69668362c6524b514a7c Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 3 Sep 2026 17:07:21 +0800 Subject: [PATCH 4/6] fix: count only unacknowledged rows as dropped An attributed row whose batch exhausted its retries is written directly by its caller, so counting it as dropped made the shutdown flush report billing loss that did not happen. --- crates/state/src/admission.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index f7e223b..f1b2df0 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -77,10 +77,10 @@ impl BillingLedger { }; next.take(&mut batch, &mut row_acks, &mut ack); } - let rows = batch.len() as u64; + let unacked = (batch.len() - row_acks.len()) as u64; let committed = Self::commit(&worker_store, &mut batch).await; if !committed { - dropped += rows; + dropped += unacked; } for tx in row_acks.drain(..) { let _ = tx.send(committed); From 2ca8fbceae9e20348ae8d4cb29e5d4f7d43ca33e Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 3 Sep 2026 18:54:09 +0800 Subject: [PATCH 5/6] perf: mock transport encodes its reply through Display again 01e6087 switched the mock transport's reply encoding from v.to_string() to serde_json::to_vec. A per-commit bisect of PR #48 on the testbed put the whole +88k instructions per 48 KB request of the big-body regression on that line (p3 4,039k vs p4 4,127k instr/req; restoring the file returns to 4,039k), while the four production groups of that commit run the same instruction stream within 0.01%. The mock transport is benchmark and demo code; production uses the HTTP transport. The remaining rps spread between builds is code placement, not added work. --- crates/engines/src/mock_transport.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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(), }) } From f03d514342d6bceaafdbe5c5d7b72d9833be34df Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 3 Sep 2026 20:28:18 +0800 Subject: [PATCH 6/6] fix: report canceled attributed ledger writes --- crates/state/src/admission.rs | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index f1b2df0..787d116 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -83,7 +83,9 @@ impl BillingLedger { dropped += unacked; } for tx in row_acks.drain(..) { - let _ = tx.send(committed); + if tx.send(committed).is_err() && !committed { + dropped += 1; + } } if let Some(tx) = ack.take() { let _ = tx.send(std::mem::take(&mut dropped)); @@ -520,6 +522,25 @@ mod tests { 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());