diff --git a/pkg/agent/endpoint/attempt.go b/pkg/agent/endpoint/attempt.go index 8bb539a1a..ff925a926 100644 --- a/pkg/agent/endpoint/attempt.go +++ b/pkg/agent/endpoint/attempt.go @@ -37,7 +37,6 @@ type attempt struct { escPath string // target is the origin-form request target, still percent-encoded target string - route *Route requestID string granted bool body io.Reader @@ -96,9 +95,6 @@ func (a *attempt) newPreamble() *livekit.AgentHttp_StreamPreamble { ClientAddr: clientAddr, Scheme: scheme, } - if a.route != nil && a.route.Template != nil { - p.Route = a.route.Template.String() - } return p } diff --git a/pkg/agent/endpoint/front.go b/pkg/agent/endpoint/front.go index 434e1e2f0..1a198efbd 100644 --- a/pkg/agent/endpoint/front.go +++ b/pkg/agent/endpoint/front.go @@ -208,15 +208,15 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } - candidates := f.params.Registry.Candidates(access.APIKey, agentName, deployment) - if len(candidates) == 0 && f.params.Fallback == nil { + tbl := f.params.Registry.table(access.APIKey, agentName, deployment) + if tbl == nil && f.params.Fallback == nil { f.writeUnavailable(w, "no workers available for deployment") return } mask := methodMask(r.Method) granted := access.Level >= AccessGranted - matched, route, partial, denied := matchDeployment(candidates, path, mask, granted) + matched, partial, denied := f.matchDeployment(tbl, path, mask, granted) // no exact match: if only the trailing-slash alternate matches a registered // route, normalize the path to that form and serve it directly (no client // redirect). The exact form is tried first, so a route registered with a @@ -224,9 +224,9 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { // the registered form. When the request must be relayed, the serving node // runs this same normalization, so no redirect is ever emitted. if len(matched) == 0 && !partial && !denied { - if alt, altEsc, ok := slashAlternatePaths(candidates, path, escPath, mask); ok { + if alt, altEsc, ok := slashAlternatePaths(tbl, path, escPath, mask); ok { path, escPath = alt, altEsc - matched, route, partial, denied = matchDeployment(candidates, path, mask, granted) + matched, partial, denied = f.matchDeployment(tbl, path, mask, granted) } } if len(matched) == 0 && f.params.Fallback != nil { @@ -240,7 +240,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { }) { return } - if len(candidates) == 0 && !denied && !partial { + if tbl == nil && !denied && !partial { f.writeUnavailable(w, "no workers available for deployment") return } @@ -268,7 +268,6 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { req: r, escPath: escPath, target: requestTarget(escPath, r.URL.RawQuery), - route: route, requestID: reqID, granted: granted, pools: f.pools, @@ -288,15 +287,18 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { attempted := make(map[*Registration]bool) for range maxAttempts { - reg := pickWorker(matched, attempted) - if reg == nil { + picked := pickWorker(matched, attempted) + if picked == nil { break } - attempted[reg] = true + attempted[picked.reg] = true a.before = bodyConsumed.Load() a.refreshTimeout() - switch f.bridge(w, a, reg) { + // the preamble is re-serialized per attempt, so this reaches only the + // worker it was set for + a.preamble.Route = picked.raw + switch f.bridge(w, a, picked.reg) { case bridgeDone: return case bridgeAbort: @@ -324,42 +326,37 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { f.writeUnavailable(w, "no worker could serve the request") } -// matchDeployment resolves a path against every candidate worker's manifest. -// A candidate joins matched when it can serve it; route stays nil when a table -// was too ambiguous to decide, so matched is the dispatch test. -func matchDeployment(candidates []*Registration, path string, mask router.Mask, granted bool) (matched []*Registration, route *Route, partial, denied bool) { - for _, reg := range candidates { - rt, res := reg.Manifest.Match(path, mask) - switch res { - case router.ResultFull: - if !granted && !rt.Public { - denied = true - continue - } - if route == nil { - route = rt - } - matched = append(matched, reg) - case router.ResultPartial: - partial = true - case router.ResultOverBudget: - // no route was decided, so its Public flag is unknown and only a grant - // can clear the request - if !granted { - denied = true - continue - } - matched = append(matched, reg) +// matchDeployment resolves a path against the deployment's merged route table. +// matched is the dispatch test: a route can be decided while nothing is left to +// serve it, and an undecidable table dispatches with no route at all. +func (f *Front) matchDeployment(tbl *routeTable, path string, mask router.Mask, granted bool) (matched []routeWorker, partial, denied bool) { + if tbl == nil { + return nil, false, false + } + var res router.Result + matched, _, res, denied = tbl.match(path, mask, granted) + switch res { + case router.ResultPartial: + partial = true + case router.ResultOverBudget: + // no route was decided, so its Public flag is unknown and only a grant + // can clear the request + if !granted { + denied = true + break + } + for _, reg := range f.params.Registry.Candidates(tbl.key.apiKey, tbl.key.agentName, tbl.key.deployment) { + matched = append(matched, routeWorker{reg: reg}) } } return } -// slashAlternatePaths looks for a candidate that serves the trailing-slash -// alternate of path, returning the decoded and escaped forms to retry. Both +// slashAlternatePaths looks for the trailing-slash alternate of path in the +// deployment's table, returning the decoded and escaped forms to retry. Both // forms come out of the same transform. -func slashAlternatePaths(candidates []*Registration, path, escPath string, mask router.Mask) (string, string, bool) { - if path == "/" { +func slashAlternatePaths(tbl *routeTable, path, escPath string, mask router.Mask) (string, string, bool) { + if tbl == nil || path == "/" { return path, escPath, false } var alt, altEsc string @@ -373,12 +370,10 @@ func slashAlternatePaths(candidates []*Registration, path, escPath string, mask } else { alt, altEsc = path+"/", escPath+"/" } - for _, reg := range candidates { - if _, res := reg.Manifest.Match(alt, mask); res == router.ResultFull { - return alt, altEsc, true - } + if !tbl.serves(alt, mask) { + return path, escPath, false } - return path, escPath, false + return alt, altEsc, true } // endpointPath is a request split into its routing components. @@ -432,20 +427,20 @@ func splitEndpointPath(u *url.URL) (endpointPath, error) { } // pickWorker chooses a worker by the power of two choices: sample two eligible -// registrations at random and take the one with fewer in-flight streams. +// declarations at random and take the one with fewer in-flight streams. // Eligible = not already attempted, has a live session, not draining. -func pickWorker(regs []*Registration, ignore map[*Registration]bool) *Registration { - var eligible []*Registration - for _, reg := range regs { - if ignore[reg] || !reg.HasSession() || reg.IsDraining() { +func pickWorker(workers []routeWorker, ignore map[*Registration]bool) *routeWorker { + var eligible []routeWorker + for _, w := range workers { + if ignore[w.reg] || !w.reg.HasSession() || w.reg.IsDraining() { continue } - eligible = append(eligible, reg) + eligible = append(eligible, w) } if len(eligible) == 0 { return nil } - return eligible[p2c(eligible, (*Registration).InflightStreams)] + return &eligible[p2c(eligible, func(w routeWorker) int { return w.reg.InflightStreams() })] } // p2c returns the index of the less-loaded of two distinct random draws from diff --git a/pkg/agent/endpoint/manifest.go b/pkg/agent/endpoint/manifest.go index 066755d7d..2946f986e 100644 --- a/pkg/agent/endpoint/manifest.go +++ b/pkg/agent/endpoint/manifest.go @@ -20,6 +20,8 @@ import ( "net/http" "strings" + "go.uber.org/zap/zapcore" + "github.com/livekit/protocol/livekit" "github.com/livekit/livekit-server/pkg/agent/endpoint/router" @@ -28,15 +30,46 @@ import ( const MaxManifestRoutes = 256 -// Manifest is a worker's route table, in declaration order. -type Manifest = router.Router[*Route] +// Manifest is one worker's validated route table, in declaration order. +type Manifest struct { + Endpoints []Endpoint + ambiguous []string +} -// Route is one validated manifest entry. -type Route struct { +// Endpoint is one validated manifest entry. Mask carries the methods exactly as +// declared, so it may name several. +type Endpoint struct { Template *router.Template + Mask router.Mask Public bool } +// Ambiguous returns the declared templates whose shape forces the matcher to +// backtrack. The result is read-only. +func (m *Manifest) Ambiguous() []string { + if m == nil { + return nil + } + return m.ambiguous +} + +// MarshalLogObject describes the manifest's shape. +func (m *Manifest) MarshalLogObject(e zapcore.ObjectEncoder) error { + if m == nil { + return nil + } + e.AddInt("routes", len(m.Endpoints)) + if len(m.ambiguous) == 0 { + return nil + } + return e.AddArray("ambiguous", zapcore.ArrayMarshalerFunc(func(a zapcore.ArrayEncoder) error { + for _, t := range m.ambiguous { + a.AppendString(t) + } + return nil + })) +} + // methodMask tags a route with the methods it serves. The set is exactly the // verbs FastAPI can route; anything else masks to 0 and matches nothing. func methodMask(method string) router.Mask { @@ -85,13 +118,12 @@ func NegotiateSettings(req *livekit.RegisterWorkerRequest) (*livekit.AgentHttp_A return &livekit.AgentHttp_AgentEndpointSettings{Protocol: negotiated}, nil } -// ParseManifest validates a registration's endpoint list and compiles it into -// a router. +// 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) } - b := router.NewBuilder[*Route]() + m := &Manifest{Endpoints: make([]Endpoint, 0, len(endpoints))} for _, ep := range endpoints { tpl, err := router.ParseTemplate(ep.GetPath()) if err != nil { @@ -109,15 +141,16 @@ func ParseManifest(endpoints []*livekit.AgentHttp_AgentEndpoint) (*Manifest, err if u != method { return nil, fmt.Errorf("endpoint %q method %q must be uppercase", ep.GetPath(), method) } - m := methodMask(u) - if m == 0 { + b := methodMask(u) + if b == 0 { return nil, fmt.Errorf("endpoint %q declares unsupported method %q", ep.GetPath(), method) } - mask |= m + mask |= b } - if err := b.Add(tpl, mask, &Route{Template: tpl, Public: ep.GetPublic()}); err != nil { - return nil, err + if tpl.Ambiguous() { + m.ambiguous = append(m.ambiguous, tpl.String()) } + m.Endpoints = append(m.Endpoints, Endpoint{Template: tpl, Mask: mask, Public: ep.GetPublic()}) } - return b.Build(), nil + return m, nil } diff --git a/pkg/agent/endpoint/manifest_test.go b/pkg/agent/endpoint/manifest_test.go index 00e81b561..d81922437 100644 --- a/pkg/agent/endpoint/manifest_test.go +++ b/pkg/agent/endpoint/manifest_test.go @@ -20,40 +20,37 @@ func ep(path string, methods []string, public bool) *livekit.AgentHttp_AgentEndp 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{ + tbl := tableOf(t, mustManifest(t, ep("/x", []string{"GET"}, true), ep("/x", []string{"POST"}, true), - }) - require.NoError(t, err) + )) - r, res := m.Match("/x", methodMask(http.MethodPost)) + _, r, res, _ := tbl.match("/x", methodMask(http.MethodPost), true) require.Equal(t, router.ResultFull, res) require.Equal(t, "/x", r.Template.String()) // the manifest carries the app's methods verbatim: FastAPI does not imply // HEAD from GET, so neither does the matcher - _, res = m.Match("/x", methodMask(http.MethodHead)) + _, _, res, _ = tbl.match("/x", methodMask(http.MethodHead), true) require.Equal(t, router.ResultPartial, res) // PARTIAL only when no route serves the method - _, res = m.Match("/x", methodMask(http.MethodDelete)) + _, _, res, _ = tbl.match("/x", methodMask(http.MethodDelete), true) require.Equal(t, router.ResultPartial, res) // an unroutable method masks to 0 - _, res = m.Match("/x", methodMask("BREW")) + _, _, res, _ = tbl.match("/x", methodMask("BREW"), true) require.Equal(t, router.ResultPartial, res) - _, res = m.Match("/nope", methodMask(http.MethodGet)) + _, _, res, _ = tbl.match("/nope", methodMask(http.MethodGet), true) require.Equal(t, router.ResultNone, res) } func TestSlashAlternatePaths(t *testing.T) { - m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ + tbl := tableOf(t, mustManifest(t, ep("/hook", []string{"POST"}, true), ep("/", []string{"POST"}, true), - }) - require.NoError(t, err) - candidates := []*Registration{{Manifest: m}} + )) cases := []struct { path, escPath string @@ -72,7 +69,7 @@ func TestSlashAlternatePaths(t *testing.T) { {"/", "/", "/", "/", false}, } for _, c := range cases { - alt, altEsc, ok := slashAlternatePaths(candidates, c.path, c.escPath, methodMask(http.MethodPost)) + alt, altEsc, ok := slashAlternatePaths(tbl, c.path, c.escPath, methodMask(http.MethodPost)) require.Equal(t, c.ok, ok, "%q %q", c.path, c.escPath) require.Equal(t, c.alt, alt, "%q %q", c.path, c.escPath) require.Equal(t, c.altEsc, altEsc, "%q %q", c.path, c.escPath) @@ -82,16 +79,16 @@ func TestSlashAlternatePaths(t *testing.T) { // Templates anchor as `\n?$`: path refuses '\n' where str accepts it, so a // decoded %0A would otherwise change which route wins. func TestManifestTrailingNewlineRouteIdentity(t *testing.T) { - m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ + tbl := tableOf(t, mustManifest(t, ep("/files/{p:path}", []string{"GET"}, false), ep("/files/{p}", []string{"GET"}, true), - }) - require.NoError(t, err) + )) - r, res := m.Match("/files/x\n", methodMask(http.MethodGet)) + matched, r, res, denied := tbl.match("/files/x\n", methodMask(http.MethodGet), false) require.Equal(t, router.ResultFull, res) require.Equal(t, "/files/{p:path}", r.Template.String()) - require.False(t, r.Public, "the private route must win, as it does on the worker") + require.Empty(t, matched) + require.True(t, denied, "the private route must win, as it does on the worker") } func TestManifestValidation(t *testing.T) { diff --git a/pkg/agent/endpoint/registry.go b/pkg/agent/endpoint/registry.go index 7e9d78c45..4a54c2129 100644 --- a/pkg/agent/endpoint/registry.go +++ b/pkg/agent/endpoint/registry.go @@ -19,6 +19,8 @@ import ( "errors" "slices" "sync" + + "github.com/livekit/protocol/logger" ) // DefaultDeployment is the URL segment that addresses workers registered with an @@ -99,6 +101,10 @@ func NewRegistration(params RegistrationParams) *Registration { } } +func (r *Registration) key() regKey { + return regKey{r.APIKey, r.AgentName, normalizeDeployment(r.Deployment)} +} + // IsDraining is false when no drain signal was supplied. func (r *Registration) IsDraining() bool { return r.draining != nil && r.draining() @@ -168,10 +174,17 @@ func (r *Registration) close() { // Registry tracks data-plane registrations on this node, keyed by // (api key, agent name, deployment). The api key is the project identity in OSS. +// +// Each key also owns a merged route table. Registrations and routes move +// together under g.lock, so a worker present in byKey always has its routes +// installed. type Registry struct { - lock sync.RWMutex - regs map[string]*Registration // by worker id - byKey map[regKey][]*Registration + logger logger.Logger + + lock sync.RWMutex + regs map[string]*Registration // by worker id + byKey map[regKey][]*Registration + tables map[regKey]*routeTable } type regKey struct { @@ -182,8 +195,10 @@ type regKey struct { func NewRegistry() *Registry { return &Registry{ - regs: make(map[string]*Registration), - byKey: make(map[regKey][]*Registration), + logger: logger.GetLogger().WithComponent("agents.endpoint"), + regs: make(map[string]*Registration), + byKey: make(map[regKey][]*Registration), + tables: make(map[regKey]*routeTable), } } @@ -192,24 +207,58 @@ func NewRegistry() *Registry { // able to strand the new epoch (its own Deregister is a no-op once replaced). // The superseded epoch's session is closed. func (g *Registry) Register(r *Registration) { - key := regKey{r.APIKey, r.AgentName, normalizeDeployment(r.Deployment)} + key := r.key() g.lock.Lock() - old := g.regs[r.WorkerID] + superseded := g.regs[r.WorkerID] + old := superseded if old != nil { - g.removeLocked(old) + g.unlinkLocked(old) + if old.key() != key { + // nothing pins agent name or deployment across epochs, so the + // retiring epoch's routes may live in another table + g.retractLocked(old) + old = nil + } } + // one transaction, so an unchanged manifest keeps its Route pointers + g.tableLocked(key).mutate(old, r) g.regs[r.WorkerID] = r g.byKey[key] = append(g.byKey[key], r) g.lock.Unlock() - if old != nil { - old.close() + if superseded != nil { + superseded.close() } } +// tableLocked returns the key's route table, creating it on first use. Callers +// hold g.lock. +func (g *Registry) tableLocked(key regKey) *routeTable { + tbl := g.tables[key] + if tbl == nil { + tbl = newRouteTable(key, g.logger) + g.tables[key] = tbl + } + return tbl +} + +// table returns a deployment's merged route table, or nil when no worker holds +// the key. +func (g *Registry) table(apiKey, agentName, deployment string) *routeTable { + g.lock.RLock() + defer g.lock.RUnlock() + return g.tables[regKey{apiKey, agentName, normalizeDeployment(deployment)}] +} + // removeLocked unlinks a registration from all indexes. Callers hold g.lock. func (g *Registry) removeLocked(r *Registration) { + g.unlinkLocked(r) + g.retractLocked(r) +} + +// unlinkLocked drops a registration from the worker-id and key indexes. +func (g *Registry) unlinkLocked(r *Registration) { delete(g.regs, r.WorkerID) - key := regKey{r.APIKey, r.AgentName, normalizeDeployment(r.Deployment)} + key := r.key() if regs := g.byKey[key]; len(regs) > 0 { if i := slices.Index(regs, r); i != -1 { regs = slices.Delete(regs, i, i+1) @@ -222,6 +271,20 @@ func (g *Registry) removeLocked(r *Registration) { } } +// retractLocked drops a registration's routes. A table exists only while it +// holds routes. +func (g *Registry) retractLocked(r *Registration) { + key := r.key() + tbl := g.tables[key] + if tbl == nil { + return + } + tbl.mutate(r, nil) + if tbl.empty() { + delete(g.tables, key) + } +} + // Deregister removes exactly this registration; it is a no-op when a newer // epoch has already superseded it. func (g *Registry) Deregister(r *Registration) { diff --git a/pkg/agent/endpoint/routes.go b/pkg/agent/endpoint/routes.go new file mode 100644 index 000000000..c146cc396 --- /dev/null +++ b/pkg/agent/endpoint/routes.go @@ -0,0 +1,257 @@ +// 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 ( + "cmp" + "slices" + "sync" + "sync/atomic" + + "github.com/livekit/protocol/logger" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/router" +) + +// routeKey identifies one route: the template shape it matches and the single +// method it serves. Templates differing only in param names share a shape, so +// /x/{a} and /x/{b} are one route. +type routeKey struct { + canonical string + method router.Mask +} + +// Route is one endpoint of a deployment and every worker that declared it. The +// compiled table holds these by pointer: a route outlives the tables built +// around it, and its worker set changes without a rebuild. +type Route struct { + Template *router.Template + Mask router.Mask + + key routeKey + // ord mirrors the merge position the published table was built from + ord uint32 + + lock sync.RWMutex + workers []routeWorker +} + +// routeWorker is one worker's declaration of a route. Declarers of one route can +// disagree about visibility, position and spelling. +type routeWorker struct { + reg *Registration + raw string // the template as this worker declared it + public bool + idx uint32 // its position in this worker's manifest +} + +// declare records one worker's claim. A registration appears at most once, and +// callers walk endpoints in declaration order, so a path and method named twice +// keeps its first position. +func (r *Route) declare(w routeWorker) { + r.lock.Lock() + defer r.lock.Unlock() + for i := range r.workers { + if r.workers[i].reg == w.reg { + return + } + } + r.workers = append(r.workers, w) +} + +// retract drops a registration's claim. Matching is by pointer: a supersede has +// both epochs live at once under one worker id. +func (r *Route) retract(reg *Registration) { + r.lock.Lock() + defer r.lock.Unlock() + r.workers = slices.DeleteFunc(r.workers, func(w routeWorker) bool { return w.reg == reg }) +} + +// position reports the route's merge position - the lowest index it holds in any +// declaring manifest - and whether anything still declares it. +func (r *Route) position() (ord uint32, live bool) { + r.lock.RLock() + defer r.lock.RUnlock() + ord = noOrder + for _, w := range r.workers { + ord = min(ord, w.idx) + } + return ord, len(r.workers) > 0 +} + +const noOrder = ^uint32(0) + +// eligible copies out the workers that may serve the request. Nothing here may +// reach into a Registration: this lock must not be held across the transport. +func (r *Route) eligible(granted bool) (workers []routeWorker, denied bool) { + r.lock.RLock() + defer r.lock.RUnlock() + for _, w := range r.workers { + if !granted && !w.public { + denied = true + continue + } + workers = append(workers, w) + } + return +} + +// routeTable is a deployment's merged route table. The tree is rebuilt only when +// the route set or its order changes, and is published by a single atomic store +// under lock. +type routeTable struct { + key regKey + logger logger.Logger + + lock sync.Mutex + routes map[routeKey]*Route + tree atomic.Pointer[router.Router[*Route]] +} + +func newRouteTable(key regKey, l logger.Logger) *routeTable { + return &routeTable{ + key: key, + logger: l.WithValues("apiKey", key.apiKey, "agentName", key.agentName, "deployment", key.deployment), + routes: make(map[routeKey]*Route), + } +} + +// mutate applies a removal and an addition as one transaction, republishing the +// tree if the merged table moved. Routes the removal empties are deleted only +// once the addition has run, so a registration replaced by an equivalent one +// keeps its Route pointers and the published tree stays live. +func (t *routeTable) mutate(remove, add *Registration) { + t.lock.Lock() + defer t.lock.Unlock() + + touched := make(map[*Route]struct{}) + if remove != nil { + t.retractLocked(remove, touched) + } + if add != nil { + t.declareLocked(add, touched) + } + + dirty := false + for r := range touched { + ord, live := r.position() + if !live { + delete(t.routes, r.key) + dirty = true + continue + } + // a departure can raise a route's position and an arrival lower it; + // either reorders the table + if ord != r.ord { + dirty = true + } + } + if dirty { + t.publishLocked() + } +} + +func (t *routeTable) declareLocked(reg *Registration, touched map[*Route]struct{}) { + if reg.Manifest == nil { + return + } + for i, ep := range reg.Manifest.Endpoints { + for mask := ep.Mask; mask != 0; mask &= mask - 1 { + key := routeKey{canonical: ep.Template.Canonical(), method: mask & -mask} + r := t.routes[key] + if r == nil { + r = &Route{Template: ep.Template, Mask: key.method, key: key, ord: noOrder} + t.routes[key] = r + } + r.declare(routeWorker{reg: reg, raw: ep.Template.String(), public: ep.Public, idx: uint32(i)}) + touched[r] = struct{}{} + } + } +} + +func (t *routeTable) retractLocked(reg *Registration, touched map[*Route]struct{}) { + for _, r := range t.routes { + r.retract(reg) + touched[r] = struct{}{} + } +} + +// publishLocked compiles the merged table and swaps it in. The builder assigns +// each route an index from its position in the order added, so routes must be +// added in merge order. +func (t *routeTable) publishLocked() { + ordered := make([]*Route, 0, len(t.routes)) + for _, r := range t.routes { + ordered = append(ordered, r) + } + for _, r := range ordered { + r.ord, _ = r.position() + } + // routes sharing a position are separated by the tie-break alone, so the + // comparison must be total and reach no further than the key + slices.SortFunc(ordered, func(a, b *Route) int { + if c := cmp.Compare(a.ord, b.ord); c != 0 { + return c + } + if c := cmp.Compare(a.key.canonical, b.key.canonical); c != 0 { + return c + } + return cmp.Compare(a.key.method, b.key.method) + }) + + b := router.NewBuilder[*Route]() + for i, r := range ordered { + if err := b.Add(r.Template, r.Mask, r); err != nil { + // MaxManifestRoutes caps a worker, not the merged table; what + // survives is the lowest stretch of the merge order + t.logger.Warnw("agent endpoint route table truncated", err, + "routes", len(ordered), "installed", i) + break + } + } + t.tree.Store(b.Build()) +} + +func (t *routeTable) empty() bool { + t.lock.Lock() + defer t.lock.Unlock() + return len(t.routes) == 0 +} + +// match resolves a path against the merged table. matched is the dispatch test: +// a route reached through a tree loaded before its last worker left resolves to +// an empty set. +func (t *routeTable) match(path string, mask router.Mask, granted bool) (matched []routeWorker, route *Route, res router.Result, denied bool) { + tree := t.tree.Load() + if tree == nil { + return nil, nil, router.ResultNone, false + } + rt, res := tree.Match(path, mask) + if res != router.ResultFull { + return nil, nil, res, false + } + matched, denied = rt.eligible(granted) + return matched, rt, res, denied +} + +// serves reports whether the table has a route for this exact path and method. +func (t *routeTable) serves(path string, mask router.Mask) bool { + tree := t.tree.Load() + if tree == nil { + return false + } + _, res := tree.Match(path, mask) + return res == router.ResultFull +} diff --git a/pkg/agent/endpoint/routes_test.go b/pkg/agent/endpoint/routes_test.go new file mode 100644 index 000000000..0ff87f14c --- /dev/null +++ b/pkg/agent/endpoint/routes_test.go @@ -0,0 +1,354 @@ +// Copyright 2026 LiveKit, Inc. + +package endpoint + +import ( + "fmt" + "net/http" + "slices" + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/protocol/livekit" + + "github.com/livekit/livekit-server/pkg/agent/endpoint/router" +) + +func mustManifest(t *testing.T, eps ...*livekit.AgentHttp_AgentEndpoint) *Manifest { + t.Helper() + m, err := ParseManifest(eps) + require.NoError(t, err) + return m +} + +func regOf(workerID string, m *Manifest) *Registration { + return NewRegistration(RegistrationParams{ + WorkerID: workerID, APIKey: "proj", AgentName: "a", Deployment: "d", + Manifest: m, Session: &fakeSession{}, + }) +} + +// tableOf registers one worker per manifest under a single key and returns the +// deployment's merged table. +func tableOf(t *testing.T, manifests ...*Manifest) *routeTable { + t.Helper() + g := NewRegistry() + for i, m := range manifests { + g.Register(regOf(fmt.Sprintf("w%d", i), m)) + } + tbl := g.table("proj", "a", "d") + require.NotNil(t, tbl) + return tbl +} + +// serving names the workers a request would be dispatched to, sorted. +func serving(tbl *routeTable, path, method string, granted bool) []string { + matched, _, _, _ := tbl.match(path, methodMask(method), granted) + ids := make([]string, 0, len(matched)) + for _, w := range matched { + ids = append(ids, w.reg.WorkerID) + } + slices.Sort(ids) + return ids +} + +// A mixed fleet serves the union of its routes, and a route only ever +// dispatches to the workers that declared it. +func TestRouteTableUnion(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, ep("/health", []string{"GET"}, true), ep("/v1", []string{"GET"}, true)), + mustManifest(t, ep("/health", []string{"GET"}, true), ep("/v2", []string{"GET"}, true)), + ) + + require.Equal(t, []string{"w0", "w1"}, serving(tbl, "/health", http.MethodGet, true)) + require.Equal(t, []string{"w0"}, serving(tbl, "/v1", http.MethodGet, true)) + require.Equal(t, []string{"w1"}, serving(tbl, "/v2", http.MethodGet, true)) + require.Empty(t, serving(tbl, "/v3", http.MethodGet, true)) +} + +// A worker leaving takes only the routes nothing else declares. +func TestRouteTableDeparture(t *testing.T) { + g := NewRegistry() + old := regOf("w0", mustManifest(t, ep("/health", []string{"GET"}, true), ep("/v1", []string{"GET"}, true))) + g.Register(old) + g.Register(regOf("w1", mustManifest(t, ep("/health", []string{"GET"}, true), ep("/v2", []string{"GET"}, true)))) + tbl := g.table("proj", "a", "d") + + g.Deregister(old) + + require.Equal(t, []string{"w1"}, serving(tbl, "/health", http.MethodGet, true)) + require.Empty(t, serving(tbl, "/v1", http.MethodGet, true), "the departing worker's own route is gone") + require.Equal(t, []string{"w1"}, serving(tbl, "/v2", http.MethodGet, true)) + + // the last worker out drops the table with it + g.Deregister(g.Candidates("proj", "a", "d")[0]) + require.Nil(t, g.table("proj", "a", "d")) +} + +// Declaration order is the whole priority rule, and a homogeneous fleet keeps +// the order its manifest declared: a literal registered ahead of a param that +// would also match it must still win. +func TestRouteTableMergeOrder(t *testing.T) { + m := mustManifest(t, ep("/users/me", []string{"GET"}, true), ep("/users/{id}", []string{"GET"}, true)) + tbl := tableOf(t, m, m) + + _, r, res, _ := tbl.match("/users/me", methodMask(http.MethodGet), true) + require.Equal(t, router.ResultFull, res) + require.Equal(t, "/users/me", r.Template.String()) + require.Equal(t, []string{"w0", "w1"}, serving(tbl, "/users/me", http.MethodGet, true)) + + _, r, _, _ = tbl.match("/users/7", methodMask(http.MethodGet), true) + require.Equal(t, "/users/{id}", r.Template.String()) +} + +// Workers may disagree about a route's visibility mid-rollout. An ungranted +// caller is narrowed to the workers that declared it public. +func TestRouteTablePerWorkerPublic(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, ep("/x", []string{"GET"}, true)), + mustManifest(t, ep("/x", []string{"GET"}, false)), + ) + + require.Equal(t, []string{"w0", "w1"}, serving(tbl, "/x", http.MethodGet, true)) + require.Equal(t, []string{"w0"}, serving(tbl, "/x", http.MethodGet, false)) + + matched, _, _, denied := tbl.match("/x", methodMask(http.MethodGet), false) + require.Len(t, matched, 1) + require.True(t, denied, "the private declaration is reported, not silently dropped") + + // no public declarer at all denies outright + private := tableOf(t, mustManifest(t, ep("/y", []string{"GET"}, false))) + matched, _, _, denied = private.match("/y", methodMask(http.MethodGet), false) + require.Empty(t, matched) + require.True(t, denied) +} + +// One route per method: workers declaring different verbs on one path each +// serve their own, and a verb nobody declared is partial. +func TestRouteTableMethodSplit(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, ep("/thing", []string{"GET"}, true)), + mustManifest(t, ep("/thing", []string{"POST", "DELETE"}, true)), + ) + + require.Equal(t, []string{"w0"}, serving(tbl, "/thing", http.MethodGet, true)) + require.Equal(t, []string{"w1"}, serving(tbl, "/thing", http.MethodPost, true)) + require.Equal(t, []string{"w1"}, serving(tbl, "/thing", http.MethodDelete, true)) + + _, _, res, _ := tbl.match("/thing", methodMask(http.MethodPut), true) + require.Equal(t, router.ResultPartial, res) +} + +// Param names are invisible to the matcher, so two spellings of one shape are +// one route serving both workers, each carrying its own spelling. +func TestRouteTableCanonicalMerge(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, ep("/items/{id}", []string{"GET"}, true)), + mustManifest(t, ep("/items/{key}", []string{"GET"}, true)), + ) + + require.Len(t, tbl.routes, 1, "one shape is one route") + matched, _, res, _ := tbl.match("/items/7", methodMask(http.MethodGet), true) + require.Equal(t, router.ResultFull, res) + require.Len(t, matched, 2) + + spelling := map[string]string{} + for _, w := range matched { + spelling[w.reg.WorkerID] = w.raw + } + require.Equal(t, map[string]string{"w0": "/items/{id}", "w1": "/items/{key}"}, spelling) + + // a different shape stays a different route + require.Len(t, tableOf(t, + mustManifest(t, ep("/items/{id}", []string{"GET"}, true)), + mustManifest(t, ep("/items/{id:int}", []string{"GET"}, true)), + ).routes, 2) +} + +// A literal brace is not a param, and must not collide with one. +func TestRouteTableLiteralBraceDistinctFromParam(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, ep("/x/{a}", []string{"GET"}, true)), + mustManifest(t, ep("/x/{:str}", []string{"GET"}, true)), + ) + require.Len(t, tbl.routes, 2) +} + +// A manifest may name one path and method twice; the worker joins the route +// once. +func TestRouteTableDuplicateDeclaration(t *testing.T) { + tbl := tableOf(t, mustManifest(t, + ep("/x", []string{"GET", "POST"}, true), + ep("/x", []string{"GET"}, true), + )) + require.Equal(t, []string{"w0"}, serving(tbl, "/x", http.MethodGet, true)) +} + +// The published tree holds routes by pointer, so a reconnect that rebuilds the +// same route set keeps those pointers and stays routable. +func TestRouteTableSupersedeStaysRoutable(t *testing.T) { + g := NewRegistry() + m := mustManifest(t, ep("/x", []string{"GET"}, true)) + g.Register(regOf("w0", m)) + tbl := g.table("proj", "a", "d") + before := tbl.tree.Load() + + g.Register(regOf("w0", m)) + + require.Equal(t, []string{"w0"}, serving(tbl, "/x", http.MethodGet, true), + "a reconnect with an unchanged manifest stays routable") + require.Same(t, before, tbl.tree.Load(), "an unchanged route set needs no rebuild") +} + +// A reconnect may land on a different agent name or deployment. The retiring +// epoch's routes live in the old table and must be retracted there. +func TestRouteTableSupersedeAcrossKeys(t *testing.T) { + g := NewRegistry() + m := mustManifest(t, ep("/x", []string{"GET"}, true)) + g.Register(regOf("w0", m)) + + g.Register(NewRegistration(RegistrationParams{ + WorkerID: "w0", APIKey: "proj", AgentName: "a", Deployment: "other", + Manifest: m, Session: &fakeSession{}, + })) + + require.Nil(t, g.table("proj", "a", "d"), "the old deployment's table is gone") + require.Equal(t, []string{"w0"}, serving(g.table("proj", "a", "other"), "/x", http.MethodGet, true)) +} + +// The tree is rebuilt exactly when the merged route set or its order moves. +func TestRouteTableRebuildTrigger(t *testing.T) { + g := NewRegistry() + m := mustManifest(t, ep("/x", []string{"GET"}, true)) + g.Register(regOf("w0", m)) + tbl := g.table("proj", "a", "d") + + before := tbl.tree.Load() + g.Register(regOf("w1", m)) + require.Same(t, before, tbl.tree.Load(), "a worker joining an existing route changes no route") + + g.Register(regOf("w2", mustManifest(t, ep("/x", []string{"GET"}, true), ep("/y", []string{"GET"}, true)))) + require.NotSame(t, before, tbl.tree.Load(), "a new route rebuilds") +} + +// A departure can raise a route's merge position and an arrival lower it; +// either reorders the table. +func TestRouteTableOrderFollowsDeparture(t *testing.T) { + g := NewRegistry() + first := regOf("w0", mustManifest(t, + ep("/a/{x}", []string{"GET"}, true), + ep("/a/b", []string{"GET"}, true), + )) + g.Register(first) + g.Register(regOf("w1", mustManifest(t, + ep("/a/b", []string{"GET"}, true), + ep("/a/{x}", []string{"GET"}, true), + ))) + tbl := g.table("proj", "a", "d") + + // both routes hold position 0 - w0 declares the param first, w1 the literal + // - so the tie-break decides, and the param shadows the literal + _, r, _, _ := tbl.match("/a/b", methodMask(http.MethodGet), true) + require.Equal(t, "/a/{x}", r.Template.String()) + + // w0 leaving raises the param to w1's position 1, putting the literal ahead + // of it + g.Deregister(first) + _, r, _, _ = tbl.match("/a/b", methodMask(http.MethodGet), true) + require.Equal(t, "/a/b", r.Template.String()) +} + +// Ambiguity is a property of the merged tree: one worker's shape can carry the +// whole deployment's matcher over its step budget, and an undecided route has no +// Public flag left to clear an ungranted request. +func TestRouteTableAmbiguityIsContagious(t *testing.T) { + long := "/" + strings.Repeat("a", 512) + + open := tableOf(t, mustManifest(t, ep("/{p:path}", []string{"GET"}, true))) + _, _, res, _ := open.match(long, methodMask(http.MethodGet), false) + require.Equal(t, router.ResultFull, res, "alone, the catch-all decides") + + merged := tableOf(t, + mustManifest(t, ep("/{a}{b}{c}{d}x", []string{"GET"}, true)), + mustManifest(t, ep("/{p:path}", []string{"GET"}, true)), + ) + _, _, res, _ = merged.match(long, methodMask(http.MethodGet), false) + require.Equal(t, router.ResultOverBudget, res) +} + +// KNOWN LIMITATION: the merge takes each route's lowest position across the +// fleet, which does not preserve any one worker's internal order. A worker can +// be handed a request the merged table resolved to a public route while its own +// table routes that path to a private one. An order-preserving merge of the +// declaration chains would hold the property. +func TestRouteTableMergeOrderInversion(t *testing.T) { + tbl := tableOf(t, + mustManifest(t, + ep("/x/{z:int}", []string{"GET"}, false), + ep("/x/{b}", []string{"GET"}, true), + ), + mustManifest(t, ep("/x/{b}", []string{"GET"}, true)), + ) + + _, r, res, _ := tbl.match("/x/7", methodMask(http.MethodGet), false) + require.Equal(t, router.ResultFull, res) + require.Equal(t, "/x/{b}", r.Template.String()) + require.Equal(t, []string{"w0", "w1"}, serving(tbl, "/x/7", http.MethodGet, false), + "w0 is reachable anonymously while its own table routes /x/7 to a private route") +} + +// Registrations churn while requests resolve; run under -race. +func TestRouteTableConcurrentChurn(t *testing.T) { + g := NewRegistry() + stable := regOf("stable", mustManifest(t, ep("/x", []string{"GET"}, true))) + g.Register(stable) + + var wg sync.WaitGroup + stop := make(chan struct{}) + for i := range 4 { + wg.Add(1) + go func() { + defer wg.Done() + m := mustManifest(t, ep("/x", []string{"GET"}, true), ep(fmt.Sprintf("/w%d", i), []string{"GET"}, true)) + for n := 0; ; n++ { + select { + case <-stop: + return + default: + } + r := regOf(fmt.Sprintf("w%d", i), m) + g.Register(r) + g.Deregister(r) + } + }() + } + for range 4 { + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + if tbl := g.table("proj", "a", "d"); tbl != nil { + tbl.match("/x", methodMask(http.MethodGet), true) + tbl.serves("/x/", methodMask(http.MethodGet)) + } + } + }() + } + for range 2000 { + if tbl := g.table("proj", "a", "d"); tbl != nil { + tbl.match("/x", methodMask(http.MethodGet), false) + } + } + close(stop) + wg.Wait() + + require.Contains(t, serving(g.table("proj", "a", "d"), "/x", http.MethodGet, true), "stable") +}