From 89b951556db5ccd15c9abe570d7fbf503e62bd66 Mon Sep 17 00:00:00 2001
From: "Jonathan A. Sternberg" <jonathan.sternberg@docker.com>
Date: Tue, 24 Sep 2024 15:16:48 -0500
Subject: [PATCH] protobuf: add marshaling benchmarks for some protobuf
 messages

The benchmarks that are added are related to certain areas that are
considered hotspots for performance where many messages are marshaled.

Signed-off-by: Jonathan A. Sternberg <jonathan.sternberg@docker.com>
---
 api/services/control/control_bench_test.go | 118 +++++++++++++++++++++
 cache/contenthash/checksum_bench_test.go   |  69 ++++++++++++
 2 files changed, 187 insertions(+)
 create mode 100644 api/services/control/control_bench_test.go
 create mode 100644 cache/contenthash/checksum_bench_test.go

diff --git a/api/services/control/control_bench_test.go b/api/services/control/control_bench_test.go
new file mode 100644
index 000000000000..fd386ad91954
--- /dev/null
+++ b/api/services/control/control_bench_test.go
@@ -0,0 +1,118 @@
+package moby_buildkit_v1 //nolint:revive
+
+import (
+	"testing"
+	"time"
+
+	digest "github.com/opencontainers/go-digest"
+	"github.com/stretchr/testify/require"
+)
+
+// Buf is used to prevent the benchmark from being optimized away.
+var Buf []byte
+
+func BenchmarkMarshalVertex(b *testing.B) {
+	v := sampleVertex()
+	for i := 0; i < b.N; i++ {
+		var err error
+		Buf, err = v.Marshal()
+		require.NoError(b, err)
+	}
+}
+
+func BenchmarkMarshalVertexStatus(b *testing.B) {
+	v := sampleVertexStatus()
+	for i := 0; i < b.N; i++ {
+		var err error
+		Buf, err = v.Marshal()
+		require.NoError(b, err)
+	}
+}
+
+func BenchmarkMarshalVertexLog(b *testing.B) {
+	v := sampleVertexLog()
+	for i := 0; i < b.N; i++ {
+		var err error
+		Buf, err = v.Marshal()
+		require.NoError(b, err)
+	}
+}
+
+var VertexOutput Vertex
+
+func BenchmarkUnmarshalVertex(b *testing.B) {
+	v := sampleVertex()
+	buf, err := v.Marshal()
+	require.NoError(b, err)
+
+	for i := 0; i < b.N; i++ {
+		err := VertexOutput.Unmarshal(buf)
+		require.NoError(b, err)
+	}
+}
+
+var VertexStatusOutput VertexStatus
+
+func BenchmarkUnmarshalVertexStatus(b *testing.B) {
+	v := sampleVertexStatus()
+	buf, err := v.Marshal()
+	require.NoError(b, err)
+
+	for i := 0; i < b.N; i++ {
+		err := VertexStatusOutput.Unmarshal(buf)
+		require.NoError(b, err)
+	}
+}
+
+var VertexLogOutput VertexLog
+
+func BenchmarkUnmarshalVertexLog(b *testing.B) {
+	v := sampleVertexLog()
+	buf, err := v.Marshal()
+	require.NoError(b, err)
+
+	for i := 0; i < b.N; i++ {
+		err := VertexLogOutput.Unmarshal(buf)
+		require.NoError(b, err)
+	}
+}
+
+func sampleVertex() *Vertex {
+	now := time.Now()
+	started := now.Add(-time.Minute)
+	return &Vertex{
+		Digest: digest.FromString("abc"),
+		Inputs: []digest.Digest{
+			digest.FromString("dep1"),
+			digest.FromString("dep2"),
+		},
+		Name:      "abc",
+		Started:   &started,
+		Completed: &now,
+	}
+}
+
+func sampleVertexStatus() *VertexStatus {
+	now := time.Now()
+	started := now.Add(-time.Minute)
+	return &VertexStatus{
+		ID:        "abc",
+		Vertex:    digest.FromString("abc"),
+		Name:      "abc",
+		Current:   1024,
+		Total:     1024,
+		Timestamp: now,
+		Started:   &started,
+		Completed: &now,
+	}
+}
+
+func sampleVertexLog() *VertexLog {
+	now := time.Now()
+	return &VertexLog{
+		Vertex:    digest.FromString("abc"),
+		Timestamp: now,
+		Stream:    1,
+		Msg:       []byte("this is a log message"),
+	}
+}
diff --git a/cache/contenthash/checksum_bench_test.go b/cache/contenthash/checksum_bench_test.go
new file mode 100644
index 000000000000..affe0e03ae89
--- /dev/null
+++ b/cache/contenthash/checksum_bench_test.go
@@ -0,0 +1,69 @@
+package contenthash
+
+import (
+	"testing"
+
+	digest "github.com/opencontainers/go-digest"
+	"github.com/stretchr/testify/require"
+)
+
+// Buf is used to prevent the benchmark from being optimized away.
+var Buf []byte
+
+func BenchmarkMarshalCacheRecords(b *testing.B) {
+	v := sampleCacheRecords()
+	for i := 0; i < b.N; i++ {
+		var err error
+		Buf, err = v.Marshal()
+		require.NoError(b, err)
+	}
+}
+
+var CacheRecordsOutput CacheRecords
+
+func BenchmarkUnmarshalCacheRecords(b *testing.B) {
+	v := sampleCacheRecords()
+	buf, err := v.Marshal()
+	require.NoError(b, err)
+
+	for i := 0; i < b.N; i++ {
+		err := CacheRecordsOutput.Unmarshal(buf)
+		require.NoError(b, err)
+	}
+}
+
+func sampleCacheRecords() *CacheRecords {
+	return &CacheRecords{
+		Paths: []*CacheRecordWithPath{
+			{
+				Path: "/foo",
+				Record: &CacheRecord{
+					Digest: digest.FromString("/foo"),
+					Type:   CacheRecordTypeDir,
+				},
+			},
+			{
+				Path: "/foo/",
+				Record: &CacheRecord{
+					Digest: digest.FromString("/foo/"),
+					Type:   CacheRecordTypeDirHeader,
+				},
+			},
+			{
+				Path: "/foo/bar.txt",
+				Record: &CacheRecord{
+					Digest: digest.FromString("/foo/bar.txt"),
+					Type:   CacheRecordTypeFile,
+				},
+			},
+			{
+				Path: "/foo/link",
+				Record: &CacheRecord{
+					Digest:   digest.FromString("/foo/link"),
+					Type:     CacheRecordTypeSymlink,
+					Linkname: "/foo/bar.txt",
+				},
+			},
+		},
+	}
+}