fix(bindings/bench): Prevent IO from going out of scope (#5007)

aws · Jan 10, 2025 · ae3b36c · ae3b36c
1 parent a5d4e67
commit ae3b36c
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 14 deletions.
diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml
@@ -1,6 +1,8 @@
 name: Benchmarking
 
 on:
+  pull_request:
+    branches: [main]
   push:
     branches: [main]
   schedule:
@@ -29,23 +31,25 @@ jobs:
           pip3 install "boto3[crt]"
 
       - name: Generate
-        working-directory: bindings/rust
+        working-directory: bindings/rust/extended
         run: ./generate.sh --skip-tests
 
       - name: Benchmark
-        working-directory: bindings/rust/bench
+        working-directory: bindings/rust/standard/bench
         run: cargo criterion --message-format json > criterion_output.log
 
       - name: Configure AWS Credentials
+        if: github.event_name == 'push'
         uses: aws-actions/configure-aws-credentials@v4.0.2
         with:
           role-to-assume: arn:aws:iam::024603541914:role/GitHubOIDCRole
           role-session-name: s2ntlsghabenchsession
           aws-region: us-west-2
 
       - name: Emit CloudWatch metrics
+        if: github.event_name == 'push'
         run: |
           python3 .github/bin/criterion_to_cloudwatch.py \
-            --criterion_output_path bindings/rust/bench/criterion_output.log \
+            --criterion_output_path bindings/rust/standard/bench/criterion_output.log \
             --namespace s2n-tls-bench \
             --platform ${{ runner.os }}-${{ runner.arch }}
diff --git a/bindings/rust/standard/bench/benches/resumption.rs b/bindings/rust/standard/bench/benches/resumption.rs
@@ -45,19 +45,18 @@ where
         bench_group.bench_function(format!("{:?}-{}", handshake, T::name()), |b| {
             b.iter_batched_ref(
                 || {
-                    let pair = TlsConnPair::<T, T>::new_bench_pair(
+                    let mut pair = TlsConnPair::<T, T>::new_bench_pair(
                         CryptoConfig::new(CipherSuite::default(), KXGroup::default(), sig_type),
                         handshake,
                     )
                     .unwrap();
-                    let (mut c, s) = pair.split();
-                    c.handshake().unwrap();
-                    s
+                    pair.client_mut().handshake().unwrap();
+                    pair
                 },
-                |server| {
+                |pair| {
                     // this represents the work that the server does during the
                     // first RTT
-                    server.handshake().unwrap()
+                    pair.server_mut().handshake().unwrap();
                 },
                 BatchSize::SmallInput,
             )

diff --git a/bindings/rust/standard/bench/src/harness/mod.rs b/bindings/rust/standard/bench/src/harness/mod.rs
@@ -260,9 +260,20 @@ where
         Self { client, server, io }
     }
 
-    /// Take back ownership of individual connections in the TlsConnPair
-    pub fn split(self) -> (C, S) {
-        (self.client, self.server)
+    pub fn client(&self) -> &C {
+        &self.client
+    }
+
+    pub fn client_mut(&mut self) -> &mut C {
+        &mut self.client
+    }
+
+    pub fn server(&self) -> &S {
+        &self.server
+    }
+
+    pub fn server_mut(&mut self) -> &mut S {
+        &mut self.server
     }
 
     /// Run handshake on connections
@@ -386,8 +397,7 @@ mod tests {
             TlsConnPair::<C, S>::new_bench_pair(CryptoConfig::default(), HandshakeType::Resumption)
                 .unwrap();
         conn_pair.handshake().unwrap();
-        let (_, server) = conn_pair.split();
-        assert!(server.resumed_connection());
+        assert!(conn_pair.server().resumed_connection());
     }
 
     #[test]