fix: close billing and batch shutdown races

crates/engines/src/mock_transport.rs

@@ -55,11 +55,9 @@ impl MockTransport {
     }
 
     fn ok_json(v: Value) -> GResult<UpstreamResponse> {
-        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(),
         })
     }

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));

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<bool>,
+    ) {
         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>,
+) -> bool {
+    changed.is_err() || *shutdown.borrow()
+}
+
+async fn pause_or_stop(
+    shutdown: &mut tokio::sync::watch::Receiver<bool>,
+    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,

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,19 +203,24 @@ async fn main() -> anyhow::Result<()> {
         listener,
         router.into_make_service_with_connect_info::<std::net::SocketAddr>(),
     )
-    .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();
     rollup_task.abort();
     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(())
 }

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<LedgerWrite>, 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<dyn Store>, batch: &mut Vec<BillingRecord>) -> 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<oneshot::Sender<bool>>),
     Flush(oneshot::Sender<u64>),
 }
 
 impl LedgerWrite {
-    fn take(self, batch: &mut Vec<BillingRecord>, ack: &mut Option<oneshot::Sender<u64>>) {
+    fn take(
+        self,
+        batch: &mut Vec<BillingRecord>,
+        row_acks: &mut Vec<oneshot::Sender<bool>>,
+        ack: &mut Option<oneshot::Sender<u64>>,
+    ) {
         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());