From 1f2fb8afcd5dc245c89a24c82922266366172f1f Mon Sep 17 00:00:00 2001
From: Teddy Tennant <teddytennant@icloud.com>
Date: Wed, 8 Jul 2026 17:11:10 -0400
Subject: [PATCH] fix(runtime): saturate FlushingPolicyState counters to
 prevent add-overflow panic (#1359)

`FlushingPolicyState::register` accumulated its two u32 counters with `+=`
on every staged `Bytes` allocation. Flushing is advisory, so between the
consumer's flush/reset calls the counters can grow without bound; on a large
workload `bytes_size += bytes.len() as u32` eventually exceeds `u32::MAX` and
panics with "attempt to add with overflow" (release builds silently wrap and
corrupt the flush decision).

Use `saturating_add` for both counters, and convert the length with
`u32::try_from(..).unwrap_or(u32::MAX)` so a single >= 4 GiB allocation
saturates instead of truncating. Because `should_flush` compares with `>=`,
saturating at `u32::MAX` never changes the flush decision. No public API
change.

Fixes #1359
---
 .../memory_management/drop_queue/policy.rs    | 20 +++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/crates/cubecl-server/src/memory_management/drop_queue/policy.rs b/crates/cubecl-server/src/memory_management/drop_queue/policy.rs
index 9129e73aad..f8a05f71ad 100644
--- a/crates/cubecl-server/src/memory_management/drop_queue/policy.rs
+++ b/crates/cubecl-server/src/memory_management/drop_queue/policy.rs
@@ -32,8 +32,10 @@ pub(crate) struct FlushingPolicyState {
 impl FlushingPolicyState {
     /// Record a newly staged [`Bytes`] allocation.
     pub(crate) fn register(&mut self, bytes: &Bytes) {
-        self.bytes_count += 1;
-        self.bytes_size += bytes.len() as u32;
+        self.bytes_count = self.bytes_count.saturating_add(1);
+        self.bytes_size = self
+            .bytes_size
+            .saturating_add(u32::try_from(bytes.len()).unwrap_or(u32::MAX));
     }
 
     /// Reset all counters, typically called after a flush.
@@ -99,6 +101,20 @@ mod policy_tests {
         assert!(s.should_flush(&policy()));
     }
 
+    #[test]
+    fn register_saturates_instead_of_overflowing() {
+        // Regression test for #1359: staging allocations without an intervening
+        // reset must not panic with "attempt to add with overflow".
+        let mut s = FlushingPolicyState {
+            bytes_count: u32::MAX,
+            bytes_size: u32::MAX - 1,
+        };
+        s.register(&Bytes::from_elems(vec![0u8; 8]));
+        assert_eq!(s.bytes_count, u32::MAX);
+        assert_eq!(s.bytes_size, u32::MAX);
+        assert!(s.should_flush(&policy()));
+    }
+
     #[test]
     fn reset_clears_state() {
         let mut s = state();