From 845c29f6b7e3d71537c143b542c5d8c7badd3837 Mon Sep 17 00:00:00 2001
From: Rob Findley <rfindley@google.com>
Date: Tue, 16 Sep 2025 18:56:28 +0000
Subject: [PATCH] mcp: add a benchmark for the MemoryEventStore; remove
 validation

Add a new benchmark measuring the append and purge performance of the
MemoryEventStore. This benchmark revealed that the store is orders of
magnitude slower than it should be due to conservative validation
(hugely so: 300KB/s vs 568MB/s). Turn off this validation by default.

For #190
---
 mcp/event.go      |  4 ++--
 mcp/event_test.go | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/mcp/event.go b/mcp/event.go
index 0dd8734..d309c4e 100644
--- a/mcp/event.go
+++ b/mcp/event.go
@@ -23,8 +23,8 @@ import (
 )
 
 // If true, MemoryEventStore will do frequent validation to check invariants, slowing it down.
-// Remove when we're confident in the code.
-const validateMemoryEventStore = true
+// Enable for debugging.
+const validateMemoryEventStore = false
 
 // An Event is a server-sent event.
 // See https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#fields.
diff --git a/mcp/event_test.go b/mcp/event_test.go
index 601e830..20808c7 100644
--- a/mcp/event_test.go
+++ b/mcp/event_test.go
@@ -10,6 +10,7 @@ import (
 	"slices"
 	"strings"
 	"testing"
+	"time"
 )
 
 func TestScanEvents(t *testing.T) {
@@ -252,3 +253,46 @@ func TestMemoryEventStoreAfter(t *testing.T) {
 		})
 	}
 }
+
+func BenchmarkMemoryEventStore(b *testing.B) {
+	// Benchmark with various settings for event store size, number of session,
+	// and payload size.
+	//
+	// Assume a small number of streams per session, which is probably realistic.
+	tests := []struct {
+		name     string
+		limit    int
+		sessions int
+		datasize int
+	}{
+		{"1KB", 1024, 1, 16},
+		{"1MB", 1024 * 1024, 10, 16},
+		{"10MB", 10 * 1024 * 1024, 100, 16},
+		{"10MB_big", 10 * 1024 * 1024, 1000, 128},
+	}
+
+	for _, test := range tests {
+		b.Run(test.name, func(b *testing.B) {
+			store := NewMemoryEventStore(nil)
+			store.SetMaxBytes(test.limit)
+			ctx := context.Background()
+			sessionIDs := make([]string, test.sessions)
+			streamIDs := make([][3]StreamID, test.sessions)
+			for i := range sessionIDs {
+				sessionIDs[i] = fmt.Sprint(i)
+				for j := range 3 {
+					streamIDs[i][j] = StreamID(randText())
+				}
+			}
+			payload := make([]byte, test.datasize)
+			start := time.Now()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				sessionID := sessionIDs[i%len(sessionIDs)]
+				streamID := streamIDs[i%len(sessionIDs)][i%3]
+				store.Append(ctx, sessionID, streamID, payload)
+			}
+			b.ReportMetric(float64(test.datasize)*float64(b.N)/time.Since(start).Seconds(), "bytes/s")
+		})
+	}
+}