Optimize telemetry event publishing

server/lib/events/event.go

@@ -12,6 +12,15 @@ import (
 // maxS2RecordBytes is the maximum record size for the S2 event pipeline (1 MB).
 const maxS2RecordBytes = 1_000_000
 
+// Below this conservative size estimate, JSON overhead cannot push a record
+// near the S2 limit, so Publish can avoid a full marshal on the hot path.
+const truncateMarshalThreshold = maxS2RecordBytes - 64*1024
+
+const (
+	jsonEnvelopeOverhead      = 512
+	jsonMetadataEntryOverhead = 16
+)
+
 const (
 	Console     = oapi.TelemetryEventCategory("console")
 	Network     = oapi.TelemetryEventCategory("network")
@@ -97,8 +106,11 @@ type Envelope struct {
 
 // truncateIfNeeded marshals env and returns the (possibly truncated) envelope.
 // If the envelope still exceeds maxS2RecordBytes after nulling data (e.g. huge
-// source.metadata), it is returned as-is, callers must handle nil data.
+// source.metadata), it is returned as-is with data set to null.
 func truncateIfNeeded(env Envelope) (Envelope, []byte) {
+	if estimatedEnvelopeBytes(env) < truncateMarshalThreshold {
+		return env, nil
+	}
 	data, err := json.Marshal(env)
 	if err != nil {
 		return env, nil
@@ -113,7 +125,27 @@ func truncateIfNeeded(env Envelope) (Envelope, []byte) {
 		return env, nil
 	}
 	if len(data) > maxS2RecordBytes {
-		slog.Warn("truncateIfNeeded: envelope exceeds limit even without data", "seq", env.Seq, "size", len(data))
+		slog.Warn("truncateIfNeeded: envelope exceeds limit even without data", "size", len(data))
 	}
 	return env, data
 }
+
+func estimatedEnvelopeBytes(env Envelope) int {
+	n := jsonEnvelopeOverhead + len(env.Event.Data)
+	n += maxEscapedJSONLen(env.Event.Type)
+	n += maxEscapedJSONLen(string(env.Event.Category))
+	n += maxEscapedJSONLen(string(env.Event.Source.Kind))
+	if env.Event.Source.Event != nil {
+		n += maxEscapedJSONLen(*env.Event.Source.Event)
+	}
+	if env.Event.Source.Metadata != nil {
+		for k, v := range *env.Event.Source.Metadata {
+			n += jsonMetadataEntryOverhead + maxEscapedJSONLen(k) + maxEscapedJSONLen(v)
+		}
+	}
+	return n
+}
+
+func maxEscapedJSONLen(s string) int {
+	return len(s) * 6
+}

server/lib/events/events_test.go

@@ -3,6 +3,7 @@ package events
 import (
 	"context"
 	"encoding/json"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -118,6 +119,44 @@ func TestEventOmitEmpty(t *testing.T) {
 	assert.NotContains(t, s, `"event"`)
 }
 
+func TestTruncateIfNeededChecksLargeMetadata(t *testing.T) {
+	metadata := map[string]string{"large": strings.Repeat("x", maxS2RecordBytes)}
+	env := Envelope{
+		Event: Event{
+			Type:     "console.log",
+			Category: Console,
+			Source:   oapi.BrowserEventSource{Kind: oapi.Cdp, Metadata: &metadata},
+			Data:     json.RawMessage(`{"message":"hello"}`),
+		},
+	}
+
+	truncated, data := truncateIfNeeded(env)
+
+	require.NotNil(t, data)
+	assert.True(t, truncated.Event.Truncated)
+	assert.Equal(t, json.RawMessage("null"), truncated.Event.Data)
+	assert.Greater(t, len(data), maxS2RecordBytes)
+}
+
+func TestTruncateIfNeededChecksEscapedMetadata(t *testing.T) {
+	metadata := map[string]string{"escaped": strings.Repeat("<", maxS2RecordBytes/5)}
+	env := Envelope{
+		Event: Event{
+			Type:     "console.log",
+			Category: Console,
+			Source:   oapi.BrowserEventSource{Kind: oapi.Cdp, Metadata: &metadata},
+			Data:     json.RawMessage(`{"message":"hello"}`),
+		},
+	}
+
+	truncated, data := truncateIfNeeded(env)
+
+	require.NotNil(t, data)
+	assert.True(t, truncated.Event.Truncated)
+	assert.Equal(t, json.RawMessage("null"), truncated.Event.Data)
+	assert.Greater(t, len(data), maxS2RecordBytes)
+}
+
 func mkEnv(seq uint64, ev Event) Envelope {
 	return Envelope{Seq: seq, Event: ev}
 }
@@ -133,6 +172,74 @@ func newTestRingBuffer(t *testing.T, capacity int) *ringBuffer {
 	return rb
 }
 
+func TestEventStreamPublishAssignsSeq(t *testing.T) {
+	es, err := NewEventStream(EventStreamConfig{RingCapacity: 10})
+	require.NoError(t, err)
+	reader := es.NewReader(0)
+
+	first := es.Publish(Envelope{Event: cdpEvent("console.log", Console)})
+	second := es.Publish(Envelope{Event: cdpEvent("network.request", Network)})
+
+	assert.Equal(t, uint64(1), first.Seq)
+	assert.Equal(t, uint64(2), second.Seq)
+	assert.Equal(t, uint64(2), es.Seq())
+
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	assert.Equal(t, uint64(1), readEnvelope(t, reader, ctx).Seq)
+	assert.Equal(t, uint64(2), readEnvelope(t, reader, ctx).Seq)
+}
+
+func TestEventStreamPublishTruncatesWithAssignedSeq(t *testing.T) {
+	es, err := NewEventStream(EventStreamConfig{RingCapacity: 10})
+	require.NoError(t, err)
+	nextSeq := uint64(1_000_000_000_000_000_000)
+	es.ring.latestSeq = nextSeq - 1
+
+	env := Envelope{
+		Event: Event{
+			Type:     "console.log",
+			Category: Console,
+			Source:   oapi.BrowserEventSource{Kind: oapi.Cdp},
+			Data:     json.RawMessage(`""`),
+		},
+	}
+	seqZeroLen := marshaledEnvelopeLen(t, env)
+	env.Seq = nextSeq
+	nextSeqLen := marshaledEnvelopeLen(t, env)
+	payloadLen := maxS2RecordBytes - seqZeroLen
+	payload := json.RawMessage(`"` + strings.Repeat("x", payloadLen) + `"`)
+	seqZeroEnv := Envelope{Event: Event{
+		Type:     "console.log",
+		Category: Console,
+		Source:   oapi.BrowserEventSource{Kind: oapi.Cdp},
+		Data:     payload,
+	}}
+	nextSeqEnv := seqZeroEnv
+	nextSeqEnv.Seq = nextSeq
+	require.LessOrEqual(t, marshaledEnvelopeLen(t, seqZeroEnv), maxS2RecordBytes)
+	require.Greater(t, marshaledEnvelopeLen(t, nextSeqEnv), maxS2RecordBytes)
+	require.Greater(t, nextSeqLen-seqZeroLen, 0)
+
+	published := es.Publish(Envelope{Event: Event{
+		Type:     "console.log",
+		Category: Console,
+		Source:   oapi.BrowserEventSource{Kind: oapi.Cdp},
+		Data:     payload,
+	}})
+
+	assert.Equal(t, nextSeq, published.Seq)
+	assert.True(t, published.Event.Truncated)
+	assert.Equal(t, json.RawMessage("null"), published.Event.Data)
+}
+
+func marshaledEnvelopeLen(t *testing.T, env Envelope) int {
+	t.Helper()
+	data, err := json.Marshal(env)
+	require.NoError(t, err)
+	return len(data)
+}
+
 // TestRingBuffer: publish 3 envelopes; reader reads all 3 in order
 func TestRingBuffer(t *testing.T) {
 	rb := newTestRingBuffer(t, 10)

server/lib/events/eventstream.go

@@ -2,14 +2,11 @@ package events
 
 import (
 	"fmt"
-	"sync"
 )
 
-// EventStream is the process-lifetime event bus. It owns the ring buffer and
-// sequence counter, which outlive individual capture sessions.
+// EventStream is the process-lifetime event bus. Its ring buffer and sequence
+// counter outlive individual capture sessions.
 type EventStream struct {
-	mu   sync.Mutex
-	seq  uint64
 	ring *ringBuffer
 }
 
@@ -29,14 +26,7 @@ func NewEventStream(cfg EventStreamConfig) (*EventStream, error) {
 // Publish assigns a monotonically increasing seq to env, truncates oversized
 // payloads, and pushes it to the ring buffer.
 func (es *EventStream) Publish(env Envelope) Envelope {
-	es.mu.Lock()
-	es.seq++
-	env.Seq = es.seq
-	es.mu.Unlock()
-
-	env, _ = truncateIfNeeded(env)
-	es.ring.publish(env)
-	return env
+	return es.ring.publishNext(env)
 }
 
 // NewReader returns a Reader positioned after afterSeq. Pass 0 to start from
@@ -47,7 +37,5 @@ func (es *EventStream) NewReader(afterSeq uint64) *Reader {
 
 // Seq returns the sequence number of the last published event.
 func (es *EventStream) Seq() uint64 {
-	es.mu.Lock()
-	defer es.mu.Unlock()
-	return es.seq
+	return es.ring.seq()
 }

server/lib/events/eventstream_bench_test.go

@@ -0,0 +1,49 @@
+package events
+
+import (
+	"encoding/json"
+	"testing"
+
+	oapi "github.com/kernel/kernel-images/server/lib/oapi"
+)
+
+func benchmarkEvent() Event {
+	return Event{
+		Ts:       123456789,
+		Type:     "console.log",
+		Category: Console,
+		Source:   oapi.BrowserEventSource{Kind: oapi.Cdp},
+		Data:     json.RawMessage(`{"message":"hello","level":"log"}`),
+	}
+}
+
+func BenchmarkEventStreamPublish(b *testing.B) {
+	es, err := NewEventStream(EventStreamConfig{RingCapacity: 1024})
+	if err != nil {
+		b.Fatal(err)
+	}
+	ev := benchmarkEvent()
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		es.Publish(Envelope{Event: ev})
+	}
+}
+
+func BenchmarkEventStreamPublishRead(b *testing.B) {
+	es, err := NewEventStream(EventStreamConfig{RingCapacity: 1024})
+	if err != nil {
+		b.Fatal(err)
+	}
+	reader := es.NewReader(0)
+	ev := benchmarkEvent()
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		es.Publish(Envelope{Event: ev})
+		res, ok := reader.TryRead()
+		if !ok || res.Envelope == nil {
+			b.Fatal("expected envelope")
+		}
+	}
+}