Account for CPU time spent on handshakes in bach sims

quic/s2n-quic-core/Cargo.toml

@@ -14,7 +14,7 @@ exclude = ["corpus.tar.gz"]
 default = ["alloc", "std"]
 alloc = ["atomic-waker", "bytes", "crossbeam-utils", "s2n-codec/alloc"]
 std = ["alloc", "once_cell"]
-testing = ["std", "generator", "s2n-codec/testing", "checked-counters", "insta", "futures-test"]
+testing = ["std", "generator", "s2n-codec/testing", "checked-counters", "insta", "futures-test", "bach"]
 generator = ["bolero-generator"]
 checked-counters = []
 branch-tracing = ["tracing"]
@@ -47,6 +47,7 @@ tracing = { version = "0.1", default-features = false, optional = true }
 zerocopy = { version = "0.8", features = ["derive"] }
 futures-test = { version = "0.3", optional = true } # For testing Waker interactions
 once_cell = { version = "1", optional = true }
+bach = { version = "0.1.0", optional = true }
 
 [dev-dependencies]
 bolero = "0.13"

quic/s2n-quic-core/src/io/event_loop.rs

@@ -97,6 +97,9 @@ where
 
             let select = cooldown.wrap(select);
 
+            #[cfg(feature = "testing")]
+            bach_cpu::assert_zero_cpu();
+
             let select::Outcome {
                 rx_result,
                 tx_result,
@@ -109,6 +112,9 @@ where
                 return;
             };
 
+            #[cfg(feature = "testing")]
+            bach_cpu::take_cpu().await;
+
             // notify the application that we woke up and why
             let wakeup_timestamp = clock.get_time();
             {
@@ -126,10 +132,16 @@ where
 
             match rx_result {
                 Some(Ok(())) => {
+                    #[cfg(feature = "testing")]
+                    bach_cpu::assert_zero_cpu();
+
                     // we received some packets. give them to the endpoint.
                     rx.queue(|queue| {
                         endpoint.receive(queue, &clock);
                     });
+
+                    #[cfg(feature = "testing")]
+                    bach_cpu::take_cpu().await;
                 }
                 Some(Err(error)) => {
                     // The RX provider has encountered an error. shut down the event loop
@@ -160,11 +172,20 @@ where
                 }
             }
 
+            #[cfg(feature = "testing")]
+            bach_cpu::assert_zero_cpu();
+
             // Let the endpoint transmit, if possible
             tx.queue(|queue| {
                 endpoint.transmit(queue, &clock);
             });
 
+            #[cfg(feature = "testing")]
+            bach_cpu::take_cpu().await;
+
+            #[cfg(feature = "testing")]
+            bach_cpu::assert_zero_cpu();
+
             // Get the next expiration from the endpoint and update the timer
             let timeout = endpoint.timeout();
             if let Some(timeout) = timeout {
@@ -187,3 +208,55 @@ where
         }
     }
 }
+
+/// This allows various parts of s2n-quic to "spend" CPU cycles within bach simulations
+/// deterministically. The goal is to allow simulating (especially) handshakes accurately, which
+/// incur significant CPU cycles and as such delay processing subsequent packets. It's inaccurate
+/// to model this as network delay.
+mod bach_cpu {
+    #[cfg(feature = "testing")]
+    use core::cell::Cell;
+    use core::time::Duration;
+
+    // CPU today is attributed within the event loop, which is at least today always single
+    // threaded, and we never yield while there's still unspent CPU.
+    //
+    // FIXME: I *think* an alternative to this is to wire up an event or pseudo-event that s2n-quic
+    // itself would subscribe to -- that would be a bit less plumbing, but the crypto code doesn't
+    // directly publish events today so it wouldn't be quite enough either.
+    #[cfg(feature = "testing")]
+    thread_local! {
+        static CPU_SPENT: Cell<Duration> = const { Cell::new(Duration::ZERO) };
+    }
+
+    #[inline]
+    pub fn attribute_cpu(time: Duration) {
+        #[cfg(feature = "testing")]
+        {
+            CPU_SPENT.with(|c| {
+                let old = c.get();
+                let new = old + time;
+                c.set(new);
+            });
+        }
+    }
+
+    #[cfg(feature = "testing")]
+    pub(super) async fn take_cpu() {
+        // Make sure assert_zero_cpu works in all cfg(testing), not just with bach.
+        let taken = CPU_SPENT.take();
+
+        if !bach::is_active() {
+            return;
+        }
+
+        bach::time::sleep(taken).await;
+    }
+
+    #[cfg(feature = "testing")]
+    pub(super) fn assert_zero_cpu() {
+        assert_eq!(CPU_SPENT.get(), Duration::ZERO);
+    }
+}
+
+pub use bach_cpu::attribute_cpu;

quic/s2n-quic-core/src/io/rx.rs

@@ -12,7 +12,7 @@ pub trait Rx: Sized {
     // TODO make this generic over lifetime
     // See https://github.com/aws/s2n-quic/issues/1742
     type Queue: Queue<Handle = Self::PathHandle>;
-    type Error;
+    type Error: Send;
 
     /// Returns a future that yields after a packet is ready to be received
     #[inline]

quic/s2n-quic-core/src/io/tx.rs

@@ -15,7 +15,7 @@ pub trait Tx: Sized {
     // TODO make this generic over lifetime
     // See https://github.com/aws/s2n-quic/issues/1742
     type Queue: Queue<Handle = Self::PathHandle>;
-    type Error;
+    type Error: Send;
 
     /// Returns a future that yields after a packet is ready to be transmitted
     #[inline]

quic/s2n-quic-tls/src/callback.rs

@@ -360,6 +360,16 @@ where
     fn on_read(&mut self, data: &mut [u8]) -> usize {
         let max_len = Some(data.len());
 
+        // This is a semi-random number. However, it happens to work out OK to approximate
+        // spend during handshakes. On one side of a connection a handshake typically costs about
+        // 0.5-1ms of CPU. Most of that is driven by the peer sending information that the local
+        // endpoint acts on (e.g., verifying signatures).
+        //
+        // In practice this was chosen to make s2n-quic-sim simulate an uncontended mTLS handshake
+        // as taking 2ms (in combination with the transition edge adding some extra cost), which is
+        // fairly close to what we see in one scenario with real handshakes.
+        s2n_quic_core::io::event_loop::attribute_cpu(core::time::Duration::from_micros(100));
+
         let chunk = match self.state.rx_phase {
             HandshakePhase::Initial => self.context.receive_initial(max_len),
             HandshakePhase::Handshake => self.context.receive_handshake(max_len),
@@ -407,6 +417,9 @@ enum HandshakePhase {
 
 impl HandshakePhase {
     fn transition(&mut self) {
+        // See comment in `on_read` for value and why this exists.
+        s2n_quic_core::io::event_loop::attribute_cpu(core::time::Duration::from_micros(100));
+
         *self = match self {
             Self::Initial => Self::Handshake,
             _ => Self::Application,