From b0b04f3fdc984d1a5333b6271be931a9dcb16358 Mon Sep 17 00:00:00 2001
From: Matheus Degiovani <opensource@matheusd.com>
Date: Thu, 1 Aug 2024 14:55:24 -0300
Subject: [PATCH] main: cleanup marshalling and unmarshalling benchmarks

This adds comments to the Marshal/Unmarshal benchmarks to make it clear
their goals and caveats.

It also splits the marshalling benchmark into two: one with and one
without message reuse.
---
 integration_test.go | 94 +++++++++++++++++++++++++++++++++++++++------
 1 file changed, 83 insertions(+), 11 deletions(-)

diff --git a/integration_test.go b/integration_test.go
index e2401d94..73d8f820 100644
--- a/integration_test.go
+++ b/integration_test.go
@@ -1725,43 +1725,112 @@ func unsafeBytesToString(b []byte) string {
 	return *(*string)(unsafe.Pointer(&b))
 }
 
+// BenchmarkMarshal benchmarks marshalling a message using the simplest API
+// available to go-capnp users.
 func BenchmarkMarshal(b *testing.B) {
 	r := rand.New(rand.NewSource(12345))
 	data := make([]*A, 1000)
 	for i := range data {
 		data[i] = generateA(r)
 	}
-	arena := make([]byte, 0, 512)
 	b.ReportAllocs()
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		a := data[r.Intn(len(data))]
-		msg, seg, _ := capnp.NewMessage(capnp.SingleSegment(arena[:0]))
-		root, _ := air.NewRootBenchmarkA(seg)
+		msg, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
+		if err != nil {
+			b.Fatal(err)
+		}
+		root, err := air.NewRootBenchmarkA(seg)
+		if err != nil {
+			b.Fatal(err)
+		}
+		a.fill(root)
+		_, err = msg.Marshal()
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkMarshal_ReuseMsg benchmarks marshalling while reusing the message
+// object.
+func BenchmarkMarshal_ReuseMsg(b *testing.B) {
+	r := rand.New(rand.NewSource(12345))
+	data := make([]*A, 1000)
+	for i := range data {
+		data[i] = generateA(r)
+	}
+	msg, _, err := capnp.NewMessage(capnp.SingleSegment(nil))
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		a := data[r.Intn(len(data))]
+		seg, err := msg.Reset(msg.Arena)
+		if err != nil {
+			b.Fatal(err)
+		}
+		root, err := air.NewRootBenchmarkA(seg)
+		if err != nil {
+			b.Fatal(err)
+		}
 		a.fill(root)
-		msg.Marshal()
+		_, err = msg.Marshal()
+		if err != nil {
+			b.Fatal(err)
+		}
 	}
 }
 
+// BenchmarkUnmarshal benchmarks unmarshalling using the simplest API possible
+// available to go-capnp users.
 func BenchmarkUnmarshal(b *testing.B) {
 	r := rand.New(rand.NewSource(12345))
-	data := make([][]byte, 1000)
+	type testCase struct {
+		a    A
+		data []byte
+	}
+
+	data := make([]testCase, 1000)
 	for i := range data {
 		a := generateA(r)
 		msg, seg, _ := capnp.NewMessage(capnp.SingleSegment(nil))
 		root, _ := air.NewRootBenchmarkA(seg)
 		a.fill(root)
-		data[i], _ = msg.Marshal()
+		buf, err := msg.Marshal()
+		if err != nil {
+			b.Fatal(err)
+		}
+		data[i].data, data[i].a = buf, *a
 	}
+
 	b.ReportAllocs()
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		msg, _ := capnp.Unmarshal(data[r.Intn(len(data))])
-		a, _ := air.ReadRootBenchmarkA(msg)
-		unmarshalA(a)
+		testIdx := r.Intn(len(data))
+		msg, err := capnp.Unmarshal(data[testIdx].data)
+		if err != nil {
+			b.Fatal(err)
+		}
+		a, err := air.ReadRootBenchmarkA(msg)
+		if err != nil {
+			b.Fatal(err)
+		}
+		gotA := unmarshalA(a)
+		if gotA != data[testIdx].a {
+			b.Fatal("unexpected unmarshalled data")
+		}
 	}
 }
 
+// BenchmarkUnmarshal_Reuse benchmarks unmarshalling by reusing the msg object
+// and directly using the test data buffer as an arena.
+//
+// NOTE: this bypasses framing done by capnp marshalling, thus expects the
+// caller to frame the message appropriately.
 func BenchmarkUnmarshal_Reuse(b *testing.B) {
 	type testCase struct {
 		a    A
@@ -1786,10 +1855,13 @@ func BenchmarkUnmarshal_Reuse(b *testing.B) {
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		testIdx := r.Intn(len(data))
-		*ta = testArena(data[testIdx].data[8:])
+		*ta = testArena(data[testIdx].data[8:]) // [8:] to skip header
 		msg.Release()
 		msg.Arena = ta
-		a, _ := air.ReadRootBenchmarkA(msg)
+		a, err := air.ReadRootBenchmarkA(msg)
+		if err != nil {
+			b.Fatal(err)
+		}
 		gotA := unmarshalA(a)
 		if gotA != data[testIdx].a {
 			b.Fatal("unexpected unmarshalled data")