jleung
f85c3ef781
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Lx9FNLG2CqeBbVGGd9bFsP
100 lines
2.6 KiB
Go
100 lines
2.6 KiB
Go
package api
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/danielgtaylor/huma/v2"
|
|
)
|
|
|
|
// admission keeps metric-store query load bounded: a per-client cap on
|
|
// free-form queries (so one dashboard cannot starve the rest) and a global
|
|
// semaphore on every store query (so the UI cannot crowd out the poll path).
|
|
// A request over its client cap is refused at once; one waiting for a global
|
|
// slot waits up to wait, then is refused. Both are 429 problems.
|
|
type admission struct {
|
|
global chan struct{}
|
|
perClient int
|
|
wait time.Duration
|
|
|
|
mu sync.Mutex
|
|
active map[string]int
|
|
}
|
|
|
|
func newAdmission(perClient, global int, wait time.Duration) *admission {
|
|
return &admission{global: make(chan struct{}, global), perClient: perClient, wait: wait, active: map[string]int{}}
|
|
}
|
|
|
|
// clientKey carries the caller's identity from the middleware to the handlers.
|
|
// Until authentication exists (M6) it is the remote address without the port.
|
|
type clientKey struct{}
|
|
|
|
func clientKeyMiddleware(ctx huma.Context, next func(huma.Context)) {
|
|
host, _, err := net.SplitHostPort(ctx.RemoteAddr())
|
|
if err != nil {
|
|
host = ctx.RemoteAddr()
|
|
}
|
|
next(huma.WithValue(ctx, clientKey{}, host))
|
|
}
|
|
|
|
func clientOf(ctx context.Context) string {
|
|
c, _ := ctx.Value(clientKey{}).(string)
|
|
return c
|
|
}
|
|
|
|
func tooMany(msg string) error {
|
|
return huma.Error429TooManyRequests(msg)
|
|
}
|
|
|
|
// acquire admits one query. freeForm queries also count against the client's cap.
|
|
// The returned release must be called exactly once.
|
|
func (a *admission) acquire(ctx context.Context, freeForm bool) (release func(), err error) {
|
|
client := clientOf(ctx)
|
|
if freeForm && !a.takeClient(client) {
|
|
return nil, tooMany("too many concurrent queries from " + client + "; the limit is " + itoa(a.perClient))
|
|
}
|
|
wctx, cancel := context.WithTimeout(ctx, a.wait)
|
|
defer cancel()
|
|
select {
|
|
case a.global <- struct{}{}:
|
|
case <-wctx.Done():
|
|
if freeForm {
|
|
a.dropClient(client)
|
|
}
|
|
return nil, tooMany("the metric store is busy; try again shortly")
|
|
}
|
|
return func() {
|
|
<-a.global
|
|
if freeForm {
|
|
a.dropClient(client)
|
|
}
|
|
}, nil
|
|
}
|
|
|
|
func (a *admission) takeClient(client string) bool {
|
|
a.mu.Lock()
|
|
defer a.mu.Unlock()
|
|
if a.active[client] >= a.perClient {
|
|
return false
|
|
}
|
|
a.active[client]++
|
|
return true
|
|
}
|
|
|
|
func (a *admission) dropClient(client string) {
|
|
a.mu.Lock()
|
|
defer a.mu.Unlock()
|
|
if a.active[client] <= 1 {
|
|
delete(a.active, client)
|
|
return
|
|
}
|
|
a.active[client]--
|
|
}
|
|
|
|
// admit takes a query slot for the request in ctx.
|
|
func (s *server) admit(ctx context.Context, freeForm bool) (func(), error) {
|
|
return s.adm.acquire(ctx, freeForm)
|
|
}
|