fix(s2n-quic-dc-metrics): Avoid panicking in Counter::increment with >u32::MAX count (#2995)

Author: Mark-Simulacrum

dc/s2n-quic-dc-metrics/src/counter.rs

@@ -52,9 +52,19 @@ impl Counter {
         }
     }
 
-    #[track_caller]
     pub fn increment(&self, count: u64) {
-        assert!(count <= u32::MAX as u64);
+        // If we get a particularly large count, directly serialize it into the underlying
+        // aggregate. We expect this to be rare, so acquiring the lock across all counters should
+        // be relatively cheap.
+        //
+        // It's possible to shift the exact threshold if we wanted to (e.g., by reducing the
+        // counter index bits, or having a more complicated serialization scheme), but we don't
+        // expect this to matter much in practice. If the event recorded is this large there's
+        // probably a good deal of compute needed to produce it in the first place.
+        if count > (u32::MAX as u64) {
+            self.channels.lock_aggregate()[self.counter as usize].value += count;
+            return;
+        }
         self.channels
             .send_event(((self.counter as u64) << 32) | count);
     }
@@ -76,3 +86,18 @@ fn basic() {
 
     assert_eq!(registry.take_current_metrics_line(), "a=5,b=9");
 }
+
+#[test]
+fn check_u64_max() {
+    let registry = crate::Registry::new();
+    let a = registry.register_counter(String::from("a"), None);
+    let b = registry.register_counter(String::from("b"), None);
+
+    a.increment(u64::MAX);
+    b.increment(u32::MAX as u64);
+
+    assert_eq!(
+        registry.take_current_metrics_line(),
+        format!("a={},b={}", u64::MAX, u32::MAX)
+    );
+}

dc/s2n-quic-dc-metrics/src/rseq.rs

@@ -9,7 +9,7 @@ use std::{
     ptr::NonNull,
     sync::{
         atomic::{AtomicBool, AtomicPtr, AtomicU64, Ordering},
-        Mutex,
+        Mutex, MutexGuard,
     },
 };
 
@@ -249,6 +249,10 @@ impl<T: Absorb> Channels<T> {
         self.empty_pages.push(page);
     }
 
+    pub(crate) fn lock_aggregate(&self) -> MutexGuard<'_, Vec<T>> {
+        self.aggregate.lock().expect("propagate panic")
+    }
+
     #[cfg(not(target_os = "linux"))]
     pub(crate) fn send_event(&self, event: u64) {
         self.fallback_push(event)