From fc5c0e967deba1181b45d3f3ff6b5fa1e1c76586 Mon Sep 17 00:00:00 2001 From: Admin Bot Date: Sun, 27 Sep 2026 18:21:41 +0800 Subject: [PATCH] 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 --- cmd/gateway/main.go | 6 + internal/config/config.go | 16 +++ internal/config/loader.go | 25 ++-- internal/config/models_test.go | 138 ++++++++++++++++++- internal/logging/logger.go | 8 ++ internal/observability/metrics.go | 46 +++++++ internal/observability/prometheus.go | 8 ++ internal/proxy/proxy.go | 120 ++++++++++++++++- internal/proxy/proxy_test.go | 84 ++++++++++++ internal/proxy/router.go | 1 + internal/proxy/scheduler.go | 155 ++++++++++++++++++++++ internal/proxy/scheduler_test.go | 126 ++++++++++++++++++ internal/server/router.go | 11 ++ internal/serviceadapter/memory_adapter.go | 67 ++++++++++ internal/serviceadapter/router.go | 11 +- internal/serviceadapter/router_test.go | 56 ++++++++ 16 files changed, 861 insertions(+), 17 deletions(-) create mode 100644 internal/proxy/scheduler.go create mode 100644 internal/proxy/scheduler_test.go create mode 100644 internal/serviceadapter/memory_adapter.go diff --git a/cmd/gateway/main.go b/cmd/gateway/main.go index ba826d5..15610ef 100644 --- a/cmd/gateway/main.go +++ b/cmd/gateway/main.go @@ -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) diff --git a/internal/config/config.go b/internal/config/config.go index 89e50ac..2c09268 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -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" diff --git a/internal/config/loader.go b/internal/config/loader.go index e0d9ed3..a21375c 100644 --- a/internal/config/loader.go +++ b/internal/config/loader.go @@ -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, diff --git a/internal/config/models_test.go b/internal/config/models_test.go index 8a5da54..a362d74 100644 --- a/internal/config/models_test.go +++ b/internal/config/models_test.go @@ -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 + }()) +} diff --git a/internal/logging/logger.go b/internal/logging/logger.go index 2043612..04487e5 100644 --- a/internal/logging/logger.go +++ b/internal/logging/logger.go @@ -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')) +} diff --git a/internal/observability/metrics.go b/internal/observability/metrics.go index 608d1ad..caa6e7a 100644 --- a/internal/observability/metrics.go +++ b/internal/observability/metrics.go @@ -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 } diff --git a/internal/observability/prometheus.go b/internal/observability/prometheus.go index 6c2d040..b308e98 100644 --- a/internal/observability/prometheus.go +++ b/internal/observability/prometheus.go @@ -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" } diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 66fddc5..aeac194 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -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. diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index d04e307..30b0e4b 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -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) + } + }) + } +} diff --git a/internal/proxy/router.go b/internal/proxy/router.go index da95efa..544ff3a 100644 --- a/internal/proxy/router.go +++ b/internal/proxy/router.go @@ -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) { diff --git a/internal/proxy/scheduler.go b/internal/proxy/scheduler.go new file mode 100644 index 0000000..743f928 --- /dev/null +++ b/internal/proxy/scheduler.go @@ -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) +} diff --git a/internal/proxy/scheduler_test.go b/internal/proxy/scheduler_test.go new file mode 100644 index 0000000..f72386d --- /dev/null +++ b/internal/proxy/scheduler_test.go @@ -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"]) + } +} diff --git a/internal/server/router.go b/internal/server/router.go index c3fcd0f..4601714 100644 --- a/internal/server/router.go +++ b/internal/server/router.go @@ -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 diff --git a/internal/serviceadapter/memory_adapter.go b/internal/serviceadapter/memory_adapter.go new file mode 100644 index 0000000..9fdd95e --- /dev/null +++ b/internal/serviceadapter/memory_adapter.go @@ -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"}, + }}, + }, + }, + } +} diff --git a/internal/serviceadapter/router.go b/internal/serviceadapter/router.go index 4c090d3..a4c0d99 100644 --- a/internal/serviceadapter/router.go +++ b/internal/serviceadapter/router.go @@ -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: :read for GET/HEAD, :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", diff --git a/internal/serviceadapter/router_test.go b/internal/serviceadapter/router_test.go index 35ac7a7..b0016ca 100644 --- a/internal/serviceadapter/router_test.go +++ b/internal/serviceadapter/router_test.go @@ -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) + } +}