mirror of
https://github.com/livekit/livekit.git
synced 2026-09-13 00:25:35 +00:00
agent: HTTP endpoints data plane
Serve worker-declared FastAPI routes at /agents/{deployment}/{path} without
any worker-side listener: workers dial a fixed pool of wires speaking
AgentHttp.Frame, the server opens multiplexed streams carrying one opaque
HTTP/1.1 exchange each (two-level credit flow control, prioritized write
scheduler, attach epoch fencing). The front does starlette-exact manifest
matching with per-endpoint public access, typed 401/404/405, a retry table,
and SSE/WebSocket passthrough; a pluggable fallback hook lets multi-node
deployments resolve misses elsewhere. Includes a conformance client, a
manual sidecar, and the acceptance suite.
The auth middleware now reads access_token from the query string only:
FormValue consumed the bodies of proxied url-encoded POSTs.
This commit is contained in:
@@ -0,0 +1,111 @@
|
||||
// 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.
|
||||
|
||||
// agent-endpoint-client is the reference sidecar for the agent HTTP endpoints
|
||||
// data plane: it registers a manifest against a livekit-server and bridges
|
||||
// tunnel streams to any local HTTP server, standing in for the SDK's tunnel
|
||||
// client until it lands.
|
||||
//
|
||||
// Example, against a dev server:
|
||||
//
|
||||
// livekit-server --dev &
|
||||
// python app.py # any local HTTP server on :8080
|
||||
// agent-endpoint-client -url ws://localhost:7880/agent -api-key devkey \
|
||||
// -api-secret secret -deployment production -target 127.0.0.1:8080 \
|
||||
// -route "GET /json" -route "public GET /sse"
|
||||
// curl http://localhost:7880/agents/production/json
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strings"
|
||||
"syscall"
|
||||
|
||||
"github.com/livekit/protocol/livekit"
|
||||
"github.com/livekit/protocol/logger"
|
||||
|
||||
"github.com/livekit/livekit-server/pkg/agent/endpoint/client"
|
||||
)
|
||||
|
||||
type routeFlags []string
|
||||
|
||||
func (r *routeFlags) String() string { return strings.Join(*r, ",") }
|
||||
func (r *routeFlags) Set(v string) error { *r = append(*r, v); return nil }
|
||||
|
||||
func main() {
|
||||
var (
|
||||
url = flag.String("url", "ws://localhost:7880/agent", "livekit-server /agent URL")
|
||||
apiKey = flag.String("api-key", "devkey", "API key")
|
||||
apiSecret = flag.String("api-secret", "secret", "API secret")
|
||||
agentName = flag.String("agent-name", "endpoint-sidecar", "agent name")
|
||||
deployment = flag.String("deployment", "", "deployment name (empty = default)")
|
||||
target = flag.String("target", "127.0.0.1:8080", "local HTTP server to bridge into")
|
||||
routes routeFlags
|
||||
)
|
||||
flag.Var(&routes, "route", `route to expose, e.g. "GET /json", "public POST /sms", repeatable`)
|
||||
flag.Parse()
|
||||
|
||||
if len(routes) == 0 {
|
||||
fmt.Fprintln(os.Stderr, "at least one -route is required")
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
var endpoints []*livekit.AgentHttp_AgentEndpoint
|
||||
for _, r := range routes {
|
||||
parts := strings.Fields(r)
|
||||
public := false
|
||||
if len(parts) > 0 && parts[0] == "public" {
|
||||
public = true
|
||||
parts = parts[1:]
|
||||
}
|
||||
if len(parts) != 2 {
|
||||
fmt.Fprintf(os.Stderr, "invalid -route %q, want \"[public] METHOD /path\"\n", r)
|
||||
os.Exit(1)
|
||||
}
|
||||
endpoints = append(endpoints, &livekit.AgentHttp_AgentEndpoint{
|
||||
Path: parts[1],
|
||||
Methods: []string{strings.ToUpper(parts[0])},
|
||||
Public: public,
|
||||
})
|
||||
}
|
||||
|
||||
logger.InitFromConfig(&logger.Config{Level: "info"}, "agent-endpoint-client")
|
||||
|
||||
w := client.New(client.Config{
|
||||
ServerURL: *url,
|
||||
APIKey: *apiKey,
|
||||
APISecret: *apiSecret,
|
||||
AgentName: *agentName,
|
||||
Deployment: *deployment,
|
||||
Endpoints: endpoints,
|
||||
TargetAddr: *target,
|
||||
Logger: logger.GetLogger(),
|
||||
})
|
||||
|
||||
ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer cancel()
|
||||
|
||||
if err := w.Start(ctx); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "start failed:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
logger.Infow("endpoint sidecar attached", "workerID", w.WorkerID(), "routes", len(endpoints))
|
||||
|
||||
<-ctx.Done()
|
||||
w.Close()
|
||||
}
|
||||
@@ -21,7 +21,7 @@ require (
|
||||
github.com/jxskiss/base62 v1.1.0
|
||||
github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731
|
||||
github.com/livekit/mediatransportutil v0.0.0-20260608063931-a3417d38cda0
|
||||
github.com/livekit/protocol v1.50.5-0.20260814120900-8b1ab81c7d00
|
||||
github.com/livekit/protocol v1.50.5-0.20260820011432-2480d740da5e
|
||||
github.com/livekit/psrpc v0.7.3
|
||||
github.com/mackerelio/go-osstat v0.2.8
|
||||
github.com/magefile/mage v1.17.2
|
||||
|
||||
@@ -166,8 +166,8 @@ github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731 h1:9x+U2HGLrSw5AT
|
||||
github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731/go.mod h1:Rs3MhFwutWhGwmY1VQsygw28z5bWcnEYmS1OG9OxjOQ=
|
||||
github.com/livekit/mediatransportutil v0.0.0-20260608063931-a3417d38cda0 h1:XHNNzebIKZRkLimla/hFGrAIX5EMWHctrgt3hLw7s+I=
|
||||
github.com/livekit/mediatransportutil v0.0.0-20260608063931-a3417d38cda0/go.mod h1:o8CFmAdrVwzJNOCsQCLUzXRjokkufNshnQHOe4fRaqU=
|
||||
github.com/livekit/protocol v1.50.5-0.20260814120900-8b1ab81c7d00 h1:g4Rdg7gAqPX/CrV7ptAERPtDX1YbjpqAF67G/6pd2Zs=
|
||||
github.com/livekit/protocol v1.50.5-0.20260814120900-8b1ab81c7d00/go.mod h1:/kYxa0dlTuH981LaBFHG/Swyr969d0+2+/6Lm7fFc34=
|
||||
github.com/livekit/protocol v1.50.5-0.20260820011432-2480d740da5e h1:AjYyvyLEC+rB2HKit2RbwMnqUwujH2b7OSUaRUcYbVA=
|
||||
github.com/livekit/protocol v1.50.5-0.20260820011432-2480d740da5e/go.mod h1:/kYxa0dlTuH981LaBFHG/Swyr969d0+2+/6Lm7fFc34=
|
||||
github.com/livekit/psrpc v0.7.3 h1:bekuZt/ZQzg8+/M8G6G5jq7bvV9fAKdPHSOZeTwrIIc=
|
||||
github.com/livekit/psrpc v0.7.3/go.mod h1:rAI+m2+/cb4x9RXhLRtUx5ZwdfjjXOl4zi46IjEetaw=
|
||||
github.com/livekit/webrtc-pion/v4 v4.2.18-warp.1 h1:fH+v4W+NFp9FfPzON6FaUFNmazGcctaAhb2P+Ksf+1s=
|
||||
|
||||
+3
-3
@@ -62,9 +62,9 @@ type JobRequest struct {
|
||||
// only set for participant jobs
|
||||
Participant *livekit.ParticipantInfo
|
||||
Metadata string
|
||||
AgentName string
|
||||
Deployment string
|
||||
Attributes map[string]string
|
||||
AgentName string
|
||||
Deployment string
|
||||
Attributes map[string]string
|
||||
}
|
||||
|
||||
type agentClient struct {
|
||||
|
||||
@@ -6,4 +6,20 @@ type Config struct {
|
||||
EnableUserDataRecording bool `yaml:"enable_user_data_recording"`
|
||||
EnableUserDataRedaction bool `yaml:"enable_user_data_redaction"`
|
||||
TargetLoad float32 `yaml:"target_load,omitempty"`
|
||||
|
||||
// agent HTTP endpoints data plane; zero values take the endpoint package
|
||||
// defaults
|
||||
Endpoints EndpointsConfig `yaml:"endpoints,omitempty"`
|
||||
}
|
||||
|
||||
type EndpointsConfig struct {
|
||||
// Disabled turns off the /agents/{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"`
|
||||
}
|
||||
|
||||
@@ -0,0 +1,799 @@
|
||||
// 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 client is the reference worker-side implementation of the agent HTTP
|
||||
// endpoints data plane: it registers a manifest over the control connection,
|
||||
// attaches the fixed wire pool, and bridges each stream to a local HTTP server.
|
||||
// It doubles as the protocol conformance harness and as a manual sidecar for
|
||||
// non-SDK workers.
|
||||
package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/url"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/livekit/protocol/auth"
|
||||
"github.com/livekit/protocol/livekit"
|
||||
"github.com/livekit/protocol/logger"
|
||||
"github.com/livekit/protocol/utils/guid"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
// ServerURL is the ws(s):// URL of the /agent endpoint
|
||||
ServerURL string
|
||||
APIKey string
|
||||
APISecret string
|
||||
|
||||
AgentName string
|
||||
Deployment string
|
||||
Endpoints []*livekit.AgentHttp_AgentEndpoint
|
||||
|
||||
// TargetAddr is the host:port of the local HTTP server streams bridge into
|
||||
TargetAddr string
|
||||
|
||||
Logger logger.Logger
|
||||
}
|
||||
|
||||
// Worker is one registration epoch: a control connection plus the fixed
|
||||
// data-connection pool.
|
||||
type Worker struct {
|
||||
cfg Config
|
||||
instanceID string
|
||||
|
||||
mu sync.Mutex
|
||||
workerID string
|
||||
settings *livekit.AgentHttp_AgentEndpointSettings
|
||||
control *wsConn
|
||||
dataConns []*dataConn
|
||||
closed bool
|
||||
registered chan struct{}
|
||||
statusSeq uint64
|
||||
}
|
||||
|
||||
func New(cfg Config) *Worker {
|
||||
if cfg.Logger == nil {
|
||||
cfg.Logger = logger.GetLogger()
|
||||
}
|
||||
return &Worker{
|
||||
cfg: cfg,
|
||||
instanceID: guid.New("AEI_"),
|
||||
registered: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Start registers and attaches the pool; it returns once the data plane is up.
|
||||
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)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
w.mu.Lock()
|
||||
w.control = control
|
||||
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,
|
||||
},
|
||||
}}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
go w.controlLoop(ctx, control)
|
||||
|
||||
select {
|
||||
case <-w.registered:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(10 * time.Second):
|
||||
return errors.New("registration timeout")
|
||||
}
|
||||
|
||||
w.mu.Lock()
|
||||
settings := w.settings
|
||||
workerID := w.workerID
|
||||
w.mu.Unlock()
|
||||
if settings == nil {
|
||||
return errors.New("server did not negotiate endpoint settings")
|
||||
}
|
||||
|
||||
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)
|
||||
w.mu.Unlock()
|
||||
}
|
||||
|
||||
// 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) {
|
||||
at := auth.NewAccessToken(w.cfg.APIKey, w.cfg.APISecret).
|
||||
SetVideoGrant(&auth.VideoGrant{Agent: true}).
|
||||
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()
|
||||
}
|
||||
@@ -0,0 +1,375 @@
|
||||
// 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)
|
||||
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)
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,579 @@
|
||||
// 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 (
|
||||
"bufio"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/rand/v2"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/livekit/protocol/livekit"
|
||||
"github.com/livekit/protocol/logger"
|
||||
"github.com/livekit/protocol/utils/guid"
|
||||
)
|
||||
|
||||
const (
|
||||
// PathPrefix is the public route namespace: /agents/{deployment}/{path...}
|
||||
PathPrefix = "/agents/"
|
||||
|
||||
// responseHeadTimeout bounds the wait for the worker's response head. Bodies
|
||||
// (SSE, long streams) are unbounded; the head never legitimately takes this
|
||||
// long.
|
||||
responseHeadTimeout = 90 * time.Second
|
||||
|
||||
// maxAttempts bounds worker retries per request
|
||||
maxAttempts = 3
|
||||
)
|
||||
|
||||
// ScopeResolver maps an inbound request to its project scope. It returns the
|
||||
// api key the request is authorized for (empty when unauthenticated) - the
|
||||
// service layer implements it from validated grants.
|
||||
type ScopeResolver func(r *http.Request) (apiKey string, authenticated bool)
|
||||
|
||||
type Front struct {
|
||||
registry *Registry
|
||||
resolveScope ScopeResolver
|
||||
logger logger.Logger
|
||||
|
||||
// fallback is consulted when nothing local can serve the request (no
|
||||
// candidates, no route match, or every match without capacity); a
|
||||
// multi-node deployment plugs its resolve-and-relay here. nil means local
|
||||
// misses are final.
|
||||
fallback Fallback
|
||||
// see WithSingleKeyFallback
|
||||
singleKeyFallback bool
|
||||
}
|
||||
|
||||
func NewFront(registry *Registry, resolveScope ScopeResolver, log logger.Logger) *Front {
|
||||
return &Front{
|
||||
registry: registry,
|
||||
resolveScope: resolveScope,
|
||||
logger: log.WithComponent("agents.endpoint"),
|
||||
}
|
||||
}
|
||||
|
||||
// FallbackRequest describes a request nothing local could serve. The request
|
||||
// body is untouched when the fallback runs.
|
||||
type FallbackRequest struct {
|
||||
// Scope is the project identity the front resolved (api key in OSS)
|
||||
Scope string
|
||||
Authenticated bool
|
||||
Deployment string
|
||||
// Path within the deployment, '/'-rooted
|
||||
Path string
|
||||
WebSocket bool
|
||||
}
|
||||
|
||||
// Fallback serves a request elsewhere (e.g. a multi-node relay); it reports
|
||||
// whether a response was written. Returning false falls back to the local
|
||||
// status mapping.
|
||||
type Fallback func(w http.ResponseWriter, r *http.Request, req *FallbackRequest) bool
|
||||
|
||||
// WithFallback installs the miss handler consulted when nothing local can
|
||||
// serve a request.
|
||||
func (f *Front) WithFallback(fb Fallback) *Front {
|
||||
f.fallback = fb
|
||||
return f
|
||||
}
|
||||
|
||||
// WithSingleKeyFallback resolves unauthenticated requests to the registry's
|
||||
// single api key when the scope resolver yields none. Self-hosted convenience
|
||||
// only: a multi-tenant front must never guess a scope from what happens to be
|
||||
// registered.
|
||||
func (f *Front) WithSingleKeyFallback() *Front {
|
||||
f.singleKeyFallback = true
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
rest, ok := strings.CutPrefix(r.URL.Path, PathPrefix)
|
||||
if !ok {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
deployment, path, found := strings.Cut(rest, "/")
|
||||
if !found {
|
||||
path = ""
|
||||
}
|
||||
path = "/" + path
|
||||
if deployment == "" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
|
||||
isWS := isWebSocketUpgrade(r)
|
||||
|
||||
apiKey, authenticated := f.resolveScope(r)
|
||||
if apiKey == "" && f.singleKeyFallback {
|
||||
// unauthenticated: OSS serves public routes when the worker fleet
|
||||
// belongs to a single key
|
||||
apiKey, _ = f.registry.SingleAPIKey()
|
||||
}
|
||||
if apiKey == "" {
|
||||
w.Header().Set("WWW-Authenticate", "Bearer")
|
||||
http.Error(w, "authentication required", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
|
||||
candidates := f.registry.Candidates(apiKey, deployment)
|
||||
if len(candidates) == 0 && f.fallback == nil {
|
||||
w.Header().Set("Retry-After", "1")
|
||||
http.Error(w, "no workers available for deployment", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
|
||||
// manifest match across the deployment's workers: FULL wins; PARTIAL only
|
||||
// yields 405 when nothing matches fully; slash-redirect mirrors FastAPI
|
||||
var matched []*Registration
|
||||
var route *Route
|
||||
partial := false
|
||||
restricted := false
|
||||
for _, reg := range candidates {
|
||||
rt, res := reg.Manifest.Match(path, r.Method, isWS)
|
||||
switch res {
|
||||
case MatchFull:
|
||||
if !authenticated && !rt.Public {
|
||||
restricted = true
|
||||
continue
|
||||
}
|
||||
if route == nil {
|
||||
route = rt
|
||||
}
|
||||
matched = append(matched, reg)
|
||||
case MatchPartial:
|
||||
partial = true
|
||||
}
|
||||
}
|
||||
if route == nil && f.fallback != nil {
|
||||
// nothing local can serve: hand off before the local status mapping
|
||||
if f.fallback(w, r, &FallbackRequest{
|
||||
Scope: apiKey, Authenticated: authenticated,
|
||||
Deployment: deployment, Path: path, WebSocket: isWS,
|
||||
}) {
|
||||
return
|
||||
}
|
||||
if len(candidates) == 0 && !restricted && !partial {
|
||||
w.Header().Set("Retry-After", "1")
|
||||
http.Error(w, "no workers available for deployment", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
}
|
||||
if route == nil {
|
||||
switch {
|
||||
case restricted:
|
||||
w.Header().Set("WWW-Authenticate", "Bearer")
|
||||
http.Error(w, "authentication required", http.StatusUnauthorized)
|
||||
case partial:
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
default:
|
||||
for _, reg := range candidates {
|
||||
if alt, ok := reg.Manifest.RedirectSlashes(path, r.Method, isWS); ok {
|
||||
u := *r.URL
|
||||
u.Path = PathPrefix + deployment + alt
|
||||
http.Redirect(w, r, u.String(), http.StatusTemporaryRedirect)
|
||||
return
|
||||
}
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
bodyConsumed := int64(0)
|
||||
countingBody := &countingReader{r: r.Body, n: &bodyConsumed}
|
||||
|
||||
attempted := make(map[*Registration]bool)
|
||||
for attempt := 0; attempt < maxAttempts; attempt++ {
|
||||
reg := pickWeighted(matched, attempted)
|
||||
if reg == nil {
|
||||
break
|
||||
}
|
||||
attempted[reg] = true
|
||||
|
||||
done, retryable := f.bridge(w, r, reg, path, countingBody, bodyConsumed, isWS)
|
||||
if done || !retryable {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// the route matched locally but nothing served it (matches draining or
|
||||
// conn-less, or every attempt failed before writing): the fallback may hold
|
||||
// capacity elsewhere. Safe exactly while no request bytes were consumed -
|
||||
// reaching this point implies it, since consuming attempts are never
|
||||
// retryable.
|
||||
if bodyConsumed == 0 && f.fallback != nil {
|
||||
if f.fallback(w, r, &FallbackRequest{
|
||||
Scope: apiKey, Authenticated: authenticated,
|
||||
Deployment: deployment, Path: path, WebSocket: isWS,
|
||||
}) {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
w.Header().Set("Retry-After", "1")
|
||||
http.Error(w, "no worker could serve the request", http.StatusServiceUnavailable)
|
||||
}
|
||||
|
||||
// pickWeighted is capacity-weighted random over non-draining workers with
|
||||
// attached data connections.
|
||||
func pickWeighted(regs []*Registration, ignore map[*Registration]bool) *Registration {
|
||||
var sum float32
|
||||
weights := make([]float32, len(regs))
|
||||
for i, reg := range regs {
|
||||
if ignore[reg] || reg.AttachedConns() == 0 {
|
||||
continue
|
||||
}
|
||||
if reg.Draining != nil && reg.Draining() {
|
||||
continue
|
||||
}
|
||||
w := float32(1)
|
||||
if reg.Load != nil {
|
||||
w = max(0.01, 1-reg.Load())
|
||||
}
|
||||
weights[i] = w
|
||||
sum += w
|
||||
}
|
||||
if sum == 0 {
|
||||
return nil
|
||||
}
|
||||
target := rand.Float32() * sum
|
||||
for i, reg := range regs {
|
||||
if target -= weights[i]; weights[i] > 0 && target <= 0 {
|
||||
return reg
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// bridge runs one attempt against one worker. done means a response (or abort)
|
||||
// reached the client; retryable reports whether another attempt is safe per the
|
||||
// retry table: idempotent/bodyless until any response byte arrived,
|
||||
// anything on HSR_REFUSED, nothing once bytes were consumed otherwise.
|
||||
func (f *Front) bridge(
|
||||
w http.ResponseWriter,
|
||||
r *http.Request,
|
||||
reg *Registration,
|
||||
path string,
|
||||
body io.Reader,
|
||||
bodyConsumedBefore int64,
|
||||
isWS bool,
|
||||
) (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,
|
||||
})
|
||||
if err != nil {
|
||||
return false, true
|
||||
}
|
||||
defer stream.Close()
|
||||
|
||||
ctx := r.Context()
|
||||
stop := context.AfterFunc(ctx, func() {
|
||||
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "client disconnected")
|
||||
})
|
||||
defer stop()
|
||||
|
||||
// serialize the request into the stream concurrently with response reading:
|
||||
// directions are independent (full duplex within the stream)
|
||||
outReq := f.outboundRequest(r, path, body, isWS)
|
||||
writeErrCh := make(chan error, 1)
|
||||
if isWS {
|
||||
// upgrade requests have no body: write the head inline and do NOT
|
||||
// half-close - the client->worker direction carries the session, and
|
||||
// the upgrade pump must be the stream's only writer
|
||||
if err := outReq.Write(stream); err != nil {
|
||||
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "request write failed")
|
||||
return false, true
|
||||
}
|
||||
writeErrCh <- nil
|
||||
} else {
|
||||
go func() {
|
||||
err := outReq.Write(stream)
|
||||
if err == nil {
|
||||
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")
|
||||
}
|
||||
writeErrCh <- err
|
||||
}()
|
||||
}
|
||||
|
||||
counted := &countingReader{r: stream, n: new(int64)}
|
||||
br := bufio.NewReader(counted)
|
||||
|
||||
resp, err := f.readResponseHead(w, br, outReq, stream)
|
||||
if err != nil {
|
||||
retryable = f.classifyRetry(r, stream, *counted.n, bodyConsumedBefore, err)
|
||||
if !retryable {
|
||||
f.logger.Warnw("agent endpoint request failed", err,
|
||||
"workerID", reg.WorkerID, "path", path)
|
||||
http.Error(w, "bad gateway", http.StatusBadGateway)
|
||||
return true, false
|
||||
}
|
||||
// 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")
|
||||
<-writeErrCh
|
||||
return false, true
|
||||
}
|
||||
|
||||
// a response byte arrived: from here every failure is surfaced, never retried
|
||||
if resp.StatusCode == http.StatusSwitchingProtocols {
|
||||
f.bridgeUpgrade(w, resp, br, stream)
|
||||
return true, false
|
||||
}
|
||||
|
||||
copyResponseHeaders(w.Header(), resp)
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
|
||||
rc := http.NewResponseController(w)
|
||||
buf := make([]byte, 32<<10)
|
||||
for {
|
||||
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")
|
||||
return true, false
|
||||
}
|
||||
_ = rc.Flush()
|
||||
}
|
||||
if rerr == io.EOF {
|
||||
break
|
||||
}
|
||||
if rerr != nil {
|
||||
// never expose a clean-looking short body
|
||||
select {
|
||||
case werr := <-writeErrCh:
|
||||
f.logger.Debugw("request write result after response failure", "error", werr)
|
||||
default:
|
||||
}
|
||||
panic(http.ErrAbortHandler)
|
||||
}
|
||||
}
|
||||
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) {
|
||||
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")
|
||||
case <-headCh:
|
||||
}
|
||||
}()
|
||||
defer close(headCh)
|
||||
|
||||
for {
|
||||
resp, err := http.ReadResponse(br, outReq)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode >= 100 && resp.StatusCode < 200 && resp.StatusCode != http.StatusSwitchingProtocols {
|
||||
// informational: relay and keep reading
|
||||
for k, vv := range resp.Header {
|
||||
for _, v := range vv {
|
||||
w.Header().Add(k, v)
|
||||
}
|
||||
}
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
clear(w.Header())
|
||||
continue
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
}
|
||||
|
||||
// bridgeUpgrade hijacks the client connection after a 101 and pumps raw bytes in
|
||||
// both directions; the stream carries the rest of the WebSocket session.
|
||||
func (f *Front) bridgeUpgrade(w http.ResponseWriter, resp *http.Response, br *bufio.Reader, stream *Stream) {
|
||||
hj, ok := w.(http.Hijacker)
|
||||
if !ok {
|
||||
// e.g. HTTP/2 client conns cannot be upgraded
|
||||
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "client does not support upgrade")
|
||||
http.Error(w, "upgrade not supported on this connection", http.StatusBadGateway)
|
||||
return
|
||||
}
|
||||
clientConn, clientRW, err := hj.Hijack()
|
||||
if err != nil {
|
||||
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "hijack failed")
|
||||
return
|
||||
}
|
||||
defer clientConn.Close()
|
||||
|
||||
if err := resp.Write(clientRW); err != nil {
|
||||
return
|
||||
}
|
||||
if err := clientRW.Flush(); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
errCh := make(chan error, 2)
|
||||
go func() {
|
||||
// worker -> client, including bytes the bufio reader already buffered
|
||||
_, err := io.Copy(clientConn, br)
|
||||
errCh <- err
|
||||
}()
|
||||
go func() {
|
||||
// client -> worker
|
||||
buf := make([]byte, 32<<10)
|
||||
for {
|
||||
n, rerr := clientRW.Read(buf)
|
||||
if n > 0 {
|
||||
if _, werr := stream.Write(buf[:n]); werr != nil {
|
||||
errCh <- werr
|
||||
return
|
||||
}
|
||||
}
|
||||
if rerr != nil {
|
||||
_ = stream.CloseWrite()
|
||||
errCh <- rerr
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
<-errCh
|
||||
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "upgrade session ended")
|
||||
}
|
||||
|
||||
// classifyRetry implements the retry table.
|
||||
func (f *Front) classifyRetry(r *http.Request, stream *Stream, responseBytes, bodyConsumedBefore int64, err error) bool {
|
||||
if responseBytes > 0 || stream.BytesRead() > 0 {
|
||||
return false
|
||||
}
|
||||
if stream.Refused() {
|
||||
// the worker proved non-dispatch; safe for any method, but only when the
|
||||
// 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
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// outboundRequest builds the request serialized into the stream: the path the
|
||||
// worker's router sees (deployment prefix stripped), hop-by-hop headers removed,
|
||||
// forwarding headers appended.
|
||||
func (f *Front) outboundRequest(r *http.Request, path string, body io.Reader, isWS bool) *http.Request {
|
||||
out := r.Clone(r.Context())
|
||||
out.RequestURI = ""
|
||||
out.URL = &url.URL{Path: path, RawQuery: r.URL.RawQuery}
|
||||
out.Host = r.Host
|
||||
out.Body = io.NopCloser(body)
|
||||
// one exchange per stream: closing the worker-local app connection after the
|
||||
// response is what lets the opaque pump observe the end of the exchange and
|
||||
// free the stream slot
|
||||
out.Close = !isWS
|
||||
|
||||
removeHopByHopHeaders(out.Header, isWS)
|
||||
out.Header.Del("Expect") // the front owns 100-continue semantics client-side
|
||||
|
||||
if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
|
||||
prior := out.Header.Get("X-Forwarded-For")
|
||||
if prior != "" {
|
||||
out.Header.Set("X-Forwarded-For", prior+", "+host)
|
||||
} else {
|
||||
out.Header.Set("X-Forwarded-For", host)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// hop-by-hop headers per RFC 9110; Connection-nominated headers are dropped too.
|
||||
// For WebSocket upgrades Connection/Upgrade survive so the worker-side bridge
|
||||
// sees a real upgrade request.
|
||||
func removeHopByHopHeaders(h http.Header, isWS bool) {
|
||||
for _, f := range h.Values("Connection") {
|
||||
for _, sf := range strings.Split(f, ",") {
|
||||
if sf = strings.TrimSpace(sf); sf != "" && !strings.EqualFold(sf, "upgrade") {
|
||||
h.Del(sf)
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, k := range []string{
|
||||
"Keep-Alive", "Proxy-Authenticate", "Proxy-Authorization",
|
||||
"Te", "Trailer", "Transfer-Encoding",
|
||||
} {
|
||||
h.Del(k)
|
||||
}
|
||||
if !isWS {
|
||||
h.Del("Connection")
|
||||
h.Del("Upgrade")
|
||||
} else {
|
||||
h.Set("Connection", "Upgrade")
|
||||
}
|
||||
}
|
||||
|
||||
func copyResponseHeaders(dst http.Header, resp *http.Response) {
|
||||
for k, vv := range resp.Header {
|
||||
for _, v := range vv {
|
||||
dst.Add(k, v)
|
||||
}
|
||||
}
|
||||
removeHopByHopHeaders(dst, false)
|
||||
if resp.ContentLength >= 0 && dst.Get("Content-Length") == "" {
|
||||
dst.Set("Content-Length", fmt.Sprintf("%d", resp.ContentLength))
|
||||
}
|
||||
}
|
||||
|
||||
type countingReader struct {
|
||||
r io.Reader
|
||||
n *int64
|
||||
}
|
||||
|
||||
func (c *countingReader) Read(p []byte) (int, error) {
|
||||
n, err := c.r.Read(p)
|
||||
*c.n += int64(n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
func isWebSocketUpgrade(r *http.Request) bool {
|
||||
return strings.EqualFold(r.Header.Get("Upgrade"), "websocket") &&
|
||||
httpHeaderContainsToken(r.Header, "Connection", "upgrade")
|
||||
}
|
||||
|
||||
func httpHeaderContainsToken(h http.Header, name, token string) bool {
|
||||
for _, v := range h.Values(name) {
|
||||
for _, f := range strings.Split(v, ",") {
|
||||
if strings.EqualFold(strings.TrimSpace(f), token) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
// 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 (
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/livekit/protocol/livekit"
|
||||
)
|
||||
|
||||
const MaxManifestRoutes = 256
|
||||
|
||||
// Route is one validated manifest entry.
|
||||
type Route struct {
|
||||
Template *Template
|
||||
Methods []string // uppercase; empty for websocket routes
|
||||
Kind livekit.AgentHttp_AgentEndpointKind
|
||||
Public bool
|
||||
}
|
||||
|
||||
// Manifest is a worker's ordered route table.
|
||||
type Manifest struct {
|
||||
routes []Route
|
||||
}
|
||||
|
||||
// MatchResult mirrors starlette's Match enum: a FULL match selects the route, a
|
||||
// PARTIAL match (path matched, method didn't) yields 405 only after the whole
|
||||
// table has been scanned, so a later route with the right method still wins.
|
||||
type MatchResult int
|
||||
|
||||
const (
|
||||
MatchNone MatchResult = iota
|
||||
MatchPartial
|
||||
MatchFull
|
||||
)
|
||||
|
||||
// ParseManifest validates a registration's endpoint list.
|
||||
func ParseManifest(endpoints []*livekit.AgentHttp_AgentEndpoint) (*Manifest, error) {
|
||||
if len(endpoints) > MaxManifestRoutes {
|
||||
return nil, fmt.Errorf("manifest exceeds %d routes", MaxManifestRoutes)
|
||||
}
|
||||
m := &Manifest{routes: make([]Route, 0, len(endpoints))}
|
||||
for _, ep := range endpoints {
|
||||
tpl, err := ParseTemplate(ep.GetPath())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var methods []string
|
||||
switch ep.GetKind() {
|
||||
case livekit.AgentHttp_AEK_HTTP:
|
||||
if len(ep.GetMethods()) == 0 {
|
||||
return nil, fmt.Errorf("endpoint %q declares no methods", ep.GetPath())
|
||||
}
|
||||
for _, method := range ep.GetMethods() {
|
||||
u := strings.ToUpper(method)
|
||||
if u != method {
|
||||
return nil, fmt.Errorf("endpoint %q method %q must be uppercase", ep.GetPath(), method)
|
||||
}
|
||||
methods = append(methods, u)
|
||||
}
|
||||
case livekit.AgentHttp_AEK_WEBSOCKET:
|
||||
if len(ep.GetMethods()) != 0 {
|
||||
return nil, fmt.Errorf("websocket endpoint %q must not declare methods", ep.GetPath())
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("endpoint %q has unsupported kind %s", ep.GetPath(), ep.GetKind())
|
||||
}
|
||||
m.routes = append(m.routes, Route{
|
||||
Template: tpl,
|
||||
Methods: methods,
|
||||
Kind: ep.GetKind(),
|
||||
Public: ep.GetPublic(),
|
||||
})
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// Match resolves a request path against the table. websocket selects the
|
||||
// websocket route class (upgrade requests); otherwise the HTTP class with
|
||||
// method matching.
|
||||
func (m *Manifest) Match(path, method string, websocket bool) (*Route, MatchResult) {
|
||||
partial := false
|
||||
for i := range m.routes {
|
||||
r := &m.routes[i]
|
||||
if websocket != (r.Kind == livekit.AgentHttp_AEK_WEBSOCKET) {
|
||||
continue
|
||||
}
|
||||
if !r.Template.Match(path) {
|
||||
continue
|
||||
}
|
||||
if websocket || slices.Contains(r.Methods, method) {
|
||||
return r, MatchFull
|
||||
}
|
||||
partial = true
|
||||
}
|
||||
if partial {
|
||||
return nil, MatchPartial
|
||||
}
|
||||
return nil, MatchNone
|
||||
}
|
||||
|
||||
// RedirectSlashes reports whether the alternate-slash form of path would match,
|
||||
// mirroring starlette's redirect_slashes default (FastAPI 307s /x/ to /x and
|
||||
// vice versa when only the alternate form matches).
|
||||
func (m *Manifest) RedirectSlashes(path, method string, websocket bool) (string, bool) {
|
||||
if path == "/" {
|
||||
return "", false
|
||||
}
|
||||
var alt string
|
||||
if strings.HasSuffix(path, "/") {
|
||||
alt = strings.TrimSuffix(path, "/")
|
||||
} else {
|
||||
alt = path + "/"
|
||||
}
|
||||
if _, res := m.Match(alt, method, websocket); res == MatchFull {
|
||||
return alt, true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
@@ -0,0 +1,409 @@
|
||||
// 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 (
|
||||
"crypto/subtle"
|
||||
"errors"
|
||||
"slices"
|
||||
"sync"
|
||||
|
||||
"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")
|
||||
)
|
||||
|
||||
func normalizeDeployment(d string) string {
|
||||
if d == "" {
|
||||
return DefaultDeployment
|
||||
}
|
||||
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).
|
||||
type Registration struct {
|
||||
WorkerID string
|
||||
InstanceID string
|
||||
APIKey string
|
||||
Deployment string
|
||||
Manifest *Manifest
|
||||
Settings Settings
|
||||
Logger logger.Logger
|
||||
|
||||
// Load and Draining are provided by the control-plane layer that owns the
|
||||
// worker (reported load rides UpdateWorkerStatus on the control connection).
|
||||
Load func() float32
|
||||
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
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
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 {
|
||||
return nil, ErrAttachRejected
|
||||
}
|
||||
if instanceID != r.InstanceID {
|
||||
return nil, ErrAttachRejected
|
||||
}
|
||||
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.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 {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return len(r.conns)
|
||||
}
|
||||
|
||||
func (r *Registration) close() {
|
||||
r.mu.Lock()
|
||||
if r.closed {
|
||||
r.mu.Unlock()
|
||||
return
|
||||
}
|
||||
r.closed = true
|
||||
conns := slices.Clone(r.conns)
|
||||
r.conns = nil
|
||||
r.mu.Unlock()
|
||||
|
||||
for _, c := range conns {
|
||||
c.Close(ErrConnClosed)
|
||||
}
|
||||
}
|
||||
|
||||
// Registry tracks data-plane registrations on this node, keyed by
|
||||
// (api key, deployment). The api key is the project identity in OSS.
|
||||
type Registry struct {
|
||||
mu sync.Mutex
|
||||
regs map[string]*Registration // by worker id
|
||||
byKey map[regKey][]*Registration
|
||||
|
||||
// scope hooks fire when a scope (api key / project) gains its first or
|
||||
// loses its last registration; the Remote uses them to (de)register the
|
||||
// resolve topic on the bus
|
||||
scopeCounts map[string]int
|
||||
onScopeActive func(scope string)
|
||||
onScopeIdle func(scope string)
|
||||
|
||||
// hookMu serializes scope hook invocations. Transitions are decided against
|
||||
// the registry's CURRENT state under hookMu, never from values computed
|
||||
// earlier: a reconnect's activation racing the old connection's idle must
|
||||
// not leave a live scope without its resolve topic.
|
||||
hookMu sync.Mutex
|
||||
scopeNotified map[string]bool
|
||||
}
|
||||
|
||||
type regKey struct {
|
||||
apiKey string
|
||||
deployment string
|
||||
}
|
||||
|
||||
func NewRegistry() *Registry {
|
||||
return &Registry{
|
||||
regs: make(map[string]*Registration),
|
||||
byKey: make(map[regKey][]*Registration),
|
||||
scopeCounts: make(map[string]int),
|
||||
scopeNotified: make(map[string]bool),
|
||||
}
|
||||
}
|
||||
|
||||
// notifyScope reconciles the scope's hook state with the registry's current
|
||||
// truth. Racing register/deregister notifications converge instead of
|
||||
// interleaving.
|
||||
func (g *Registry) notifyScope(scope string) {
|
||||
g.hookMu.Lock()
|
||||
defer g.hookMu.Unlock()
|
||||
g.mu.Lock()
|
||||
active := g.scopeCounts[scope] > 0
|
||||
activeHook, idleHook := g.onScopeActive, g.onScopeIdle
|
||||
g.mu.Unlock()
|
||||
if activeHook == nil && idleHook == nil {
|
||||
return
|
||||
}
|
||||
if g.scopeNotified[scope] == active {
|
||||
return
|
||||
}
|
||||
if active {
|
||||
g.scopeNotified[scope] = true
|
||||
activeHook(scope)
|
||||
} else {
|
||||
delete(g.scopeNotified, scope)
|
||||
idleHook(scope)
|
||||
}
|
||||
}
|
||||
|
||||
// SetScopeHooks installs the scope activation callbacks, replaying currently
|
||||
// active scopes so a resolver layered on after registrations still registers
|
||||
// its topics. Multi-node deployments use this to answer endpoint resolves only
|
||||
// while they hold registrations for a scope.
|
||||
func (g *Registry) SetScopeHooks(active, idle func(scope string)) {
|
||||
g.hookMu.Lock()
|
||||
defer g.hookMu.Unlock()
|
||||
g.mu.Lock()
|
||||
g.onScopeActive = active
|
||||
g.onScopeIdle = idle
|
||||
scopes := make([]string, 0, len(g.scopeCounts))
|
||||
for scope, n := range g.scopeCounts {
|
||||
if n > 0 {
|
||||
scopes = append(scopes, scope)
|
||||
}
|
||||
}
|
||||
g.mu.Unlock()
|
||||
if active != nil {
|
||||
for _, s := range scopes {
|
||||
if !g.scopeNotified[s] {
|
||||
g.scopeNotified[s] = true
|
||||
active(s)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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.
|
||||
func (g *Registry) Register(r *Registration) error {
|
||||
key := regKey{r.APIKey, normalizeDeployment(r.Deployment)}
|
||||
g.mu.Lock()
|
||||
old := g.regs[r.WorkerID]
|
||||
if old != nil {
|
||||
g.removeLocked(old)
|
||||
}
|
||||
g.regs[r.WorkerID] = r
|
||||
g.byKey[key] = append(g.byKey[key], r)
|
||||
g.scopeCounts[r.APIKey]++
|
||||
g.mu.Unlock()
|
||||
g.notifyScope(r.APIKey)
|
||||
if old != nil {
|
||||
if old.APIKey != r.APIKey {
|
||||
g.notifyScope(old.APIKey)
|
||||
}
|
||||
old.close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// 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)}
|
||||
if regs := g.byKey[key]; len(regs) > 0 {
|
||||
if i := slices.Index(regs, r); i != -1 {
|
||||
regs = slices.Delete(regs, i, i+1)
|
||||
}
|
||||
if len(regs) == 0 {
|
||||
delete(g.byKey, key)
|
||||
} else {
|
||||
g.byKey[key] = regs
|
||||
}
|
||||
}
|
||||
g.scopeCounts[r.APIKey]--
|
||||
if g.scopeCounts[r.APIKey] == 0 {
|
||||
delete(g.scopeCounts, r.APIKey)
|
||||
}
|
||||
}
|
||||
|
||||
// Deregister removes exactly this registration; it is a no-op when a newer
|
||||
// epoch has already superseded it.
|
||||
func (g *Registry) Deregister(r *Registration) {
|
||||
g.mu.Lock()
|
||||
if g.regs[r.WorkerID] != r {
|
||||
g.mu.Unlock()
|
||||
return
|
||||
}
|
||||
g.removeLocked(r)
|
||||
g.mu.Unlock()
|
||||
g.notifyScope(r.APIKey)
|
||||
r.close()
|
||||
}
|
||||
|
||||
// ValidateAttach checks an attach without adopting the connection, so the
|
||||
// attach response can be written before stream frames may flow.
|
||||
func (g *Registry) ValidateAttach(workerID, instanceID, token string) error {
|
||||
g.mu.Lock()
|
||||
r, ok := g.regs[workerID]
|
||||
g.mu.Unlock()
|
||||
if !ok {
|
||||
return ErrUnknownWorker
|
||||
}
|
||||
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.
|
||||
func (g *Registry) BeginAttach(workerID, instanceID, token string) (*AttachTicket, error) {
|
||||
g.mu.Lock()
|
||||
r, ok := g.regs[workerID]
|
||||
g.mu.Unlock()
|
||||
if !ok {
|
||||
return nil, ErrUnknownWorker
|
||||
}
|
||||
return r.beginAttach(instanceID, token)
|
||||
}
|
||||
|
||||
// Candidates returns the registrations for (api key, deployment segment).
|
||||
func (g *Registry) Candidates(apiKey, deployment string) []*Registration {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
return slices.Clone(g.byKey[regKey{apiKey, normalizeDeployment(deployment)}])
|
||||
}
|
||||
|
||||
// SingleAPIKey returns the api key when every registration shares one - the OSS
|
||||
// resolution for unauthenticated requests to public endpoints. ok is false when
|
||||
// zero or multiple keys are present.
|
||||
func (g *Registry) SingleAPIKey() (string, bool) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
var key string
|
||||
for _, r := range g.regs {
|
||||
if key == "" {
|
||||
key = r.APIKey
|
||||
} else if key != r.APIKey {
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
return key, key != ""
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package endpoint
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/livekit/protocol/livekit"
|
||||
"github.com/livekit/protocol/logger"
|
||||
)
|
||||
|
||||
// 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.
|
||||
func TestRegistrySupersede(t *testing.T) {
|
||||
g := NewRegistry()
|
||||
manifest, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{
|
||||
Path: "/x", Methods: []string{"GET"}, Public: true,
|
||||
}})
|
||||
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(),
|
||||
}
|
||||
}
|
||||
|
||||
oldReg := mk("i-1")
|
||||
require.NoError(t, g.Register(oldReg))
|
||||
newReg := mk("i-2")
|
||||
require.NoError(t, g.Register(newReg))
|
||||
|
||||
require.Equal(t, []*Registration{newReg}, g.Candidates("key", "production"))
|
||||
require.Error(t, g.ValidateAttach("AW_1", "i-1", "ATT_i-1"), "old epoch must not attach")
|
||||
require.NoError(t, g.ValidateAttach("AW_1", "i-2", "ATT_i-2"))
|
||||
|
||||
// the old control connection tears down after the new one registered
|
||||
g.Deregister(oldReg)
|
||||
require.Equal(t, []*Registration{newReg}, g.Candidates("key", "production"),
|
||||
"the retiring epoch must not deregister its successor")
|
||||
|
||||
g.Deregister(newReg)
|
||||
require.Empty(t, g.Candidates("key", "production"))
|
||||
}
|
||||
@@ -0,0 +1,239 @@
|
||||
// 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.
|
||||
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
|
||||
}
|
||||
|
||||
type streamQueue struct {
|
||||
chunks [][]byte
|
||||
eof bool // send an EOF frame after the last chunk drains
|
||||
}
|
||||
|
||||
func newScheduler(wire WireConn) *scheduler {
|
||||
s := &scheduler{
|
||||
wire: wire,
|
||||
queues: make(map[uint32]*streamQueue),
|
||||
}
|
||||
s.cond = sync.NewCond(&s.mu)
|
||||
go s.writeLoop()
|
||||
return s
|
||||
}
|
||||
|
||||
// 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 []byte) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for !s.closed && s.queuedBytes+int64(len(chunk)) > connBufferBudget && s.queuedBytes > 0 {
|
||||
s.cond.Wait()
|
||||
}
|
||||
if s.closed {
|
||||
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))
|
||||
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))
|
||||
}
|
||||
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.
|
||||
func (s *scheduler) next() (*livekit.AgentHttp_Frame, bool) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for {
|
||||
if s.closed {
|
||||
return nil, false
|
||||
}
|
||||
if len(s.control) > 0 {
|
||||
f := s.control[0]
|
||||
s.control = s.control[1:]
|
||||
return f, true
|
||||
}
|
||||
if len(s.rr) > 0 {
|
||||
if s.rrIdx >= len(s.rr) {
|
||||
s.rrIdx = 0
|
||||
}
|
||||
id := s.rr[s.rrIdx]
|
||||
q := s.queues[id]
|
||||
var f *livekit.AgentHttp_Frame
|
||||
if len(q.chunks) > 0 {
|
||||
chunk := q.chunks[0]
|
||||
q.chunks = q.chunks[1:]
|
||||
s.queuedBytes -= int64(len(chunk))
|
||||
f = &livekit.AgentHttp_Frame{
|
||||
StreamId: id,
|
||||
Message: &livekit.AgentHttp_Frame_Data{Data: chunk},
|
||||
}
|
||||
}
|
||||
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, 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, true
|
||||
}
|
||||
s.cond.Wait()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *scheduler) writeLoop() {
|
||||
for {
|
||||
f, ok := s.next()
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
_ = s.wire.SetWriteDeadline(time.Now().Add(writeStallTimeout))
|
||||
if err := s.wire.WriteFrame(f); err != nil {
|
||||
s.close(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,341 @@
|
||||
// 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
|
||||
}
|
||||
|
||||
chunk := make([]byte, n)
|
||||
copy(chunk, p[:n])
|
||||
if err := s.conn.sched.enqueueData(s.id, chunk); err != nil {
|
||||
return total, err
|
||||
}
|
||||
s.conn.noteActivity(n)
|
||||
p = p[n:]
|
||||
total += int(n)
|
||||
}
|
||||
return total, nil
|
||||
}
|
||||
|
||||
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()
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
// 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 (
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Template is a compiled starlette-style path template. The compilation mirrors
|
||||
// starlette's compile_path exactly: {name} or {name:convertor} segments become the
|
||||
// convertor's pattern, everything else is matched literally, anchored on both ends.
|
||||
// Workers declare templates through FastAPI routes, so any divergence from
|
||||
// starlette's semantics would make the server route requests the worker-side
|
||||
// router then refuses.
|
||||
type Template struct {
|
||||
raw string
|
||||
re *regexp.Regexp
|
||||
}
|
||||
|
||||
// convertor patterns copied verbatim from starlette's convertors.py. The uuid
|
||||
// pattern deliberately makes every hyphen optional, so 32 bare hex characters
|
||||
// match too.
|
||||
var convertorPatterns = map[string]string{
|
||||
"str": `[^/]+`,
|
||||
"path": `.*`,
|
||||
"int": `[0-9]+`,
|
||||
"float": `[0-9]+(?:\.[0-9]+)?`,
|
||||
"uuid": `[0-9a-fA-F]{8}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{12}`,
|
||||
}
|
||||
|
||||
// starlette's PARAM_REGEX
|
||||
var paramRegex = regexp.MustCompile(`\{([a-zA-Z_][a-zA-Z0-9_]*)(:[a-zA-Z_][a-zA-Z0-9_]*)?\}`)
|
||||
|
||||
// ParseTemplate compiles a starlette path template. Custom convertors are
|
||||
// rejected: only the five built-ins may travel over the wire.
|
||||
func ParseTemplate(path string) (*Template, error) {
|
||||
if !strings.HasPrefix(path, "/") {
|
||||
return nil, fmt.Errorf("path template must start with '/': %q", path)
|
||||
}
|
||||
|
||||
var pattern strings.Builder
|
||||
pattern.WriteString("^")
|
||||
|
||||
idx := 0
|
||||
seen := map[string]bool{}
|
||||
for _, m := range paramRegex.FindAllStringSubmatchIndex(path, -1) {
|
||||
start, end := m[0], m[1]
|
||||
name := path[m[2]:m[3]]
|
||||
convertor := "str"
|
||||
if m[4] != -1 {
|
||||
convertor = path[m[4]+1 : m[5]] // skip the ':'
|
||||
}
|
||||
convPattern, ok := convertorPatterns[convertor]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown path convertor %q in template %q", convertor, path)
|
||||
}
|
||||
if seen[name] {
|
||||
return nil, fmt.Errorf("duplicated param name %q in template %q", name, path)
|
||||
}
|
||||
seen[name] = true
|
||||
|
||||
pattern.WriteString(regexp.QuoteMeta(path[idx:start]))
|
||||
pattern.WriteString("(?:")
|
||||
pattern.WriteString(convPattern)
|
||||
pattern.WriteString(")")
|
||||
idx = end
|
||||
}
|
||||
pattern.WriteString(regexp.QuoteMeta(path[idx:]))
|
||||
pattern.WriteString("$")
|
||||
|
||||
re, err := regexp.Compile(pattern.String())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid path template %q: %w", path, err)
|
||||
}
|
||||
return &Template{raw: path, re: re}, nil
|
||||
}
|
||||
|
||||
func (t *Template) Match(path string) bool {
|
||||
return t.re.MatchString(path)
|
||||
}
|
||||
|
||||
func (t *Template) String() string {
|
||||
return t.raw
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
package endpoint
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/livekit/protocol/livekit"
|
||||
)
|
||||
|
||||
func TestTemplateStarletteSemantics(t *testing.T) {
|
||||
cases := []struct {
|
||||
template string
|
||||
path string
|
||||
match bool
|
||||
}{
|
||||
{"/token", "/token", true},
|
||||
{"/token", "/token/", false},
|
||||
{"/token", "/Token", false},
|
||||
{"/users/{id}", "/users/42", true},
|
||||
{"/users/{id}", "/users/42/posts", false},
|
||||
{"/users/{id}", "/users/", false},
|
||||
{"/users/{id:int}", "/users/42", true},
|
||||
{"/users/{id:int}", "/users/4x2", false},
|
||||
{"/files/{p:path}", "/files/a/b/c.txt", true},
|
||||
{"/files/{p:path}", "/files/", true},
|
||||
{"/price/{v:float}", "/price/1.25", true},
|
||||
{"/price/{v:float}", "/price/1.", false},
|
||||
{"/obj/{u:uuid}", "/obj/123e4567-e89b-12d3-a456-426614174000", true},
|
||||
// starlette's uuid convertor makes every hyphen optional
|
||||
{"/obj/{u:uuid}", "/obj/123e4567e89b12d3a456426614174000", true},
|
||||
{"/obj/{u:uuid}", "/obj/123e4567", false},
|
||||
{"/a/{x}/b/{y}", "/a/1/b/2", true},
|
||||
{"/a/{x}/b/{y}", "/a/1/c/2", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
tpl, err := ParseTemplate(c.template)
|
||||
require.NoError(t, err, c.template)
|
||||
require.Equal(t, c.match, tpl.Match(c.path), "%s vs %s", c.template, c.path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTemplateRejectsCustomConvertors(t *testing.T) {
|
||||
_, err := ParseTemplate("/x/{id:slug}")
|
||||
require.Error(t, err)
|
||||
_, err = ParseTemplate("/x/{a}/{a}")
|
||||
require.Error(t, err)
|
||||
_, err = ParseTemplate("no-slash")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func ep(path string, methods []string, public bool) *livekit.AgentHttp_AgentEndpoint {
|
||||
return &livekit.AgentHttp_AgentEndpoint{Path: path, Methods: methods, Public: public}
|
||||
}
|
||||
|
||||
func TestManifestFullPartialSemantics(t *testing.T) {
|
||||
// POST /x registered after GET /x must still serve POSTs (starlette scans
|
||||
// for a FULL match before settling for the PARTIAL 405)
|
||||
m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{
|
||||
ep("/x", []string{"GET"}, true),
|
||||
ep("/x", []string{"POST"}, true),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
r, res := m.Match("/x", http.MethodPost, false)
|
||||
require.Equal(t, MatchFull, res)
|
||||
require.Contains(t, r.Methods, "POST")
|
||||
|
||||
// the manifest carries the app's methods verbatim: FastAPI does not imply
|
||||
// HEAD from GET, so neither does the matcher
|
||||
_, res = m.Match("/x", http.MethodHead, false)
|
||||
require.Equal(t, MatchPartial, res)
|
||||
|
||||
// PARTIAL only when no route serves the method
|
||||
_, res = m.Match("/x", http.MethodDelete, false)
|
||||
require.Equal(t, MatchPartial, res)
|
||||
|
||||
_, res = m.Match("/nope", http.MethodGet, false)
|
||||
require.Equal(t, MatchNone, res)
|
||||
}
|
||||
|
||||
func TestManifestRedirectSlashes(t *testing.T) {
|
||||
m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{
|
||||
ep("/hook", []string{"POST"}, true),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
alt, ok := m.RedirectSlashes("/hook/", http.MethodPost, false)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "/hook", alt)
|
||||
|
||||
_, ok = m.RedirectSlashes("/other/", http.MethodPost, false)
|
||||
require.False(t, ok)
|
||||
}
|
||||
|
||||
func TestManifestWebSocketRoutes(t *testing.T) {
|
||||
m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{
|
||||
{Path: "/ws", Kind: livekit.AgentHttp_AEK_WEBSOCKET, Public: true},
|
||||
ep("/http", []string{"GET"}, true),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, res := m.Match("/ws", http.MethodGet, true)
|
||||
require.Equal(t, MatchFull, res)
|
||||
_, res = m.Match("/ws", http.MethodGet, false)
|
||||
require.Equal(t, MatchNone, res)
|
||||
_, res = m.Match("/http", http.MethodGet, true)
|
||||
require.Equal(t, MatchNone, res)
|
||||
}
|
||||
|
||||
func TestManifestValidation(t *testing.T) {
|
||||
_, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x", nil, false)})
|
||||
require.Error(t, err, "http endpoint without methods")
|
||||
|
||||
_, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x", []string{"get"}, false)})
|
||||
require.Error(t, err, "lowercase method")
|
||||
|
||||
_, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{
|
||||
{Path: "/ws", Kind: livekit.AgentHttp_AEK_WEBSOCKET, Methods: []string{"GET"}},
|
||||
})
|
||||
require.Error(t, err, "websocket route with methods")
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
// 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
|
||||
}
|
||||
|
||||
func (s Settings) Proto() *livekit.AgentHttp_AgentEndpointSettings {
|
||||
return &livekit.AgentHttp_AgentEndpointSettings{
|
||||
Protocol: s.Protocol,
|
||||
AttachToken: s.AttachToken,
|
||||
DataConnectionCount: s.DataConnCount,
|
||||
}
|
||||
}
|
||||
+59
-7
@@ -155,6 +155,12 @@ type WorkerRegistration struct {
|
||||
Permissions *livekit.ParticipantPermission
|
||||
ClientIP string
|
||||
Deployment string
|
||||
|
||||
// agent HTTP endpoints data plane
|
||||
Endpoints []*livekit.AgentHttp_AgentEndpoint
|
||||
InstanceID string
|
||||
EndpointProtocol uint32
|
||||
EndpointSettings *livekit.AgentHttp_AgentEndpointSettings
|
||||
}
|
||||
|
||||
func MakeWorkerRegistration() WorkerRegistration {
|
||||
@@ -166,10 +172,17 @@ func MakeWorkerRegistration() WorkerRegistration {
|
||||
|
||||
var _ WorkerSignalHandler = (*WorkerRegisterer)(nil)
|
||||
|
||||
// EndpointSettingsFunc validates a registration's endpoint manifest and returns the
|
||||
// negotiated data-plane settings (including the attach token). Returning an error
|
||||
// fails the registration. It runs only when the registration declares endpoints AND
|
||||
// a data-plane protocol version; a nil func rejects such registrations.
|
||||
type EndpointSettingsFunc func(req *livekit.RegisterWorkerRequest) (*livekit.AgentHttp_AgentEndpointSettings, error)
|
||||
|
||||
type WorkerRegisterer struct {
|
||||
WorkerPingHandler
|
||||
serverInfo *livekit.ServerInfo
|
||||
deadline time.Time
|
||||
serverInfo *livekit.ServerInfo
|
||||
deadline time.Time
|
||||
endpointSettings EndpointSettingsFunc
|
||||
|
||||
registration WorkerRegistration
|
||||
registered bool
|
||||
@@ -184,6 +197,13 @@ func NewWorkerRegisterer(conn SignalConn, serverInfo *livekit.ServerInfo, base W
|
||||
}
|
||||
}
|
||||
|
||||
// WithEndpointSettings enables the HTTP endpoints data plane for registrations that
|
||||
// declare endpoints.
|
||||
func (h *WorkerRegisterer) WithEndpointSettings(f EndpointSettingsFunc) *WorkerRegisterer {
|
||||
h.endpointSettings = f
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *WorkerRegisterer) Deadline() time.Time {
|
||||
return h.deadline
|
||||
}
|
||||
@@ -221,13 +241,28 @@ func (h *WorkerRegisterer) HandleRegister(req *livekit.RegisterWorkerRequest) er
|
||||
h.registration.JobType = req.GetType()
|
||||
h.registration.Permissions = permissions
|
||||
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")
|
||||
}
|
||||
settings, err := h.endpointSettings(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
h.registration.EndpointSettings = settings
|
||||
}
|
||||
h.registered = true
|
||||
|
||||
_, err := h.conn.WriteServerMessage(&livekit.ServerMessage{
|
||||
Message: &livekit.ServerMessage_Register{
|
||||
Register: &livekit.RegisterWorkerResponse{
|
||||
WorkerId: h.registration.ID,
|
||||
ServerInfo: h.serverInfo,
|
||||
WorkerId: h.registration.ID,
|
||||
ServerInfo: h.serverInfo,
|
||||
EndpointSettings: h.registration.EndpointSettings,
|
||||
},
|
||||
},
|
||||
})
|
||||
@@ -248,9 +283,11 @@ type Worker struct {
|
||||
cancel context.CancelFunc
|
||||
closed chan struct{}
|
||||
|
||||
mu sync.Mutex
|
||||
load float32
|
||||
status livekit.WorkerStatus
|
||||
mu sync.Mutex
|
||||
load float32
|
||||
status livekit.WorkerStatus
|
||||
statusSeq uint64
|
||||
draining bool
|
||||
|
||||
runningJobs map[livekit.JobID]*livekit.Job
|
||||
availability map[livekit.JobID]chan *livekit.AvailabilityResponse
|
||||
@@ -581,15 +618,30 @@ func (w *Worker) HandleUpdateWorker(update *livekit.UpdateWorkerStatus) error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
// status updates may interleave across connections; never regress newer state
|
||||
if seq := update.GetSeq(); seq != 0 {
|
||||
if seq <= w.statusSeq {
|
||||
return nil
|
||||
}
|
||||
w.statusSeq = seq
|
||||
}
|
||||
|
||||
if status := update.Status; status != nil && w.status != *status {
|
||||
w.status = *status
|
||||
w.Logger().Debugw("worker status changed", "status", w.status)
|
||||
}
|
||||
w.load = update.GetLoad()
|
||||
w.draining = update.GetDraining()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Worker) Draining() bool {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
return w.draining
|
||||
}
|
||||
|
||||
func (w *Worker) HandleMigrateJob(req *livekit.MigrateJobRequest) error {
|
||||
// TODO(theomonnom): On OSS this is not implemented
|
||||
// We could maybe just move a specific job to another worker
|
||||
|
||||
@@ -0,0 +1,502 @@
|
||||
// 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 service_test
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/livekit/livekit-server/pkg/agent"
|
||||
"github.com/livekit/livekit-server/pkg/agent/endpoint"
|
||||
"github.com/livekit/livekit-server/pkg/agent/endpoint/client"
|
||||
"github.com/livekit/livekit-server/pkg/config"
|
||||
"github.com/livekit/livekit-server/pkg/routing"
|
||||
"github.com/livekit/livekit-server/pkg/service"
|
||||
"github.com/livekit/protocol/auth"
|
||||
"github.com/livekit/protocol/livekit"
|
||||
"github.com/livekit/psrpc"
|
||||
)
|
||||
|
||||
const (
|
||||
testKey = "test"
|
||||
testSecret = "verysecretsecret"
|
||||
)
|
||||
|
||||
type endpointStack struct {
|
||||
t *testing.T
|
||||
ts *httptest.Server
|
||||
svc *service.AgentService
|
||||
}
|
||||
|
||||
func newEndpointStack(t *testing.T, endpointsCfg agent.EndpointsConfig) *endpointStack {
|
||||
localNode, err := routing.NewLocalNode(nil)
|
||||
require.NoError(t, err)
|
||||
keyProvider := auth.NewSimpleKeyProvider(testKey, testSecret)
|
||||
|
||||
svc, err := service.NewAgentService(
|
||||
&config.Config{
|
||||
Region: "test",
|
||||
Keys: map[string]string{testKey: testSecret},
|
||||
Agents: agent.Config{TargetLoad: agent.DefaultTargetLoad, Endpoints: endpointsCfg},
|
||||
},
|
||||
localNode,
|
||||
psrpc.NewLocalMessageBus(),
|
||||
keyProvider,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle("/agent", svc)
|
||||
mux.Handle(endpoint.PathPrefix, svc.EndpointFront())
|
||||
|
||||
authMW := service.NewAPIKeyAuthMiddleware(keyProvider)
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
authMW.ServeHTTP(w, r, mux.ServeHTTP)
|
||||
}))
|
||||
t.Cleanup(ts.Close)
|
||||
t.Cleanup(func() { svc.DrainConnections(time.Millisecond, true) })
|
||||
|
||||
return &endpointStack{t: t, ts: ts, svc: svc}
|
||||
}
|
||||
|
||||
func (s *endpointStack) wsURL() string {
|
||||
return "ws" + strings.TrimPrefix(s.ts.URL, "http") + "/agent"
|
||||
}
|
||||
|
||||
func (s *endpointStack) startWorker(target string, deployment string, endpoints []*livekit.AgentHttp_AgentEndpoint) *client.Worker {
|
||||
w := client.New(client.Config{
|
||||
ServerURL: s.wsURL(),
|
||||
APIKey: testKey,
|
||||
APISecret: testSecret,
|
||||
AgentName: "test-agent",
|
||||
Deployment: deployment,
|
||||
Endpoints: endpoints,
|
||||
TargetAddr: strings.TrimPrefix(target, "http://"),
|
||||
})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
require.NoError(s.t, w.Start(ctx))
|
||||
s.t.Cleanup(w.Close)
|
||||
return w
|
||||
}
|
||||
|
||||
func (s *endpointStack) clientToken(t *testing.T) string {
|
||||
at := auth.NewAccessToken(testKey, testSecret).SetVideoGrant(&auth.VideoGrant{RoomJoin: true, Room: "x"})
|
||||
tok, err := at.ToJWT()
|
||||
require.NoError(t, err)
|
||||
return tok
|
||||
}
|
||||
|
||||
func httpEP(path string, methods []string, public bool) *livekit.AgentHttp_AgentEndpoint {
|
||||
return &livekit.AgentHttp_AgentEndpoint{Path: path, Methods: methods, Public: public}
|
||||
}
|
||||
|
||||
// newTargetApp is the local app the worker bridges into; it never listens on a
|
||||
// port reachable through the stack, only via the tunnel.
|
||||
func newTargetApp(t *testing.T, mux *http.ServeMux) *httptest.Server {
|
||||
app := httptest.NewServer(mux)
|
||||
t.Cleanup(app.Close)
|
||||
return app
|
||||
}
|
||||
|
||||
func TestAgentEndpointsCorrectnessGate(t *testing.T) {
|
||||
bigDown := make([]byte, 4<<20)
|
||||
_, _ = rand.Read(bigDown)
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /json", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"ok":true}`))
|
||||
})
|
||||
mux.HandleFunc("POST /upload", func(w http.ResponseWriter, r *http.Request) {
|
||||
sum := sha256.New()
|
||||
n, err := io.Copy(sum, r.Body)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), 500)
|
||||
return
|
||||
}
|
||||
fmt.Fprintf(w, "%d:%x", n, sum.Sum(nil))
|
||||
})
|
||||
mux.HandleFunc("GET /big", func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write(bigDown)
|
||||
})
|
||||
mux.HandleFunc("GET /sse", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
f := w.(http.Flusher)
|
||||
for i := 0; i < 5; i++ {
|
||||
fmt.Fprintf(w, "data: event-%d\n\n", i)
|
||||
f.Flush()
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
}
|
||||
})
|
||||
app := newTargetApp(t, mux)
|
||||
|
||||
stack := newEndpointStack(t, agent.EndpointsConfig{})
|
||||
stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
|
||||
httpEP("/json", []string{"GET"}, true),
|
||||
httpEP("/upload", []string{"POST"}, true),
|
||||
httpEP("/big", []string{"GET"}, true),
|
||||
httpEP("/sse", []string{"GET"}, true),
|
||||
})
|
||||
|
||||
base := stack.ts.URL + "/agents/production"
|
||||
|
||||
t.Run("json round trip", func(t *testing.T) {
|
||||
resp, err := http.Get(base + "/json")
|
||||
require.NoError(t, err)
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
require.Equal(t, 200, resp.StatusCode)
|
||||
require.JSONEq(t, `{"ok":true}`, string(body))
|
||||
})
|
||||
|
||||
t.Run("upload byte exact", func(t *testing.T) {
|
||||
up := make([]byte, 8<<20)
|
||||
_, _ = rand.Read(up)
|
||||
resp, err := http.Post(base+"/upload", "application/octet-stream", bytes.NewReader(up))
|
||||
require.NoError(t, err)
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
require.Equal(t, 200, resp.StatusCode)
|
||||
require.Equal(t, fmt.Sprintf("%d:%x", len(up), sha256.Sum256(up)), string(body))
|
||||
})
|
||||
|
||||
t.Run("download byte exact", func(t *testing.T) {
|
||||
resp, err := http.Get(base + "/big")
|
||||
require.NoError(t, err)
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
require.NoError(t, err)
|
||||
require.True(t, bytes.Equal(bigDown, body))
|
||||
})
|
||||
|
||||
t.Run("sse events arrive one at a time", func(t *testing.T) {
|
||||
start := time.Now()
|
||||
resp, err := http.Get(base + "/sse")
|
||||
require.NoError(t, err)
|
||||
defer resp.Body.Close()
|
||||
br := bufio.NewReader(resp.Body)
|
||||
var arrivals []time.Duration
|
||||
for {
|
||||
line, err := br.ReadString('\n')
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
arrivals = append(arrivals, time.Since(start))
|
||||
}
|
||||
}
|
||||
require.Len(t, arrivals, 5)
|
||||
// incremental delivery: the first event arrives well before the last is
|
||||
// even written (5 x 150ms); buffering the whole body would collapse gaps
|
||||
require.Less(t, arrivals[0], 450*time.Millisecond)
|
||||
require.Greater(t, arrivals[4]-arrivals[0], 300*time.Millisecond)
|
||||
})
|
||||
|
||||
t.Run("32 concurrent requests", func(t *testing.T) {
|
||||
errCh := make(chan error, 32)
|
||||
for i := 0; i < 32; i++ {
|
||||
go func() {
|
||||
resp, err := http.Get(base + "/json")
|
||||
if err == nil {
|
||||
io.Copy(io.Discard, resp.Body)
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode != 200 {
|
||||
err = fmt.Errorf("status %d", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
errCh <- err
|
||||
}()
|
||||
}
|
||||
for i := 0; i < 32; i++ {
|
||||
require.NoError(t, <-errCh)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestAgentEndpointsStatusMapping(t *testing.T) {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(200) })
|
||||
app := newTargetApp(t, mux)
|
||||
|
||||
stack := newEndpointStack(t, agent.EndpointsConfig{})
|
||||
stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
|
||||
httpEP("/hook", []string{"POST"}, true),
|
||||
httpEP("/private", []string{"GET"}, false),
|
||||
})
|
||||
|
||||
base := stack.ts.URL + "/agents/production"
|
||||
|
||||
t.Run("404 unknown path", func(t *testing.T) {
|
||||
resp, _ := http.Get(base + "/nope")
|
||||
resp.Body.Close()
|
||||
require.Equal(t, 404, resp.StatusCode)
|
||||
})
|
||||
t.Run("405 wrong method", func(t *testing.T) {
|
||||
resp, _ := http.Get(base + "/hook")
|
||||
resp.Body.Close()
|
||||
require.Equal(t, 405, resp.StatusCode)
|
||||
})
|
||||
t.Run("401 non-public without token", func(t *testing.T) {
|
||||
resp, _ := http.Get(base + "/private")
|
||||
resp.Body.Close()
|
||||
require.Equal(t, 401, resp.StatusCode)
|
||||
})
|
||||
t.Run("200 non-public with token", func(t *testing.T) {
|
||||
req, _ := http.NewRequest("GET", base+"/private", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+stack.clientToken(t))
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
require.NoError(t, err)
|
||||
resp.Body.Close()
|
||||
require.Equal(t, 200, resp.StatusCode)
|
||||
})
|
||||
t.Run("503 unknown deployment", func(t *testing.T) {
|
||||
resp, _ := http.Get(stack.ts.URL + "/agents/staging/hook")
|
||||
resp.Body.Close()
|
||||
require.Equal(t, 503, resp.StatusCode)
|
||||
})
|
||||
t.Run("307 slash redirect", func(t *testing.T) {
|
||||
c := &http.Client{CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
}}
|
||||
resp, err := c.Post(base+"/hook/", "text/plain", nil)
|
||||
require.NoError(t, err)
|
||||
resp.Body.Close()
|
||||
require.Equal(t, 307, resp.StatusCode)
|
||||
require.Equal(t, "/agents/production/hook", resp.Header.Get("Location"))
|
||||
})
|
||||
}
|
||||
|
||||
func TestAgentEndpointsHOL(t *testing.T) {
|
||||
// one data conn forces every stream onto the same socket: the credit windows
|
||||
// and the write scheduler are the only things standing between a stalled
|
||||
// reader and its siblings
|
||||
blocked := make(chan struct{})
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /drip", func(w http.ResponseWriter, r *http.Request) {
|
||||
f := w.(http.Flusher)
|
||||
buf := make([]byte, 64<<10)
|
||||
for {
|
||||
if _, err := w.Write(buf); err != nil {
|
||||
return
|
||||
}
|
||||
f.Flush()
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
}
|
||||
})
|
||||
mux.HandleFunc("GET /quick", func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte("ok"))
|
||||
})
|
||||
app := newTargetApp(t, mux)
|
||||
|
||||
stack := newEndpointStack(t, agent.EndpointsConfig{DataConnCount: 1})
|
||||
stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
|
||||
httpEP("/drip", []string{"GET"}, true),
|
||||
httpEP("/quick", []string{"GET"}, true),
|
||||
})
|
||||
base := stack.ts.URL + "/agents/production"
|
||||
|
||||
// a stalled client: open /drip, read a little, then stop reading entirely
|
||||
resp, err := http.Get(base + "/drip")
|
||||
require.NoError(t, err)
|
||||
defer resp.Body.Close()
|
||||
small := make([]byte, 1024)
|
||||
_, err = io.ReadFull(resp.Body, small)
|
||||
require.NoError(t, err)
|
||||
// do not read further; the stream's credit window fills and stays full
|
||||
close(blocked)
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// sibling streams on the SAME connection must proceed at full speed
|
||||
for i := 0; i < 5; i++ {
|
||||
start := time.Now()
|
||||
q, err := http.Get(base + "/quick")
|
||||
require.NoError(t, err)
|
||||
body, _ := io.ReadAll(q.Body)
|
||||
q.Body.Close()
|
||||
require.Equal(t, "ok", string(body))
|
||||
require.Less(t, time.Since(start), 2*time.Second,
|
||||
"sibling stream stalled behind a blocked heavy stream")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentEndpointsRetrySafety(t *testing.T) {
|
||||
var hits atomic.Int32
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /json", func(w http.ResponseWriter, r *http.Request) {
|
||||
hits.Add(1)
|
||||
_, _ = w.Write([]byte("ok"))
|
||||
})
|
||||
app := newTargetApp(t, mux)
|
||||
|
||||
// a worker whose local app is unreachable REFUSES streams; the front must
|
||||
// retry on the healthy worker exactly once
|
||||
stack := newEndpointStack(t, agent.EndpointsConfig{})
|
||||
deadTarget := "127.0.0.1:1" // nothing listens
|
||||
eps := []*livekit.AgentHttp_AgentEndpoint{httpEP("/json", []string{"GET"}, true)}
|
||||
|
||||
broken := client.New(client.Config{
|
||||
ServerURL: stack.wsURL(), APIKey: testKey, APISecret: testSecret,
|
||||
AgentName: "test-agent", Deployment: "production",
|
||||
Endpoints: eps, TargetAddr: deadTarget,
|
||||
})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
require.NoError(t, broken.Start(ctx))
|
||||
t.Cleanup(broken.Close)
|
||||
|
||||
stack.startWorker(app.URL, "production", eps)
|
||||
|
||||
// run enough requests that both workers get picked first sometimes
|
||||
okCount := 0
|
||||
for i := 0; i < 12; i++ {
|
||||
resp, err := http.Get(stack.ts.URL + "/agents/production/json")
|
||||
require.NoError(t, err)
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode == 200 && string(body) == "ok" {
|
||||
okCount++
|
||||
}
|
||||
}
|
||||
require.Equal(t, 12, okCount, "REFUSED streams must fall through to the healthy worker")
|
||||
require.EqualValues(t, 12, hits.Load())
|
||||
}
|
||||
|
||||
func TestAgentEndpointsTruncationAborts(t *testing.T) {
|
||||
// a worker connection dying mid-response must abort the client connection,
|
||||
// never expose a clean-looking short body
|
||||
mux := http.NewServeMux()
|
||||
release := make(chan struct{})
|
||||
mux.HandleFunc("GET /partial", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Length", "1000000")
|
||||
_, _ = w.Write(make([]byte, 1000))
|
||||
w.(http.Flusher).Flush()
|
||||
<-release
|
||||
})
|
||||
app := newTargetApp(t, mux)
|
||||
|
||||
stack := newEndpointStack(t, agent.EndpointsConfig{})
|
||||
w := stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
|
||||
httpEP("/partial", []string{"GET"}, true),
|
||||
})
|
||||
|
||||
resp, err := http.Get(stack.ts.URL + "/agents/production/partial")
|
||||
require.NoError(t, err)
|
||||
defer resp.Body.Close()
|
||||
|
||||
head := make([]byte, 1000)
|
||||
_, err = io.ReadFull(resp.Body, head)
|
||||
require.NoError(t, err)
|
||||
|
||||
w.Close() // kill the worker mid-stream
|
||||
close(release)
|
||||
|
||||
_, err = io.ReadAll(resp.Body)
|
||||
require.Error(t, err, "truncated response must not read as clean EOF")
|
||||
}
|
||||
|
||||
func TestAgentEndpointsNoLocalListenerContract(t *testing.T) {
|
||||
// the front never routes undeclared paths: the worker-local health/info
|
||||
// routes are unreachable by construction
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte("local")) })
|
||||
app := newTargetApp(t, mux)
|
||||
|
||||
stack := newEndpointStack(t, agent.EndpointsConfig{})
|
||||
stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
|
||||
httpEP("/declared", []string{"GET"}, true),
|
||||
})
|
||||
|
||||
for _, path := range []string{"/", "/worker", "/undeclared"} {
|
||||
resp, err := http.Get(stack.ts.URL + "/agents/production" + path)
|
||||
require.NoError(t, err)
|
||||
resp.Body.Close()
|
||||
require.Equal(t, 404, resp.StatusCode, path)
|
||||
}
|
||||
}
|
||||
|
||||
// dial helper kept for upgrade tests once the SDK lands; avoids unused imports
|
||||
var _ = net.Dialer{}
|
||||
|
||||
// the upgrade exchange and the raw bidirectional session both ride one stream;
|
||||
// this guards the front's hijack path and the no-half-close rule for upgrades
|
||||
func TestAgentEndpointsWebSocketUpgrade(t *testing.T) {
|
||||
upgrader := websocket.Upgrader{}
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
|
||||
c, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer c.Close()
|
||||
for {
|
||||
mt, msg, err := c.ReadMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if err := c.WriteMessage(mt, append([]byte("echo:"), msg...)); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
app := newTargetApp(t, mux)
|
||||
|
||||
stack := newEndpointStack(t, agent.EndpointsConfig{})
|
||||
stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
|
||||
{Path: "/ws", Kind: livekit.AgentHttp_AEK_WEBSOCKET, Public: true},
|
||||
})
|
||||
|
||||
wsURL := "ws" + strings.TrimPrefix(stack.ts.URL, "http") + "/agents/production/ws"
|
||||
c, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
||||
require.NoError(t, err)
|
||||
if resp != nil {
|
||||
resp.Body.Close()
|
||||
}
|
||||
defer c.Close()
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
payload := fmt.Sprintf("msg-%d", i)
|
||||
require.NoError(t, c.WriteMessage(websocket.TextMessage, []byte(payload)))
|
||||
_, echoed, err := c.ReadMessage()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "echo:"+payload, string(echoed))
|
||||
}
|
||||
|
||||
// a larger frame exercises credit flow through the raw pump
|
||||
big := bytes.Repeat([]byte("x"), 256<<10)
|
||||
require.NoError(t, c.WriteMessage(websocket.BinaryMessage, big))
|
||||
_, echoed, err := c.ReadMessage()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, len(big)+5, len(echoed))
|
||||
}
|
||||
+253
-5
@@ -26,9 +26,11 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/types/known/emptypb"
|
||||
|
||||
"github.com/livekit/livekit-server/pkg/agent"
|
||||
"github.com/livekit/livekit-server/pkg/agent/endpoint"
|
||||
"github.com/livekit/livekit-server/pkg/config"
|
||||
"github.com/livekit/livekit-server/pkg/routing"
|
||||
"github.com/livekit/livekit-server/pkg/rtc"
|
||||
@@ -112,8 +114,11 @@ func DispatchAgentWorkerSignal(c agent.SignalConn, h agent.WorkerSignalHandler,
|
||||
return true
|
||||
}
|
||||
|
||||
func HandshakeAgentWorker(c agent.SignalConn, serverInfo *livekit.ServerInfo, registration agent.WorkerRegistration, l logger.Logger) (r agent.WorkerRegistration, ok bool) {
|
||||
func HandshakeAgentWorker(c agent.SignalConn, serverInfo *livekit.ServerInfo, registration agent.WorkerRegistration, l logger.Logger, opts ...func(*agent.WorkerRegisterer)) (r agent.WorkerRegistration, ok bool) {
|
||||
wr := agent.NewWorkerRegisterer(c, serverInfo, registration)
|
||||
for _, opt := range opts {
|
||||
opt(wr)
|
||||
}
|
||||
if err := c.SetReadDeadline(wr.Deadline()); err != nil {
|
||||
return
|
||||
}
|
||||
@@ -138,8 +143,12 @@ type AgentService struct {
|
||||
|
||||
type AgentHandler struct {
|
||||
agentServer rpc.AgentInternalServer
|
||||
mu sync.Mutex
|
||||
logger logger.Logger
|
||||
// the server's only configured api key, when exactly one exists: the
|
||||
// unauthenticated scope for public endpoints regardless of which node holds
|
||||
// the workers
|
||||
singleAPIKey string
|
||||
mu sync.Mutex
|
||||
logger logger.Logger
|
||||
|
||||
serverInfo *livekit.ServerInfo
|
||||
workers map[string]*agent.Worker
|
||||
@@ -147,6 +156,9 @@ type AgentHandler struct {
|
||||
keyProvider auth.KeyProvider
|
||||
targetLoad float32
|
||||
|
||||
endpointRegistry *endpoint.Registry
|
||||
endpointsConfig agent.EndpointsConfig
|
||||
|
||||
namespaceWorkers map[workerKey][]*agent.Worker
|
||||
roomKeyCount int
|
||||
publisherKeyCount int
|
||||
@@ -199,10 +211,34 @@ func NewAgentService(
|
||||
agent.PublisherAgentTopic,
|
||||
agent.ParticipantAgentTopic,
|
||||
)
|
||||
s.AgentHandler.endpointsConfig = conf.Agents.Endpoints
|
||||
if len(conf.Keys) == 1 {
|
||||
for key := range conf.Keys {
|
||||
s.AgentHandler.singleAPIKey = key
|
||||
}
|
||||
}
|
||||
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// EndpointFront is the /agents/{deployment}/{path...} handler backed by this
|
||||
// node's attached workers. Project scope comes from validated grants when a
|
||||
// token is present; unauthenticated requests reach public endpoints only.
|
||||
func (s *AgentService) EndpointFront() http.Handler {
|
||||
front := endpoint.NewFront(s.endpointRegistry, func(r *http.Request) (string, bool) {
|
||||
if claims := GetGrants(r.Context()); claims != nil {
|
||||
return GetAPIKey(r.Context()), true
|
||||
}
|
||||
// unauthenticated: with a single configured key the scope is
|
||||
// unambiguous even when this node holds no registrations (multi-node)
|
||||
return s.singleAPIKey, false
|
||||
}, s.logger)
|
||||
front.WithSingleKeyFallback()
|
||||
return front
|
||||
}
|
||||
|
||||
func (s *AgentService) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
attach := r.URL.Query().Get("attach") != ""
|
||||
if conn, registration, ok := s.upgrader.Upgrade(w, r, nil); ok {
|
||||
// bound the size of a single signalling frame so an oversized message is
|
||||
// rejected by the transport before being fully buffered in memory. This
|
||||
@@ -211,13 +247,157 @@ func (s *AgentService) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
if s.signalMessageSizeLimit > 0 {
|
||||
conn.SetReadLimit(s.signalMessageSizeLimit)
|
||||
}
|
||||
|
||||
if attach {
|
||||
// data wire: no registration handshake, no signal loop; the wire
|
||||
// speaks AgentHttp.Frame exclusively and the endpoint mux owns it
|
||||
// after a successful attach
|
||||
HandleEndpointAttach(s.endpointRegistry, NewEndpointWireConn(conn), s.wireParams())
|
||||
return
|
||||
}
|
||||
|
||||
sigConn := NewWSSignalConnection(conn, s.signalMessageSizeLimit)
|
||||
defer sigConn.Close()
|
||||
|
||||
s.HandleConnection(r.Context(), sigConn, registration)
|
||||
}
|
||||
}
|
||||
|
||||
// endpointWireIdleTimeout is the read-side liveness bound on adopted wires:
|
||||
// the SDK pings every 30s at the websocket level, so a wire silent for this
|
||||
// long is dead and must release its pool slot instead of occupying it until a
|
||||
// write happens to trip the write-stall deadline.
|
||||
const endpointWireIdleTimeout = 2 * time.Minute
|
||||
|
||||
// endpointWireConn frames AgentHttp.Frame over a raw websocket: one frame per
|
||||
// binary message.
|
||||
type endpointWireConn struct {
|
||||
ws *websocket.Conn
|
||||
writeMu sync.Mutex
|
||||
}
|
||||
|
||||
// NewEndpointWireConn wraps an upgraded agent websocket as a data-plane wire.
|
||||
func NewEndpointWireConn(ws *websocket.Conn) endpoint.WireConn {
|
||||
c := &endpointWireConn{ws: ws}
|
||||
ws.SetPingHandler(func(m string) error {
|
||||
_ = ws.SetReadDeadline(time.Now().Add(endpointWireIdleTimeout))
|
||||
return ws.WriteControl(websocket.PongMessage, []byte(m), time.Now().Add(10*time.Second))
|
||||
})
|
||||
return c
|
||||
}
|
||||
|
||||
func (c *endpointWireConn) WriteFrame(f *livekit.AgentHttp_Frame) error {
|
||||
b, err := proto.Marshal(f)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.writeMu.Lock()
|
||||
defer c.writeMu.Unlock()
|
||||
return c.ws.WriteMessage(websocket.BinaryMessage, b)
|
||||
}
|
||||
|
||||
func (c *endpointWireConn) ReadFrame() (*livekit.AgentHttp_Frame, error) {
|
||||
for {
|
||||
_ = c.ws.SetReadDeadline(time.Now().Add(endpointWireIdleTimeout))
|
||||
mt, b, err := c.ws.ReadMessage()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if mt != websocket.BinaryMessage {
|
||||
continue
|
||||
}
|
||||
f := &livekit.AgentHttp_Frame{}
|
||||
if err := proto.Unmarshal(b, f); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return f, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (c *endpointWireConn) SetWriteDeadline(t time.Time) error {
|
||||
return c.ws.SetWriteDeadline(t)
|
||||
}
|
||||
|
||||
func (c *endpointWireConn) SetReadDeadline(t time.Time) error {
|
||||
return c.ws.SetReadDeadline(t)
|
||||
}
|
||||
|
||||
func (c *endpointWireConn) Close() error {
|
||||
return c.ws.Close()
|
||||
}
|
||||
|
||||
// wireParams are this node's wire-level flow-control parameters, announced to
|
||||
// the worker in each attach response.
|
||||
func (s *AgentService) wireParams() endpoint.WireParams {
|
||||
p := endpoint.WireParams{
|
||||
CreditWindow: s.endpointsConfig.CreditWindow,
|
||||
ConnectionWindow: s.endpointsConfig.ConnectionWindow,
|
||||
MaxFrameSize: s.endpointsConfig.MaxFrameSize,
|
||||
MaxStreamsPerConn: s.endpointsConfig.MaxStreamsPerConn,
|
||||
}.WithDefaults()
|
||||
// data frames ride the same websocket read limit as signalling: keep the
|
||||
// frame size safely under it so a full frame can never kill the wire
|
||||
if lim := s.signalMessageSizeLimit; lim > 0 && int64(p.MaxFrameSize) > lim/2 {
|
||||
p.MaxFrameSize = uint32(lim / 2)
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// HandleEndpointAttach adopts a worker-dialed data wire: the first frame is
|
||||
// Attach on stream 0, validated against the registration's epoch and token; the
|
||||
// response carries this node's wire parameters. Shared by OSS and cloud
|
||||
// servers.
|
||||
func HandleEndpointAttach(registry *endpoint.Registry, wire endpoint.WireConn, params endpoint.WireParams) {
|
||||
if err := wire.SetReadDeadline(time.Now().Add(agent.RegisterTimeout)); err != nil {
|
||||
_ = wire.Close()
|
||||
return
|
||||
}
|
||||
f, err := wire.ReadFrame()
|
||||
if err != nil {
|
||||
_ = wire.Close()
|
||||
return
|
||||
}
|
||||
att, ok := f.Message.(*livekit.AgentHttp_Frame_Attach)
|
||||
if !ok || f.StreamId != 0 {
|
||||
_ = wire.Close()
|
||||
return
|
||||
}
|
||||
|
||||
effective := params.WithDefaults()
|
||||
respond := func(errStr string) error {
|
||||
resp := &livekit.AgentHttp_AttachDataConnectionResponse{Error: errStr}
|
||||
if errStr == "" {
|
||||
resp.CreditWindow = effective.CreditWindow
|
||||
resp.ConnectionWindow = effective.ConnectionWindow
|
||||
resp.MaxFrameSize = effective.MaxFrameSize
|
||||
resp.MaxStreamsPerConn = effective.MaxStreamsPerConn
|
||||
}
|
||||
// the only wire write outside the scheduler: bound it too
|
||||
_ = wire.SetWriteDeadline(time.Now().Add(agent.RegisterTimeout))
|
||||
return wire.WriteFrame(&livekit.AgentHttp_Frame{
|
||||
Message: &livekit.AgentHttp_Frame_AttachResponse{AttachResponse: resp},
|
||||
})
|
||||
}
|
||||
|
||||
a := att.Attach
|
||||
// the slot is reserved before the ack is written (never a success ack for a
|
||||
// wire that then loses the cap race) and adopted only after it, so the
|
||||
// worker cannot observe stream frames ahead of the attach outcome
|
||||
ticket, err := registry.BeginAttach(a.GetWorkerId(), a.GetInstanceId(), a.GetAttachToken())
|
||||
if err != nil {
|
||||
_ = respond(err.Error())
|
||||
_ = wire.Close()
|
||||
return
|
||||
}
|
||||
if err := respond(""); err != nil {
|
||||
ticket.Abort()
|
||||
_ = wire.Close()
|
||||
return
|
||||
}
|
||||
if !ticket.Complete(wire, effective) {
|
||||
_ = wire.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func NewAgentHandler(
|
||||
agentServer rpc.AgentInternalServer,
|
||||
keyProvider auth.KeyProvider,
|
||||
@@ -240,11 +420,37 @@ func NewAgentHandler(
|
||||
roomTopic: roomTopic,
|
||||
publisherTopic: publisherTopic,
|
||||
participantTopic: participantTopic,
|
||||
endpointRegistry: endpoint.NewRegistry(),
|
||||
}
|
||||
}
|
||||
|
||||
// endpointSettings validates a registration's manifest and negotiates the
|
||||
// data-plane settings, minting the epoch's attach token.
|
||||
func (h *AgentHandler) endpointSettings(req *livekit.RegisterWorkerRequest) (*livekit.AgentHttp_AgentEndpointSettings, error) {
|
||||
if h.endpointsConfig.Disabled {
|
||||
return nil, errors.New("agent HTTP endpoints are disabled on this server")
|
||||
}
|
||||
if _, err := endpoint.ParseManifest(req.GetEndpoints()); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if req.GetInstanceId() == "" {
|
||||
return nil, errors.New("registrations with endpoints require an instance_id")
|
||||
}
|
||||
settings := endpoint.Settings{
|
||||
Protocol: endpoint.CurrentProtocol,
|
||||
AttachToken: endpoint.NewAttachToken(),
|
||||
DataConnCount: h.endpointsConfig.DataConnCount,
|
||||
}
|
||||
if settings.DataConnCount == 0 {
|
||||
settings.DataConnCount = endpoint.DefaultDataConnCount
|
||||
}
|
||||
return settings.Proto(), nil
|
||||
}
|
||||
|
||||
func (h *AgentHandler) HandleConnection(ctx context.Context, conn agent.SignalConn, registration agent.WorkerRegistration) {
|
||||
registration, ok := HandshakeAgentWorker(conn, h.serverInfo, registration, h.logger)
|
||||
registration, ok := HandshakeAgentWorker(conn, h.serverInfo, registration, h.logger, func(wr *agent.WorkerRegisterer) {
|
||||
wr.WithEndpointSettings(h.endpointSettings)
|
||||
})
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
@@ -255,15 +461,57 @@ func (h *AgentHandler) HandleConnection(ctx context.Context, conn agent.SignalCo
|
||||
worker := agent.NewWorker(registration, apiKey, apiSecret, conn, h.logger)
|
||||
h.registerWorker(worker)
|
||||
|
||||
endpointReg := h.registerEndpoints(worker)
|
||||
|
||||
handlerWorker := &agentHandlerWorker{h, worker}
|
||||
for ok := true; ok; {
|
||||
ok = DispatchAgentWorkerSignal(conn, handlerWorker, worker.Logger())
|
||||
}
|
||||
|
||||
if endpointReg != nil {
|
||||
h.endpointRegistry.Deregister(endpointReg)
|
||||
}
|
||||
h.deregisterWorker(worker)
|
||||
worker.Close()
|
||||
}
|
||||
|
||||
// registerEndpoints adopts a worker's endpoint manifest into the data-plane
|
||||
// registry. The registration lives exactly as long as the control connection.
|
||||
func (h *AgentHandler) registerEndpoints(w *agent.Worker) *endpoint.Registration {
|
||||
settings := w.EndpointSettings
|
||||
if settings == nil {
|
||||
return nil
|
||||
}
|
||||
manifest, err := endpoint.ParseManifest(w.Endpoints)
|
||||
if err != nil {
|
||||
// validated during the handshake; a failure here is a programming error
|
||||
w.Logger().Errorw("endpoint manifest failed to re-parse", err)
|
||||
return nil
|
||||
}
|
||||
reg := &endpoint.Registration{
|
||||
WorkerID: w.ID,
|
||||
InstanceID: w.InstanceID,
|
||||
APIKey: w.APIKey(),
|
||||
Deployment: w.Deployment,
|
||||
Manifest: manifest,
|
||||
Settings: endpoint.Settings{
|
||||
Protocol: settings.GetProtocol(),
|
||||
AttachToken: settings.GetAttachToken(),
|
||||
DataConnCount: settings.GetDataConnectionCount(),
|
||||
},
|
||||
Logger: w.Logger(),
|
||||
Load: w.Load,
|
||||
Draining: w.Draining,
|
||||
}
|
||||
if err := h.endpointRegistry.Register(reg); err != nil {
|
||||
w.Logger().Errorw("failed to register endpoints", err)
|
||||
return nil
|
||||
}
|
||||
w.Logger().Infow("endpoints registered",
|
||||
"deployment", w.Deployment, "routes", len(w.Endpoints))
|
||||
return reg
|
||||
}
|
||||
|
||||
func (h *AgentHandler) registerWorker(w *agent.Worker) {
|
||||
h.mu.Lock()
|
||||
|
||||
|
||||
+6
-2
@@ -76,8 +76,12 @@ func (m *APIKeyAuthMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request,
|
||||
|
||||
authToken = authHeader[len(bearerPrefix):]
|
||||
} else {
|
||||
// attempt to find from request header
|
||||
authToken = r.FormValue(accessTokenParam)
|
||||
// attempt to find from the query string. FormValue would also parse
|
||||
// url-encoded POST bodies, consuming the body of any request that gets
|
||||
// proxied further (agent HTTP endpoints)
|
||||
if r.URL != nil {
|
||||
authToken = r.URL.Query().Get(accessTokenParam)
|
||||
}
|
||||
}
|
||||
|
||||
if authToken != "" {
|
||||
|
||||
@@ -40,6 +40,7 @@ import (
|
||||
"github.com/livekit/protocol/logger"
|
||||
"github.com/livekit/protocol/utils/xtwirp"
|
||||
|
||||
"github.com/livekit/livekit-server/pkg/agent/endpoint"
|
||||
"github.com/livekit/livekit-server/pkg/config"
|
||||
"github.com/livekit/livekit-server/pkg/routing"
|
||||
"github.com/livekit/livekit-server/version"
|
||||
@@ -150,6 +151,7 @@ func NewLivekitServer(conf *config.Config,
|
||||
rtcService.SetupRoutes(mux)
|
||||
whipService.SetupRoutes(mux)
|
||||
mux.Handle("/agent", agentService)
|
||||
mux.Handle(endpoint.PathPrefix, agentService.EndpointFront())
|
||||
mux.HandleFunc("/", s.defaultHandler)
|
||||
|
||||
s.httpServer = &http.Server{
|
||||
|
||||
Reference in New Issue
Block a user