diff --git a/brontide/bench_test.go b/brontide/bench_test.go index 214d61f1f..0605c54f5 100644 --- a/brontide/bench_test.go +++ b/brontide/bench_test.go @@ -2,6 +2,7 @@ package brontide import ( "bytes" + "io" "math" "math/rand" "testing" @@ -63,3 +64,47 @@ func BenchmarkReadHeaderAndBody(t *testing.B) { } require.NoError(t, benchErr) } + +// BenchmarkWriteMessage benchmarks the performance of writing a maximum-sized +// message and flushing it to an io.Discard to measure the allocation and CPU +// overhead of the encryption and writing logic. +func BenchmarkWriteMessage(b *testing.B) { + localConn, remoteConn, err := establishTestConnection(b) + require.NoError(b, err, "unable to establish test connection: %v", err) + + noiseLocalConn, ok := localConn.(*Conn) + require.True(b, ok, "expected *Conn type for localConn") + + // Create the largest possible message we can write (MaxUint16 bytes). + // This is the maximum message size allowed by the protocol. + const maxMsgSize = math.MaxUint16 + largeMsg := bytes.Repeat([]byte("a"), maxMsgSize) + + // Use io.Discard to simulate writing to a network connection that + // continuously accepts data without needing resets. + discard := io.Discard + + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + // Write our massive message, then call flush to actually write + // the encrypted message This simulates a full write operation + // to a network. + err := noiseLocalConn.noise.WriteMessage(largeMsg) + if err != nil { + b.Fatalf("WriteMessage failed: %v", err) + } + _, err = noiseLocalConn.noise.Flush(discard) + if err != nil { + b.Fatalf("Flush failed: %v", err) + } + } + + // We'll make sure to clean up the connections at the end of the + // benchmark. + b.Cleanup(func() { + localConn.Close() + remoteConn.Close() + }) +}