Optimize telemetry event publishing

Author: IlyaasK

server/lib/events/event.go

@@ -12,6 +12,10 @@ import (
 // maxS2RecordBytes is the maximum record size for the S2 event pipeline (1 MB).
 const maxS2RecordBytes = 1_000_000
 
+// Below this estimated size, normal JSON overhead cannot push a record near the
+// S2 limit, so Publish can avoid a full JSON marshal on the hot path.
+const truncateMarshalThreshold = maxS2RecordBytes - 64*1024
+
 const (
 	Console     = oapi.TelemetryEventCategory("console")
 	Network     = oapi.TelemetryEventCategory("network")
@@ -97,8 +101,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 +120,20 @@ 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 := len(env.Event.Data) + len(env.Event.Type) + len(env.Event.Category) + len(env.Event.Source.Kind)
+	if env.Event.Source.Event != nil {
+		n += len(*env.Event.Source.Event)
+	}
+	if env.Event.Source.Metadata != nil {
+		for k, v := range *env.Event.Source.Metadata {
+			n += len(k) + len(v)
+		}
+	}
+	return n
+}

server/lib/events/events_test.go

@@ -3,6 +3,7 @@ package events
 import (
 	"context"
 	"encoding/json"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -118,6 +119,25 @@ 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 mkEnv(seq uint64, ev Event) Envelope {
 	return Envelope{Seq: seq, Event: ev}
 }
@@ -133,6 +153,24 @@ 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)
+}
+
 // 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,8 @@ 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 +38,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")
+		}
+	}
+}

server/lib/events/ringbuffer.go

@@ -13,7 +13,8 @@ type ringBuffer struct {
 	buf        []Envelope
 	cap        uint64
 	latestSeq  uint64        // highest envelope.Seq published
-	readerWake chan struct{} // closed-and-replaced on each Publish to wake blocked readers
+	readerWake chan struct{} // closed-and-replaced when blocked readers need a wakeup
+	waiters    int
 }
 
 func newRingBuffer(capacity int) (*ringBuffer, error) {
@@ -36,21 +37,55 @@ func (rb *ringBuffer) reset() {
 		rb.buf[i] = Envelope{}
 	}
 	rb.latestSeq = 0
-	old := rb.readerWake
-	rb.readerWake = make(chan struct{})
+	var old chan struct{}
+	if rb.waiters > 0 {
+		old = rb.readerWake
+		rb.readerWake = make(chan struct{})
+	}
 	rb.mu.Unlock()
-	close(old)
+	if old != nil {
+		close(old)
+	}
 }
 
-// publish adds an envelope to the ring, evicting the oldest on overflow.
-func (rb *ringBuffer) publish(env Envelope) {
-	rb.mu.Lock()
+func (rb *ringBuffer) publishLocked(env Envelope) chan struct{} {
 	rb.buf[env.Seq%rb.cap] = env
 	rb.latestSeq = env.Seq
+	if rb.waiters == 0 {
+		return nil
+	}
 	old := rb.readerWake
 	rb.readerWake = make(chan struct{})
+	return old
+}
+
+func (rb *ringBuffer) closeWake(old chan struct{}) {
+	if old != nil {
+		close(old)
+	}
+}
+
+// publish adds an envelope to the ring, evicting the oldest on overflow.
+func (rb *ringBuffer) publish(env Envelope) {
+	rb.mu.Lock()
+	old := rb.publishLocked(env)
+	rb.mu.Unlock()
+	rb.closeWake(old)
+}
+
+func (rb *ringBuffer) publishNext(env Envelope) Envelope {
+	rb.mu.Lock()
+	env.Seq = rb.latestSeq + 1
+	old := rb.publishLocked(env)
 	rb.mu.Unlock()
-	close(old)
+	rb.closeWake(old)
+	return env
+}
+
+func (rb *ringBuffer) seq() uint64 {
+	rb.mu.RLock()
+	defer rb.mu.RUnlock()
+	return rb.latestSeq
 }
 
 func (rb *ringBuffer) oldestSeq() uint64 {
@@ -108,7 +143,7 @@ func (r *Reader) TryRead() (ReadResult, bool) {
 // Read blocks until the next envelope is available or ctx is cancelled.
 func (r *Reader) Read(ctx context.Context) (ReadResult, error) {
 	for {
-		r.rb.mu.RLock()
+		r.rb.mu.Lock()
 		wake := r.rb.readerWake
 		latest := r.rb.latestSeq
 		oldest := r.rb.oldestSeq()
@@ -117,35 +152,49 @@ func (r *Reader) Read(ctx context.Context) (ReadResult, error) {
 			// Buffer is empty (or was just reset). Reset reader position
 			// so it starts from the beginning when new data arrives.
 			r.nextSeq = 1
-			r.rb.mu.RUnlock()
-			select {
-			case <-ctx.Done():
+			r.rb.waiters++
+			r.rb.mu.Unlock()
+			err := waitForWake(ctx, wake)
+			r.rb.mu.Lock()
+			r.rb.waiters--
+			r.rb.mu.Unlock()
+			if err != nil {
 				return ReadResult{}, ctx.Err()
-			case <-wake:
-				continue
 			}
+			continue
 		}
 
 		if r.nextSeq < oldest {
 			dropped := oldest - r.nextSeq
 			r.nextSeq = oldest
-			r.rb.mu.RUnlock()
+			r.rb.mu.Unlock()
 			return ReadResult{Dropped: dropped}, nil
 		}
 
 		if r.nextSeq <= latest {
 			env := r.rb.buf[r.nextSeq%r.rb.cap]
 			r.nextSeq++
-			r.rb.mu.RUnlock()
+			r.rb.mu.Unlock()
 			return ReadResult{Envelope: &env}, nil
 		}
 
-		r.rb.mu.RUnlock()
-
-		select {
-		case <-ctx.Done():
+		r.rb.waiters++
+		r.rb.mu.Unlock()
+		err := waitForWake(ctx, wake)
+		r.rb.mu.Lock()
+		r.rb.waiters--
+		r.rb.mu.Unlock()
+		if err != nil {
 			return ReadResult{}, ctx.Err()
-		case <-wake:
 		}
 	}
 }
+
+func waitForWake(ctx context.Context, wake <-chan struct{}) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-wake:
+		return nil
+	}
+}

server/lib/telemetry/telemetry.go

@@ -31,6 +31,7 @@ type TelemetrySession struct {
 	id              string
 	sessionStartSeq uint64
 	categories      map[oapi.TelemetryEventCategory]struct{}
+	sessionMetadata *map[string]string
 	appliedAt       time.Time
 }
 
@@ -68,6 +69,8 @@ func (s *TelemetrySession) Start(telemetrySessionID string, cfg TelemetryConfig)
 	s.sessionStartSeq = s.es.Seq()
 	s.appliedAt = time.Now()
 	s.categories = categorySet(cfg.Categories)
+	metadata := map[string]string{"telemetry_session_id": telemetrySessionID}
+	s.sessionMetadata = &metadata
 }
 
 // publishLocked stamps telemetry_session_id into ev.Source.Metadata and forwards to the bus.
@@ -77,10 +80,10 @@ func (s *TelemetrySession) publishLocked(ev events.Event) events.Envelope {
 		ev.Ts = time.Now().UnixMicro()
 	}
 	if ev.Source.Metadata == nil {
-		m := make(map[string]string)
-		ev.Source.Metadata = &m
+		ev.Source.Metadata = s.sessionMetadata
+	} else {
+		(*ev.Source.Metadata)["telemetry_session_id"] = s.id
 	}
-	(*ev.Source.Metadata)["telemetry_session_id"] = s.id
 	return s.es.Publish(events.Envelope{Event: ev})
 }
 
@@ -176,5 +179,6 @@ func (s *TelemetrySession) Stop() {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	s.id = ""
+	s.sessionMetadata = nil
 	s.appliedAt = time.Time{}
 }