From 2f45e54e70fcec4de0d29e0f3a7f0321563bf323 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Th=C3=A9o=20Monnom?= <theo.monnom@outlook.com>
Date: Fri, 21 Aug 2026 19:02:11 -0700
Subject: [PATCH] agent endpoints: pool the per-request bridge buffers

- reuse the 32KiB response-copy buffer and the response-head reader from
  sync.Pools; ~27% fewer bytes allocated per small request.
- reuse a per-conn marshal buffer in WriteFrame instead of allocating one per
  frame; ~19% fewer bytes on a 1MiB upload. Safe because stream.Write copies
  the payload and gorilla WriteMessage does not retain it.
---
 pkg/agent/endpoint/front.go | 20 ++++++++++++++++++--
 pkg/service/agentservice.go | 12 +++++++++---
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/pkg/agent/endpoint/front.go b/pkg/agent/endpoint/front.go
index b4dbf5b6b..e136297a5 100644
--- a/pkg/agent/endpoint/front.go
+++ b/pkg/agent/endpoint/front.go
@@ -25,6 +25,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/livekit/protocol/livekit"
@@ -63,6 +64,14 @@ const (
 	maxAttempts = 3
 )
 
+// Per-request scratch is pooled: a served request would otherwise allocate a
+// 32KiB response-copy buffer and a response-head reader every time, and at high
+// request rates that dominates the front's garbage.
+var (
+	copyBufferPool   = sync.Pool{New: func() any { b := make([]byte, 32<<10); return &b }}
+	responseReadPool = sync.Pool{New: func() any { return bufio.NewReaderSize(nil, 4<<10) }}
+)
+
 // APIKeyResolver maps an inbound request to the api key it is authorized for
 // (empty when unauthenticated) - the service layer implements it from validated
 // grants.
@@ -377,7 +386,12 @@ func (f *Front) bridge(
 	}
 
 	counted := &countingReader{r: stream, n: new(int64)}
-	br := bufio.NewReader(counted)
+	br := responseReadPool.Get().(*bufio.Reader)
+	br.Reset(counted)
+	// bridge returns only after the response (or the hijacked upgrade session)
+	// is fully drained, so the reader is free to recycle here; Reset(nil) drops
+	// the stream reference so the pool never pins a dead conn.
+	defer func() { br.Reset(nil); responseReadPool.Put(br) }()
 
 	resp, err := f.readResponseHead(w, br, outReq, stream)
 	if err != nil {
@@ -405,7 +419,9 @@ func (f *Front) bridge(
 	w.WriteHeader(resp.StatusCode)
 
 	rc := http.NewResponseController(w)
-	buf := make([]byte, 32<<10)
+	bufp := copyBufferPool.Get().(*[]byte)
+	buf := *bufp
+	defer copyBufferPool.Put(bufp)
 	for {
 		n, rerr := resp.Body.Read(buf)
 		if n > 0 {
diff --git a/pkg/service/agentservice.go b/pkg/service/agentservice.go
index 8d9ec2419..7802051d4 100644
--- a/pkg/service/agentservice.go
+++ b/pkg/service/agentservice.go
@@ -275,6 +275,9 @@ const endpointWireIdleTimeout = 2 * time.Minute
 type endpointWireConn struct {
 	ws      *websocket.Conn
 	writeMu sync.Mutex
+	// scratch holds the last frame's marshalled bytes, reused across writes so a
+	// high-rate stream does not allocate a buffer per frame. Guarded by writeMu.
+	scratch []byte
 }
 
 // NewEndpointWireConn wraps an upgraded agent websocket as a data-plane wire.
@@ -288,12 +291,15 @@ func NewEndpointWireConn(ws *websocket.Conn) endpoint.WireConn {
 }
 
 func (c *endpointWireConn) WriteFrame(f *livekit.AgentHttp_Frame) error {
-	b, err := proto.Marshal(f)
+	c.writeMu.Lock()
+	defer c.writeMu.Unlock()
+	// marshal into the reused buffer and write it before returning; gorilla's
+	// WriteMessage does not retain the payload, so the next frame can reuse it.
+	b, err := proto.MarshalOptions{}.MarshalAppend(c.scratch[:0], f)
 	if err != nil {
 		return err
 	}
-	c.writeMu.Lock()
-	defer c.writeMu.Unlock()
+	c.scratch = b
 	return c.ws.WriteMessage(websocket.BinaryMessage, b)
 }