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.
This commit is contained in:
Théo Monnom
2026-08-21 19:02:11 -07:00
parent 4a9af5cc3d
commit 2f45e54e70
2 changed files with 27 additions and 5 deletions
+18 -2
View File
@@ -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 {
+9 -3
View File
@@ -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)
}