dandri/livekit
1
0 Fork 0
mirror of https://github.com/livekit/livekit.git synced 2026-10-23 08:42:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

agent endpoints: WebTransport data plane

Browse Source 
Replace the capsule mux / credit / attach-pool machinery with one WebTransport
(QUIC) session per worker: a control stream carries registration and status,
and each HTTP exchange rides its own QUIC stream (native mux, per-stream flow
control, half-close, and reset). Encode the HTTP method into the cross-node
route filter so candidate selection is method-aware. Keep the legacy WebSocket
control transport for backward compatibility.
This commit is contained in:
Théo Monnom committed 2026-08-30 18:59:30 -07:00
1 parent 1652c0d889
commit 751d85c9bf
25 files changed
+1373 -2485
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

pkg/agent/config.go
+14 -7
View File
@@ -13,13 +13,20 @@ type Config struct {
}
type EndpointsConfig struct {
// Disabled turns off the /agents/{deployment}/... front and rejects
// registrations that declare endpoints.
// Disabled turns off the /agents/{agent_name}/{deployment}/... front and
// rejects registrations that declare endpoints.
Disabled bool `yaml:"disabled,omitempty"`
DataConnCount uint32 `yaml:"data_conn_count,omitempty"`
CreditWindow uint32 `yaml:"credit_window,omitempty"`
ConnectionWindow uint32 `yaml:"connection_window,omitempty"`
MaxFrameSize uint32 `yaml:"max_frame_size,omitempty"`
MaxStreamsPerConn uint32 `yaml:"max_streams_per_conn,omitempty"`
// WebTransportPort is the UDP port for the worker WebTransport listener that
// carries the control stream and HTTP exchanges. 0 disables the listener, so
// no worker can serve endpoints (the front then only relays, in cloud).
WebTransportPort uint32 `yaml:"webtransport_port,omitempty"`
// TLSCertFile / TLSKeyFile is the certificate the WebTransport (QUIC) listener
// presents. QUIC has no plaintext mode, so a listener needs one; in dev mode a
// self-signed cert is generated when these are unset.
TLSCertFile string `yaml:"tls_cert_file,omitempty"`
TLSKeyFile string `yaml:"tls_key_file,omitempty"`
// MaxStreams is the soft per-session concurrent-stream cap used for capacity
// weighting; 0 takes the endpoint package default.
MaxStreams uint32 `yaml:"max_streams,omitempty"`
}
pkg/agent/endpoint/conformance/worker.go
+156 -697
View File
@@ -13,25 +13,27 @@
// limitations under the License.
// Package conformance is a reference worker-side implementation of the agent
// HTTP endpoints data plane, used as test infrastructure: it registers a
// manifest over the control connection, attaches the fixed wire pool, and
// bridges each stream to a local HTTP server. The acceptance suite drives it to
// HTTP endpoints data plane over WebTransport, used as test infrastructure: it
// opens one WebTransport session to /agent, registers its manifest on the
// control stream, then serves each node-opened stream by bridging the opaque
// HTTP exchange to a local HTTP server. The acceptance suite drives it to
// exercise the server end to end in Go, without a Python SDK worker.
package conformance
import (
"context"
"errors"
"crypto/tls"
"fmt"
"io"
"net"
"net/url"
"net/http"
"sync"
"time"
"github.com/gorilla/websocket"
"google.golang.org/protobuf/proto"
"github.com/quic-go/quic-go/http3"
"github.com/quic-go/webtransport-go"
"github.com/livekit/livekit-server/pkg/agent/endpoint"
"github.com/livekit/protocol/auth"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
@@ -39,7 +41,7 @@ import (
)
type Config struct {
// ServerURL is the ws(s):// URL of the /agent endpoint
// ServerURL is the https:// URL of the /agent endpoint (WebTransport).
ServerURL string
APIKey string
APISecret string
@@ -48,26 +50,26 @@ type Config struct {
Deployment string
Endpoints []*livekit.AgentHttp_AgentEndpoint
// TargetAddr is the host:port of the local HTTP server streams bridge into
// TargetAddr is the host:port of the local HTTP server streams bridge into.
TargetAddr string
// Insecure skips TLS verification (self-signed certs in tests).
Insecure bool
Logger logger.Logger
}
// Worker is one registration epoch: a control connection plus the fixed
// data-connection pool.
// Worker is one registration epoch: a single WebTransport session carrying the
// control stream and every HTTP-exchange stream.
type Worker struct {
cfg Config
instanceID string
mu sync.Mutex
sess *webtransport.Session
workerID string
settings *livekit.AgentHttp_AgentEndpointSettings
control *wsConn
dataConns []*dataConn
closed bool
registered chan struct{}
statusSeq uint64
}
func New(cfg Config) *Worker {
@@ -81,68 +83,171 @@ func New(cfg Config) *Worker {
}
}
// Start registers and attaches the pool; it returns once the data plane is up.
// Start dials the WebTransport session, registers on the control stream, and
// begins serving node-opened streams. It returns once registration completes.
func (w *Worker) Start(ctx context.Context) error {
token, err := w.mintToken()
if err != nil {
return err
}
control, err := dialWS(ctx, w.cfg.ServerURL, token, false)
tr := &webtransport.Transport{
TLSClientConfig: &tls.Config{
InsecureSkipVerify: w.cfg.Insecure, //nolint:gosec // test/self-signed only
NextProtos: []string{http3.NextProtoH3},
},
}
hdr := http.Header{}
hdr.Set("Authorization", "Bearer "+token)
_, sess, err := tr.Dial(ctx, w.cfg.ServerURL, hdr)
if err != nil {
return err
return fmt.Errorf("dial webtransport: %w", err)
}
w.mu.Lock()
w.control = control
w.sess = sess
w.mu.Unlock()
if err := control.writeWorker(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Register{
Register: &livekit.RegisterWorkerRequest{
Type: livekit.JobType_JT_ROOM,
AgentName: w.cfg.AgentName,
Version: "endpoint-conformance-client",
PingInterval: 30,
Deployment: w.cfg.Deployment,
Endpoints: w.cfg.Endpoints,
InstanceId: w.instanceID,
EndpointProtocol: 1,
// control stream: the worker opens it first, then registers
control, err := sess.OpenStreamSync(ctx)
if err != nil {
return fmt.Errorf("open control stream: %w", err)
}
if err := endpoint.WriteControlMessage(control, &livekit.WorkerMessage{
Message: &livekit.WorkerMessage_Register{
Register: &livekit.RegisterWorkerRequest{
Type: livekit.JobType_JT_ROOM,
AgentName: w.cfg.AgentName,
Version: "endpoint-conformance-client",
PingInterval: 30,
Deployment: w.cfg.Deployment,
Endpoints: w.cfg.Endpoints,
InstanceId: w.instanceID,
EndpointProtocol: 1,
},
},
}}); err != nil {
return err
}); err != nil {
return fmt.Errorf("write register: %w", err)
}
go w.controlLoop(ctx, control)
var resp livekit.ServerMessage
if err := endpoint.ReadControlMessage(control, &resp); err != nil {
return fmt.Errorf("read register response: %w", err)
}
reg := resp.GetRegister()
if reg == nil {
return fmt.Errorf("expected register response, got %T", resp.GetMessage())
}
w.mu.Lock()
w.workerID = reg.GetWorkerId()
close(w.registered)
w.mu.Unlock()
go w.controlLoop(control)
// the serve loop lives as long as the session, not the dial context (which
// the caller may cancel as soon as Start returns)
go w.serveLoop(sess)
return nil
}
// controlLoop drains further control messages (availability requests are
// declined; the conformance worker takes no jobs).
func (w *Worker) controlLoop(control *webtransport.Stream) {
for {
var msg livekit.ServerMessage
if err := endpoint.ReadControlMessage(control, &msg); err != nil {
return
}
if a := msg.GetAvailability(); a != nil {
_ = endpoint.WriteControlMessage(control, &livekit.WorkerMessage{
Message: &livekit.WorkerMessage_Availability{
Availability: &livekit.AvailabilityResponse{
JobId: a.GetJob().GetId(),
Available: false,
},
},
})
}
}
}
// serveLoop accepts node-opened streams (one HTTP exchange each) and bridges
// them to the local target. It runs until the session ends.
func (w *Worker) serveLoop(sess *webtransport.Session) {
ctx := sess.Context()
for {
stream, err := sess.AcceptStream(ctx)
if err != nil {
return
}
go w.serve(stream)
}
}
// serve bridges one HTTP exchange: the opaque request bytes on the stream are
// piped to a fresh TCP connection to the target, and the target's response
// bytes are piped back. Half-closes are propagated in both directions so
// streaming responses (SSE) flush incrementally.
func (w *Worker) serve(stream *webtransport.Stream) {
tcp, err := net.Dial("tcp", w.cfg.TargetAddr)
if err != nil {
// the request was never dispatched: refuse so the front may retry
stream.CancelWrite(endpoint.StreamCodeRefused)
stream.CancelRead(endpoint.StreamCodeRefused)
return
}
var wg sync.WaitGroup
wg.Add(2)
// request: stream -> target, then half-close the target's write side
go func() {
defer wg.Done()
_, _ = io.Copy(tcp, stream)
if c, ok := tcp.(*net.TCPConn); ok {
_ = c.CloseWrite()
}
}()
// response: target -> stream, then FIN the stream's response side
go func() {
defer wg.Done()
_, _ = io.Copy(stream, tcp)
_ = stream.Close()
}()
wg.Wait()
_ = tcp.Close()
}
// WaitRegistered blocks until registration completes or ctx is done.
func (w *Worker) WaitRegistered(ctx context.Context) error {
select {
case <-w.registered:
return nil
case <-ctx.Done():
return ctx.Err()
case <-time.After(10 * time.Second):
return errors.New("registration timeout")
return fmt.Errorf("timed out waiting for registration")
}
}
// WorkerID returns the id assigned by the server (valid after registration).
func (w *Worker) WorkerID() string {
w.mu.Lock()
settings := w.settings
workerID := w.workerID
w.mu.Unlock()
if settings == nil {
return errors.New("server did not negotiate endpoint settings")
}
defer w.mu.Unlock()
return w.workerID
}
for i := uint32(0); i < settings.GetDataConnectionCount(); i++ {
dc, err := w.attachDataConn(ctx, token, workerID, settings)
if err != nil {
return fmt.Errorf("attach %d: %w", i, err)
}
w.mu.Lock()
w.dataConns = append(w.dataConns, dc)
// Close tears down the session.
func (w *Worker) Close() {
w.mu.Lock()
if w.closed {
w.mu.Unlock()
return
}
w.closed = true
sess := w.sess
w.mu.Unlock()
if sess != nil {
_ = sess.CloseWithError(0, "worker closed")
}
// report availability so the front's load weighting sees the worker
_ = w.UpdateStatus(livekit.WorkerStatus_WS_AVAILABLE, 0, false)
return nil
}
func (w *Worker) mintToken() (string, error) {
@@ -151,649 +256,3 @@ func (w *Worker) mintToken() (string, error) {
SetValidFor(24 * time.Hour)
return at.ToJWT()
}
func (w *Worker) WorkerID() string {
w.mu.Lock()
defer w.mu.Unlock()
return w.workerID
}
func (w *Worker) UpdateStatus(status livekit.WorkerStatus, load float32, draining bool) error {
w.mu.Lock()
w.statusSeq++
seq := w.statusSeq
control := w.control
w.mu.Unlock()
if control == nil {
return errors.New("not started")
}
return control.writeWorker(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_UpdateWorker{
UpdateWorker: &livekit.UpdateWorkerStatus{
Status: &status,
Load: load,
Draining: draining,
Seq: seq,
},
}})
}
func (w *Worker) Close() {
w.mu.Lock()
if w.closed {
w.mu.Unlock()
return
}
w.closed = true
control := w.control
conns := w.dataConns
w.mu.Unlock()
if control != nil {
_ = control.close()
}
for _, dc := range conns {
dc.close(errors.New("worker closed"))
}
}
func (w *Worker) controlLoop(ctx context.Context, control *wsConn) {
for {
msg, err := control.readServer()
if err != nil {
return
}
switch m := msg.Message.(type) {
case *livekit.ServerMessage_Register:
w.mu.Lock()
w.workerID = m.Register.GetWorkerId()
w.settings = m.Register.GetEndpointSettings()
w.mu.Unlock()
close(w.registered)
case *livekit.ServerMessage_Availability:
// decline room jobs: the conformance client only serves endpoints
_ = control.writeWorker(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Availability{
Availability: &livekit.AvailabilityResponse{
JobId: m.Availability.GetJob().GetId(),
Available: false,
},
}})
case *livekit.ServerMessage_GoAway:
w.cfg.Logger.Infow("server draining, closing", "reason", m.GoAway.GetReason())
w.Close()
return
case *livekit.ServerMessage_Pong:
default:
}
}
}
// --- data connections ---
type dataConn struct {
ws *wsConn
params *livekit.AgentHttp_AttachDataConnectionResponse
target string
log logger.Logger
writeMu sync.Mutex // client side serializes writes; correctness over throughput
mu sync.Mutex
streams map[uint32]*clientStream
closed bool
// connection-level send window (response bytes toward the server), refilled
// by server credit frames on stream 0
sendMu sync.Mutex
sendCond *sync.Cond
connSendCredit int64
// connection-level receive accounting (request bytes from the server)
recvMu sync.Mutex
connRecvUnacked int64
}
func (w *Worker) attachDataConn(ctx context.Context, token, workerID string, settings *livekit.AgentHttp_AgentEndpointSettings) (*dataConn, error) {
ws, err := dialWS(ctx, w.cfg.ServerURL, token, true)
if err != nil {
return nil, err
}
if err := ws.writeFrame(&livekit.AgentHttp_Frame{
Message: &livekit.AgentHttp_Frame_Attach{Attach: &livekit.AgentHttp_AttachDataConnection{
WorkerId: workerID,
InstanceId: w.instanceID,
AttachToken: settings.GetAttachToken(),
}},
}); err != nil {
_ = ws.close()
return nil, err
}
f, err := ws.readFrame()
if err != nil {
_ = ws.close()
return nil, err
}
resp, ok := f.Message.(*livekit.AgentHttp_Frame_AttachResponse)
if !ok {
_ = ws.close()
return nil, errors.New("expected attach response")
}
if e := resp.AttachResponse.GetError(); e != "" {
_ = ws.close()
return nil, fmt.Errorf("attach rejected: %s", e)
}
dc := &dataConn{
ws: ws,
params: resp.AttachResponse,
target: w.cfg.TargetAddr,
log: w.cfg.Logger,
streams: make(map[uint32]*clientStream),
connSendCredit: int64(resp.AttachResponse.GetConnectionWindow()),
}
dc.sendCond = sync.NewCond(&dc.sendMu)
ws.enableWireLiveness()
go dc.readLoop()
go dc.pingLoop()
return dc, nil
}
// pingLoop keeps the wire's liveness visible in both directions: the server
// refreshes its read deadline on our pings, and our pong-refreshed deadline
// detects a dead server.
func (dc *dataConn) pingLoop() {
t := time.NewTicker(30 * time.Second)
defer t.Stop()
for range t.C {
dc.mu.Lock()
closed := dc.closed
dc.mu.Unlock()
if closed {
return
}
if err := dc.ws.ping(); err != nil {
return
}
}
}
// writeFrame serializes all frames FIFO: unlike the server's scheduler there is
// no control/data priority here, so a credit can queue behind sibling data. Fine
// for a conformance client; bounded by the per-write deadline.
func (dc *dataConn) writeFrame(f *livekit.AgentHttp_Frame) error {
dc.writeMu.Lock()
defer dc.writeMu.Unlock()
return dc.ws.writeFrame(f)
}
// reserveConnSend takes up to want bytes from the shared window.
func (dc *dataConn) reserveConnSend(want int64, failed func() bool) (int64, error) {
dc.sendMu.Lock()
defer dc.sendMu.Unlock()
for dc.connSendCredit <= 0 {
dc.mu.Lock()
closed := dc.closed
dc.mu.Unlock()
if closed || failed() {
return 0, errors.New("connection or stream closed")
}
dc.sendCond.Wait()
}
n := want
if n > dc.connSendCredit {
n = dc.connSendCredit
}
dc.connSendCredit -= n
return n, nil
}
// connConsumed replenishes the shared receive window once request bytes were
// consumed, threshold-acked at half the window on stream 0.
func (dc *dataConn) connConsumed(n int64) {
if n <= 0 {
return
}
dc.recvMu.Lock()
dc.connRecvUnacked += n
var credit int64
if dc.connRecvUnacked >= int64(dc.params.GetConnectionWindow())/2 {
credit = dc.connRecvUnacked
dc.connRecvUnacked = 0
}
dc.recvMu.Unlock()
if credit > 0 {
_ = dc.writeFrame(&livekit.AgentHttp_Frame{
StreamId: 0,
Message: &livekit.AgentHttp_Frame_Credit{Credit: uint32(credit)},
})
}
}
func (dc *dataConn) close(err error) {
dc.mu.Lock()
if dc.closed {
dc.mu.Unlock()
return
}
dc.closed = true
streams := make([]*clientStream, 0, len(dc.streams))
for _, s := range dc.streams {
streams = append(streams, s)
}
dc.streams = map[uint32]*clientStream{}
dc.mu.Unlock()
for _, s := range streams {
s.fail(err)
}
dc.sendMu.Lock()
dc.sendCond.Broadcast()
dc.sendMu.Unlock()
_ = dc.ws.close()
}
func (dc *dataConn) readLoop() {
for {
f, err := dc.ws.readFrame()
if err != nil {
dc.close(err)
return
}
switch m := f.Message.(type) {
case *livekit.AgentHttp_Frame_Open:
dc.handleOpen(f.StreamId, m.Open)
case *livekit.AgentHttp_Frame_Data:
delivered := false
dc.withStream(f.StreamId, func(s *clientStream) {
s.onData(m.Data)
delivered = true
})
if !delivered {
// unknown stream: protocol contract says ignore, but the bytes
// consumed the shared window and must be credited back
dc.connConsumed(int64(len(m.Data)))
}
case *livekit.AgentHttp_Frame_Eof:
dc.withStream(f.StreamId, func(s *clientStream) {
s.onEOF()
})
case *livekit.AgentHttp_Frame_Reset_:
dc.withStream(f.StreamId, func(s *clientStream) {
s.fail(fmt.Errorf("reset: %s", m.Reset_.GetError()))
})
case *livekit.AgentHttp_Frame_Credit:
if f.StreamId == 0 {
dc.sendMu.Lock()
dc.connSendCredit += int64(m.Credit)
dc.sendCond.Broadcast()
dc.sendMu.Unlock()
} else {
dc.withStream(f.StreamId, func(s *clientStream) {
s.onCredit(m.Credit)
})
}
default:
// attach frames after the handshake are a protocol violation
dc.close(errors.New("unexpected frame on data connection"))
return
}
}
}
func (dc *dataConn) withStream(id uint32, f func(*clientStream)) {
dc.mu.Lock()
s := dc.streams[id]
dc.mu.Unlock()
if s != nil {
f(s)
}
}
func (dc *dataConn) handleOpen(id uint32, open *livekit.AgentHttp_HttpStreamOpen) {
s := &clientStream{
id: id,
dc: dc,
sendCredit: int64(dc.params.GetCreditWindow()),
failCh: make(chan struct{}),
}
s.sendCond = sync.NewCond(&s.sendMu)
s.recvCond = sync.NewCond(&s.recvMu)
_ = open
dc.mu.Lock()
if dc.closed || uint32(len(dc.streams)) >= dc.params.GetMaxStreamsPerConn() {
dc.mu.Unlock()
_ = dc.writeFrame(resetFrame(id, livekit.AgentHttp_HSR_REFUSED, "no stream capacity"))
return
}
dc.streams[id] = s
dc.mu.Unlock()
go s.run()
}
func resetFrame(id uint32, code livekit.AgentHttp_HttpStreamResetCode, reason string) *livekit.AgentHttp_Frame {
return &livekit.AgentHttp_Frame{
StreamId: id,
Message: &livekit.AgentHttp_Frame_Reset_{
Reset_: &livekit.AgentHttp_HttpStreamReset{Code: code, Error: reason},
},
}
}
// clientStream bridges one stream to one TCP connection to the local HTTP
// server. The transport parses nothing.
type clientStream struct {
id uint32
dc *dataConn
sendMu sync.Mutex
sendCond *sync.Cond
sendCredit int64
sendFailed bool
// recv side: a cond-guarded buffer, never a blocking channel - the wire
// read loop must not stall behind one slow local app (bytes in flight are
// already bounded by the credit window)
recvMu sync.Mutex
recvCond *sync.Cond
recvBuf [][]byte
recvEOF bool
recvDone bool // removed or failed: arriving bytes are settled immediately
recvUnacked int64
failCh chan struct{}
failOnce sync.Once
err error
}
func (s *clientStream) fail(err error) {
s.failOnce.Do(func() {
s.recvMu.Lock()
s.err = err
s.recvDone = true
// bytes the app will never read are settled against the shared window
for _, p := range s.recvBuf {
s.dc.connConsumed(int64(len(p)))
}
s.recvBuf = nil
s.recvCond.Broadcast()
s.recvMu.Unlock()
s.sendMu.Lock()
s.sendFailed = true
s.sendCond.Broadcast()
s.sendMu.Unlock()
close(s.failCh)
})
}
func (s *clientStream) onData(payload []byte) {
s.recvMu.Lock()
if s.recvDone {
s.recvMu.Unlock()
// never consumed by the app: release the shared window immediately
s.dc.connConsumed(int64(len(payload)))
return
}
s.recvBuf = append(s.recvBuf, payload)
s.recvCond.Broadcast()
s.recvMu.Unlock()
}
func (s *clientStream) onEOF() {
s.recvMu.Lock()
s.recvEOF = true
s.recvCond.Broadcast()
s.recvMu.Unlock()
}
func (s *clientStream) onCredit(inc uint32) {
s.sendMu.Lock()
s.sendCredit += int64(inc)
s.sendCond.Broadcast()
s.sendMu.Unlock()
}
func (s *clientStream) remove() {
// mark done BEFORE unlinking so a delivery racing the removal settles the
// shared window instead of vanishing into an orphaned buffer
s.recvMu.Lock()
s.recvDone = true
for _, p := range s.recvBuf {
s.dc.connConsumed(int64(len(p)))
}
s.recvBuf = nil
s.recvCond.Broadcast()
s.recvMu.Unlock()
s.dc.mu.Lock()
delete(s.dc.streams, s.id)
s.dc.mu.Unlock()
}
// nextRecv blocks for the next request chunk; ok=false means EOF (done=false)
// or stream failure (done=true).
func (s *clientStream) nextRecv() (payload []byte, ok bool, failed bool) {
s.recvMu.Lock()
defer s.recvMu.Unlock()
for len(s.recvBuf) == 0 && !s.recvEOF && !s.recvDone {
s.recvCond.Wait()
}
if len(s.recvBuf) > 0 {
payload = s.recvBuf[0]
s.recvBuf = s.recvBuf[1:]
return payload, true, false
}
if s.recvDone {
return nil, false, true
}
return nil, false, false
}
// run executes the pump: dial the local app, copy stream->app and app->stream
// with two-level credit accounting on both directions. This is deliberately the
// report's five-line tunnel plus flow control.
func (s *clientStream) run() {
defer s.remove()
app, err := net.DialTimeout("tcp", s.dc.target, 10*time.Second)
if err != nil {
// nothing was dispatched: REFUSED tells the server a retry is safe
_ = s.dc.writeFrame(resetFrame(s.id, livekit.AgentHttp_HSR_REFUSED, "local app unreachable"))
return
}
defer app.Close()
done := make(chan struct{}, 2)
// stream -> app (request bytes)
go func() {
defer func() { done <- struct{}{} }()
for {
payload, ok, failed := s.nextRecv()
if failed {
return
}
if !ok {
// EOF after all buffered chunks drained
if tc, ok := app.(*net.TCPConn); ok {
_ = tc.CloseWrite()
}
return
}
// enforce, consume, and replenish both windows as bytes reach the app
if _, err := app.Write(payload); err != nil {
s.fail(err)
return
}
s.dc.connConsumed(int64(len(payload)))
s.recvUnacked += int64(len(payload))
if s.recvUnacked >= int64(s.dc.params.GetCreditWindow())/2 {
_ = s.dc.writeFrame(&livekit.AgentHttp_Frame{
StreamId: s.id,
Message: &livekit.AgentHttp_Frame_Credit{Credit: uint32(s.recvUnacked)},
})
s.recvUnacked = 0
}
}
}()
// app -> stream (response bytes), under the stream window and the wire's
// shared connection window
go func() {
defer func() { done <- struct{}{} }()
failed := func() bool {
select {
case <-s.failCh:
return true
default:
return false
}
}
buf := make([]byte, int(s.dc.params.GetMaxFrameSize()))
for {
n, rerr := app.Read(buf)
if n > 0 {
remaining := buf[:n]
for len(remaining) > 0 {
s.sendMu.Lock()
for s.sendCredit <= 0 && !s.sendFailed {
s.sendCond.Wait()
}
if s.sendFailed {
s.sendMu.Unlock()
return
}
want := int64(len(remaining))
if want > s.sendCredit {
want = s.sendCredit
}
s.sendMu.Unlock()
c, err := s.dc.reserveConnSend(want, failed)
if err != nil {
return
}
s.sendMu.Lock()
s.sendCredit -= c
s.sendMu.Unlock()
chunk := make([]byte, c)
copy(chunk, remaining[:c])
if err := s.dc.writeFrame(&livekit.AgentHttp_Frame{
StreamId: s.id,
Message: &livekit.AgentHttp_Frame_Data{Data: chunk},
}); err != nil {
s.fail(err)
return
}
remaining = remaining[c:]
}
}
if rerr != nil {
if rerr != io.EOF {
s.fail(rerr)
}
_ = s.dc.writeFrame(&livekit.AgentHttp_Frame{
StreamId: s.id,
Message: &livekit.AgentHttp_Frame_Eof{Eof: &livekit.AgentHttp_HttpStreamEof{}},
})
return
}
}
}()
<-done
<-done
}
// --- websocket transport ---
type wsConn struct {
ws *websocket.Conn
mu sync.Mutex
}
func dialWS(ctx context.Context, serverURL, token string, attach bool) (*wsConn, error) {
u, err := url.Parse(serverURL)
if err != nil {
return nil, err
}
q := u.Query()
q.Set("access_token", token)
if attach {
q.Set("attach", "1")
}
u.RawQuery = q.Encode()
ws, _, err := websocket.DefaultDialer.DialContext(ctx, u.String(), nil)
if err != nil {
return nil, err
}
return &wsConn{ws: ws}, nil
}
func (c *wsConn) writeWorker(msg *livekit.WorkerMessage) error {
payload, err := proto.Marshal(msg)
if err != nil {
return err
}
c.mu.Lock()
defer c.mu.Unlock()
_ = c.ws.SetWriteDeadline(time.Now().Add(30 * time.Second))
return c.ws.WriteMessage(websocket.BinaryMessage, payload)
}
func (c *wsConn) readServer() (*livekit.ServerMessage, error) {
_, payload, err := c.ws.ReadMessage()
if err != nil {
return nil, err
}
msg := &livekit.ServerMessage{}
if err := proto.Unmarshal(payload, msg); err != nil {
return nil, err
}
return msg, nil
}
func (c *wsConn) writeFrame(f *livekit.AgentHttp_Frame) error {
payload, err := proto.Marshal(f)
if err != nil {
return err
}
c.mu.Lock()
defer c.mu.Unlock()
_ = c.ws.SetWriteDeadline(time.Now().Add(30 * time.Second))
return c.ws.WriteMessage(websocket.BinaryMessage, payload)
}
func (c *wsConn) readFrame() (*livekit.AgentHttp_Frame, error) {
_ = c.ws.SetReadDeadline(time.Now().Add(wireIdleTimeout))
_, payload, err := c.ws.ReadMessage()
if err != nil {
return nil, err
}
f := &livekit.AgentHttp_Frame{}
if err := proto.Unmarshal(payload, f); err != nil {
return nil, err
}
return f, nil
}
const wireIdleTimeout = 2 * time.Minute
// enableWireLiveness arms a read deadline refreshed by any inbound traffic and
// by pongs to our pings.
func (c *wsConn) enableWireLiveness() {
_ = c.ws.SetReadDeadline(time.Now().Add(wireIdleTimeout))
c.ws.SetPongHandler(func(string) error {
return c.ws.SetReadDeadline(time.Now().Add(wireIdleTimeout))
})
}
func (c *wsConn) ping() error {
return c.ws.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second))
}
func (c *wsConn) close() error {
return c.ws.Close()
}
pkg/agent/endpoint/conn.go
-386
View File
@@ -1,386 +0,0 @@
// Copyright 2026 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package endpoint
import (
"sync"
"sync/atomic"
"time"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
)
// DataConn is one adopted worker wire carrying multiplexed streams. Only the
// server opens streams; ids are odd and scoped to the wire (even ids are
// reserved for worker-opened streams). Flow control is two-level: per-stream
// windows plus a shared connection window (stream 0 credit), so one stalled
// consumer can neither starve its siblings nor pin unbounded memory.
type DataConn struct {
wire WireConn
params WireParams
logger logger.Logger
sched *scheduler
mu sync.Mutex
streams map[uint32]*Stream
nextID uint32
closed bool
closeErr error
onClose func(*DataConn)
// connection-level send window (request bytes toward the worker), refilled
// by worker credit frames on stream 0
sendMu sync.Mutex
sendCond *sync.Cond
connSendCredit int64
// connection-level receive accounting (response bytes from the worker):
// recvConnAvail is what the worker may still put in flight, replenished on
// consumption via connConsumed
recvMu sync.Mutex
recvConnAvail int64
recvConnUnacked int64
// activity is an exponentially decayed byte counter used by stream
// placement: a connection busy moving a heavy transfer should not receive
// new streams while a lighter sibling exists.
activity atomic.Int64
lastDecay atomic.Int64 // unix nanos
decayHalfLife time.Duration
}
func NewDataConn(wire WireConn, params WireParams, onClose func(*DataConn), log logger.Logger) *DataConn {
params = params.WithDefaults()
c := &DataConn{
wire: wire,
params: params,
logger: log,
streams: make(map[uint32]*Stream),
nextID: 1, // odd ids; even reserved for worker-opened streams
onClose: onClose,
connSendCredit: int64(params.ConnectionWindow),
recvConnAvail: int64(params.ConnectionWindow),
decayHalfLife: 500 * time.Millisecond,
}
c.sendCond = sync.NewCond(&c.sendMu)
c.lastDecay.Store(time.Now().UnixNano())
c.sched = newScheduler(wire, int(params.MaxFrameSize))
go c.readLoop()
return c
}
// OpenStream opens a stream and sends its OPEN frame. It fails fast when the
// per-wire stream cap is reached; the caller picks another wire or worker.
func (c *DataConn) OpenStream(meta *livekit.AgentHttp_HttpStreamOpen) (*Stream, error) {
c.mu.Lock()
if c.closed {
err := c.closeErr
c.mu.Unlock()
if err == nil {
err = ErrConnClosed
}
return nil, err
}
if uint32(len(c.streams)) >= c.params.MaxStreamsPerConn {
c.mu.Unlock()
return nil, ErrTooManyStreams
}
id := c.nextID
c.nextID += 2
s := newStream(id, c, int64(c.params.CreditWindow))
c.streams[id] = s
c.mu.Unlock()
if err := c.sched.enqueueControl(&livekit.AgentHttp_Frame{
StreamId: id,
Message: &livekit.AgentHttp_Frame_Open{Open: meta},
}); err != nil {
c.removeStream(id)
return nil, err
}
return s, nil
}
// OpenStreams reports open streams (active and parked).
func (c *DataConn) OpenStreams() int {
c.mu.Lock()
defer c.mu.Unlock()
return len(c.streams)
}
// SpareStreams reports how many more streams this conn can accept before its
// per-conn cap; never negative.
func (c *DataConn) SpareStreams() int {
c.mu.Lock()
defer c.mu.Unlock()
if n := int(c.params.MaxStreamsPerConn) - len(c.streams); n > 0 {
return n
}
return 0
}
// HasCapacity reports whether a new stream may be opened.
func (c *DataConn) HasCapacity() bool {
c.mu.Lock()
defer c.mu.Unlock()
return !c.closed && uint32(len(c.streams)) < c.params.MaxStreamsPerConn
}
// --- connection-level send window ---
// reserveConnSend blocks until part of the shared window is available and
// reserves up to want bytes. cancelled lets a caller whose stream died stop
// waiting for credit that may never come.
func (c *DataConn) reserveConnSend(want int64, cancelled func() bool) (int64, error) {
c.sendMu.Lock()
defer c.sendMu.Unlock()
for c.connSendCredit <= 0 {
if c.closedErr() != nil {
return 0, c.closedErr()
}
if cancelled() {
return 0, ErrStreamClosed
}
c.sendCond.Wait()
}
if err := c.closedErr(); err != nil {
return 0, err
}
n := want
if n > c.connSendCredit {
n = c.connSendCredit
}
c.connSendCredit -= n
return n, nil
}
func (c *DataConn) returnConnSend(n int64) {
if n <= 0 {
return
}
c.sendMu.Lock()
c.connSendCredit += n
c.sendCond.Broadcast()
c.sendMu.Unlock()
}
func (c *DataConn) addConnSend(increment uint32) {
c.sendMu.Lock()
c.connSendCredit += int64(increment)
c.sendCond.Broadcast()
c.sendMu.Unlock()
}
// wakeSendWaiters unblocks reserveConnSend callers so a reset stream's writer
// can observe its cancellation.
func (c *DataConn) wakeSendWaiters() {
c.sendMu.Lock()
c.sendCond.Broadcast()
c.sendMu.Unlock()
}
func (c *DataConn) closedErr() error {
c.mu.Lock()
defer c.mu.Unlock()
if !c.closed {
return nil
}
if c.closeErr != nil {
return c.closeErr
}
return ErrConnClosed
}
// --- connection-level receive window ---
// consumeConnRecv accounts an arriving payload against the shared window.
func (c *DataConn) consumeConnRecv(n int64) error {
c.recvMu.Lock()
defer c.recvMu.Unlock()
if n > c.recvConnAvail {
return ErrProtocol
}
c.recvConnAvail -= n
return nil
}
// connConsumed replenishes the shared window once bytes are consumed (read by
// the client, or drained by a torn-down stream), threshold-acked at half the
// window on stream 0.
func (c *DataConn) connConsumed(n int64) {
if n <= 0 {
return
}
c.recvMu.Lock()
c.recvConnUnacked += n
var credit int64
if c.recvConnUnacked >= int64(c.params.ConnectionWindow)/2 {
credit = c.recvConnUnacked
c.recvConnUnacked = 0
c.recvConnAvail += credit
}
c.recvMu.Unlock()
if credit > 0 {
_ = c.sched.enqueueControl(&livekit.AgentHttp_Frame{
StreamId: 0,
Message: &livekit.AgentHttp_Frame_Credit{Credit: uint32(credit)},
})
}
}
// noteActivity feeds the placement score: every payload byte in either
// direction counts, decayed over time so idle (parked) streams stop weighing.
func (c *DataConn) noteActivity(n int64) {
c.decay()
c.activity.Add(n)
}
func (c *DataConn) decay() {
now := time.Now().UnixNano()
last := c.lastDecay.Load()
elapsed := time.Duration(now - last)
if elapsed < c.decayHalfLife {
return
}
if !c.lastDecay.CompareAndSwap(last, now) {
return
}
halvings := int(elapsed / c.decayHalfLife)
v := c.activity.Load()
for i := 0; i < halvings && v != 0; i++ {
v /= 2
}
c.activity.Store(v)
}
// Score is the placement weight: recent payload activity plus queued backlog.
// Lower is lighter. Heavy returns true when the connection is actively moving a
// large transfer and should not be co-located onto.
func (c *DataConn) Score() (score int64, heavy bool) {
c.decay()
a := c.activity.Load()
c.sched.mu.Lock()
q := c.sched.queuedBytes
c.sched.mu.Unlock()
score = a + q + int64(c.OpenStreams())*1024
heavy = a+q >= int64(c.params.CreditWindow)/2
return score, heavy
}
func (c *DataConn) Close(err error) {
c.mu.Lock()
if c.closed {
c.mu.Unlock()
return
}
c.closed = true
c.closeErr = err
streams := make([]*Stream, 0, len(c.streams))
for _, s := range c.streams {
streams = append(streams, s)
}
c.streams = map[uint32]*Stream{}
onClose := c.onClose
c.mu.Unlock()
if err == nil {
err = ErrConnClosed
}
for _, s := range streams {
s.mu.Lock()
if s.err == nil {
s.err = err
}
s.closed = true
s.cond.Broadcast()
s.mu.Unlock()
}
c.wakeSendWaiters()
c.sched.close(err)
if onClose != nil {
onClose(c)
}
}
func (c *DataConn) removeStream(id uint32) {
c.mu.Lock()
delete(c.streams, id)
c.mu.Unlock()
}
func (c *DataConn) stream(id uint32) *Stream {
c.mu.Lock()
defer c.mu.Unlock()
return c.streams[id]
}
func (c *DataConn) readLoop() {
for {
f, err := c.wire.ReadFrame()
if err != nil {
c.Close(err)
return
}
switch m := f.Message.(type) {
case *livekit.AgentHttp_Frame_Data:
if err := c.consumeConnRecv(int64(len(m.Data))); err != nil {
c.logger.Warnw("data plane protocol violation", err)
c.Close(err)
return
}
s := c.stream(f.StreamId)
if s == nil {
// reset locally; the payload will never be read
c.connConsumed(int64(len(m.Data)))
continue
}
c.noteActivity(int64(len(m.Data)))
if err := s.onData(m.Data); err != nil {
if err == errStreamGone {
c.connConsumed(int64(len(m.Data)))
continue
}
c.logger.Warnw("data plane protocol violation", err)
c.Close(err)
return
}
case *livekit.AgentHttp_Frame_Eof:
if s := c.stream(f.StreamId); s != nil {
s.onEOF()
}
case *livekit.AgentHttp_Frame_Reset_:
if s := c.stream(f.StreamId); s != nil {
s.onReset(m.Reset_.GetCode(), m.Reset_.GetError())
c.removeStream(f.StreamId)
c.sched.dropStream(f.StreamId)
c.wakeSendWaiters()
}
case *livekit.AgentHttp_Frame_Credit:
if f.StreamId == 0 {
c.addConnSend(m.Credit)
} else if s := c.stream(f.StreamId); s != nil {
s.onCredit(m.Credit)
}
default:
// attach handshakes ended before adoption; open only flows toward
// the worker
c.logger.Warnw("unexpected frame on data connection", ErrProtocol, "frame", logger.Proto(f))
c.Close(ErrProtocol)
return
}
}
}
pkg/agent/endpoint/control.go
+64
View File
@@ -0,0 +1,64 @@
// Copyright 2026 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package endpoint
import (
"encoding/binary"
"fmt"
"io"
"google.golang.org/protobuf/proto"
)
// MaxControlMessageSize bounds one framed control message so a malformed length
// prefix cannot make the reader allocate without bound.
const MaxControlMessageSize = 1 << 20
// WriteControlMessage writes a length-delimited protobuf on the control stream:
// a 4-byte big-endian length followed by the marshaled message. The control
// stream carries the same WorkerMessage/ServerMessage exchange the WebSocket
// control connection used; only the framing (QUIC stream instead of WS message)
// differs.
func WriteControlMessage(w io.Writer, m proto.Message) error {
b, err := proto.Marshal(m)
if err != nil {
return err
}
var hdr [4]byte
binary.BigEndian.PutUint32(hdr[:], uint32(len(b)))
if _, err := w.Write(hdr[:]); err != nil {
return err
}
_, err = w.Write(b)
return err
}
// ReadControlMessage reads one length-delimited protobuf written by
// WriteControlMessage into m.
func ReadControlMessage(r io.Reader, m proto.Message) error {
var hdr [4]byte
if _, err := io.ReadFull(r, hdr[:]); err != nil {
return err
}
n := binary.BigEndian.Uint32(hdr[:])
if n > MaxControlMessageSize {
return fmt.Errorf("control message too large: %d bytes", n)
}
b := make([]byte, n)
if _, err := io.ReadFull(r, b); err != nil {
return err
}
return proto.Unmarshal(b, m)
}
pkg/agent/endpoint/front.go
+33 -37
View File
@@ -17,7 +17,6 @@ package endpoint
import (
"bufio"
"context"
"errors"
"fmt"
"io"
"math/rand/v2"
@@ -28,9 +27,7 @@ import (
"sync"
"time"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
"github.com/livekit/protocol/utils/guid"
)
// HeaderEndpointMiss marks a response the front produced ITSELF - a routing
@@ -52,7 +49,7 @@ const (
)
const (
// PathPrefix is the public route namespace: /agents/{deployment}/{path...}
// PathPrefix is the public route namespace: /agents/{agent_name}/{deployment}/{path...}
PathPrefix = "/agents/"
// responseHeadTimeout bounds the wait for the worker's response head. Bodies
@@ -109,7 +106,12 @@ type FallbackRequest struct {
// unauthenticated)
APIKey string
Authenticated bool
Deployment string
AgentName string
Deployment string
// Method is the request's HTTP method, so a multi-node layer can select
// candidates method-aware (a node serving the path under a different method
// is not a candidate).
Method string
// Path within the deployment, '/'-rooted
Path string
}
@@ -160,6 +162,11 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) {
http.NotFound(w, r)
return
}
agentName, rest, found := strings.Cut(rest, "/")
if !found || agentName == "" {
http.NotFound(w, r)
return
}
deployment, path, found := strings.Cut(rest, "/")
if !found {
path = ""
@@ -182,7 +189,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) {
return
}
candidates := f.registry.Candidates(apiKey, deployment)
candidates := f.registry.Candidates(apiKey, agentName, deployment)
if len(candidates) == 0 && f.fallback == nil {
w.Header().Set("Retry-After", "1")
f.writeMiss(w, http.StatusServiceUnavailable, MissUnavailable, "no workers available for deployment")
@@ -215,7 +222,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// nothing local can serve: hand off before the local status mapping
if f.fallback(w, r, &FallbackRequest{
APIKey: apiKey, Authenticated: authenticated,
Deployment: deployment, Path: path,
AgentName: agentName, Deployment: deployment, Method: r.Method, Path: path,
}) {
return
}
@@ -236,7 +243,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) {
for _, reg := range candidates {
if alt, ok := reg.Manifest.RedirectSlashes(path, r.Method); ok {
u := *r.URL
u.Path = PathPrefix + deployment + alt
u.Path = PathPrefix + agentName + "/" + deployment + alt
http.Redirect(w, r, u.String(), http.StatusTemporaryRedirect)
return
}
@@ -271,7 +278,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if bodyConsumed == 0 && f.fallback != nil {
if f.fallback(w, r, &FallbackRequest{
APIKey: apiKey, Authenticated: authenticated,
Deployment: deployment, Path: path,
AgentName: agentName, Deployment: deployment, Method: r.Method, Path: path,
}) {
return
}
@@ -285,13 +292,12 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// registrations at random and take the one with fewer in-flight streams (least
// outstanding requests). This approximates optimal load spreading without global
// coordination or the herding of exact least-loaded, and - unlike a
// self-reported load - the in-flight count is observed here on this node's own
// data conns, so it is accurate for adopted (satellite) registrations too.
// Eligible = not already attempted, has attached conns, not draining.
// self-reported load - the in-flight count is observed here from the worker's
// own session. Eligible = not already attempted, has a live session, not draining.
func pickWorker(regs []*Registration, ignore map[*Registration]bool) *Registration {
eligible := regs[:0:0]
for _, reg := range regs {
if ignore[reg] || reg.AttachedConns() == 0 {
if ignore[reg] || !reg.HasSession() {
continue
}
if reg.Draining != nil && reg.Draining() {
@@ -335,23 +341,15 @@ func (f *Front) bridge(
body io.Reader,
bodyConsumedBefore int64,
) (done bool, retryable bool) {
conn := reg.PickConn()
if conn == nil {
return false, true // no capacity here; try another worker
}
stream, err := conn.OpenStream(&livekit.AgentHttp_HttpStreamOpen{
RequestId: guid.New("AER_"),
ClientAddr: r.RemoteAddr,
})
ctx := r.Context()
stream, err := reg.OpenStream(ctx)
if err != nil {
return false, true
return false, true // no session/capacity here; try another worker
}
defer stream.Close()
ctx := r.Context()
stop := context.AfterFunc(ctx, func() {
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "client disconnected")
stream.Reset(ResetCancel, "client disconnected")
})
defer stop()
@@ -365,7 +363,7 @@ func (f *Front) bridge(
err = stream.CloseWrite()
} else {
// fail fast: the worker is waiting for bytes that will never come
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "request write failed")
stream.Reset(ResetCancel, "request write failed")
}
writeErrCh <- err
}()
@@ -389,7 +387,7 @@ func (f *Front) bridge(
}
// join the request writer before another attempt touches the shared
// body reader (retries are bodyless per the table, so this is prompt)
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "retrying elsewhere")
stream.Reset(ResetCancel, "retrying elsewhere")
<-writeErrCh
return false, true
}
@@ -406,7 +404,7 @@ func (f *Front) bridge(
n, rerr := resp.Body.Read(buf)
if n > 0 {
if _, werr := w.Write(buf[:n]); werr != nil {
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "client write failed")
stream.Reset(ResetCancel, "client write failed")
return true, false
}
_ = rc.Flush()
@@ -427,16 +425,18 @@ func (f *Front) bridge(
return true, false
}
// readResponseHead reads the worker's response head, relaying informational
// responses (1xx except 101) to the client.
func (f *Front) readResponseHead(w http.ResponseWriter, br *bufio.Reader, outReq *http.Request, stream *Stream) (*http.Response, error) {
// readResponseHead reads the worker's response head, relaying 1xx informational
// responses to the client and returning the final head. A 101 is not treated as
// informational (there is no protocol upgrade over an opaque request stream): it
// is returned as-is as the final response rather than looped past.
func (f *Front) readResponseHead(w http.ResponseWriter, br *bufio.Reader, outReq *http.Request, stream Stream) (*http.Response, error) {
deadline := time.NewTimer(responseHeadTimeout)
defer deadline.Stop()
headCh := make(chan struct{})
go func() {
select {
case <-deadline.C:
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "response head timeout")
stream.Reset(ResetCancel, "response head timeout")
case <-headCh:
}
}()
@@ -463,7 +463,7 @@ func (f *Front) readResponseHead(w http.ResponseWriter, br *bufio.Reader, outReq
}
// classifyRetry implements the retry table.
func (f *Front) classifyRetry(r *http.Request, stream *Stream, responseBytes, bodyConsumedBefore int64, err error) bool {
func (f *Front) classifyRetry(r *http.Request, stream Stream, responseBytes, bodyConsumedBefore int64, err error) bool {
if responseBytes > 0 || stream.BytesRead() > 0 {
return false
}
@@ -472,9 +472,6 @@ func (f *Front) classifyRetry(r *http.Request, stream *Stream, responseBytes, bo
// request body can be replayed (nothing consumed yet)
return bodyConsumedBefore == 0 && r.ContentLength == 0
}
if errors.Is(err, ErrStreamRefused) {
return bodyConsumedBefore == 0 && r.ContentLength == 0
}
switch r.Method {
case http.MethodGet, http.MethodHead, http.MethodOptions:
return r.ContentLength == 0
@@ -549,4 +546,3 @@ func (c *countingReader) Read(p []byte) (int, error) {
*c.n += int64(n)
return n, err
}
pkg/agent/endpoint/manifest.go
+19
View File
@@ -16,6 +16,7 @@ package endpoint
import (
"fmt"
"net/http"
"slices"
"strings"
@@ -24,6 +25,21 @@ import (
const MaxManifestRoutes = 256
// allowedMethods is the set of HTTP methods a route may declare: exactly the
// verbs FastAPI (starlette) can route. Registration rejects anything else, so a
// typo ("GTE") or a garbage token can't sit in a manifest silently never
// matching. Widen this only if the worker side ever routes beyond FastAPI.
var allowedMethods = map[string]struct{}{
http.MethodGet: {},
http.MethodHead: {},
http.MethodPost: {},
http.MethodPut: {},
http.MethodPatch: {},
http.MethodDelete: {},
http.MethodOptions: {},
http.MethodTrace: {},
}
// Route is one validated manifest entry.
type Route struct {
Template *Template
@@ -77,6 +93,9 @@ func ParseManifest(endpoints []*livekit.AgentHttp_AgentEndpoint) (*Manifest, err
if u != method {
return nil, fmt.Errorf("endpoint %q method %q must be uppercase", ep.GetPath(), method)
}
if _, ok := allowedMethods[u]; !ok {
return nil, fmt.Errorf("endpoint %q declares unsupported method %q", ep.GetPath(), method)
}
methods = append(methods, u)
}
m.routes = append(m.routes, Route{
pkg/agent/endpoint/registry.go
+72 -218
View File
@@ -15,29 +15,25 @@
package endpoint
import (
"crypto/subtle"
"context"
"errors"
"slices"
"sync"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
"github.com/livekit/protocol/utils/guid"
)
// DefaultDeployment is the URL segment that addresses workers registered with an
// empty deployment name (self-hosted workers typically set none).
const DefaultDeployment = "default"
var (
ErrAttachRejected = errors.New("attach rejected")
ErrUnknownWorker = errors.New("unknown worker registration")
// ErrWrongEpoch means a registration exists for the worker but the wire's
// credentials name a different epoch: a multi-node adopter re-validates
// against the registration holder and refreshes, since the holder is the
// authority on the current epoch.
ErrWrongEpoch = errors.New("wrong registration epoch")
)
// DefaultMaxStreams is the soft per-session concurrency cap used only for
// capacity weighting; QUIC's own stream limit is the hard bound.
const DefaultMaxStreams = 256
// ErrNoSession means the registration has no live data-plane session, so no
// stream can be opened toward the worker.
var ErrNoSession = errors.New("registration has no data-plane session")
func normalizeDeployment(d string) string {
if d == "" {
@@ -46,187 +42,84 @@ func normalizeDeployment(d string) string {
return d
}
// Registration is one worker's data-plane state: its manifest, negotiated
// settings, and attached data connections. It lives exactly
// as long as the worker's control connection (epoch fencing: a reconnecting
// worker forms a new registration with a fresh attach token; connections of the
// old epoch die with it).
// Registration is one worker's data-plane state: its manifest and the single
// WebTransport session that carries both its control stream and the HTTP
// exchanges the node opens toward it. It lives exactly as long as that session
// (epoch fencing: a reconnecting worker forms a new registration, and the old
// session dies with it).
type Registration struct {
WorkerID string
InstanceID string
APIKey string
AgentName string
Deployment string
Manifest *Manifest
// Endpoints is the raw manifest as declared, kept so a multi-node layer can
// replicate the route set for remote path matching without re-deriving it.
Endpoints []*livekit.AgentHttp_AgentEndpoint
Settings Settings
Logger logger.Logger
// Draining is provided by the control-plane layer that owns the worker; a
// draining worker takes no new streams. Worker selection uses live in-flight
// streams, not a reported load, so no load hook is needed here.
// streams, not a reported load.
Draining func() bool
// pendingAttaches counts slots reserved by validated-but-unadopted wires so
// two racing attaches at the cap cannot both be acked
pendingAttaches int
mu sync.Mutex
conns []*DataConn
closed bool
mu sync.Mutex
session Session
closed bool
}
func (r *Registration) validateAttach(instanceID, token string) error {
if subtle.ConstantTimeCompare([]byte(token), []byte(r.Settings.AttachToken)) != 1 {
return ErrAttachRejected
}
if instanceID != r.InstanceID {
return ErrAttachRejected
}
// SetSession attaches the worker's live data-plane session. One session per
// worker: the WebTransport session that also carries its control stream.
func (r *Registration) SetSession(s Session) {
r.mu.Lock()
defer r.mu.Unlock()
if r.closed {
return ErrAttachRejected
}
if len(r.conns) >= int(r.Settings.DataConnCount) {
return ErrAttachRejected
}
return nil
}
// beginAttach validates the wire's credentials and atomically reserves a pool
// slot, so the attach ack can be written before adoption without two racing
// wires both being acked at the cap.
func (r *Registration) beginAttach(instanceID, token string) (*AttachTicket, error) {
if subtle.ConstantTimeCompare([]byte(token), []byte(r.Settings.AttachToken)) != 1 ||
instanceID != r.InstanceID {
return nil, ErrWrongEpoch
}
r.mu.Lock()
defer r.mu.Unlock()
if r.closed {
return nil, ErrAttachRejected
}
// the pool is fixed size: tolerate redials, but never more live conns than
// negotiated
if len(r.conns)+r.pendingAttaches >= int(r.Settings.DataConnCount) {
return nil, ErrAttachRejected
}
r.pendingAttaches++
return &AttachTicket{r: r}, nil
}
// AttachTicket is a reserved pool slot: Complete adopts the wire into it, Abort
// releases it.
type AttachTicket struct {
r *Registration
used bool
}
// Complete adopts the wire; it reports false when the registration closed while
// the ack was in flight (the caller must close the wire).
func (t *AttachTicket) Complete(wire WireConn, params WireParams) bool {
r := t.r
r.mu.Lock()
defer r.mu.Unlock()
if t.used {
return false
}
t.used = true
r.pendingAttaches--
if r.closed {
return false
}
conn := NewDataConn(wire, params, r.removeConn, r.Logger)
r.conns = append(r.conns, conn)
return true
}
func (t *AttachTicket) Abort() {
r := t.r
r.mu.Lock()
defer r.mu.Unlock()
if t.used {
return
}
t.used = true
r.pendingAttaches--
}
func (r *Registration) removeConn(c *DataConn) {
r.mu.Lock()
if i := slices.Index(r.conns, c); i != -1 {
r.conns = slices.Delete(r.conns, i, i+1)
}
r.session = s
r.mu.Unlock()
}
// PickConn places a new stream: the lightest non-heavy connection with capacity,
// falling back to spread-evenly when everything is heavy.
func (r *Registration) PickConn() *DataConn {
r.mu.Lock()
conns := slices.Clone(r.conns)
r.mu.Unlock()
var best, bestHeavy *DataConn
var bestScore, bestHeavyScore int64
for _, c := range conns {
if !c.HasCapacity() {
continue
}
score, heavy := c.Score()
if heavy {
if bestHeavy == nil || score < bestHeavyScore {
bestHeavy, bestHeavyScore = c, score
}
continue
}
if best == nil || score < bestScore {
best, bestScore = c, score
}
}
if best != nil {
return best
}
return bestHeavy
}
// AttachedConns reports live data connections.
func (r *Registration) AttachedConns() int {
func (r *Registration) getSession() Session {
r.mu.Lock()
defer r.mu.Unlock()
return len(r.conns)
if r.closed {
return nil
}
return r.session
}
// InflightStreams reports the open streams across this registration's data
// conns - the requests currently being served from this node. It is the
// least-outstanding-requests signal for worker selection: a per-node view (each
// node counts only the streams it opened), which is exactly what a node needs
// to spread its own traffic, and it needs no load reported by the worker.
// OpenStream opens a data-plane stream toward the worker for one HTTP exchange.
func (r *Registration) OpenStream(ctx context.Context) (Stream, error) {
s := r.getSession()
if s == nil {
return nil, ErrNoSession
}
return s.OpenStream(ctx)
}
// HasSession reports whether the worker has a live data-plane session, i.e. can
// serve HTTP requests right now.
func (r *Registration) HasSession() bool {
return r.getSession() != nil
}
// InflightStreams reports the open streams on this worker's session - the
// least-outstanding-requests signal for worker selection.
func (r *Registration) InflightStreams() int {
r.mu.Lock()
defer r.mu.Unlock()
n := 0
for _, c := range r.conns {
n += c.OpenStreams()
s := r.getSession()
if s == nil {
return 0
}
return n
return s.OpenStreams()
}
// SpareStreams reports the total spare stream capacity across this
// registration's data conns on this node - how many more concurrent requests it
// can accept. It is the node's live, locally-observed serving headroom for the
// worker; being measured on the conns this node holds, it is accurate for
// adopted (satellite) registrations too, which report no load.
// SpareStreams reports remaining stream capacity - the node's live serving
// headroom for the worker, used to weight node selection.
func (r *Registration) SpareStreams() int {
r.mu.Lock()
defer r.mu.Unlock()
n := 0
for _, c := range r.conns {
n += c.SpareStreams()
s := r.getSession()
if s == nil {
return 0
}
return n
if spare := s.MaxStreams() - s.OpenStreams(); spare > 0 {
return spare
}
return 0
}
func (r *Registration) close() {
@@ -236,17 +129,16 @@ func (r *Registration) close() {
return
}
r.closed = true
conns := slices.Clone(r.conns)
r.conns = nil
s := r.session
r.session = nil
r.mu.Unlock()
for _, c := range conns {
c.Close(ErrConnClosed)
if s != nil {
s.Close("registration closed")
}
}
// Registry tracks data-plane registrations on this node, keyed by
// (api key, deployment). The api key is the project identity in OSS.
// (api key, agent name, deployment). The api key is the project identity in OSS.
type Registry struct {
mu sync.Mutex
regs map[string]*Registration // by worker id
@@ -255,6 +147,7 @@ type Registry struct {
type regKey struct {
apiKey string
agentName string
deployment string
}
@@ -265,17 +158,12 @@ func NewRegistry() *Registry {
}
}
// NewAttachToken mints a registration's attach token.
func NewAttachToken() string {
return guid.New("ATT_")
}
// Register adopts a registration. A worker id already present is superseded:
// worker ids are stable across reconnects, and the retiring control connection
// must not be able to strand the new epoch (its own Deregister is a no-op once
// replaced). The superseded epoch's connections die with it.
// Register records a registration. A worker id already present is superseded:
// worker ids are stable across reconnects, and the retiring session must not be
// able to strand the new epoch (its own Deregister is a no-op once replaced).
// The superseded epoch's session is closed.
func (g *Registry) Register(r *Registration) error {
key := regKey{r.APIKey, normalizeDeployment(r.Deployment)}
key := regKey{r.APIKey, r.AgentName, normalizeDeployment(r.Deployment)}
g.mu.Lock()
old := g.regs[r.WorkerID]
if old != nil {
@@ -293,7 +181,7 @@ func (g *Registry) Register(r *Registration) error {
// removeLocked unlinks a registration from all indexes. Callers hold g.mu.
func (g *Registry) removeLocked(r *Registration) {
delete(g.regs, r.WorkerID)
key := regKey{r.APIKey, normalizeDeployment(r.Deployment)}
key := regKey{r.APIKey, r.AgentName, normalizeDeployment(r.Deployment)}
if regs := g.byKey[key]; len(regs) > 0 {
if i := slices.Index(regs, r); i != -1 {
regs = slices.Delete(regs, i, i+1)
@@ -319,45 +207,11 @@ func (g *Registry) Deregister(r *Registration) {
r.close()
}
// ValidateAttach checks an attach without adopting the connection, so the
// attach response can be written before stream frames may flow. apiKey is the
// project the connecting grant is authorized for; it must match the
// registration's, so a leaked attach token is useless without a grant for the
// worker's own project (defense in depth over the token secret).
func (g *Registry) ValidateAttach(workerID, instanceID, apiKey, token string) error {
g.mu.Lock()
r, ok := g.regs[workerID]
g.mu.Unlock()
if !ok {
return ErrUnknownWorker
}
if apiKey != r.APIKey {
return ErrAttachRejected // wrong project: terminal, never adopt-retried
}
return r.validateAttach(instanceID, token)
}
// BeginAttach validates an attach and reserves a pool slot; the caller writes
// the ack and then Completes (or Aborts) the ticket. apiKey binds the attach to
// the connecting grant's project (see ValidateAttach).
func (g *Registry) BeginAttach(workerID, instanceID, apiKey, token string) (*AttachTicket, error) {
g.mu.Lock()
r, ok := g.regs[workerID]
g.mu.Unlock()
if !ok {
return nil, ErrUnknownWorker
}
if apiKey != r.APIKey {
return nil, ErrAttachRejected // wrong project: terminal, never adopt-retried
}
return r.beginAttach(instanceID, token)
}
// Candidates returns the registrations for (api key, deployment segment).
func (g *Registry) Candidates(apiKey, deployment string) []*Registration {
// Candidates returns the registrations for (api key, agent name, deployment segment).
func (g *Registry) Candidates(apiKey, agentName, deployment string) []*Registration {
g.mu.Lock()
defer g.mu.Unlock()
return slices.Clone(g.byKey[regKey{apiKey, normalizeDeployment(deployment)}])
return slices.Clone(g.byKey[regKey{apiKey, agentName, normalizeDeployment(deployment)}])
}
// SingleAPIKey returns the api key when every registration shares one - the OSS
pkg/agent/endpoint/registry_test.go
+59 -19
View File
@@ -1,17 +1,31 @@
// Copyright 2026 LiveKit, Inc.
package endpoint
import (
"context"
"testing"
"github.com/stretchr/testify/require"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
)
// fakeSession is a no-op Session for registry tests: it records whether it was
// closed so supersede/deregister can be asserted.
type fakeSession struct {
open int
closed bool
}
func (s *fakeSession) OpenStream(context.Context) (Stream, error) { return nil, ErrNoSession }
func (s *fakeSession) OpenStreams() int { return s.open }
func (s *fakeSession) MaxStreams() int { return DefaultMaxStreams }
func (s *fakeSession) Close(string) { s.closed = true }
// worker ids are stable across reconnects: a re-registration must supersede the
// old epoch, and the retiring connection's Deregister must not strand the new
// one.
// old epoch and close its session, and the retiring session's Deregister must
// not strand the new one.
func TestRegistrySupersede(t *testing.T) {
g := NewRegistry()
manifest, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{
@@ -19,32 +33,58 @@ func TestRegistrySupersede(t *testing.T) {
}})
require.NoError(t, err)
mk := func(instance string) *Registration {
return &Registration{
WorkerID: "AW_1", InstanceID: instance, APIKey: "key",
Deployment: "production", Manifest: manifest,
Settings: Settings{DataConnCount: 1, AttachToken: "ATT_" + instance},
Logger: logger.GetLogger(),
mk := func() (*Registration, *fakeSession) {
r := &Registration{
WorkerID: "AW_1", APIKey: "key",
AgentName: "agent", Deployment: "production", Manifest: manifest,
}
s := &fakeSession{}
r.SetSession(s)
return r, s
}
oldReg := mk("i-1")
oldReg, oldSess := mk()
require.NoError(t, g.Register(oldReg))
newReg := mk("i-2")
newReg, newSess := mk()
require.NoError(t, g.Register(newReg))
require.Equal(t, []*Registration{newReg}, g.Candidates("key", "production"))
require.Error(t, g.ValidateAttach("AW_1", "i-1", "key", "ATT_i-1"), "old epoch must not attach")
require.NoError(t, g.ValidateAttach("AW_1", "i-2", "key", "ATT_i-2"))
// a grant for another project cannot attach even with the right token
require.ErrorIs(t, g.ValidateAttach("AW_1", "i-2", "other", "ATT_i-2"), ErrAttachRejected,
"wrong-project grant must not attach")
require.Equal(t, []*Registration{newReg}, g.Candidates("key", "agent", "production"))
require.True(t, oldSess.closed, "superseded epoch's session must be closed")
require.False(t, newSess.closed)
// the old control connection tears down after the new one registered
g.Deregister(oldReg)
require.Equal(t, []*Registration{newReg}, g.Candidates("key", "production"),
require.Equal(t, []*Registration{newReg}, g.Candidates("key", "agent", "production"),
"the retiring epoch must not deregister its successor")
g.Deregister(newReg)
require.Empty(t, g.Candidates("key", "production"))
require.Empty(t, g.Candidates("key", "agent", "production"))
require.True(t, newSess.closed, "deregistered session must be closed")
}
// Candidates keys on (apiKey, agentName, deployment): a request for one agent
// must not see another agent's workers in the same project/deployment.
func TestRegistryAgentScoping(t *testing.T) {
g := NewRegistry()
manifest, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{
Path: "/x", Methods: []string{"GET"}, Public: true,
}})
require.NoError(t, err)
mk := func(workerID, agentName string) *Registration {
r := &Registration{
WorkerID: workerID, APIKey: "key",
AgentName: agentName, Deployment: "production", Manifest: manifest,
}
r.SetSession(&fakeSession{})
return r
}
a := mk("AW_a", "alpha")
b := mk("AW_b", "beta")
require.NoError(t, g.Register(a))
require.NoError(t, g.Register(b))
require.Equal(t, []*Registration{a}, g.Candidates("key", "alpha", "production"))
require.Equal(t, []*Registration{b}, g.Candidates("key", "beta", "production"))
require.Empty(t, g.Candidates("key", "gamma", "production"))
}
pkg/agent/endpoint/routematch.go
+36 -20
View File
@@ -22,18 +22,20 @@ import (
// Route matching for the multi-node edge. Templates are canonicalized by
// replacing each path parameter with a typed wildcard token ({id:int} -> the
// int token, {name} or {name:str} -> the str token, {rest:path} -> the greedy
// glob), then indexed by ALL prefixes plus the full route as a terminal. A
// concrete request path is matched by a left-to-right pruned walk: at each
// segment we branch over the literal value and the wildcard tokens the value is
// glob), then indexed by ALL prefixes (path-only) plus, for each declared HTTP
// method, a method-qualified terminal route key. A concrete request (method +
// path) is matched by a left-to-right pruned walk over the path: at each segment
// we branch over the literal value and the wildcard tokens the value is
// type-compatible with, keeping only branches that are a prefix of some route -
// so the cost is linear in the path depth, never 2^segments, and the surviving
// set never exceeds the route count.
// set never exceeds the route count - then the terminal check is method-aware,
// so a node that serves the path under a different method is not a candidate.
//
// This is a SHAPE match only: the exact convertor (int/uuid/...) and the method
// are re-checked by the worker's own router, which stays authoritative. A shape
// match that the worker then rejects costs one wasted relay, never a wrong
// route; a false positive from a probabilistic backing (see FilterIndex) is the
// same bounded cost.
// This is a SHAPE match only: the exact convertor (int/uuid/...) is re-checked
// by the worker's own router, which stays authoritative (the method is now part
// of the match, but the worker still enforces it). A shape match the worker then
// rejects costs one wasted relay, never a wrong route; a false positive from a
// probabilistic backing (see FilterIndex) is the same bounded cost.
// canonical wildcard tokens; the control-char prefix keeps them from colliding
// with any real path segment
@@ -158,6 +160,15 @@ func extendKey(prefix, tok string) string {
return prefix + segSep + tok
}
// RouteKey qualifies a canonical path terminal with its HTTP method, so a
// terminal is a route only for the method(s) the worker declared. Methods are
// uppercase ASCII (allowlisted at registration) and canonical tokens start with
// a control char, so method and path can't collide. Prefix keys stay path-only
// (the pruned walk is method-agnostic until the terminal check).
func RouteKey(method, canonicalRoute string) string {
return method + segSep + canonicalRoute
}
// PrefixIndex is the backing an edge matches against: whether a canonical prefix
// key is a prefix of some route, and whether it is a complete route. An exact
// map-backed index (ExactIndex) gives precise answers; a cuckoo-filter backing
@@ -167,10 +178,12 @@ type PrefixIndex interface {
IsRoute(key string) bool
}
// Matches reports whether the request path matches any route in the index, via
// the pruned walk. A path deeper than maxMatchDepth returns true so the caller
// relays and lets the worker decide (never a spurious 404 on a deep path).
func Matches(idx PrefixIndex, path string) bool {
// Matches reports whether the request (method + path) matches any route in the
// index, via the pruned walk. The path walk is method-agnostic (prefix keys are
// path-only); only the terminal check is method-qualified. A path deeper than
// maxMatchDepth returns true so the caller relays and lets the worker decide
// (never a spurious 404 on a deep path).
func Matches(idx PrefixIndex, method, path string) bool {
segs := splitSegments(path)
if len(segs) > maxMatchDepth {
return true
@@ -180,7 +193,7 @@ func Matches(idx PrefixIndex, path string) bool {
for _, seg := range segs {
// a glob at any currently-live prefix eats this segment and the rest
for p := range live {
if idx.IsRoute(extendKey(p, tokGlob)) {
if idx.IsRoute(RouteKey(method, extendKey(p, tokGlob))) {
return true
}
}
@@ -198,10 +211,10 @@ func Matches(idx PrefixIndex, path string) bool {
}
live = next
}
// consumed every segment: a live prefix that is a complete route matches,
// and a glob at a live prefix matches the empty remainder too
// consumed every segment: a live prefix that is a complete route for this
// method matches, and a glob at a live prefix matches the empty remainder too
for p := range live {
if idx.IsRoute(p) || idx.IsRoute(extendKey(p, tokGlob)) {
if idx.IsRoute(RouteKey(method, p)) || idx.IsRoute(RouteKey(method, extendKey(p, tokGlob))) {
return true
}
}
@@ -224,8 +237,9 @@ func NewExactIndex() *ExactIndex {
}
}
// Add canonicalizes a template and inserts its prefixes and terminal.
func (x *ExactIndex) Add(path string) error {
// Add canonicalizes a template and inserts its path prefixes plus a
// method-qualified terminal per declared method.
func (x *ExactIndex) Add(methods []string, path string) error {
segs, err := canonicalizeTemplate(path)
if err != nil {
return err
@@ -235,7 +249,9 @@ func (x *ExactIndex) Add(path string) error {
key = extendKey(key, s)
x.prefixes[key] = struct{}{}
}
x.routes[key] = struct{}{}
for _, m := range methods {
x.routes[RouteKey(m, key)] = struct{}{}
}
return nil
}
pkg/agent/endpoint/routematch_test.go
+76 -44
View File
@@ -3,49 +3,57 @@
package endpoint
import (
"net/http"
"testing"
"github.com/stretchr/testify/require"
)
// idx builds an ExactIndex whose routes all serve GET; the path-shape tests
// below match with GET via the matches helper. Method-awareness has its own test.
func idx(t *testing.T, routes ...string) *ExactIndex {
x := NewExactIndex()
for _, r := range routes {
require.NoError(t, x.Add(r), r)
require.NoError(t, x.Add([]string{http.MethodGet}, r), r)
}
return x
}
// matches is a GET-method shorthand so the path-shape assertions read cleanly.
func matches(idx PrefixIndex, path string) bool {
return Matches(idx, http.MethodGet, path)
}
func TestRouteMatchStaticAndParams(t *testing.T) {
x := idx(t, "/sms", "/orders/{id}/items", "/orders/{id}/list", "/health")
// static
require.True(t, Matches(x, "/sms"))
require.True(t, Matches(x, "/health"))
require.False(t, Matches(x, "/unknown"))
require.True(t, matches(x, "/sms"))
require.True(t, matches(x, "/health"))
require.False(t, matches(x, "/unknown"))
// single-segment params, and the two sub-routes are distinguished
require.True(t, Matches(x, "/orders/42/items"))
require.True(t, Matches(x, "/orders/abc/list"))
require.False(t, Matches(x, "/orders/42/refund")) // neither sub-route exists
require.False(t, Matches(x, "/orders/42")) // prefix only, not a full route
require.False(t, Matches(x, "/orders/42/items/x"))
require.True(t, matches(x, "/orders/42/items"))
require.True(t, matches(x, "/orders/abc/list"))
require.False(t, matches(x, "/orders/42/refund")) // neither sub-route exists
require.False(t, matches(x, "/orders/42")) // prefix only, not a full route
require.False(t, matches(x, "/orders/42/items/x"))
// trailing slash matches the same shape (worker 307s)
require.True(t, Matches(x, "/orders/42/items/"))
require.True(t, matches(x, "/orders/42/items/"))
}
func TestRouteMatchTypedParams(t *testing.T) {
x := idx(t, "/orders/{id:int}", "/u/{u:uuid}")
require.True(t, Matches(x, "/orders/42")) // int
require.False(t, Matches(x, "/orders/abc")) // not an int -> edge miss, no relay
require.False(t, Matches(x, "/orders/3.14")) // float is not int
require.True(t, matches(x, "/orders/42")) // int
require.False(t, matches(x, "/orders/abc")) // not an int -> edge miss, no relay
require.False(t, matches(x, "/orders/3.14")) // float is not int
uuid := "550e8400-e29b-41d4-a716-446655440000"
require.True(t, Matches(x, "/u/"+uuid))
require.True(t, Matches(x, "/u/"+"550e8400e29b41d4a716446655440000")) // hyphenless
require.False(t, Matches(x, "/u/not-a-uuid"))
require.True(t, matches(x, "/u/"+uuid))
require.True(t, matches(x, "/u/"+"550e8400e29b41d4a716446655440000")) // hyphenless
require.False(t, matches(x, "/u/not-a-uuid"))
}
func TestRouteMatchUUIDvsStr(t *testing.T) {
@@ -53,45 +61,45 @@ func TestRouteMatchUUIDvsStr(t *testing.T) {
uuid := "550e8400-e29b-41d4-a716-446655440000"
strOnly := idx(t, "/x/{v}") // str
require.True(t, Matches(strOnly, "/x/"+uuid)) // uuid value matches a str route
require.True(t, Matches(strOnly, "/x/anything"))
require.True(t, matches(strOnly, "/x/"+uuid)) // uuid value matches a str route
require.True(t, matches(strOnly, "/x/anything"))
uuidOnly := idx(t, "/x/{v:uuid}")
require.True(t, Matches(uuidOnly, "/x/"+uuid))
require.False(t, Matches(uuidOnly, "/x/anything")) // non-uuid never matches a uuid route
require.True(t, matches(uuidOnly, "/x/"+uuid))
require.False(t, matches(uuidOnly, "/x/anything")) // non-uuid never matches a uuid route
}
func TestRouteMatchGlob(t *testing.T) {
x := idx(t, "/files/{rest:path}", "/static")
require.True(t, Matches(x, "/files/a"))
require.True(t, Matches(x, "/files/a/b/c/d.txt")) // spans segments
require.True(t, matches(x, "/files/a"))
require.True(t, matches(x, "/files/a/b/c/d.txt")) // spans segments
// /files with no trailing slash is a tolerated FALSE POSITIVE (Starlette
// requires the separating slash); the worker returns the real status
require.True(t, Matches(x, "/files"))
require.False(t, Matches(x, "/other/a/b"))
require.True(t, Matches(x, "/static"))
require.True(t, matches(x, "/files"))
require.False(t, matches(x, "/other/a/b"))
require.True(t, matches(x, "/static"))
}
// a path convertor that is not the last segment must still match (the glob is
// treated as terminal, over-approximating the suffix)
func TestRouteMatchNonTerminalGlob(t *testing.T) {
x := idx(t, "/files/{rest:path}/edit")
require.True(t, Matches(x, "/files/a/b/edit"))
require.True(t, Matches(x, "/files/a/edit"))
require.True(t, matches(x, "/files/a/b/edit"))
require.True(t, matches(x, "/files/a/edit"))
// suffix is over-approximated: a path under /files matches even without
// /edit (the worker's router returns the real 404)
require.True(t, Matches(x, "/files/a/b"))
require.False(t, Matches(x, "/other/a/edit"))
require.True(t, matches(x, "/files/a/b"))
require.False(t, matches(x, "/other/a/edit"))
}
// a path convertor mixed with literal text in a segment spans slashes and must
// not be narrowed to a single-segment str
func TestRouteMatchMixedGlob(t *testing.T) {
x := idx(t, "/files/pre{rest:path}")
require.True(t, Matches(x, "/files/prea"))
require.True(t, Matches(x, "/files/prea/b/c")) // spans segments (would drop under str)
require.False(t, Matches(x, "/other/x"))
require.True(t, matches(x, "/files/prea"))
require.True(t, matches(x, "/files/prea/b/c")) // spans segments (would drop under str)
require.False(t, matches(x, "/other/x"))
}
func TestRouteMatchHeterogeneousShapes(t *testing.T) {
@@ -99,18 +107,41 @@ func TestRouteMatchHeterogeneousShapes(t *testing.T) {
ints := idx(t, "/item/{id:int}")
strs := idx(t, "/item/{slug}")
require.True(t, Matches(ints, "/item/42"))
require.False(t, Matches(ints, "/item/foo")) // routes only to the str worker
require.True(t, Matches(strs, "/item/foo"))
require.True(t, Matches(strs, "/item/42")) // str accepts a numeric slug too
require.True(t, matches(ints, "/item/42"))
require.False(t, matches(ints, "/item/foo")) // routes only to the str worker
require.True(t, matches(strs, "/item/foo"))
require.True(t, matches(strs, "/item/42")) // str accepts a numeric slug too
}
func TestRouteMatchLiteralBeatsWildcardCoexist(t *testing.T) {
// a literal segment and a param at the same position coexist
x := idx(t, "/users/me", "/users/{id:int}")
require.True(t, Matches(x, "/users/me")) // literal
require.True(t, Matches(x, "/users/42")) // param
require.False(t, Matches(x, "/users/abc")) // neither: not "me", not an int
require.True(t, matches(x, "/users/me")) // literal
require.True(t, matches(x, "/users/42")) // param
require.False(t, matches(x, "/users/abc")) // neither: not "me", not an int
}
// the terminal is method-qualified: the SAME path shape served under different
// methods (e.g. two workers, POST vs PUT) matches only for the declared method,
// so the cross-node presence filter is method-aware.
func TestRouteMatchMethodAware(t *testing.T) {
postOnly := NewExactIndex()
require.NoError(t, postOnly.Add([]string{http.MethodPost}, "/test"))
putOnly := NewExactIndex()
require.NoError(t, putOnly.Add([]string{http.MethodPut}, "/test"))
require.True(t, Matches(postOnly, http.MethodPost, "/test"))
require.False(t, Matches(postOnly, http.MethodPut, "/test"))
require.True(t, Matches(putOnly, http.MethodPut, "/test"))
require.False(t, Matches(putOnly, http.MethodPost, "/test"))
// a route declaring several methods matches each, and only those
multi := NewExactIndex()
require.NoError(t, multi.Add([]string{http.MethodGet, http.MethodHead}, "/thing/{id:int}"))
require.True(t, Matches(multi, http.MethodGet, "/thing/5"))
require.True(t, Matches(multi, http.MethodHead, "/thing/5"))
require.False(t, Matches(multi, http.MethodDelete, "/thing/5"))
require.False(t, Matches(multi, http.MethodGet, "/thing/abc")) // still shape-checked
}
func TestRouteMatchDepthCapFallsBackToRelay(t *testing.T) {
@@ -119,16 +150,17 @@ func TestRouteMatchDepthCapFallsBackToRelay(t *testing.T) {
for i := 0; i < maxMatchDepth+5; i++ {
deep += "x/"
}
require.True(t, Matches(x, deep)) // over cap -> relay and let the worker decide
require.True(t, matches(x, deep)) // over cap -> relay and let the worker decide
}
func TestRouteDepthRejectedAtAdd(t *testing.T) {
// depth enforcement lives in the manifest layer (ParseManifest), not in the
// index: Add canonicalizes any depth without error, so the two layers stay
// decoupled.
func TestRouteDepthNotEnforcedAtAdd(t *testing.T) {
x := NewExactIndex()
long := "/"
for i := 0; i < MaxRouteDepth+2; i++ {
long += "s/"
}
// canonicalization succeeds (depth enforcement lives in the manifest layer);
// this test just documents that very deep templates still canonicalize
require.NoError(t, x.Add(long))
require.NoError(t, x.Add([]string{http.MethodGet}, long))
}
pkg/agent/endpoint/scheduler.go
-296
View File
@@ -1,296 +0,0 @@
// Copyright 2026 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package endpoint
import (
"slices"
"sync"
"time"
"github.com/livekit/protocol/livekit"
)
// scheduler serializes all writes on one wire through a single writer goroutine
// with two classes: control frames (OPEN/CREDIT/RESET) preempt DATA, and DATA
// drains one chunk per stream in round-robin order; a stream's EOF rides its
// data queue so it can never overtake the stream's own bytes. Without the class
// split, credits queue behind a deep data backlog and the peer's send windows
// never refill; without the round-robin, one heavy
// stream starves its siblings; without the budget, many streams times a full
// credit window pin unbounded memory; without the write deadline, a peer that
// stops reading freezes cancellations too.
//
// Locking rule: the wire write happens OUTSIDE the scheduler lock. Holding the
// lock across a blocking write would let one stalled connection stop enqueues
// from every stream sharing it.
// chunkPoolMinSize is the payload size at or above which a frame's buffer is
// drawn from the recycle pool. Below it (webhook-sized bodies, headers) a
// right-sized make is cheaper than borrowing a full MaxFrameSize buffer, so
// small/average requests never pay the pooled buffer's footprint.
const chunkPoolMinSize = 16 << 10
type scheduler struct {
wire WireConn
mu sync.Mutex
cond *sync.Cond
closed bool
err error
control []*livekit.AgentHttp_Frame
queues map[uint32]*streamQueue
rr []uint32 // stream ids with queued data, round-robin order
rrIdx int
queuedBytes int64 // aggregate queued DATA payload bytes, capped by connBufferBudget
// chunkPool recycles MaxFrameSize payload buffers for large frames (>=
// chunkPoolMinSize), so a big upload does not allocate one per frame. A
// buffer is drawn in Stream.Write, held in a streamQueue until its frame is
// written, then returned (in writeLoop after WriteFrame, or in dropStream).
chunkPool sync.Pool
}
type streamQueue struct {
chunks []queuedChunk
eof bool // send an EOF frame after the last chunk drains
}
// queuedChunk is one pending DATA payload. data is the bytes to send; bp is the
// pooled buffer backing them and is non-nil only for pooled (large) chunks,
// which are the ones recycled once their frame is written.
type queuedChunk struct {
data []byte
bp *[]byte
}
func newScheduler(wire WireConn, maxFrameSize int) *scheduler {
s := &scheduler{
wire: wire,
queues: make(map[uint32]*streamQueue),
chunkPool: sync.Pool{New: func() any {
b := make([]byte, maxFrameSize)
return &b
}},
}
s.cond = sync.NewCond(&s.mu)
go s.writeLoop()
return s
}
// newChunk copies p into a queued chunk, drawing a pooled MaxFrameSize buffer
// for large payloads and a right-sized make for small ones (so webhook-sized
// requests never borrow a full pooled buffer). The bytes are copied because the
// scheduler is async: the caller's buffer is reused before the frame is written.
func (s *scheduler) newChunk(p []byte) queuedChunk {
if len(p) >= chunkPoolMinSize {
bp := s.chunkPool.Get().(*[]byte)
n := copy(*bp, p)
return queuedChunk{data: (*bp)[:n], bp: bp}
}
data := make([]byte, len(p))
copy(data, p)
return queuedChunk{data: data}
}
// recycle returns a chunk's pooled buffer if it has one.
func (s *scheduler) recycle(c queuedChunk) {
if c.bp != nil {
s.chunkPool.Put(c.bp)
}
}
// enqueueControl never blocks: control frames are small and bounded by the
// number of live streams.
func (s *scheduler) enqueueControl(f *livekit.AgentHttp_Frame) error {
s.mu.Lock()
defer s.mu.Unlock()
if s.closed {
return s.errLocked()
}
s.control = append(s.control, f)
s.cond.Broadcast()
return nil
}
// enqueueData blocks while the connection buffer budget is exhausted. Stream
// and connection credit are enforced by the caller (Stream.Write); the budget
// is the local memory cap on top of them.
func (s *scheduler) enqueueData(streamID uint32, chunk queuedChunk) error {
s.mu.Lock()
defer s.mu.Unlock()
for !s.closed && s.queuedBytes+int64(len(chunk.data)) > connBufferBudget && s.queuedBytes > 0 {
s.cond.Wait()
}
if s.closed {
s.recycle(chunk) // never enqueued; return its buffer to the pool
return s.errLocked()
}
q := s.queues[streamID]
if q == nil {
q = &streamQueue{}
s.queues[streamID] = q
}
if len(q.chunks) == 0 && !q.eof {
s.rr = append(s.rr, streamID)
}
q.chunks = append(q.chunks, chunk)
s.queuedBytes += int64(len(chunk.data))
s.cond.Broadcast()
return nil
}
// enqueueEOF marks the stream's send side done; the EOF frame is emitted after
// its queued chunks, preserving stream order.
func (s *scheduler) enqueueEOF(streamID uint32) error {
s.mu.Lock()
defer s.mu.Unlock()
if s.closed {
return s.errLocked()
}
q := s.queues[streamID]
if q == nil {
q = &streamQueue{}
s.queues[streamID] = q
s.rr = append(s.rr, streamID)
} else if len(q.chunks) == 0 && !q.eof {
s.rr = append(s.rr, streamID)
}
q.eof = true
s.cond.Broadcast()
return nil
}
// dropStream discards any queued data for a reset stream. A writer blocked on
// the budget can still enqueue one late chunk after the reset; receivers
// tolerate frames for unknown stream ids by protocol contract.
func (s *scheduler) dropStream(streamID uint32) {
s.mu.Lock()
defer s.mu.Unlock()
if q, ok := s.queues[streamID]; ok {
for _, c := range q.chunks {
s.queuedBytes -= int64(len(c.data))
s.recycle(c)
}
delete(s.queues, streamID)
if i := slices.Index(s.rr, streamID); i != -1 {
s.rr = slices.Delete(s.rr, i, i+1)
if s.rrIdx > i {
s.rrIdx--
}
}
s.cond.Broadcast()
}
}
func (s *scheduler) close(err error) {
s.mu.Lock()
if s.closed {
s.mu.Unlock()
return
}
s.closed = true
s.err = err
s.cond.Broadcast()
s.mu.Unlock()
_ = s.wire.Close()
}
func (s *scheduler) errLocked() error {
if s.err != nil {
return s.err
}
return ErrConnClosed
}
// next pops the next frame to write: all pending control first, then one DATA
// chunk (or the EOF sentinel) from the round-robin ring. recycle is the pooled
// buffer backing a DATA frame, to be returned once the frame is written (nil
// for control/EOF/unpooled frames).
func (s *scheduler) next() (f *livekit.AgentHttp_Frame, recycle *[]byte, ok bool) {
s.mu.Lock()
defer s.mu.Unlock()
for {
if s.closed {
return nil, nil, false
}
if len(s.control) > 0 {
f = s.control[0]
s.control = s.control[1:]
return f, nil, true
}
if len(s.rr) > 0 {
if s.rrIdx >= len(s.rr) {
s.rrIdx = 0
}
id := s.rr[s.rrIdx]
q := s.queues[id]
if len(q.chunks) > 0 {
chunk := q.chunks[0]
q.chunks = q.chunks[1:]
s.queuedBytes -= int64(len(chunk.data))
recycle = chunk.bp
f = &livekit.AgentHttp_Frame{
StreamId: id,
Message: &livekit.AgentHttp_Frame_Data{Data: chunk.data},
}
}
if len(q.chunks) == 0 {
if q.eof {
if f == nil {
f = &livekit.AgentHttp_Frame{
StreamId: id,
Message: &livekit.AgentHttp_Frame_Eof{Eof: &livekit.AgentHttp_HttpStreamEof{}},
}
} else {
// keep the stream in the ring so the EOF drains next round
s.rrIdx++
s.cond.Broadcast()
return f, recycle, true
}
}
delete(s.queues, id)
s.rr = slices.Delete(s.rr, s.rrIdx, s.rrIdx+1)
} else {
s.rrIdx++
}
s.cond.Broadcast() // budget freed
return f, recycle, true
}
s.cond.Wait()
}
}
func (s *scheduler) writeLoop() {
for {
f, recycle, ok := s.next()
if !ok {
return
}
_ = s.wire.SetWriteDeadline(time.Now().Add(writeStallTimeout))
err := s.wire.WriteFrame(f)
// WriteFrame marshalled the payload into its own buffer and gorilla does
// not retain it, so the pooled buffer is free to recycle here whether or
// not the write succeeded.
if recycle != nil {
s.chunkPool.Put(recycle)
}
if err != nil {
s.close(err)
return
}
}
}
pkg/agent/endpoint/session.go
+78
View File
@@ -0,0 +1,78 @@
// Copyright 2026 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package endpoint
import (
"context"
"io"
)
// CurrentProtocol is the data-plane protocol version this server implements,
// negotiated in RegisterWorkerResponse.
const CurrentProtocol uint32 = 1
// ResetCode is why a stream was aborted.
type ResetCode int
const (
// ResetCancel: the node/client abandoned the exchange (client disconnect,
// timeout, retry elsewhere).
ResetCancel ResetCode = iota
// ResetRefused: the worker aborted before dispatching the request to the
// application, so the request was never applied and is safe to retry.
ResetRefused
)
// Session is the node's handle to one worker's data plane: a single
// WebTransport (QUIC) session that also carries the worker's control stream.
// The node opens a fresh Stream per HTTP exchange; QUIC provides the
// multiplexing and per-stream flow control, so there is no capsule layer.
type Session interface {
// OpenStream opens a bidirectional stream toward the worker for one HTTP
// exchange. It blocks until a stream can be opened (QUIC stream-limit
// backpressure) or ctx is done.
OpenStream(ctx context.Context) (Stream, error)
// OpenStreams reports the streams currently open on this session - the
// least-outstanding-requests signal for worker selection.
OpenStreams() int
// MaxStreams is the soft cap on concurrent streams, used to derive spare
// capacity for node-selection weighting.
MaxStreams() int
// Close tears the session down.
Close(reason string)
}
// Stream is one HTTP exchange over a Session: opaque HTTP/1.1 request bytes are
// written toward the worker and opaque response bytes are read back. CloseWrite
// half-closes the request side (the worker then sees EOF); Reset aborts both
// directions.
type Stream interface {
io.Reader
io.Writer
// CloseWrite half-closes the send side once the request is fully written
// (QUIC stream FIN).
CloseWrite() error
// Reset aborts the stream in both directions (QUIC RESET_STREAM /
// STOP_SENDING) with the given code.
Reset(code ResetCode, reason string)
// Close releases the stream after a completed exchange.
Close() error
// BytesRead reports response bytes consumed so far; the retry boundary is
// "no response byte arrived".
BytesRead() int64
// Refused reports whether the worker aborted the stream with ResetRefused,
// i.e. the request was never dispatched and is safe to retry elsewhere.
Refused() bool
}
pkg/agent/endpoint/stream.go
-366
View File
@@ -1,366 +0,0 @@
// Copyright 2026 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package endpoint
import (
"errors"
"fmt"
"io"
"sync"
"github.com/livekit/protocol/livekit"
)
// Stream is the server side of one multiplexed stream: one opaque HTTP
// exchange. Write carries request bytes toward the worker under the
// peer-granted stream window AND the wire's shared connection window; Read
// consumes response bytes and replenishes both windows as they are consumed, so
// a slow HTTP client suspends the worker's send without touching sibling
// streams.
type Stream struct {
id uint32
conn *DataConn
mu sync.Mutex
cond *sync.Cond
// send side (request bytes)
sendCredit int64
writeClosed bool
// recv side (response bytes)
recvChunks [][]byte
recvOffset int // read offset into recvChunks[0]
recvWindow int64 // bytes the peer may still send; enforced, not advisory
recvUnacked int64 // consumed bytes not yet credited back
recvEOF bool
refused bool
err error
closed bool
bytesRead int64
bytesWritten int64
}
func newStream(id uint32, conn *DataConn, window int64) *Stream {
s := &Stream{
id: id,
conn: conn,
sendCredit: window,
recvWindow: window,
}
s.cond = sync.NewCond(&s.mu)
return s
}
func (s *Stream) ID() uint32 { return s.id }
// BytesRead reports response bytes consumed so far; the retry boundary is
// "no response byte arrived".
func (s *Stream) BytesRead() int64 {
s.mu.Lock()
defer s.mu.Unlock()
return s.bytesRead
}
// Refused reports whether the worker reset the stream with HSR_REFUSED, i.e.
// the request was never dispatched and is safe to retry elsewhere.
func (s *Stream) Refused() bool {
s.mu.Lock()
defer s.mu.Unlock()
return s.refused
}
// Write sends request bytes toward the worker, blocking on the stream window,
// the wire's shared connection window and the local buffer budget.
func (s *Stream) Write(p []byte) (int, error) {
cancelled := func() bool {
s.mu.Lock()
defer s.mu.Unlock()
return s.closed || s.err != nil || s.writeClosed
}
total := 0
for len(p) > 0 {
s.mu.Lock()
for s.sendCredit <= 0 && s.err == nil && !s.closed && !s.writeClosed {
s.cond.Wait()
}
if err := s.writeErrLocked(); err != nil {
s.mu.Unlock()
return total, err
}
want := int64(len(p))
if want > s.sendCredit {
want = s.sendCredit
}
if max := int64(s.conn.params.MaxFrameSize); want > max {
want = max
}
s.mu.Unlock()
// the shared window is reserved after the stream window and outside
// s.mu (the cancelled callback re-locks it)
n, err := s.conn.reserveConnSend(want, cancelled)
if err != nil {
return total, err
}
s.mu.Lock()
if err := s.writeErrLocked(); err != nil {
s.mu.Unlock()
s.conn.returnConnSend(n)
return total, err
}
var overshoot int64
if n > s.sendCredit {
// single-writer streams never hit this; stay safe regardless
overshoot = n - s.sendCredit
n = s.sendCredit
}
s.sendCredit -= n
s.bytesWritten += n
s.mu.Unlock()
s.conn.returnConnSend(overshoot)
if n == 0 {
// a concurrent writer drained the stream window between the peek
// and the reservation: never emit an empty frame
continue
}
if err := s.conn.sched.enqueueData(s.id, s.conn.sched.newChunk(p[:n])); err != nil {
return total, err
}
s.conn.noteActivity(n)
p = p[n:]
total += int(n)
}
return total, nil
}
// ReadFrom streams r to the worker in MaxFrameSize frames. Without it,
// http.Request.Write wraps this stream in a 4KiB bufio.Writer (Stream is not an
// io.ByteWriter) and a large body would be refragmented into ~4KiB frames;
// bufio.Writer.ReadFrom delegates here once its buffer drains, so the body
// flows at the negotiated frame size instead. Implements io.ReaderFrom. Write
// still handles headers and small bodies (bufio never delegates for those).
func (s *Stream) ReadFrom(r io.Reader) (int64, error) {
buf := make([]byte, s.conn.params.MaxFrameSize)
var total int64
for {
n, rerr := r.Read(buf)
if n > 0 {
w, werr := s.Write(buf[:n])
total += int64(w)
if werr != nil {
return total, werr
}
}
if rerr != nil {
if rerr == io.EOF {
return total, nil
}
return total, rerr
}
}
}
func (s *Stream) writeErrLocked() error {
switch {
case s.err != nil:
return s.err
case s.closed:
return ErrStreamClosed
case s.writeClosed:
return fmt.Errorf("%w: write after CloseWrite", ErrStreamClosed)
}
return nil
}
// CloseWrite half-closes the send side; the worker sees EOF after the queued
// request bytes drain.
func (s *Stream) CloseWrite() error {
s.mu.Lock()
if s.writeClosed || s.closed {
s.mu.Unlock()
return nil
}
s.writeClosed = true
s.mu.Unlock()
return s.conn.sched.enqueueEOF(s.id)
}
// Read consumes response bytes. Stream and connection credit are granted back
// to the worker on consumption (threshold acking at half the respective
// window) - receipt alone never replenishes either window.
func (s *Stream) Read(p []byte) (int, error) {
s.mu.Lock()
for len(s.recvChunks) == 0 && !s.recvEOF && s.err == nil && !s.closed {
s.cond.Wait()
}
if len(s.recvChunks) == 0 {
err := s.err
if err == nil && s.recvEOF {
err = io.EOF
} else if err == nil {
err = ErrStreamClosed
}
s.mu.Unlock()
return 0, err
}
chunk := s.recvChunks[0]
n := copy(p, chunk[s.recvOffset:])
s.recvOffset += n
if s.recvOffset == len(chunk) {
s.recvChunks = s.recvChunks[1:]
s.recvOffset = 0
}
s.bytesRead += int64(n)
s.recvUnacked += int64(n)
var credit int64
if s.recvUnacked >= int64(s.conn.params.CreditWindow)/2 {
credit = s.recvUnacked
s.recvUnacked = 0
s.recvWindow += credit
}
s.mu.Unlock()
if credit > 0 {
_ = s.conn.sched.enqueueControl(&livekit.AgentHttp_Frame{
StreamId: s.id,
Message: &livekit.AgentHttp_Frame_Credit{Credit: uint32(credit)},
})
}
s.conn.connConsumed(int64(n))
s.conn.noteActivity(int64(n))
return n, nil
}
// drainUnreadLocked releases response bytes that arrived but will never be
// Read, so a torn-down stream cannot leak the wire's connection window.
// Callers hold s.mu.
func (s *Stream) drainUnreadLocked() int64 {
var pending int64
for i, c := range s.recvChunks {
pending += int64(len(c))
if i == 0 {
pending -= int64(s.recvOffset)
}
}
s.recvChunks = nil
s.recvOffset = 0
return pending
}
// Reset aborts the stream in both directions.
func (s *Stream) Reset(code livekit.AgentHttp_HttpStreamResetCode, reason string) {
s.mu.Lock()
if s.closed {
s.mu.Unlock()
return
}
s.closed = true
if s.err == nil {
s.err = ErrStreamClosed
}
pending := s.drainUnreadLocked()
s.cond.Broadcast()
s.mu.Unlock()
s.conn.connConsumed(pending)
s.conn.wakeSendWaiters()
s.conn.sched.dropStream(s.id)
_ = s.conn.sched.enqueueControl(&livekit.AgentHttp_Frame{
StreamId: s.id,
Message: &livekit.AgentHttp_Frame_Reset_{
Reset_: &livekit.AgentHttp_HttpStreamReset{Code: code, Error: reason},
},
})
s.conn.removeStream(s.id)
}
// Close releases the stream after a completed exchange.
func (s *Stream) Close() error {
s.mu.Lock()
if s.closed {
s.mu.Unlock()
return nil
}
s.closed = true
if s.err == nil {
s.err = ErrStreamClosed
}
pending := s.drainUnreadLocked()
s.cond.Broadcast()
s.mu.Unlock()
s.conn.connConsumed(pending)
s.conn.wakeSendWaiters()
s.conn.sched.dropStream(s.id)
s.conn.removeStream(s.id)
return nil
}
// --- frame delivery from the wire read loop ---
func (s *Stream) onData(payload []byte) error {
s.mu.Lock()
defer s.mu.Unlock()
if s.closed {
// stream was reset locally; the conn window is released by the caller
return errStreamGone
}
if int64(len(payload)) > s.recvWindow {
return fmt.Errorf("%w: peer exceeded credit window on stream %d", ErrProtocol, s.id)
}
s.recvWindow -= int64(len(payload))
s.recvChunks = append(s.recvChunks, payload)
s.cond.Broadcast()
return nil
}
// errStreamGone tells the read loop the payload was not adopted and its
// connection-window share must be released immediately.
var errStreamGone = errors.New("stream gone")
func (s *Stream) onEOF() {
s.mu.Lock()
s.recvEOF = true
s.cond.Broadcast()
s.mu.Unlock()
}
func (s *Stream) onReset(code livekit.AgentHttp_HttpStreamResetCode, reason string) {
s.mu.Lock()
if s.err == nil {
if code == livekit.AgentHttp_HSR_REFUSED {
s.refused = true
s.err = ErrStreamRefused
} else {
s.err = fmt.Errorf("stream reset by worker: %s (%s)", code, reason)
}
}
pending := s.drainUnreadLocked()
s.cond.Broadcast()
s.mu.Unlock()
s.conn.connConsumed(pending)
}
func (s *Stream) onCredit(increment uint32) {
s.mu.Lock()
s.sendCredit += int64(increment)
s.cond.Broadcast()
s.mu.Unlock()
}
pkg/agent/endpoint/template.go
+2 -7
View File
@@ -27,8 +27,7 @@ import (
// starlette's semantics would make the server route requests the worker-side
// router then refuses.
type Template struct {
raw string
re *regexp.Regexp
re *regexp.Regexp
}
// convertor patterns copied verbatim from starlette's convertors.py. The uuid
@@ -86,13 +85,9 @@ func ParseTemplate(path string) (*Template, error) {
if err != nil {
return nil, fmt.Errorf("invalid path template %q: %w", path, err)
}
return &Template{raw: path, re: re}, nil
return &Template{re: re}, nil
}
func (t *Template) Match(path string) bool {
return t.re.MatchString(path)
}
func (t *Template) String() string {
return t.raw
}
pkg/agent/endpoint/template_test.go
+9
View File
@@ -101,6 +101,15 @@ func TestManifestValidation(t *testing.T) {
_, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x", []string{"get"}, false)})
require.Error(t, err, "lowercase method")
_, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x", []string{"GTE"}, false)})
require.Error(t, err, "unsupported method rejected")
// the full FastAPI/starlette verb set is accepted
_, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{
ep("/x", []string{"GET", "HEAD", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "TRACE"}, false),
})
require.NoError(t, err, "all FastAPI verbs accepted")
_, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{
{Path: "/ws", Kind: livekit.AgentHttp_AEK_TEXT, Methods: []string{"GET"}},
})
pkg/agent/endpoint/webtransport.go
+160
View File
@@ -0,0 +1,160 @@
// Copyright 2026 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package endpoint
import (
"context"
"errors"
"sync"
"sync/atomic"
"github.com/quic-go/webtransport-go"
)
// Stream reset codes on the wire, shared by the server and the Go conformance
// worker so the two never drift. QUIC carries the numeric code on
// RESET_STREAM/STOP_SENDING; a worker uses Refused to mark a request it never
// dispatched (the front then treats it as safe to retry elsewhere).
const (
StreamCodeCancel webtransport.StreamErrorCode = 0
StreamCodeRefused webtransport.StreamErrorCode = 2
)
func resetToWTCode(c ResetCode) webtransport.StreamErrorCode {
if c == ResetRefused {
return StreamCodeRefused
}
return StreamCodeCancel
}
// wtSession adapts a WebTransport session to Session: the node opens one stream
// per HTTP exchange and QUIC multiplexes them.
type wtSession struct {
sess *webtransport.Session
maxStreams int
open atomic.Int64
}
// NewWebTransportSession wraps a live WebTransport session as a data-plane
// Session. maxStreams is the soft concurrency cap used only for capacity
// weighting; QUIC's own stream limit is the hard bound.
func NewWebTransportSession(sess *webtransport.Session, maxStreams int) Session {
return &wtSession{sess: sess, maxStreams: maxStreams}
}
func (s *wtSession) OpenStream(ctx context.Context) (Stream, error) {
qs, err := s.sess.OpenStreamSync(ctx)
if err != nil {
return nil, err
}
s.open.Add(1)
return &wtStream{sess: s, qs: qs}, nil
}
func (s *wtSession) OpenStreams() int { return int(s.open.Load()) }
func (s *wtSession) MaxStreams() int { return s.maxStreams }
func (s *wtSession) Close(reason string) { _ = s.sess.CloseWithError(0, reason) }
// wtStream is one HTTP exchange over a WebTransport bidi stream.
type wtStream struct {
sess *wtSession
qs *webtransport.Stream
mu sync.Mutex
bytesRead int64
refused bool
done bool
sendClosed bool // request side FIN'd (CloseWrite) or reset
}
func (s *wtStream) Read(p []byte) (int, error) {
n, err := s.qs.Read(p)
if n > 0 {
s.mu.Lock()
s.bytesRead += int64(n)
s.mu.Unlock()
}
if err != nil {
var se *webtransport.StreamError
if errors.As(err, &se) && se.ErrorCode == StreamCodeRefused {
s.mu.Lock()
s.refused = true
s.mu.Unlock()
}
}
return n, err
}
func (s *wtStream) Write(p []byte) (int, error) { return s.qs.Write(p) }
// CloseWrite sends STREAM FIN on the request side; the worker then reads EOF.
func (s *wtStream) CloseWrite() error {
s.mu.Lock()
s.sendClosed = true
s.mu.Unlock()
return s.qs.Close()
}
func (s *wtStream) Reset(code ResetCode, _ string) {
c := resetToWTCode(code)
s.mu.Lock()
s.sendClosed = true
s.mu.Unlock()
s.qs.CancelWrite(c)
s.qs.CancelRead(c)
s.release()
}
func (s *wtStream) Close() error {
// If the request side was never cleanly FIN'd (CloseWrite) or reset, a writer
// goroutine may still be blocked in Write because the worker stopped reading:
// cancel the send side to unblock it (else the goroutine and QUIC stream leak).
// After a clean FIN we must NOT reset - that would turn a completed request
// into an abort on the wire.
s.mu.Lock()
cancelWrite := !s.sendClosed
s.sendClosed = true
s.mu.Unlock()
if cancelWrite {
s.qs.CancelWrite(StreamCodeCancel)
}
s.qs.CancelRead(StreamCodeCancel)
s.release()
return nil
}
// release decrements the session's open-stream count exactly once.
func (s *wtStream) release() {
s.mu.Lock()
if s.done {
s.mu.Unlock()
return
}
s.done = true
s.mu.Unlock()
s.sess.open.Add(-1)
}
func (s *wtStream) BytesRead() int64 {
s.mu.Lock()
defer s.mu.Unlock()
return s.bytesRead
}
func (s *wtStream) Refused() bool {
s.mu.Lock()
defer s.mu.Unlock()
return s.refused
}
pkg/agent/endpoint/webtransport_test.go
+184
View File
@@ -0,0 +1,184 @@
// Copyright 2026 LiveKit, Inc.
package endpoint_test
import (
"context"
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/tls"
"crypto/x509"
"crypto/x509/pkix"
"io"
"math/big"
"net"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/quic-go/quic-go/http3"
"github.com/quic-go/webtransport-go"
"github.com/stretchr/testify/require"
"github.com/livekit/livekit-server/pkg/agent/endpoint"
"github.com/livekit/livekit-server/pkg/agent/endpoint/conformance"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
)
func selfSignedTLS(t *testing.T) *tls.Config {
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
require.NoError(t, err)
tmpl := &x509.Certificate{
SerialNumber: big.NewInt(1),
Subject: pkix.Name{CommonName: "localhost"},
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(time.Hour),
IPAddresses: []net.IP{net.ParseIP("127.0.0.1")},
}
der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
require.NoError(t, err)
return &tls.Config{
Certificates: []tls.Certificate{{Certificate: [][]byte{der}, PrivateKey: key}},
NextProtos: []string{http3.NextProtoH3},
}
}
// startWTServer runs a WebTransport /agent server that registers each session's
// worker (read off the control stream) into reg and keeps the session alive.
func startWTServer(t *testing.T, reg *endpoint.Registry) string {
udp, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
require.NoError(t, err)
mux := http.NewServeMux()
srv := &webtransport.Server{H3: &http3.Server{TLSConfig: selfSignedTLS(t), Handler: mux}}
mux.HandleFunc("/agent", func(w http.ResponseWriter, r *http.Request) {
sess, err := srv.Upgrade(w, r)
if err != nil {
return
}
go handleSession(reg, sess)
})
go func() { _ = srv.Serve(udp) }()
t.Cleanup(func() { _ = srv.Close(); _ = udp.Close() })
return "https://" + udp.LocalAddr().String() + "/agent"
}
func handleSession(reg *endpoint.Registry, sess *webtransport.Session) {
ctx := sess.Context()
control, err := sess.AcceptStream(ctx)
if err != nil {
return
}
var msg livekit.WorkerMessage
if err := endpoint.ReadControlMessage(control, &msg); err != nil {
return
}
rw := msg.GetRegister()
if rw == nil {
return
}
manifest, err := endpoint.ParseManifest(rw.GetEndpoints())
if err != nil {
return
}
registration := &endpoint.Registration{
WorkerID: rw.GetInstanceId(),
APIKey: "test",
AgentName: rw.GetAgentName(),
Deployment: rw.GetDeployment(),
Manifest: manifest,
Endpoints: rw.GetEndpoints(),
}
registration.SetSession(endpoint.NewWebTransportSession(sess, endpoint.DefaultMaxStreams))
_ = reg.Register(registration)
_ = endpoint.WriteControlMessage(control, &livekit.ServerMessage{
Message: &livekit.ServerMessage_Register{
Register: &livekit.RegisterWorkerResponse{WorkerId: registration.WorkerID},
},
})
// keep the session alive; drain further control messages until it dies
for {
var m livekit.WorkerMessage
if err := endpoint.ReadControlMessage(control, &m); err != nil {
return
}
}
}
func TestWebTransportEndpointRoundTrip(t *testing.T) {
// the worker's local app
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/hello":
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"ok":true}`)
case "/echo":
body, _ := io.ReadAll(r.Body)
_, _ = w.Write(body)
default:
http.NotFound(w, r)
}
}))
defer target.Close()
reg := endpoint.NewRegistry()
base := startWTServer(t, reg)
front := endpoint.NewFront(reg, func(*http.Request) (string, bool) { return "", false }, logger.GetLogger()).
WithSingleKeyFallback()
ts := httptest.NewServer(front)
defer ts.Close()
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
w := conformance.New(conformance.Config{
ServerURL: base,
APIKey: "APIkey",
APISecret: "secret-that-is-long-enough-to-sign",
AgentName: "myagent",
Deployment: "production",
TargetAddr: strings.TrimPrefix(target.URL, "http://"),
Insecure: true,
Endpoints: []*livekit.AgentHttp_AgentEndpoint{
{Path: "/hello", Methods: []string{"GET"}, Public: true},
{Path: "/echo", Methods: []string{"POST"}, Public: true},
},
})
require.NoError(t, w.Start(ctx))
t.Cleanup(w.Close)
require.NoError(t, w.WaitRegistered(ctx))
agentBase := ts.URL + "/agents/myagent/production"
t.Run("GET json", func(t *testing.T) {
resp, err := http.Get(agentBase + "/hello")
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
body, _ := io.ReadAll(resp.Body)
require.JSONEq(t, `{"ok":true}`, string(body))
})
t.Run("POST echo", func(t *testing.T) {
payload := strings.Repeat("x", 64<<10)
resp, err := http.Post(agentBase+"/echo", "application/octet-stream", strings.NewReader(payload))
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
body, _ := io.ReadAll(resp.Body)
require.Equal(t, len(payload), len(body))
require.Equal(t, payload, string(body))
})
t.Run("unknown path 404", func(t *testing.T) {
resp, err := http.Get(agentBase + "/nope")
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusNotFound, resp.StatusCode)
})
}
pkg/agent/endpoint/wire.go
-122
View File
@@ -1,122 +0,0 @@
// Copyright 2026 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// Package endpoint implements the server side of the agent HTTP endpoints data
// plane: worker-dialed wires carrying multiplexed streams, each stream one
// opaque HTTP exchange, flow-controlled per stream and per wire with a
// prioritized per-wire write scheduler. Wires speak AgentHttp.Frame
// exclusively: one frame per websocket binary message.
package endpoint
import (
"errors"
"time"
"github.com/livekit/protocol/livekit"
)
// CurrentProtocol is the data-plane protocol version negotiated in
// RegisterWorkerResponse.endpoint_settings.
const CurrentProtocol = 1
var (
ErrConnClosed = errors.New("data connection closed")
ErrStreamClosed = errors.New("stream closed")
ErrStreamRefused = errors.New("stream refused by worker")
ErrTooManyStreams = errors.New("too many open streams on connection")
ErrProtocol = errors.New("data plane protocol violation")
)
// WireConn is the transport under one data wire: whole binary websocket
// messages in, one Frame each. The write deadline is what turns a peer that
// stopped reading into a dead connection instead of a stuck writer goroutine.
type WireConn interface {
WriteFrame(f *livekit.AgentHttp_Frame) error
ReadFrame() (*livekit.AgentHttp_Frame, error)
SetWriteDeadline(t time.Time) error
SetReadDeadline(t time.Time) error
Close() error
}
// Settings are the registration-level parameters negotiated on the control
// connection. Wire-level parameters ride each wire's attach response instead:
// wires may be adopted by any node, and the adopting node's parameters govern.
type Settings struct {
Protocol uint32
AttachToken string
DataConnCount uint32
}
// WireParams are one wire's flow-control parameters, chosen by the node that
// adopted it and announced in AttachDataConnectionResponse.
type WireParams struct {
CreditWindow uint32
ConnectionWindow uint32
MaxFrameSize uint32
MaxStreamsPerConn uint32
}
const (
DefaultDataConnCount = 8
DefaultCreditWindow = 1 << 20 // 1MiB per stream
DefaultConnectionWindow = 4 << 20 // shared by all streams on a wire
DefaultMaxFrameSize = 64 << 10
DefaultMaxStreamsPerConn = 64 // total open; parked streams included
// writeStallTimeout bounds a single wire write; a peer that stops reading
// kills the connection instead of freezing its sibling streams.
writeStallTimeout = 30 * time.Second
// connBufferBudget caps the aggregate queued-but-unwritten bytes per
// connection so many streams times a full window cannot pin unbounded memory.
connBufferBudget = 4 << 20
)
// WithDefaults fills zero fields with the package defaults.
func (p WireParams) WithDefaults() WireParams {
if p.CreditWindow == 0 {
p.CreditWindow = DefaultCreditWindow
}
if p.ConnectionWindow == 0 {
p.ConnectionWindow = DefaultConnectionWindow
}
if p.MaxFrameSize == 0 {
p.MaxFrameSize = DefaultMaxFrameSize
}
if p.MaxStreamsPerConn == 0 {
p.MaxStreamsPerConn = DefaultMaxStreamsPerConn
}
return p
}
// frameOverhead bounds the AgentHttp.Frame envelope around a payload (stream id,
// oneof tag, length varints); a data frame on the wire is at most
// MaxFrameSize + frameOverhead.
const frameOverhead = 1024
// MaxMessageSize is the largest websocket message a wire may legitimately carry
// (a full data frame). A reader caps its read limit at this so an oversized
// frame is rejected before it can exhaust memory, even when no signalling limit
// is configured.
func (p WireParams) MaxMessageSize() int64 {
return int64(p.MaxFrameSize) + frameOverhead
}
func (s Settings) Proto() *livekit.AgentHttp_AgentEndpointSettings {
return &livekit.AgentHttp_AgentEndpointSettings{
Protocol: s.Protocol,
AttachToken: s.AttachToken,
DataConnectionCount: s.DataConnCount,
}
}
pkg/agent/worker.go
+8 -2
View File
@@ -159,7 +159,6 @@ type WorkerRegistration struct {
// agent HTTP endpoints data plane
Endpoints []*livekit.AgentHttp_AgentEndpoint
InstanceID string
EndpointProtocol uint32
EndpointSettings *livekit.AgentHttp_AgentEndpointSettings
}
@@ -243,12 +242,19 @@ func (h *WorkerRegisterer) HandleRegister(req *livekit.RegisterWorkerRequest) er
h.registration.Deployment = req.GetDeployment()
h.registration.Endpoints = req.GetEndpoints()
h.registration.InstanceID = req.GetInstanceId()
h.registration.EndpointProtocol = req.GetEndpointProtocol()
if len(req.GetEndpoints()) > 0 && req.GetEndpointProtocol() > 0 {
if h.endpointSettings == nil {
return errors.New("agent HTTP endpoints are not supported by this server")
}
// endpoints are addressed at /agents/{agent_name}/{deployment}/...; a
// worker serving them needs a non-empty, URL-safe agent name
if req.GetAgentName() == "" {
return errors.New("agent HTTP endpoints require an agent name")
}
if err := protoagent.ValidateAgentName(req.GetAgentName()); err != nil {
return err
}
settings, err := h.endpointSettings(req)
if err != nil {
return err