From f538ee8be793fc10b896d0f5c358f04b2b48feb5 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Sat, 30 Aug 2025 14:41:41 +0100
Subject: [PATCH] brontide: add benchark for WriteMessage

In this commit, we add a benchmark for the WriteMessage method. This is
the first step in an upcoming optimizooor series of commits.

The baseline is:
```
goos: darwin
goarch: arm64
pkg: github.com/lightningnetwork/lnd/brontide
cpu: Apple M4 Max
BenchmarkWriteMessage-16    	   22736	     50667 ns/op	   73788 B/op	       5 allocs/op
BenchmarkWriteMessage-16    	   23217	     50463 ns/op	   73788 B/op	       5 allocs/op
BenchmarkWriteMessage-16    	   24241	     49941 ns/op	   73788 B/op	       5 allocs/op
BenchmarkWriteMessage-16    	   23574	     51021 ns/op	   73788 B/op	       5 allocs/op
BenchmarkWriteMessage-16    	   23784	     49926 ns/op	   73788 B/op	       5 allocs/op
BenchmarkWriteMessage-16    	   24230	     50538 ns/op	   73788 B/op	       5 allocs/op
BenchmarkWriteMessage-16    	   24058	     49971 ns/op	   73788 B/op	       5 allocs/op
BenchmarkWriteMessage-16    	   23762	     50224 ns/op	   73788 B/op	       5 allocs/op
BenchmarkWriteMessage-16    	   24266	     53034 ns/op	   73788 B/op	       5 allocs/op
BenchmarkWriteMessage-16    	   23317	     50045 ns/op	   73788 B/op	       5 allocs/op
PASS
ok  	github.com/lightningnetwork/lnd/brontide	17.433s
```
---
 brontide/bench_test.go | 45 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

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()
+	})
+}