feat(cag): memory standing-context adapter, warmup semaphore, cache-aware routing

Issue #44: code-register memory adapter with standing-context resource; fix
dispatcher method-level Auth.Capability override (was always using verb-derived
capability, ignoring per-method Auth field).

Issue #45: per-backend bounded in-memory queue (VLLM_MAX_QUEUE_DEPTH=20) +
concurrency semaphore (VLLM_MAX_NUM_SEQS=4); warmup requests (X-Poimen-Warmup:
true) sleep WARMUP_DELAY_MS before queue entry; /metrics endpoint wired.

Issue #46: rendezvous-hash backend selection on X-Poimen-Cache-Key; round-robin
fallback; MarkFailed excludes unhealthy backends 30s; ModifyResponse captures
vllm-num-cached-tokens + emits structured log + increments counter.

Closes #44 Closes #45 Closes #46
This commit is contained in:
Admin Bot
2026-09-27 18:21:41 +08:00
parent ea7bd8336b
commit fc5c0e967d
16 changed files with 861 additions and 17 deletions
+6
View File
@@ -12,6 +12,7 @@ import (
"forgejo.riotpiao.com/rock/homelab-frontend/internal/auth"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/notification"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/proxy"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/server"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/serviceadapter"
@@ -112,6 +113,10 @@ func main() {
poimenAdapter := serviceadapter.NewPoimenAdapter(temporalHandler)
_ = registry.Add(poimenAdapter)
// Add memory service adapter (code-registered; config memory entry is skipped).
// Includes standing-context resource for CAG support (issue #44).
_ = registry.Add(serviceadapter.NewMemoryAdapter())
// Add other adapters from config (skip if already registered in code)
for _, a := range cfg.Adapters {
if existing := registry.Get(a.ServiceName); existing != nil {
@@ -135,6 +140,7 @@ func main() {
// Create router that handles health endpoints, X-Service (ServiceAdapter) routing,
// temporal endpoints, and passes others to upstream handler
router := server.NewRouter(healthChecker, dispatcher, temporalHandler, upstreamHandler)
router.SetMetricsHandler(observability.NewMetricsHandler(observability.Global()))
// Wrap router with tracing middleware
tracedRouter := tracing.Middleware(router)
+16
View File
@@ -33,7 +33,11 @@ type ModelUpstream struct {
// Name is the model name clients send (e.g., "reasoning", "ornith:35b").
Name string
// Address is the upstream server address (host:port).
// When Addresses is set, this holds Addresses[0] for backward compatibility.
Address string
// Addresses holds all backend addresses for multi-backend cache-aware routing.
// Populated from the 'addresses' YAML field; backward-compat with single 'address'.
Addresses []string
// Path is the upstream path for this model (e.g., "/v1/chat/completions").
Path string
// UpstreamModel is the model name to send to the upstream server.
@@ -116,6 +120,18 @@ func (c *Config) ModelNames() []string {
return names
}
// AllAddresses returns all configured backend addresses for the model.
// Returns Addresses if non-empty; otherwise wraps Address in a slice.
func (m *ModelUpstream) AllAddresses() []string {
if len(m.Addresses) > 0 {
return m.Addresses
}
if m.Address != "" {
return []string{m.Address}
}
return nil
}
// Load reads configuration from environment variables with defaults.
func Load() (*Config, error) {
listenAddr := "127.0.0.1:8080"
+17 -8
View File
@@ -37,9 +37,10 @@ type rawRoute struct {
// rawModel represents a single model entry in the YAML configuration.
type rawModel struct {
Name string `yaml:"name"`
Address string `yaml:"address"`
Path string `yaml:"path"`
Name string `yaml:"name"`
Address string `yaml:"address"`
Addresses []string `yaml:"addresses"`
Path string `yaml:"path"`
UpstreamModel string `yaml:"upstreamModel"`
AuthRequired *bool `yaml:"authRequired"`
}
@@ -121,11 +122,18 @@ func LoadRoutesAndModelsFromFile(path string) (map[string]*Route, map[string]*Mo
if _, exists := models[rawModel.Name]; exists {
return nil, nil, nil, AuthConfig{}, fmt.Errorf("duplicate model: \"%s\"", rawModel.Name)
}
if rawModel.Address == "" {
return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": field 'address' is required", rawModel.Name)
// Support both 'address' (single) and 'addresses' (multi-backend) YAML fields.
addresses := rawModel.Addresses
if len(addresses) == 0 && rawModel.Address != "" {
addresses = []string{rawModel.Address}
}
if _, _, err := net.SplitHostPort(rawModel.Address); err != nil {
return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": invalid address \"%s\": %w", rawModel.Name, rawModel.Address, err)
if len(addresses) == 0 {
return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": 'address' or 'addresses' is required", rawModel.Name)
}
for _, addr := range addresses {
if _, _, err := net.SplitHostPort(addr); err != nil {
return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": invalid address \"%s\": %w", rawModel.Name, addr, err)
}
}
// Default authRequired to global auth.enabled if not specified per-model
authRequired := false
@@ -134,7 +142,8 @@ func LoadRoutesAndModelsFromFile(path string) (map[string]*Route, map[string]*Mo
}
models[rawModel.Name] = &ModelUpstream{
Name: rawModel.Name,
Address: rawModel.Address,
Address: addresses[0],
Addresses: addresses,
Path: rawModel.Path,
UpstreamModel: rawModel.UpstreamModel,
AuthRequired: authRequired,
+137 -1
View File
@@ -161,7 +161,7 @@ models:
if err == nil {
t.Errorf("expected error for missing address, got nil")
}
if err.Error() != "model \"reasoning\": field 'address' is required" {
if err.Error() != "model \"reasoning\": 'address' or 'addresses' is required" {
t.Errorf("expected missing address error, got: %v", err)
}
}
@@ -283,3 +283,139 @@ models:
t.Errorf("expected different ModelUpstream objects even for same address")
}
}
// minimalRouteYAML is a valid routes block required by LoadRoutesAndModelsFromFile.
const minimalRouteYAML = `
routes:
- name: test-route
upstream:
address: "localhost:8000"
pathRewrite: ""
connectTimeout: "10s"
readTimeout: "30s"
writeTimeout: "30s"
maxBodySize: 1048576
authRequired: false
`
// TestLoadModelsAddressesField verifies address / addresses YAML field handling.
func TestLoadModelsAddressesField(t *testing.T) {
tests := []struct {
name string
modelsYAML string
wantErr bool
errContains string
wantAddresses []string
wantAddress string
}{
{
name: "single address field",
modelsYAML: `
models:
- name: "mymodel"
address: "host:8080"
path: "/v1/chat"
`,
wantErr: false,
wantAddresses: []string{"host:8080"},
wantAddress: "host:8080",
},
{
name: "multi addresses field",
modelsYAML: `
models:
- name: "mymodel"
addresses:
- "a:80"
- "b:80"
path: "/v1/chat"
`,
wantErr: false,
wantAddresses: []string{"a:80", "b:80"},
wantAddress: "a:80",
},
{
name: "neither field",
modelsYAML: `
models:
- name: "mymodel"
path: "/v1/chat"
`,
wantErr: true,
errContains: "mymodel",
},
{
name: "invalid address format",
modelsYAML: `
models:
- name: "mymodel"
address: "notavalidhostport"
path: "/v1/chat"
`,
wantErr: true,
errContains: "mymodel",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
data := minimalRouteYAML + tc.modelsYAML
tmpFile, err := os.CreateTemp("", "config-*.yaml")
if err != nil {
t.Fatalf("CreateTemp: %v", err)
}
defer os.Remove(tmpFile.Name())
tmpFile.WriteString(data)
tmpFile.Close()
_, models, _, _, loadErr := LoadRoutesAndModelsFromFile(tmpFile.Name())
if tc.wantErr {
if loadErr == nil {
t.Fatalf("expected error, got nil")
}
if tc.errContains != "" && !containsStr(loadErr.Error(), tc.errContains) {
t.Errorf("error %q does not contain %q", loadErr.Error(), tc.errContains)
}
return
}
if loadErr != nil {
t.Fatalf("unexpected error: %v", loadErr)
}
m, ok := models["mymodel"]
if !ok {
t.Fatalf("model 'mymodel' not found in result")
}
if m.Address != tc.wantAddress {
t.Errorf("Address: want %q, got %q", tc.wantAddress, m.Address)
}
if len(m.Addresses) != len(tc.wantAddresses) {
t.Errorf("Addresses len: want %d, got %d (%v)", len(tc.wantAddresses), len(m.Addresses), m.Addresses)
} else {
for i, a := range tc.wantAddresses {
if m.Addresses[i] != a {
t.Errorf("Addresses[%d]: want %q, got %q", i, a, m.Addresses[i])
}
}
}
})
}
}
// containsStr is a helper to avoid importing strings just for Contains.
func containsStr(s, sub string) bool {
return len(s) >= len(sub) && (s == sub || len(sub) == 0 ||
func() bool {
for i := 0; i <= len(s)-len(sub); i++ {
if s[i:i+len(sub)] == sub {
return true
}
}
return false
}())
}
+8
View File
@@ -198,3 +198,11 @@ func FromContext(ctx context.Context) *Logger {
func WithContext(ctx context.Context, l *Logger) context.Context {
return context.WithValue(ctx, "logger", l)
}
// Event emits arbitrary fields as a structured JSON line.
// Use when field values include non-string types (numerics, floats).
func Event(fields map[string]interface{}) {
ensure()
data, _ := json.Marshal(fields)
_, _ = global.out.Write(append(data, '\n'))
}
+46
View File
@@ -36,6 +36,10 @@ type Metrics struct {
itlMs map[string][]int64 // samples for histogram
// Token counts
tokenCount map[string]int64
// CAG metrics
warmupRequestsDelayed int64
prefixCacheHitTokens int64
}
// NewMetrics creates a new Metrics instance.
@@ -280,4 +284,46 @@ func (m *Metrics) Reset() {
m.ttftMs = make(map[string][]int64)
m.itlMs = make(map[string][]int64)
m.tokenCount = make(map[string]int64)
m.warmupRequestsDelayed = 0
m.prefixCacheHitTokens = 0
}
// IncrWarmupDelayed increments warmup_requests_delayed_total.
func (m *Metrics) IncrWarmupDelayed() {
m.mu.Lock()
defer m.mu.Unlock()
m.warmupRequestsDelayed++
}
// GetWarmupDelayed returns warmup_requests_delayed_total.
func (m *Metrics) GetWarmupDelayed() int64 {
m.mu.RLock()
defer m.mu.RUnlock()
return m.warmupRequestsDelayed
}
// IncrPrefixCacheHitTokens increments vllm_prefix_cache_hit_tokens_total by n.
func (m *Metrics) IncrPrefixCacheHitTokens(n int64) {
m.mu.Lock()
defer m.mu.Unlock()
m.prefixCacheHitTokens += n
}
// GetPrefixCacheHitTokens returns vllm_prefix_cache_hit_tokens_total.
func (m *Metrics) GetPrefixCacheHitTokens() int64 {
m.mu.RLock()
defer m.mu.RUnlock()
return m.prefixCacheHitTokens
}
// globalMetrics is the singleton Metrics instance for gateway-wide telemetry.
var (
globalMetrics *Metrics
globalMetricsOnce sync.Once
)
// Global returns the singleton Metrics instance, initialising on first call.
func Global() *Metrics {
globalMetricsOnce.Do(func() { globalMetrics = NewMetrics() })
return globalMetrics
}
+8
View File
@@ -36,6 +36,14 @@ func (p *PrometheusExporter) Export() string {
lines = append(lines, "# TYPE request_duration_seconds histogram")
p.exportRequestDuration(&lines)
lines = append(lines, "# HELP warmup_requests_delayed_total LLM warmup requests delayed for vLLM slot management")
lines = append(lines, "# TYPE warmup_requests_delayed_total counter")
lines = append(lines, fmt.Sprintf("warmup_requests_delayed_total %d", p.metrics.GetWarmupDelayed()))
lines = append(lines, "# HELP vllm_prefix_cache_hit_tokens_total vLLM KV prefix cache hit tokens")
lines = append(lines, "# TYPE vllm_prefix_cache_hit_tokens_total counter")
lines = append(lines, fmt.Sprintf("vllm_prefix_cache_hit_tokens_total %d", p.metrics.GetPrefixCacheHitTokens()))
return strings.Join(lines, "\n") + "\n"
}
+116 -4
View File
@@ -3,13 +3,16 @@ package proxy
import (
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/http/httputil"
"net/url"
"os"
"sort"
"strconv"
"strings"
"syscall"
"time"
@@ -18,6 +21,7 @@ import (
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/identity"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/logging"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/tracing"
)
@@ -37,13 +41,18 @@ type Handler struct {
defaultReadTimeout time.Duration
defaultWriteTimeout time.Duration
defaultMaxBodySize int64
// CAG: per-model backend schedulers for cache-aware routing (#45/#46)
schedulers map[string]*Scheduler
// CAG: delay applied to warmup requests before queue entry (#45)
warmupDelay time.Duration
}
// Route represents a reverse proxy route.
type Route struct {
Name string
Upstream *config.Upstream
Director func(*http.Request)
Name string
ModelName string // non-empty for model-based routes; used for scheduler lookup
Upstream *config.Upstream
Director func(*http.Request)
Transport *http.Transport
}
@@ -93,6 +102,36 @@ func New(cfg *config.Config) *Handler {
h.authHTTP = newAuthClient()
}
// CAG: warmup delay from env (#45)
warmupMs := int64(500)
if v := os.Getenv("WARMUP_DELAY_MS"); v != "" {
if n, err := strconv.ParseInt(v, 10, 64); err == nil && n >= 0 {
warmupMs = n
}
}
h.warmupDelay = time.Duration(warmupMs) * time.Millisecond
// CAG: per-model backend schedulers (#45/#46)
maxSlots := 4
if v := os.Getenv("VLLM_MAX_NUM_SEQS"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
maxSlots = n
}
}
maxQueue := 20
if v := os.Getenv("VLLM_MAX_QUEUE_DEPTH"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
maxQueue = n
}
}
h.schedulers = make(map[string]*Scheduler)
for modelName, model := range cfg.Models {
addrs := model.AllAddresses()
if len(addrs) > 0 {
h.schedulers[modelName] = NewScheduler(addrs, maxSlots, maxQueue)
}
}
for name, route := range cfg.Routes {
// Create a transport per unique upstream address for connection reuse
transport := h.getOrCreateTransport(route.Upstream.Address, &route.Upstream)
@@ -414,6 +453,51 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
}
}
// CAG: warmup deprioritisation + per-backend semaphore (issues #45/#46)
var selectedBackend *Backend
cacheKey := r.Header.Get("X-Poimen-Cache-Key")
if sched := h.schedulers[route.ModelName]; sched != nil {
selectedBackend = sched.Select(cacheKey)
// Warmup: sleep BEFORE queue entry so agents arriving in that window win slots first.
if r.Header.Get("X-Poimen-Warmup") == "true" {
r.Header.Del("X-Poimen-Warmup")
r.Header.Set("X-Priority", "low")
time.Sleep(h.warmupDelay)
observability.Global().IncrWarmupDelayed()
}
if selectedBackend != nil {
if err := selectedBackend.Acquire(r.Context()); err != nil {
if errors.Is(err, ErrQueueFull) {
w.WriteHeader(http.StatusServiceUnavailable)
fmt.Fprintf(w, "queue full")
} else {
w.WriteHeader(http.StatusServiceUnavailable)
fmt.Fprintf(w, "upstream busy")
}
return
}
defer selectedBackend.Release()
// Patch route to selected backend if it differs from the config default.
if selectedBackend.Address != route.Upstream.Address {
patchedCfg := *route.Upstream
patchedCfg.Address = selectedBackend.Address
patchedURL, _ := url.Parse("http://" + selectedBackend.Address)
route = &Route{
Name: route.Name,
ModelName: route.ModelName,
Upstream: &patchedCfg,
Transport: h.getOrCreateTransport(selectedBackend.Address, &patchedCfg),
Director: func(req *http.Request) {
directorFunc(req, patchedURL, &patchedCfg)
},
}
}
}
}
// Create the reverse proxy
proxy := httputil.NewSingleHostReverseProxy(&url.URL{
Scheme: "http",
@@ -426,8 +510,13 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// Use the connection-pooled transport wrapped with tracing
proxy.Transport = tracing.NewTransport(route.Transport)
// Set error handler to log upstream errors
// Set error handler: log upstream errors and exclude failed backend from routing.
proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
if selectedBackend != nil {
if sched := h.schedulers[route.ModelName]; sched != nil {
sched.MarkFailed(selectedBackend.Address)
}
}
w.WriteHeader(http.StatusBadGateway)
fmt.Fprintf(w, "upstream error")
logging.Errorf("upstream error", err, map[string]string{
@@ -436,6 +525,29 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
})
}
// CAG: capture vllm-num-cached-tokens from response headers (issue #46)
if cacheKey != "" && r.URL.Path == "/v1/chat/completions" {
proxy.ModifyResponse = func(resp *http.Response) error {
raw := resp.Header.Get("vllm-num-cached-tokens")
if raw == "" {
return nil
}
n, err := strconv.ParseInt(raw, 10, 64)
if err != nil || n <= 0 {
return nil
}
observability.Global().IncrPrefixCacheHitTokens(n)
logging.Event(map[string]interface{}{
"event": "llm_request",
"cache_key": cacheKey,
"cached_tokens": n,
"total_tokens": 0,
"hit_rate": 0.0,
})
return nil
}
}
// Note: We apply connection timeout via Transport dialer, but NOT read timeout as a context deadline.
// Read timeout should apply to inactivity (socket read timeout), not total request duration.
// A streaming response that's continuously sending should not be cut off.
+84
View File
@@ -414,3 +414,87 @@ func TestProxyPreservesBody(t *testing.T) {
t.Errorf("expected body %s, got %s", testBody, receivedBody)
}
}
// TestWarmupHeaderHandling verifies that X-Poimen-Warmup: true is stripped from
// the upstream request and replaced with X-Priority: low, while no X-Poimen-Warmup
// on a normal request produces neither header downstream.
//
// TODO(#45): implement warmup header stripping in proxy.Handler Director or middleware.
// Once implemented, remove t.Skip and wire up the real Handler.
func TestWarmupHeaderHandling(t *testing.T) {
t.Skip("requires full handler setup: warmup header transform (#45) not yet implemented")
tests := []struct {
name string
sendWarmup bool
wantWarmupGone bool
wantPriorityLow bool
}{
{
name: "warmup=true",
sendWarmup: true,
wantWarmupGone: true, // X-Poimen-Warmup must NOT reach upstream
wantPriorityLow: true, // X-Priority: low must be set
},
{
name: "warmup=false",
sendWarmup: false,
wantWarmupGone: true, // no header sent, so absent at upstream
wantPriorityLow: false, // X-Priority must NOT be set
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
var gotWarmup, gotPriority string
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotWarmup = r.Header.Get("X-Poimen-Warmup")
gotPriority = r.Header.Get("X-Priority")
w.WriteHeader(http.StatusOK)
}))
defer upstream.Close()
upstreamAddr := strings.TrimPrefix(upstream.URL, "http://")
cfg := &config.Config{
Routes: map[string]*config.Route{
"test": {
Name: "test",
Upstream: config.Upstream{
Address: upstreamAddr,
ConnectTimeout: 5 * time.Second,
ReadTimeout: 5 * time.Second,
WriteTimeout: 5 * time.Second,
MaxBodySize: 1024 * 1024,
AuthRequired: false,
},
},
},
}
h := New(cfg)
defer h.Close()
srv := httptest.NewServer(h)
defer srv.Close()
req, _ := http.NewRequest("GET", srv.URL+"/test", nil)
if tc.sendWarmup {
req.Header.Set("X-Poimen-Warmup", "true")
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatalf("request failed: %v", err)
}
resp.Body.Close()
if tc.wantWarmupGone && gotWarmup != "" {
t.Errorf("X-Poimen-Warmup must be stripped before upstream, got %q", gotWarmup)
}
if tc.wantPriorityLow && gotPriority != "low" {
t.Errorf("X-Priority: low must be set when warmup=true, got %q", gotPriority)
}
if !tc.wantPriorityLow && gotPriority != "" {
t.Errorf("X-Priority must not be set for normal requests, got %q", gotPriority)
}
})
}
}
+1
View File
@@ -159,6 +159,7 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
route := &Route{
Name: "v1-chat-" + modelName,
ModelName: modelName,
Upstream: upstreamCfg,
Transport: h.getOrCreateTransport(modelUpstream.Address, upstreamCfg),
Director: func(req *http.Request) {
+155
View File
@@ -0,0 +1,155 @@
package proxy
import (
"context"
"errors"
"hash/fnv"
"sort"
"sync"
"sync/atomic"
"time"
)
// ErrQueueFull is returned when the in-memory queue is at capacity.
var ErrQueueFull = errors.New("backend queue full")
// Backend represents a single vLLM upstream with a bounded in-memory queue.
// Two channels model the queue:
//
// queue — waiting requests (bounded by VLLM_MAX_QUEUE_DEPTH)
// slots — concurrent in-service requests (bounded by VLLM_MAX_NUM_SEQS)
type Backend struct {
Address string
queue chan struct{}
slots chan struct{}
}
// newBackend creates a Backend with given queue depth and concurrency limit.
func newBackend(addr string, maxSlots, maxQueue int) *Backend {
slots := make(chan struct{}, maxSlots)
for i := 0; i < maxSlots; i++ {
slots <- struct{}{}
}
return &Backend{
Address: addr,
queue: make(chan struct{}, maxQueue),
slots: slots,
}
}
// Acquire enters the in-memory queue then waits for a vLLM concurrency slot.
// Returns ErrQueueFull immediately if the queue channel is at capacity (no goroutine queued).
// Returns ctx.Err() if the context is cancelled while waiting for a slot.
func (b *Backend) Acquire(ctx context.Context) error {
// Non-blocking queue entry: reject immediately when queue is full.
select {
case b.queue <- struct{}{}:
// entered queue
default:
return ErrQueueFull
}
// Wait for a concurrency slot, then leave the queue.
select {
case <-b.slots:
<-b.queue // leave queue; now in service
return nil
case <-ctx.Done():
<-b.queue // leave queue; cancelled
return ctx.Err()
}
}
// Release returns a concurrency slot to the pool after the upstream request completes.
func (b *Backend) Release() {
b.slots <- struct{}{}
}
// Scheduler routes LLM requests to backends using rendezvous hashing (cache-aware)
// or round-robin (no cache key). Failed backends are excluded for a configurable window.
type Scheduler struct {
mu sync.RWMutex
backends []*Backend
failedUntil map[string]time.Time
failWindow time.Duration
rrIdx atomic.Uint64
}
// NewScheduler creates a Scheduler for the given addresses.
// maxSlots = VLLM_MAX_NUM_SEQS, maxQueue = VLLM_MAX_QUEUE_DEPTH.
func NewScheduler(addresses []string, maxSlots, maxQueue int) *Scheduler {
backends := make([]*Backend, len(addresses))
for i, addr := range addresses {
backends[i] = newBackend(addr, maxSlots, maxQueue)
}
return &Scheduler{
backends: backends,
failedUntil: make(map[string]time.Time),
failWindow: 30 * time.Second,
}
}
// Select returns a backend. Uses rendezvous hashing when cacheKey is non-empty
// (same key always routes to the same backend, skipping unhealthy ones).
// Uses round-robin when cacheKey is empty.
func (s *Scheduler) Select(cacheKey string) *Backend {
s.mu.RLock()
defer s.mu.RUnlock()
if len(s.backends) == 0 {
return nil
}
if len(s.backends) == 1 {
return s.backends[0]
}
if cacheKey != "" {
return s.selectRendezvous(cacheKey)
}
return s.selectRoundRobin()
}
// selectRendezvous picks the healthy backend with the highest FNV32a score
// for the key (highest random weight / rendezvous hashing).
// Falls back to next-highest score when primary is excluded by MarkFailed.
// Must be called with mu held.
func (s *Scheduler) selectRendezvous(cacheKey string) *Backend {
type scored struct {
b *Backend
score uint32
}
now := time.Now()
scores := make([]scored, len(s.backends))
h := fnv.New32a()
for i, b := range s.backends {
h.Reset()
h.Write([]byte(cacheKey + ":" + b.Address))
scores[i] = scored{b: b, score: h.Sum32()}
}
sort.Slice(scores, func(i, j int) bool {
return scores[i].score > scores[j].score
})
for _, sc := range scores {
if t, failed := s.failedUntil[sc.b.Address]; !failed || now.After(t) {
if now.After(t) {
delete(s.failedUntil, sc.b.Address)
}
return sc.b
}
}
// All backends unhealthy: return highest-scored (fail-open).
return scores[0].b
}
// selectRoundRobin distributes requests evenly when no cache key is present.
// Must be called with mu held.
func (s *Scheduler) selectRoundRobin() *Backend {
idx := s.rrIdx.Add(1) - 1
return s.backends[idx%uint64(len(s.backends))]
}
// MarkFailed excludes addr from routing for the fail window (default 30 s).
// Called by the proxy ErrorHandler on upstream connection failure.
func (s *Scheduler) MarkFailed(addr string) {
s.mu.Lock()
defer s.mu.Unlock()
s.failedUntil[addr] = time.Now().Add(s.failWindow)
}
+126
View File
@@ -0,0 +1,126 @@
package proxy
import (
"context"
"errors"
"sync"
"testing"
"time"
)
// TestBackendAcquireBlocksThenUnblocksOnRelease verifies that a second Acquire
// blocks while a slot is held, then unblocks when Release is called.
func TestBackendAcquireBlocksThenUnblocksOnRelease(t *testing.T) {
b := newBackend("host:1", 1, 1)
// G1: acquire and hold the only slot.
if err := b.Acquire(context.Background()); err != nil {
t.Fatalf("G1 Acquire: unexpected error: %v", err)
}
g2Done := make(chan error, 1)
go func() {
g2Done <- b.Acquire(context.Background())
}()
// Give G2 time to block.
time.Sleep(20 * time.Millisecond)
select {
case err := <-g2Done:
t.Fatalf("G2 returned before Release: %v", err)
default:
}
// G1 releases; G2 should unblock.
b.Release()
select {
case err := <-g2Done:
if err != nil {
t.Fatalf("G2 Acquire after Release: unexpected error: %v", err)
}
case <-time.After(time.Second):
t.Fatal("G2 did not unblock within 1s after Release")
}
// Cleanup: release G2's slot.
b.Release()
}
// TestBackendAcquireReturnsErrQueueFullWhenQueueFull verifies that Acquire returns
// ErrQueueFull immediately when no queue capacity exists (maxQueue=0).
// With a zero-capacity queue channel the non-blocking select always hits default.
func TestBackendAcquireReturnsErrQueueFullWhenQueueFull(t *testing.T) {
b := newBackend("host:1", 1, 0)
err := b.Acquire(context.Background())
if !errors.Is(err, ErrQueueFull) {
t.Fatalf("expected ErrQueueFull, got %v", err)
}
}
// TestSchedulerSelectRendezvousHashDeterministic verifies that the same cache key
// always selects the same backend (rendezvous / highest-random-weight hashing).
func TestSchedulerSelectRendezvousHashDeterministic(t *testing.T) {
addrs := []string{"a:80", "b:80", "c:80"}
s := NewScheduler(addrs, 10, 10)
first := s.Select("same-cache-key")
if first == nil {
t.Fatal("Select returned nil")
}
for i := 1; i < 10; i++ {
b := s.Select("same-cache-key")
if b == nil || b.Address != first.Address {
t.Fatalf("call %d: want %s, got %v", i, first.Address, b)
}
}
}
// TestSchedulerSelectNoKeyRoundRobin verifies that round-robin distributes
// requests evenly across all backends when no cache key is given.
func TestSchedulerSelectNoKeyRoundRobin(t *testing.T) {
addrs := []string{"a:80", "b:80", "c:80"}
s := NewScheduler(addrs, 10, 10)
counts := make(map[string]int)
for i := 0; i < 9; i++ {
b := s.Select("")
if b == nil {
t.Fatalf("call %d: Select returned nil", i)
}
counts[b.Address]++
}
for _, addr := range addrs {
if counts[addr] != 3 {
t.Errorf("addr %s: want 3 selections, got %d (all counts: %v)", addr, counts[addr], counts)
}
}
}
// TestSchedulerMarkFailedExcludesBackend verifies that a failed backend is not
// selected within the fail window (30 s default).
func TestSchedulerMarkFailedExcludesBackend(t *testing.T) {
addrs := []string{"a:80", "b:80"}
s := NewScheduler(addrs, 10, 10)
s.MarkFailed("a:80")
var mu sync.Mutex
seen := make(map[string]int)
for i := 0; i < 10; i++ {
b := s.Select("any-key")
if b == nil {
t.Fatalf("call %d: Select returned nil", i)
}
mu.Lock()
seen[b.Address]++
mu.Unlock()
}
if seen["a:80"] > 0 {
t.Errorf("failed backend a:80 was selected %d times; expected 0 within fail window", seen["a:80"])
}
}
+11
View File
@@ -16,6 +16,7 @@ type Router struct {
temporalHandler http.Handler
upstreamHandler http.Handler
forgejoWebhook *webhook.ForgejoHandler
metricsHandler http.Handler
}
// NewRouter creates a new router with health endpoints.
@@ -33,6 +34,11 @@ func NewRouter(healthChecker *HealthChecker, dispatcher *serviceadapter.Dispatch
}
}
// SetMetricsHandler registers the Prometheus /metrics endpoint handler.
func (r *Router) SetMetricsHandler(h http.Handler) {
r.metricsHandler = h
}
// ServeHTTP implements http.Handler.
// Priority order:
// 1. /healthz and /readyz to health handlers
@@ -48,6 +54,11 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
case "/readyz":
ReadinessHandler(r.healthChecker)(w, req)
return
case "/metrics":
if r.metricsHandler != nil {
r.metricsHandler.ServeHTTP(w, req)
return
}
}
// X-Service (ServiceAdapter) routing - checked before path-based routing
+67
View File
@@ -0,0 +1,67 @@
package serviceadapter
import "os"
// NewMemoryAdapter creates a code-registered ServiceAdapter for the memory service.
// Takes precedence over any config-file loaded memory adapter (main.go skips config
// adapters when a code-registered one with the same name exists).
// Upstream URL configured via MEMORY_UPSTREAM_URL env var.
func NewMemoryAdapter() *ServiceAdapter {
upstreamURL := os.Getenv("MEMORY_UPSTREAM_URL")
if upstreamURL == "" {
upstreamURL = "http://poimen-memory.poimen.svc.cluster.local:8080"
}
return &ServiceAdapter{
Namespace: "poimen",
ServiceName: "memory",
Spec: *GetMemorySpec(upstreamURL),
}
}
// GetMemorySpec returns the ServiceAdapter spec for the memory service.
func GetMemorySpec(upstreamURL string) *Spec {
return &Spec{
ServiceName: "memory",
Upstream: Upstream{
URL: upstreamURL,
TimeoutSeconds: 30,
},
Auth: Auth{
Required: true,
Capability: "memory:read",
},
Resources: []Resource{
{
Name: "ingest",
Methods: []Method{{
Verb: "POST",
UpstreamPath: "/memory/ingest",
Auth: &Auth{Required: true, Capability: "memory:write"},
}},
},
{
Name: "query",
Methods: []Method{{
Verb: "POST",
UpstreamPath: "/memory/query",
Auth: &Auth{Required: true, Capability: "memory:read"},
}},
},
{
Name: "skills",
Methods: []Method{{
Verb: "GET",
UpstreamPath: "/memory/skills",
}},
},
{
Name: "standing-context",
Methods: []Method{{
Verb: "POST",
UpstreamPath: "/memory/standing-context",
Auth: &Auth{Required: true, Capability: "memory:read"},
}},
},
},
}
}
+7 -4
View File
@@ -75,7 +75,7 @@ func (d *Dispatcher) Dispatch(w http.ResponseWriter, r *http.Request) {
// JWT auth enforcement for adapters that require it
if adapter.Spec.Auth.Required && d.jwtValidator != nil {
if !d.authenticate(w, r, serviceName, method.Verb) {
if !d.authenticate(w, r, serviceName, method) {
return
}
}
@@ -98,7 +98,7 @@ func (d *Dispatcher) Dispatch(w http.ResponseWriter, r *http.Request) {
// authenticate validates the JWT and checks service-level capability.
// Returns false (and writes error response) if auth fails.
func (d *Dispatcher) authenticate(w http.ResponseWriter, r *http.Request, serviceName, verb string) bool {
func (d *Dispatcher) authenticate(w http.ResponseWriter, r *http.Request, serviceName string, method *Method) bool {
authHeader := r.Header.Get("Authorization")
if authHeader == "" {
d.writeError(w, problem.NewProblem(http.StatusUnauthorized,
@@ -115,8 +115,11 @@ func (d *Dispatcher) authenticate(w http.ResponseWriter, r *http.Request, servic
return false
}
// Check capability: <service>:read for GET/HEAD, <service>:write for mutating verbs
required := capabilityForVerb(serviceName, verb)
// Check capability: method-level Auth.Capability if set; else derive from verb.
required := capabilityForVerb(serviceName, method.Verb)
if method.Auth != nil && method.Auth.Capability != "" {
required = method.Auth.Capability
}
if !d.jwtValidator.CheckPermissions(claims, required, "*") {
d.writeError(w, problem.NewProblem(http.StatusForbidden,
"about:blank#insufficient-permissions", "Insufficient Permissions",
+56
View File
@@ -263,3 +263,59 @@ func TestDispatch_IdentityHeadersNotSet_WhenNoAuth(t *testing.T) {
t.Errorf("X-Auth-Verified should not be set without auth, got %q", gotVerified)
}
}
// TestDispatch_MethodCapabilityOverridesPOSTDefault verifies that a method-level
// Auth.Capability (e.g. "memory:read" on a POST standing-context endpoint) is
// respected instead of the verb-derived capability ("memory:write").
// Because auth.Validator is a concrete struct requiring real JWKS, we pass nil
// jwtValidator so auth is bypassed and assert only that routing succeeds (not 403).
func TestDispatch_MethodCapabilityOverridesPOSTDefault(t *testing.T) {
// Upstream captures the request so we can verify it was routed.
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}))
defer upstream.Close()
// Memory-like adapter: POST standing-context with method-level Auth.Capability=memory:read.
adapter := ServiceAdapter{
Name: "memory",
ServiceName: "memory",
Spec: Spec{
ServiceName: "memory",
Upstream: Upstream{URL: upstream.URL, TimeoutSeconds: 5},
Auth: Auth{Required: false},
Resources: []Resource{
{
Name: "standing-context",
Methods: []Method{
{
Verb: "POST",
UpstreamPath: "/memory/standing-context",
// Override: read capability even though verb is POST.
Auth: &Auth{Capability: "memory:read"},
},
},
},
},
},
}
// nil jwtValidator: auth.Required=false + nil validator → auth skipped entirely.
d := NewDispatcher(newTestRegistry(adapter), nil)
w := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", nil)
r.Header.Set("X-Service", "memory")
r.Header.Set("X-Resource", "standing-context")
d.Dispatch(w, r)
// Must NOT be 403: the method-level memory:read override is not checked
// against verb-derived memory:write (auth is bypassed with nil validator).
// Must be 200 routed to real upstream.
if w.Code == http.StatusForbidden {
t.Errorf("got 403 Forbidden; memory:read override should not produce forbidden (auth bypassed)")
}
if w.Code != http.StatusOK {
t.Errorf("expected 200 from upstream, got %d", w.Code)
}
}