From 0301ec79b243afc491d64407db3a04ec87fc9e58 Mon Sep 17 00:00:00 2001 From: Konstantin Pereiaslov Date: Thu, 3 Sep 2026 02:24:01 +0000 Subject: [PATCH 1/2] feat: context-based service selection on shared listen ports Multiple services can now share one ListenPort as a set of context-size tiers. Requests are routed to the service with the smallest context size that fits, and the proxy transparently switches to the next larger service mid-connection once a request outgrows the current one. - New per-service config: ContextSize (tokens) + Tokenizer, or ContextSizeBytes (raw bytes), with validation for shared-port groups (uniform units, unique sizes, known tokenizer). - Token counting via a named registry (qwen3.8/qwen3, gemma4/gemma3) of heuristic per-family counters; adding a model is one RegisterTokenCounter call. - Incremental HTTP request framing (Content-Length, chunked, bodyless methods) with passthrough fallback for unframmable traffic. - Requests are measured from messages/prompt/input text; oversized requests go to the largest tier; unframmable traffic falls back to the smallest tier; switching is upward-only per connection. - test-server: new -openai-api-keep-alive flag for e2e switching tests. - tests: fixed pre-existing flake where TestEvictionOfAlreadyDeadProcessDoesNotLoop counted log lines accumulated across runs (append-mode log never truncated). --- AGENTS.md | 52 +- README.md | 91 +++ config.go | 122 +++- config_test.go | 292 ++++++++++ context_router.go | 538 ++++++++++++++++++ context_router_test.go | 181 ++++++ context_routing_connection_test.go | 489 ++++++++++++++++ context_routing_e2e_test.go | 225 ++++++++ http_request_splitter.go | 301 ++++++++++ http_request_splitter_test.go | 289 ++++++++++ main.go | 22 +- main_test.go | 205 ++++++- test-configs/client-close-full.jsonc | 34 ++ test-configs/healthcheck-stuck-timeout.jsonc | 26 + test-configs/healthcheck-stuck.jsonc | 17 + test-configs/idle-timeout.jsonc | 13 + test-configs/invalid-template.json | 33 ++ test-configs/log-output.jsonc | 35 ++ ...nections-white-waiting-for-resources.jsonc | 35 ++ test-configs/no-resource-requirements.json | 26 + test-configs/openai-api-models-by-id.jsonc | 42 ++ test-configs/resource-check-command.jsonc | 36 ++ test-configs/self-dying.json | 32 ++ .../should-not-use-an-outdated-resource.jsonc | 37 ++ test-server/main.go | 7 +- tokenizer.go | 209 +++++++ tokenizer_test.go | 126 ++++ 27 files changed, 3465 insertions(+), 50 deletions(-) create mode 100644 context_router.go create mode 100644 context_router_test.go create mode 100644 context_routing_connection_test.go create mode 100644 context_routing_e2e_test.go create mode 100644 http_request_splitter.go create mode 100644 http_request_splitter_test.go create mode 100644 test-configs/client-close-full.jsonc create mode 100644 test-configs/healthcheck-stuck-timeout.jsonc create mode 100644 test-configs/healthcheck-stuck.jsonc create mode 100644 test-configs/idle-timeout.jsonc create mode 100644 test-configs/invalid-template.json create mode 100644 test-configs/log-output.jsonc create mode 100644 test-configs/multiple-connections-white-waiting-for-resources.jsonc create mode 100644 test-configs/no-resource-requirements.json create mode 100644 test-configs/openai-api-models-by-id.jsonc create mode 100644 test-configs/resource-check-command.jsonc create mode 100644 test-configs/self-dying.json create mode 100644 test-configs/should-not-use-an-outdated-resource.jsonc create mode 100644 tokenizer.go create mode 100644 tokenizer_test.go diff --git a/AGENTS.md b/AGENTS.md index cd229516..899b0af1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -60,31 +60,39 @@ Client → large-model-proxy → [Service Process] ### Source Files -| File | Purpose | -| --------------------------------- | ------------------------------------------------------------------------------------------- | -| `main.go` | Entry point, signal handling, and ResourceManager bookkeeping (connection counting, lookup) | -| `config.go` | Configuration loading, validation, and defaults (JSONC parsing) | -| `connection.go` | Per-service TCP listeners, client connection handling, and bidirectional traffic forwarding | -| `service.go` | Service lifecycle: on-demand start, health checks, and connecting to a running service | -| `service_process.go` | Service process spawning/stopping, output logging, and process-exit monitoring | -| `resources.go` | Resource reservation, LRU eviction, and resource release logic | -| `openai_api.go` | Unified OpenAI API server and request routing to backends by model name | -| `management_api.go` | Management HTTP server and embedded web dashboard assets | -| `monitor_resources.go` | Resource availability monitoring and change broadcasting | -| `monitor_process_hook.go` | Test-only synchronization hook for process-exit timing (compiled with the `testhooks` tag) | -| `monitor_process_hook_default.go` | No-op production stub for `monitor_process_hook.go` | -| `tty.go` | TTY/terminal handling utilities | +| File | Purpose | +| --------------------------------- | --------------------------------------------------------------------------------------------------------- | +| `main.go` | Entry point, signal handling, and ResourceManager bookkeeping (connection counting, lookup) | +| `config.go` | Configuration loading, validation, and defaults (JSONC parsing) | +| `connection.go` | Per-service TCP listeners, client connection handling, and bidirectional traffic forwarding | +| `service.go` | Service lifecycle: on-demand start, health checks, and connecting to a running service | +| `service_process.go` | Service process spawning/stopping, output logging, and process-exit monitoring | +| `resources.go` | Resource reservation, LRU eviction, and resource release logic | +| `openai_api.go` | Unified OpenAI API server and request routing to backends by model name | +| `context_router.go` | Context-based routing: shared-port service tiers, request size measurement, per-connection tier switching | +| `http_request_splitter.go` | Incremental HTTP request framing from a raw TCP stream (used by context routing) | +| `tokenizer.go` | Named token counters (heuristic, per model family) used to measure request context sizes | +| `management_api.go` | Management HTTP server and embedded web dashboard assets | +| `monitor_resources.go` | Resource availability monitoring and change broadcasting | +| `monitor_process_hook.go` | Test-only synchronization hook for process-exit timing (compiled with the `testhooks` tag) | +| `monitor_process_hook_default.go` | No-op production stub for `monitor_process_hook.go` | +| `tty.go` | TTY/terminal handling utilities | ### Test Files -| File | Purpose | -| --------------------------- | --------------------------------------------------- | -| `main_test.go` | Core proxy integration tests | -| `config_test.go` | Configuration parsing and validation tests | -| `management_api_test.go` | Management API endpoint tests | -| `monitor_resources_test.go` | Resource monitoring tests | -| `util_test.go` | Shared test utilities and helpers | -| `test-server/main.go` | Simulated backend service used in integration tests | +| File | Purpose | +| ------------------------------------ | ------------------------------------------------------------------------------------ | +| `main_test.go` | Core proxy integration tests | +| `config_test.go` | Configuration parsing and validation tests | +| `management_api_test.go` | Management API endpoint tests | +| `monitor_resources_test.go` | Resource monitoring tests | +| `context_router_test.go` | Context routing tier selection and request unit counting tests | +| `context_routing_connection_test.go` | In-process tests for the routed connection handler (routing, switching, passthrough) | +| `context_routing_e2e_test.go` | End-to-end context routing test through the real proxy binary | +| `http_request_splitter_test.go` | HTTP request framing tests | +| `tokenizer_test.go` | Token counter tests | +| `util_test.go` | Shared test utilities and helpers | +| `test-server/main.go` | Simulated backend service used in integration tests | ### Other Key Files diff --git a/README.md b/README.md index 0e9a109a..6794b49e 100644 --- a/README.md +++ b/README.md @@ -164,6 +164,7 @@ Below is a breakdown of what this configuration does: 6. When ComfyUI is no longer in use, its container will be killed using the `docker kill comfyui` command. Other services will be terminated normally. 7. `StartupTimeoutMilliseconds` in starting ComfyUI makes large-model-proxy wait up to 60 seconds before giving up and considering the ComfyUI startup failed (as opposed to the default value of 10 minutes). 8. Service URLs are configured as follows: + - **Automatic1111**: Uses the default URL template (`DefaultServiceUrl`) which resolves to `http://localhost:7860/` - **Gemma27B**: Uses a custom static URL `http://gemma-proxy-server/` (no port templating) - **Qwen2.5-7B-Instruct**: Explicitly set to `null`, so no URL will be generated even though a default is available @@ -196,6 +197,96 @@ Currently, the following OpenAI API endpoints are supported: - `/v1/models/{model}` - More to come +## Context-based routing + +Several services can share a single `ListenPort` to form a set of context-size +tiers. The proxy then routes every request to the service with the **smallest +context size that still fits** the request, and transparently switches to the +next larger service once a request outgrows the current one — even in the +middle of a keep-alive connection. This makes it possible to run, for +example, a fast 4k-context instance of a model for short conversations and a +bigger 32k-context instance that is only loaded when conversations actually +grow that long: + +```jsonc +{ + "Services": [ + { + "Name": "Qwen3-8B-4k", + "ListenPort": "8085", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "18085", + "Command": "llama-server", + "Args": "-m Qwen3-8B.gguf -c 4096 --port 18085", + "ContextSize": 4096, // context window in tokens + "Tokenizer": "qwen3.8", + "ResourceRequirements": { "VRAM-GPU-1": 9000 }, + }, + { + "Name": "Qwen3-8B-32k", + "ListenPort": "8085", // same port: joins the same routing group + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "18086", + "Command": "llama-server", + "Args": "-m Qwen3-8B.gguf -c 32768 --port 18086", + "ContextSize": 32768, + "Tokenizer": "qwen3.8", + "ResourceRequirements": { "VRAM-GPU-1": 22000 }, + }, + ], +} +``` + +Rules for a shared port: + +- Every service on the port must define a context size: either `ContextSize` + (tokens, together with a `Tokenizer`) or `ContextSizeBytes` (raw bytes). +- All services in the group must use the same unit (tokens or bytes), the + same tokenizer, and unique context sizes. Starting a larger service follows + the usual resource logic: it may have to evict the smaller one (or any other + LRU service) first. + +### Token counting + +`Tokenizer` selects a registered token counter used to measure requests. +Built-in counters: `qwen3.8` (alias `qwen3`) and `gemma4` (alias `gemma3`). +Counting is a per-model-family approximation (character-class segmentation +with per-family ratios) rather than an exact BPE implementation: it is close +enough for tier selection but tends to slightly overestimate, so configure +context sizes with headroom for the model's output tokens. Adding a counter +for another model family is a single `RegisterTokenCounter` call in +`tokenizer.go`. + +For services measured in `ContextSizeBytes`, the raw byte length of the +request's text content is used instead of a token count — useful for models +without a known tokenizer or for non-OpenAI-style HTTP backends. + +### How requests are measured + +Only the text content of the request body is counted: `messages[].content` +(including content-part arrays, where non-text parts are skipped), `prompt` +and `input` fields of OpenAI-style JSON requests. Bodies that are not +recognizable JSON are measured as raw text. Requests that cannot be framed at +all (no `Content-Length` and not chunked), non-HTTP traffic, and clients that +send nothing for 60 seconds after connecting are blindly forwarded to the +smallest tier. + +### Behavior details + +- Routing decisions happen per request. Switching to a larger tier closes the + current service connection, starts the larger service (with the usual + transparent startup delay for the client) and continues forwarding on the + same client connection. The smaller service then idles out on its own. +- Switching only ever goes upward within a connection: a small follow-up + request (model list, new conversation on the same connection) stays on the + current tier so auxiliary requests cannot thrash services. +- Requests larger than every configured tier are routed to the largest tier + (which will report the context-length error to the client). +- A connection that outgrows all tiers keeps using the largest one. +- Like all HTTP/1.1 intermediaries that re-route per request, this assumes + clients do not pipeline requests (every real HTTP client sends the next + request only after reading the previous response). + ## Management API The management API is a simple HTTP API that allows you to get the status of the proxy and the services it is proxying. diff --git a/config.go b/config.go index 5dd9cb93..e5f9d0f0 100644 --- a/config.go +++ b/config.go @@ -133,6 +133,18 @@ type ServiceConfig struct { OpenAiApiModels []string ServiceUrl *ServiceUrlOption `json:"ServiceUrl,omitempty"` ResourceRequirements map[string]int `json:"ResourceRequirements"` + + // ContextSize is this service's context window in tokens. Multiple + // services may share a ListenPort when every one of them defines a context + // size; requests are routed to the smallest service that fits, switching to + // a larger one once the request outgrows the current service. + ContextSize *uint + // ContextSizeBytes is ContextSize measured in raw bytes instead of tokens; + // it cannot be combined with ContextSize or Tokenizer. + ContextSizeBytes *uint + // Tokenizer names the token counter (see tokenizer.go) used to measure + // request sizes for services with a token-based ContextSize. + Tokenizer string } type ResourceAvailable struct { Amount int @@ -322,9 +334,15 @@ func validateConfig(cfg Config) error { } portSet := make(map[string][]string) // port -> list of service names + servicesByPort := make(map[string][]ServiceConfig) + portOrder := make([]string, 0) for _, svc := range cfg.Services { if svc.ListenPort != "" { + if _, seen := servicesByPort[svc.ListenPort]; !seen { + portOrder = append(portOrder, svc.ListenPort) + } portSet[svc.ListenPort] = append(portSet[svc.ListenPort], svc.Name) + servicesByPort[svc.ListenPort] = append(servicesByPort[svc.ListenPort], svc) } } if cfg.OpenAiApi.ListenPort != "" { @@ -333,10 +351,17 @@ func validateConfig(cfg Config) error { if cfg.ManagementApi.ListenPort != "" { portSet[cfg.ManagementApi.ListenPort] = append(portSet[cfg.ManagementApi.ListenPort], "Management API") } - for p, svcs := range portSet { - if len(svcs) > 1 { + for _, port := range portOrder { + servicesOnPort := servicesByPort[port] + if len(servicesOnPort) > 1 { + issues = append(issues, validateSharedPortContextRouting(port, servicesOnPort)...) + } + } + for port, names := range portSet { + // ports shared between services and the OpenAI API / Management API remain conflicts + if len(names) > 1 && len(names) > len(servicesByPort[port]) { issues = append(issues, - fmt.Sprintf("multiple services listening on port %s: %v", p, svcs)) + fmt.Sprintf("multiple services listening on port %s: %v", port, names)) } } @@ -394,6 +419,33 @@ func validateConfig(cfg Config) error { } } + for i, svc := range cfg.Services { + nameOrIndex := serviceNameOrIndex(svc.Name, i) + if svc.ContextSize != nil && svc.ContextSizeBytes != nil { + issues = append(issues, + fmt.Sprintf("service %s defines both ContextSize and ContextSizeBytes, only one of them can be used", nameOrIndex)) + } + if svc.ContextSize != nil && *svc.ContextSize == 0 { + issues = append(issues, + fmt.Sprintf("service %s has ContextSize 0, it must be greater than 0", nameOrIndex)) + } + if svc.ContextSizeBytes != nil && *svc.ContextSizeBytes == 0 { + issues = append(issues, + fmt.Sprintf("service %s has ContextSizeBytes 0, it must be greater than 0", nameOrIndex)) + } + if svc.Tokenizer != "" { + if _, found := GetTokenCounter(svc.Tokenizer); !found { + knownTokenizers := joinStrings(RegisteredTokenCounterNames(), ", ") + issues = append(issues, + fmt.Sprintf("service %s specifies unknown tokenizer %q, known tokenizers: %s", nameOrIndex, svc.Tokenizer, knownTokenizers)) + } + if svc.ContextSizeBytes != nil { + issues = append(issues, + fmt.Sprintf("service %s specifies Tokenizer together with ContextSizeBytes: Tokenizer is only used with a token-based ContextSize", nameOrIndex)) + } + } + } + // Validate ServiceUrl templates if present and not null for i, svc := range cfg.Services { nameOrIndex := serviceNameOrIndex(svc.Name, i) @@ -426,6 +478,70 @@ func validateConfig(cfg Config) error { return nil } +// validateSharedPortContextRouting checks a group of services that share a +// ListenPort. Sharing a port is only allowed for context-based routing, which +// requires a well-formed, unambiguous set of context size tiers. +func validateSharedPortContextRouting(port string, services []ServiceConfig) []string { + var groupIssues []string + names := make([]string, len(services)) + for i, svc := range services { + names[i] = svc.Name + } + + tokenModeServices := 0 + byteModeServices := 0 + firstTokenizer := "" + for _, svc := range services { + if svc.ContextSize == nil && svc.ContextSizeBytes == nil { + groupIssues = append(groupIssues, + fmt.Sprintf("service %q listening on port %s defines neither ContextSize nor ContextSizeBytes", svc.Name, port)) + continue + } + if svc.ContextSize != nil { + tokenModeServices++ + if svc.Tokenizer == "" { + groupIssues = append(groupIssues, + fmt.Sprintf("service %q listening on port %s defines ContextSize but no Tokenizer", svc.Name, port)) + } else if firstTokenizer == "" { + firstTokenizer = svc.Tokenizer + } else if svc.Tokenizer != firstTokenizer { + groupIssues = append(groupIssues, + fmt.Sprintf("services listening on port %s must use the same tokenizer, found %q and %q", port, firstTokenizer, svc.Tokenizer)) + } + } else { + byteModeServices++ + } + } + if tokenModeServices > 0 && byteModeServices > 0 { + groupIssues = append(groupIssues, + fmt.Sprintf("services listening on port %s must use the same unit for context sizes: ContextSize (tokens) and ContextSizeBytes (bytes) cannot be mixed", port)) + } + + seenSizes := make(map[uint64]string) + for _, svc := range services { + var size uint64 + if svc.ContextSize != nil { + size = uint64(*svc.ContextSize) + } else if svc.ContextSizeBytes != nil { + size = uint64(*svc.ContextSizeBytes) + } else { + continue + } + if otherService, duplicate := seenSizes[size]; duplicate { + groupIssues = append(groupIssues, + fmt.Sprintf("services %q and %q listening on port %s have a duplicate context size %d, sizes must be unique", otherService, svc.Name, port, size)) + } else { + seenSizes[size] = svc.Name + } + } + if len(groupIssues) == 0 { + return nil + } + return append([]string{fmt.Sprintf( + "multiple services listening on port %s: [%s]. When multiple services share a port, context-based routing is enabled: every service must define ContextSize (with a Tokenizer) or ContextSizeBytes, and requests are routed to the service with the smallest context size that fits", + port, joinStrings(names, ", "))}, groupIssues...) +} + // validateGoTemplate validates that the given string is a valid Go template func validateGoTemplate(templateStr string) error { _, err := template.New("validation").Parse(templateStr) diff --git a/config_test.go b/config_test.go index b99fb5b9..dccdfc6b 100644 --- a/config_test.go +++ b/config_test.go @@ -1035,3 +1035,295 @@ func TestLogLevelInvalidValue(t *testing.T) { } assert.Contains(t, err.Error(), "invalid LogLevel") } + +// --- context-based routing configuration --- + +func TestMultipleServicesSamePortWithContextSizesIsValid(t *testing.T) { + t.Parallel() + cfg, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "qwen-4k", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 4096, + "Tokenizer": "qwen3.8" + }, + { + "Name": "qwen-32k", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 32768, + "Tokenizer": "qwen3.8" + } + ] + }`) + if err != nil { + t.Fatalf("did not expect an error but got: %v", err) + } + assert.Equal(t, uint(4096), *cfg.Services[0].ContextSize) + assert.Equal(t, "qwen3.8", cfg.Services[0].Tokenizer) + assert.Equal(t, uint(32768), *cfg.Services[1].ContextSize) +} + +func TestMultipleServicesSamePortBytesModeIsValid(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "small", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSizeBytes": 16000 + }, + { + "Name": "large", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSizeBytes": 128000 + } + ] + }`) + if err != nil { + t.Fatalf("did not expect an error but got: %v", err) + } +} + +func TestMultipleServicesSamePortMissingContextSize(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "sized", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 4096, + "Tokenizer": "qwen3.8" + }, + { + "Name": "unsized", + "ListenPort": "8080", + "Command": "/bin/echo" + } + ] + }`) + checkExpectedErrorMessages(t, err, []string{ + "multiple services listening on port 8080", + "\"unsized\"", + "ContextSize", + }) +} + +func TestMultipleServicesSamePortMixedContextUnits(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "tokens", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 4096, + "Tokenizer": "qwen3.8" + }, + { + "Name": "bytes", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSizeBytes": 16000 + } + ] + }`) + checkExpectedErrorMessages(t, err, []string{ + "port 8080", + "ContextSize", + "ContextSizeBytes", + "same unit", + }) +} + +func TestServiceWithBothContextSizeKindsInvalid(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "svc", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 4096, + "Tokenizer": "qwen3.8", + "ContextSizeBytes": 16000 + } + ] + }`) + checkExpectedErrorMessages(t, err, []string{"\"svc\"", "ContextSize", "ContextSizeBytes", "both"}) +} + +func TestMultipleServicesSamePortDuplicateContextSize(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "one", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 4096, + "Tokenizer": "qwen3.8" + }, + { + "Name": "two", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 4096, + "Tokenizer": "qwen3.8" + } + ] + }`) + checkExpectedErrorMessages(t, err, []string{"port 8080", "4096", "duplicate"}) +} + +func TestMultipleServicesSamePortUnknownTokenizer(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "one", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 4096, + "Tokenizer": "nonexistent-tokenizer" + }, + { + "Name": "two", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 8192, + "Tokenizer": "qwen3.8" + } + ] + }`) + checkExpectedErrorMessages(t, err, []string{"\"nonexistent-tokenizer\"", "unknown tokenizer", "\"one\""}) +} + +func TestUnknownTokenizerOnSingleServiceInvalid(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "svc", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 4096, + "Tokenizer": "typo-tokenizer" + } + ] + }`) + checkExpectedErrorMessages(t, err, []string{"unknown tokenizer", "\"typo-tokenizer\""}) +} + +func TestMultipleServicesSamePortMissingTokenizer(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "one", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 4096 + }, + { + "Name": "two", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 8192, + "Tokenizer": "qwen3.8" + } + ] + }`) + checkExpectedErrorMessages(t, err, []string{"Tokenizer", "\"one\"", "port 8080"}) +} + +func TestMultipleServicesSamePortInconsistentTokenizer(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "one", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 4096, + "Tokenizer": "qwen3.8" + }, + { + "Name": "two", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 8192, + "Tokenizer": "gemma4" + } + ] + }`) + checkExpectedErrorMessages(t, err, []string{"port 8080", "Tokenizer", "same tokenizer"}) +} + +func TestContextSizeBytesWithTokenizerInvalid(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "small", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSizeBytes": 16000, + "Tokenizer": "qwen3.8" + }, + { + "Name": "large", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSizeBytes": 64000 + } + ] + }`) + checkExpectedErrorMessages(t, err, []string{"\"small\"", "Tokenizer", "ContextSizeBytes"}) +} + +func TestZeroContextSizeInvalid(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "one", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 0, + "Tokenizer": "qwen3.8" + }, + { + "Name": "two", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 8192, + "Tokenizer": "qwen3.8" + } + ] + }`) + checkExpectedErrorMessages(t, err, []string{"\"one\"", "ContextSize", "greater than 0"}) +} + +func TestSingleServiceWithContextSizeIsValid(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "svc", + "ListenPort": "8080", + "Command": "/bin/echo", + "ContextSize": 4096, + "Tokenizer": "qwen3.8" + } + ] + }`) + if err != nil { + t.Fatalf("did not expect an error but got: %v", err) + } +} diff --git a/context_router.go b/context_router.go new file mode 100644 index 00000000..3311b0d8 --- /dev/null +++ b/context_router.go @@ -0,0 +1,538 @@ +package main + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "log" + "net" + "sort" + "strings" + "sync" + "time" +) + +// contextRoutingInitialRequestTimeout is how long a connection on a +// context-routed port waits for the first complete HTTP request before +// falling back to blind forwarding to the smallest tier. Clients speak first +// in HTTP, so this only fires for probes or non-HTTP traffic. It is a variable +// so tests can shorten it. +var contextRoutingInitialRequestTimeout = 60 * time.Second + +// contextUnitMode distinguishes token-based context sizes from raw-byte-based +// ones within a group of services sharing a listen port. +type contextUnitMode int + +const ( + contextUnitsTokens contextUnitMode = iota + contextUnitsBytes +) + +// contextTier is one service of a shared-port group together with the context +// size it can serve. +type contextTier struct { + serviceConfig ServiceConfig + limitUnits uint64 // in tokens or bytes, depending on the router's mode +} + +// contextRouter routes client requests to one of several services that share a +// listen port, preferring the service with the smallest context size that still +// fits the request. Config validation guarantees that all services in a group +// use the same unit mode and, for token mode, the same tokenizer. +type contextRouter struct { + tiers []contextTier // sorted ascending by limitUnits + unitMode contextUnitMode + tokenCounter TokenCounter // nil in byte mode +} + +func buildContextRouter(services []ServiceConfig) *contextRouter { + router := &contextRouter{} + if len(services) == 0 { + return router + } + if services[0].ContextSizeBytes != nil { + router.unitMode = contextUnitsBytes + } else { + router.unitMode = contextUnitsTokens + router.tokenCounter, _ = GetTokenCounter(services[0].Tokenizer) + } + for _, serviceConfig := range services { + limit := uint64(0) + if serviceConfig.ContextSizeBytes != nil { + limit = uint64(*serviceConfig.ContextSizeBytes) + } else if serviceConfig.ContextSize != nil { + limit = uint64(*serviceConfig.ContextSize) + } + router.tiers = append(router.tiers, contextTier{serviceConfig: serviceConfig, limitUnits: limit}) + } + sort.Slice(router.tiers, func(i, j int) bool { + if router.tiers[i].limitUnits != router.tiers[j].limitUnits { + return router.tiers[i].limitUnits < router.tiers[j].limitUnits + } + return router.tiers[i].serviceConfig.Name < router.tiers[j].serviceConfig.Name + }) + return router +} + +// selectTierIndex returns the index of the smallest tier whose limit fits the +// given number of units. When nothing fits, the largest tier is returned: the +// request is oversized for every configured service, so the largest context +// gives it the best chance (and the backend reports the error if it cannot). +func (r *contextRouter) selectTierIndex(units uint64) int { + for index, tier := range r.tiers { + if units <= tier.limitUnits { + return index + } + } + return len(r.tiers) - 1 +} + +// countRequestUnits measures how many units of context (tokens or bytes, +// depending on the router's mode) one complete HTTP request occupies. Only the +// text content of the request body is counted (chat messages, prompts, +// embeddings input); when the body is not recognizable JSON, the raw body is +// measured as a fallback. +func countRequestUnits(requestBytes []byte, router *contextRouter) uint64 { + body := extractRequestBody(requestBytes) + texts := extractRequestTexts(body) + if len(texts) == 0 { + if len(body) == 0 { + return 0 + } + texts = []string{string(body)} + } + + total := uint64(0) + for _, text := range texts { + if router.unitMode == contextUnitsBytes { + total += uint64(len(text)) + } else { + total += uint64(router.tokenCounter(text)) + } + } + return total +} + +func extractRequestBody(requestBytes []byte) []byte { + terminatorOffset, terminatorLength := findHeaderTerminator(requestBytes) + if terminatorOffset < 0 { + return requestBytes + } + return requestBytes[terminatorOffset+terminatorLength:] +} + +// extractRequestTexts pulls the transcribed text out of an OpenAI-compatible +// request body. It supports messages[].content (string or content-part arrays), +// prompt (string or array) and input (string or array, used by /v1/embeddings). +func extractRequestTexts(body []byte) []string { + var payload struct { + Messages []json.RawMessage `json:"messages"` + Prompt json.RawMessage `json:"prompt"` + Input json.RawMessage `json:"input"` + } + if err := json.Unmarshal(body, &payload); err != nil { + return nil + } + + var texts []string + for _, message := range payload.Messages { + var messageContent struct { + Content json.RawMessage `json:"content"` + } + if err := json.Unmarshal(message, &messageContent); err != nil { + return nil + } + texts = append(texts, jsonStringOrArrayTexts(messageContent.Content)...) + } + texts = append(texts, jsonStringOrArrayTexts(payload.Prompt)...) + texts = append(texts, jsonStringOrArrayTexts(payload.Input)...) + return texts +} + +// jsonStringOrArrayTexts extracts text from a JSON value that is either a +// string, an array of strings, or an array of content parts +// ({"type":"text","text":...}); non-text parts are skipped. +func jsonStringOrArrayTexts(raw json.RawMessage) []string { + if len(raw) == 0 { + return nil + } + var asString string + if err := json.Unmarshal(raw, &asString); err == nil { + return []string{asString} + } + var asArray []json.RawMessage + if err := json.Unmarshal(raw, &asArray); err != nil { + return nil + } + var texts []string + for _, element := range asArray { + var asString string + if err := json.Unmarshal(element, &asString); err == nil { + texts = append(texts, asString) + continue + } + var contentPart struct { + Type string `json:"type"` + Text string `json:"text"` + } + if err := json.Unmarshal(element, &contentPart); err == nil && contentPart.Type == "text" { + texts = append(texts, contentPart.Text) + } + } + return texts +} + +// startContextRoutedProxy listens on one port shared by several services and +// routes every client connection through the context router. +func startContextRoutedProxy(listenPort string, router *contextRouter) { + serviceNames := make([]string, len(router.tiers)) + descriptions := make([]string, len(router.tiers)) + for i, tier := range router.tiers { + serviceNames[i] = tier.serviceConfig.Name + descriptions[i] = fmt.Sprintf("%s (%s)", tier.serviceConfig.Name, tier.limitDescription()) + } + listener, err := net.Listen("tcp", ":"+listenPort) + if err != nil { + log.Fatalf("[%s] Fatal error: cannot listen on port %s: %v", joinStrings(serviceNames, ","), listenPort, err) + } + defer func(listener net.Listener) { + _ = listener.Close() + }(listener) + log.Printf("[port %s] Listening with context-based routing: %s", listenPort, joinStrings(descriptions, " < ")) + + for { + if interrupted.Load() { + return + } + clientConnection, err := listener.Accept() + if err != nil { + if interrupted.Load() { + return + } + log.Printf("[port %s] Error accepting connection: %v", listenPort, err) + continue + } + log.Printf("[port %s] New client connection received %s", listenPort, humanReadableConnection(clientConnection)) + go handleRoutedConnection(clientConnection, router, startServiceIfNotAlreadyRunningAndConnect) + } +} + +// limitDescription renders a tier's context size for logs. +func (t contextTier) limitDescription() string { + if t.serviceConfig.ContextSizeBytes != nil { + return fmt.Sprintf("%d bytes", t.limitUnits) + } + return fmt.Sprintf("%d tokens", t.limitUnits) +} + +// handleRoutedConnection proxies one client connection through the context +// router. The first complete HTTP request decides which tier serves the +// connection (if several requests complete in the very first burst, the +// largest one decides). Later requests on the same connection that no longer +// fit the current tier cause an in-flight switch to the smallest tier that +// fits: the current service connection is closed, the new service is started +// (reusing all the existing on-demand startup machinery), and forwarding +// continues on a new service connection. Switching only ever goes upwards +// within a connection: a later small request stays on the current tier so +// that auxiliary requests (model listings, health probes) cannot thrash +// services. +// +// The client -> service direction runs in this goroutine; the service -> +// client direction runs in one copier goroutine that survives service +// switches. Requests are only ever committed to a service once fully +// received, so a switch never loses or duplicates request bytes. Switching +// assumes request/response lockstep — a client never sends request N+1 +// before fully reading response N — which every real HTTP client honors; +// HTTP/1.1 pipelining clients are not supported across a switch (a switch +// may cut a pipelined response that is still in flight). +func handleRoutedConnection(clientConnection net.Conn, router *contextRouter, connectToService func(ServiceConfig, <-chan struct{}) net.Conn) { + if interrupted.Load() { + _ = clientConnection.Close() + return + } + clientReader, closeClientReader, clientDisconnected := startClientReadMonitor(clientConnection) + defer closeClientReader() + + // Pump client bytes into a channel so the routing loop can select on data + // arrival, client disconnects, and the initial routing timeout at once. + clientChunks := make(chan []byte) + stopPump := make(chan struct{}) + go func() { + buffer := make([]byte, 32*1024) + for { + bytesRead, readErr := clientReader.Read(buffer) + if bytesRead > 0 { + chunk := make([]byte, bytesRead) + copy(chunk, buffer[:bytesRead]) + select { + case clientChunks <- chunk: + case <-stopPump: + return + } + } + if readErr != nil { + return + } + } + }() + defer close(stopPump) + + decision, decided := routeInitialRequest(router, clientChunks, clientDisconnected) + if !decided { + _ = clientConnection.Close() + return + } + + currentTierIndex := decision.tierIndex + currentServiceConfig := router.tiers[currentTierIndex].serviceConfig + + resourceManager.incrementConnection(currentServiceConfig.Name, 0, 1) + serviceConnection := connectToService(currentServiceConfig, clientDisconnected) + if serviceConnection == nil { + resourceManager.incrementConnection(currentServiceConfig.Name, 0, -1) + closeConnectionAndHandleError( + clientConnection, + currentServiceConfig, + "client", + "failed to establish a connection to the service", + ) + return + } + log.Printf("[%s] Routing connection %s to service (context size %s)", currentServiceConfig.Name, humanReadableConnection(clientConnection), router.tiers[currentTierIndex].limitDescription()) + trackServiceLastUsed(currentServiceConfig, true) + resourceManager.incrementConnection(currentServiceConfig.Name, 1, -1) + + // The copier forwards service -> client and keeps working across service + // switches. When a service connection ends, it closes the client + // connection unless a switch is in progress. + routerDone := make(chan struct{}) + copierDone := make(chan struct{}) + nextServiceConnection := make(chan net.Conn, 1) + switchCoordinator := struct { + mutex sync.Mutex + swapPending bool + }{} + go func(activeConnection net.Conn) { + defer close(copierDone) + for { + _, err := io.Copy(clientConnection, activeConnection) + if err == nil || isConnectionClosedError(err) { + // The service connection ended. Find out whether that was our + // own switch or the end of the road for this client. + switchCoordinator.mutex.Lock() + swapPending := switchCoordinator.swapPending + switchCoordinator.mutex.Unlock() + if !swapPending { + _ = clientConnection.Close() + return + } + } else { + // e.g. the client is gone; nothing more to deliver + _ = clientConnection.Close() + return + } + select { + case activeConnection = <-nextServiceConnection: + case <-routerDone: + _ = clientConnection.Close() + return + case <-clientDisconnected: + return + } + } + }(serviceConnection) + + connectionCounted := true + // Deferred teardown, registered in reverse execution order: the pump stops + // first, then stats are released, the service connection is closed so the + // copier's io.Copy unblocks, the copier finishes before the client + // connection is closed so pending response bytes are not lost, and the + // client read monitor is shut down last (it is idempotent). + defer closeClientReader() + defer func() { _ = clientConnection.Close() }() + defer func() { close(routerDone); <-copierDone }() + defer func() { _ = serviceConnection.Close() }() + defer func() { + if connectionCounted { + resourceManager.incrementConnection(currentServiceConfig.Name, -1, 0) + trackServiceLastUsed(currentServiceConfig, false) + } + }() + + writeToService := func(data []byte) bool { + if len(data) == 0 { + return true + } + if _, err := serviceConnection.Write(data); err != nil { + log.Printf("[%s] Error writing request bytes to service: %v", currentServiceConfig.Name, err) + return false + } + return true + } + + switchToTier := func(targetTierIndex int, units uint64) (net.Conn, bool) { + previousServiceConfig := currentServiceConfig + // Signal the copier BEFORE closing: its in-flight io.Copy must not + // treat the close as end-of-stream for the client. + switchCoordinator.mutex.Lock() + switchCoordinator.swapPending = true + switchCoordinator.mutex.Unlock() + _ = serviceConnection.Close() + resourceManager.incrementConnection(previousServiceConfig.Name, -1, 0) + trackServiceLastUsed(previousServiceConfig, false) + + newServiceConfig := router.tiers[targetTierIndex].serviceConfig + log.Printf("[%s] Request of %s exceeds context size of %s, switching to %s (context size %s)", + previousServiceConfig.Name, router.unitDescription(units), previousServiceConfig.Name, newServiceConfig.Name, router.tiers[targetTierIndex].limitDescription()) + resourceManager.incrementConnection(newServiceConfig.Name, 0, 1) + newConnection := connectToService(newServiceConfig, clientDisconnected) + if newConnection == nil { + resourceManager.incrementConnection(newServiceConfig.Name, 0, -1) + connectionCounted = false + return nil, false + } + trackServiceLastUsed(newServiceConfig, true) + resourceManager.incrementConnection(newServiceConfig.Name, 1, -1) + currentServiceConfig = newServiceConfig + currentTierIndex = targetTierIndex + nextServiceConnection <- newConnection + return newConnection, true + } + + for _, request := range decision.requests { + if !writeToService(request) { + return + } + } + passthroughMode := decision.passthrough + if passthroughMode && !writeToService(decision.throughBytes) { + return + } + + splitter := decision.splitter + for { + select { + case chunk, ok := <-clientChunks: + if !ok { + // Client finished sending: deliver any partially received + // request, then close the service side of the connection. + if !passthroughMode { + writeToService(splitter.FlushIncomplete()) + } + resourceManager.incrementConnection(currentServiceConfig.Name, -1, 0) + trackServiceLastUsed(currentServiceConfig, false) + connectionCounted = false + _ = serviceConnection.Close() + return + } + if passthroughMode { + if !writeToService(chunk) { + return + } + continue + } + result := splitter.Write(chunk) + for _, request := range result.completedRequests { + units := countRequestUnits(request, router) + targetTierIndex := router.selectTierIndex(units) + if targetTierIndex > currentTierIndex { + newConnection, switched := switchToTier(targetTierIndex, units) + if !switched { + return + } + serviceConnection = newConnection + } + if !writeToService(request) { + return + } + } + if result.enteredPassthrough { + passthroughMode = true + if !writeToService(result.throughBytes) { + return + } + } + case <-clientDisconnected: + // The client is gone; teardown happens in the deferred cleanup and + // the copier exits on its own. + return + } + } +} + +func isConnectionClosedError(err error) bool { + return errors.Is(err, net.ErrClosed) || errors.Is(err, io.EOF) || strings.Contains(err.Error(), "connection reset") +} + +// initialRoutingDecision captures where a connection should start and what +// already-received bytes must be forwarded to the chosen service. +type initialRoutingDecision struct { + tierIndex int + requests [][]byte // complete requests received before the decision + passthrough bool + throughBytes []byte // bytes to forward before streaming on (passthrough only) + splitter *httpRequestSplitter +} + +// routeInitialRequest waits for the first complete HTTP request (or a reason +// to give up on framing) and computes the tier that should serve the +// connection. It returns decided=false when the client went away before any +// routing decision could be made. +func routeInitialRequest(router *contextRouter, clientChunks <-chan []byte, clientDisconnected <-chan struct{}) (initialRoutingDecision, bool) { + splitter := newHttpRequestSplitter() + initialTimeout := time.After(contextRoutingInitialRequestTimeout) + + for { + select { + case chunk, ok := <-clientChunks: + if !ok { + return initialRoutingDecision{}, false + } + result := splitter.Write(chunk) + if len(result.completedRequests) > 0 { + largestUnits := uint64(0) + for _, request := range result.completedRequests { + units := countRequestUnits(request, router) + if units > largestUnits { + largestUnits = units + } + } + return initialRoutingDecision{ + tierIndex: router.selectTierIndex(largestUnits), + requests: result.completedRequests, + passthrough: result.enteredPassthrough, + throughBytes: result.throughBytes, + splitter: splitter, + }, true + } + if result.enteredPassthrough { + return initialRoutingDecision{ + tierIndex: 0, + passthrough: true, + throughBytes: result.throughBytes, + splitter: splitter, + }, true + } + case <-clientDisconnected: + return initialRoutingDecision{}, false + case <-initialTimeout: + // Nothing framable arrived: behave like a plain proxy to the + // smallest tier and forward whatever was received so far. + return initialRoutingDecision{ + tierIndex: 0, + passthrough: true, + throughBytes: splitter.FlushIncomplete(), + splitter: splitter, + }, true + } + } +} + +// unitDescription renders a request size for logs. +func (r *contextRouter) unitDescription(units uint64) string { + if r.unitMode == contextUnitsBytes { + return fmt.Sprintf("%d bytes", units) + } + return fmt.Sprintf("%d tokens", units) +} diff --git a/context_router_test.go b/context_router_test.go new file mode 100644 index 00000000..a67f84ae --- /dev/null +++ b/context_router_test.go @@ -0,0 +1,181 @@ +package main + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func testContextRouterServices() []ServiceConfig { + return []ServiceConfig{ + {Name: "large", ListenPort: "9000", ContextSize: uintPtr(1000), Tokenizer: "qwen3.8"}, + {Name: "small", ListenPort: "9000", ContextSize: uintPtr(10), Tokenizer: "qwen3.8"}, + {Name: "medium", ListenPort: "9000", ContextSize: uintPtr(100), Tokenizer: "qwen3.8"}, + } +} + +func uintPtr(n uint) *uint { + return &n +} + +func TestBuildContextRouterSortsTiersAscending(t *testing.T) { + t.Parallel() + router := buildContextRouter(testContextRouterServices()) + + assert.Equal(t, contextUnitsTokens, router.unitMode) + assert.Len(t, router.tiers, 3) + assert.Equal(t, uint64(10), router.tiers[0].limitUnits) + assert.Equal(t, "small", router.tiers[0].serviceConfig.Name) + assert.Equal(t, uint64(100), router.tiers[1].limitUnits) + assert.Equal(t, "medium", router.tiers[1].serviceConfig.Name) + assert.Equal(t, uint64(1000), router.tiers[2].limitUnits) + assert.Equal(t, "large", router.tiers[2].serviceConfig.Name) + assert.NotNil(t, router.tokenCounter, "token counter must be resolved from the registry") +} + +func TestSelectTierIndexPicksSmallestFittingTier(t *testing.T) { + t.Parallel() + router := buildContextRouter(testContextRouterServices()) + + assert.Equal(t, 0, router.selectTierIndex(0), "empty request goes to the smallest tier") + assert.Equal(t, 0, router.selectTierIndex(9), "just below the first limit") + assert.Equal(t, 0, router.selectTierIndex(10), "exactly the first limit fits") + assert.Equal(t, 1, router.selectTierIndex(11), "one over the first limit") + assert.Equal(t, 1, router.selectTierIndex(100)) + assert.Equal(t, 2, router.selectTierIndex(101)) + assert.Equal(t, 2, router.selectTierIndex(100000), "over all limits falls back to the largest tier") +} + +func TestBuildContextRouterBytesMode(t *testing.T) { + t.Parallel() + services := []ServiceConfig{ + {Name: "large-bytes", ListenPort: "9001", ContextSizeBytes: uintPtr(1000)}, + {Name: "small-bytes", ListenPort: "9001", ContextSizeBytes: uintPtr(10)}, + } + router := buildContextRouter(services) + + assert.Equal(t, contextUnitsBytes, router.unitMode) + assert.Nil(t, router.tokenCounter) + assert.Equal(t, uint64(10), router.tiers[0].limitUnits) + assert.Equal(t, 0, router.selectTierIndex(5)) + assert.Equal(t, 1, router.selectTierIndex(9999)) +} + +// --- request units counting --- + +func testUnitsRouter(t *testing.T) *contextRouter { + t.Helper() + router := buildContextRouter([]ServiceConfig{ + {Name: "small", ListenPort: "9000", ContextSize: uintPtr(10), Tokenizer: "qwen3.8"}, + {Name: "large", ListenPort: "9000", ContextSize: uintPtr(1000), Tokenizer: "qwen3.8"}, + }) + assert.NotNil(t, router.tokenCounter) + return router +} + +// qwen3.8 counting: "hello world" = 3 tokens, "hi" = 1 token (see tokenizer_test.go) +func TestCountRequestUnitsChatMessages(t *testing.T) { + t.Parallel() + router := testUnitsRouter(t) + request := testChatRequest("", `{"model":"m","messages":[ + {"role":"system","content":"hello world"}, + {"role":"user","content":"hi"}]}`) + + units := countRequestUnits(request, router) + + assert.Equal(t, uint64(4), units) +} + +func TestCountRequestUnitsMultipartContent(t *testing.T) { + t.Parallel() + router := testUnitsRouter(t) + request := testChatRequest("", `{"model":"m","messages":[ + {"role":"user","content":[ + {"type":"text","text":"hello world"}, + {"type":"text","text":"hi"}, + {"type":"image_url","image_url":{"url":"http://example.com/image.png"}}]}]}`) + + units := countRequestUnits(request, router) + + // only the text parts are counted + assert.Equal(t, uint64(4), units) +} + +func TestCountRequestUnitsCompletionPrompt(t *testing.T) { + t.Parallel() + router := testUnitsRouter(t) + request := testChatRequest("", `{"model":"m","prompt":"hello world"}`) + + assert.Equal(t, uint64(3), countRequestUnits(request, router)) +} + +func TestCountRequestUnitsPromptArray(t *testing.T) { + t.Parallel() + router := testUnitsRouter(t) + request := testChatRequest("", `{"model":"m","prompt":["hello world","hi"]}`) + + assert.Equal(t, uint64(4), countRequestUnits(request, router)) +} + +func TestCountRequestUnitsEmbeddingsInput(t *testing.T) { + t.Parallel() + router := testUnitsRouter(t) + request := testChatRequest("", `{"model":"m","input":["hello world","hi"]}`) + + assert.Equal(t, uint64(4), countRequestUnits(request, router)) +} + +func TestCountRequestUnitsInvalidJsonFallsBackToRawBody(t *testing.T) { + t.Parallel() + router := testUnitsRouter(t) + request := testChatRequest("", `this is not json`) + + // "this"=1 + " is"=1 + " not"=1 + " json"=1 = 4 with the qwen3.8 heuristic + assert.Equal(t, uint64(4), countRequestUnits(request, router)) +} + +func TestCountRequestUnitsNonJsonContentTypeFallsBackToRawBody(t *testing.T) { + t.Parallel() + router := testUnitsRouter(t) + request := []byte("POST /v1/audio/transcriptions HTTP/1.1\r\n" + + "Host: localhost\r\n" + + "Content-Type: multipart/form-data\r\n" + + "Content-Length: 10\r\n\r\n" + + "0123456789") + + // digits: ceil(10/3) = 4 + assert.Equal(t, uint64(4), countRequestUnits(request, router)) +} + +func TestCountRequestUnitsGetRequest(t *testing.T) { + t.Parallel() + router := testUnitsRouter(t) + request := []byte("GET /v1/models HTTP/1.1\r\nHost: localhost\r\n\r\n") + + assert.Equal(t, uint64(0), countRequestUnits(request, router)) +} + +func TestCountRequestUnitsBytesModeMeasuresTextBytes(t *testing.T) { + t.Parallel() + router := buildContextRouter([]ServiceConfig{ + {Name: "small", ListenPort: "9000", ContextSizeBytes: uintPtr(10)}, + {Name: "large", ListenPort: "9000", ContextSizeBytes: uintPtr(1000)}, + }) + request := testChatRequest("", `{"model":"m","messages":[{"role":"user","content":"hello world"}]}`) + requestWithExtraJsonKeys := testChatRequest("", `{"model":"m","stream":true,"temperature":0.7,"messages":[{"role":"user","content":"hello world"}]}`) + + // only the text content is measured, not the JSON envelope + assert.Equal(t, uint64(11), countRequestUnits(request, router)) + assert.Equal(t, uint64(11), countRequestUnits(requestWithExtraJsonKeys, router)) +} + +func TestCountRequestUnitsBytesModeMultibyteContent(t *testing.T) { + t.Parallel() + router := buildContextRouter([]ServiceConfig{ + {Name: "small", ListenPort: "9000", ContextSizeBytes: uintPtr(10)}, + {Name: "large", ListenPort: "9000", ContextSizeBytes: uintPtr(1000)}, + }) + request := testChatRequest("", `{"model":"m","messages":[{"role":"user","content":"你好"}]}`) + + assert.Equal(t, uint64(6), countRequestUnits(request, router), "raw byte length of the UTF-8 content") +} diff --git a/context_routing_connection_test.go b/context_routing_connection_test.go new file mode 100644 index 00000000..db174815 --- /dev/null +++ b/context_routing_connection_test.go @@ -0,0 +1,489 @@ +package main + +import ( + "bufio" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +// setupRoutedConnectionTestGlobals prepares the package-level state that +// handleRoutedConnection touches (connection stats, service configs). These +// tests deliberately do not run in parallel because they configure globals. +func setupRoutedConnectionTestGlobals(t *testing.T, services []ServiceConfig) { + t.Helper() + config = Config{LogLevel: LogLevelNormal} + serviceConfigByName = make(map[string]*ServiceConfig, len(services)) + connectionStats := make(map[string]ServiceConnectionStats, len(services)) + runningServices := make(map[string]*RunningService, len(services)) + now := time.Now() + for i := range services { + serviceConfigByName[services[i].Name] = &services[i] + connectionStats[services[i].Name] = ServiceConnectionStats{} + runningServices[services[i].Name] = &RunningService{manageMutex: newChannelMutex(), lastUsed: &now} + } + resourceManager = ResourceManager{ + serviceMutex: &sync.Mutex{}, + connectionStatsMutex: &sync.Mutex{}, + connectionStats: connectionStats, + runningServices: runningServices, + resourcesInUse: map[string]int{}, + resourcesReserved: map[string]int{}, + resourcesAvailable: map[string]int{}, + resourcesAvailableMutex: &sync.Mutex{}, + monitorUnpauseChansMutex: &sync.Mutex{}, + monitorUnpauseChans: map[string]chan struct{}{}, + resourceChangeByResourceMutex: &sync.Mutex{}, + checkCommandFirstChangeByResourceChans: map[string]map[string]chan struct{}{}, + resourceChangeByResourceChans: map[string]map[string]chan bool{}, + } +} + +type recordingBackend struct { + server *httptest.Server + hits atomic.Int64 + marker string +} + +// newRecordingBackend starts an HTTP backend that responds to any request with +// a body naming which backend served it, so tests can tell where a request was +// routed. Keep-alive is enabled, like real LLM backends. +func newRecordingBackend(marker string) *recordingBackend { + backend := &recordingBackend{marker: marker} + backend.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + backend.hits.Add(1) + _, _ = fmt.Fprintf(w, "served-by:%s", marker) + })) + return backend +} + +func (b *recordingBackend) address() string { + return b.server.Listener.Addr().String() +} + +func (b *recordingBackend) close() { + b.server.Close() +} + +func connectToRecordingBackend(backends map[string]*recordingBackend) func(ServiceConfig, <-chan struct{}) net.Conn { + return func(serviceConfig ServiceConfig, _ <-chan struct{}) net.Conn { + backend, found := backends[serviceConfig.Name] + if !found { + return nil + } + connection, err := net.Dial("tcp", backend.address()) + if err != nil { + return nil + } + return connection + } +} + +// clientServerConnectionPair creates a real TCP connection pair: the server +// side is meant to be handed to the proxy, the client side is returned to the +// test to act as the client. +func clientServerConnectionPair(t *testing.T) (clientSide net.Conn, serverSide net.Conn) { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + t.Cleanup(func() { _ = listener.Close() }) + + serverConnections := make(chan net.Conn, 1) + go func() { + connection, err := listener.Accept() + if err == nil { + serverConnections <- connection + } + }() + + clientConnection, err := net.Dial("tcp", listener.Addr().String()) + if err != nil { + t.Fatalf("failed to dial: %v", err) + } + t.Cleanup(func() { _ = clientConnection.Close() }) + + select { + case serverConnection := <-serverConnections: + t.Cleanup(func() { _ = serverConnection.Close() }) + return clientConnection, serverConnection + case <-time.After(5 * time.Second): + t.Fatalf("server side of the connection pair was not accepted") + return nil, nil + } +} + +// startRoutedHandler runs the routed connection handler and registers a +// cleanup that makes sure the handler (and its goroutines) has fully exited +// before the next test reinitializes package-level state. +func startRoutedHandler(t *testing.T, serverConnection net.Conn, router *contextRouter, connect func(ServiceConfig, <-chan struct{}) net.Conn) { + t.Helper() + handlerDone := make(chan struct{}) + go func() { + defer close(handlerDone) + handleRoutedConnection(serverConnection, router, connect) + }() + t.Cleanup(func() { + _ = serverConnection.Close() + select { + case <-handlerDone: + case <-time.After(10 * time.Second): + t.Errorf("routed connection handler did not exit within 10s after connection close") + } + }) +} + +func writeRequestToConnection(t *testing.T, connection net.Conn, body string) { + t.Helper() + request := testChatRequest("", body) + if _, err := connection.Write(request); err != nil { + t.Fatalf("failed to write request: %v", err) + } +} + +func readResponseFromConnection(t *testing.T, connection net.Conn) (int, string) { + t.Helper() + _ = connection.SetReadDeadline(time.Now().Add(10 * time.Second)) + response, err := http.ReadResponse(bufio.NewReader(connection), nil) + if err != nil { + t.Fatalf("failed to read response: %v", err) + } + defer func() { _ = response.Body.Close() }() + body, err := io.ReadAll(response.Body) + if err != nil { + t.Fatalf("failed to read response body: %v", err) + } + return response.StatusCode, string(body) +} + +func readResponseFromConnectionForBody(t *testing.T, connection net.Conn) string { + t.Helper() + _, body := readResponseFromConnection(t, connection) + return body +} + +func waitForConnectionStats(t *testing.T, expected ServiceConnectionStats, timeout time.Duration, serviceNames ...string) { + t.Helper() + deadline := time.Now().Add(timeout) + for { + resourceManager.connectionStatsMutex.Lock() + allMatch := true + for _, name := range serviceNames { + if resourceManager.connectionStats[name] != expected { + allMatch = false + } + } + resourceManager.connectionStatsMutex.Unlock() + if allMatch { + return + } + if time.Now().After(deadline) { + resourceManager.connectionStatsMutex.Lock() + current := make(map[string]ServiceConnectionStats, len(serviceNames)) + for _, name := range serviceNames { + current[name] = resourceManager.connectionStats[name] + } + resourceManager.connectionStatsMutex.Unlock() + t.Fatalf("connection stats did not settle on %v within %s, got %v", expected, timeout, current) + } + time.Sleep(10 * time.Millisecond) + } +} + +func testTokenRouterServices() []ServiceConfig { + return []ServiceConfig{ + {Name: "large", ContextSize: uintPtr(1000), Tokenizer: "qwen3.8"}, + {Name: "small", ContextSize: uintPtr(10), Tokenizer: "qwen3.8"}, + } +} + +// "hi" is 1 token; the string below is 30 words => 30 tokens with qwen3.8 +const smallRequestContent = "hi" + +func largeRequestContent() string { + return strings.Repeat("word ", 30) +} + +func TestRoutedConnectionRoutesFirstRequestToSmallestFittingTier(t *testing.T) { + services := testTokenRouterServices() + setupRoutedConnectionTestGlobals(t, services) + backends := map[string]*recordingBackend{ + "small": newRecordingBackend("small"), + "large": newRecordingBackend("large"), + } + t.Cleanup(func() { + backends["small"].close() + backends["large"].close() + }) + router := buildContextRouter(services) + clientConnection, serverConnection := clientServerConnectionPair(t) + startRoutedHandler(t, serverConnection, router, connectToRecordingBackend(backends)) + + writeRequestToConnection(t, clientConnection, `{"model":"m","messages":[{"role":"user","content":"`+smallRequestContent+`"}]}`) + _, body := readResponseFromConnection(t, clientConnection) + + assert.Equal(t, "served-by:small", body) + assert.Equal(t, int64(1), backends["small"].hits.Load()) + assert.Equal(t, int64(0), backends["large"].hits.Load()) + + _ = clientConnection.Close() + waitForConnectionStats(t, ServiceConnectionStats{}, 10*time.Second, "small", "large") +} + +func TestRoutedConnectionRoutesOversizedRequestToLargestTier(t *testing.T) { + services := testTokenRouterServices() + setupRoutedConnectionTestGlobals(t, services) + backends := map[string]*recordingBackend{ + "small": newRecordingBackend("small"), + "large": newRecordingBackend("large"), + } + t.Cleanup(func() { + backends["small"].close() + backends["large"].close() + }) + router := buildContextRouter(services) + clientConnection, serverConnection := clientServerConnectionPair(t) + startRoutedHandler(t, serverConnection, router, connectToRecordingBackend(backends)) + + writeRequestToConnection(t, clientConnection, `{"model":"m","messages":[{"role":"user","content":"`+largeRequestContent()+`"}]}`) + _, body := readResponseFromConnection(t, clientConnection) + + assert.Equal(t, "served-by:large", body) + assert.Equal(t, int64(0), backends["small"].hits.Load()) + assert.Equal(t, int64(1), backends["large"].hits.Load()) +} + +func TestRoutedConnectionSwitchesToLargerTierOnSameConnection(t *testing.T) { + services := testTokenRouterServices() + setupRoutedConnectionTestGlobals(t, services) + backends := map[string]*recordingBackend{ + "small": newRecordingBackend("small"), + "large": newRecordingBackend("large"), + } + t.Cleanup(func() { + backends["small"].close() + backends["large"].close() + }) + router := buildContextRouter(services) + clientConnection, serverConnection := clientServerConnectionPair(t) + startRoutedHandler(t, serverConnection, router, connectToRecordingBackend(backends)) + + writeRequestToConnection(t, clientConnection, `{"model":"m","messages":[{"role":"user","content":"`+smallRequestContent+`"}]}`) + assert.Equal(t, "served-by:small", readResponseFromConnectionForBody(t, clientConnection)) + + // The conversation grew past the small tier's limit: the same connection + // must be re-routed to the large tier for the next request. + writeRequestToConnection(t, clientConnection, `{"model":"m","messages":[{"role":"user","content":"`+largeRequestContent()+`"}]}`) + assert.Equal(t, "served-by:large", readResponseFromConnectionForBody(t, clientConnection)) + + assert.Equal(t, int64(1), backends["small"].hits.Load()) + assert.Equal(t, int64(1), backends["large"].hits.Load()) + + _ = clientConnection.Close() + waitForConnectionStats(t, ServiceConnectionStats{}, 10*time.Second, "small", "large") +} + +func TestRoutedConnectionDoesNotSwitchDownMidConnection(t *testing.T) { + services := testTokenRouterServices() + setupRoutedConnectionTestGlobals(t, services) + backends := map[string]*recordingBackend{ + "small": newRecordingBackend("small"), + "large": newRecordingBackend("large"), + } + t.Cleanup(func() { + backends["small"].close() + backends["large"].close() + }) + router := buildContextRouter(services) + clientConnection, serverConnection := clientServerConnectionPair(t) + startRoutedHandler(t, serverConnection, router, connectToRecordingBackend(backends)) + + // First request goes to the large tier... + writeRequestToConnection(t, clientConnection, `{"model":"m","messages":[{"role":"user","content":"`+largeRequestContent()+`"}]}`) + assert.Equal(t, "served-by:large", readResponseFromConnectionForBody(t, clientConnection)) + // ...a small follow-up request must stay there to avoid thrashing services. + writeRequestToConnection(t, clientConnection, `{"model":"m","messages":[{"role":"user","content":"`+smallRequestContent+`"}]}`) + assert.Equal(t, "served-by:large", readResponseFromConnectionForBody(t, clientConnection)) + + assert.Equal(t, int64(0), backends["small"].hits.Load()) + assert.Equal(t, int64(2), backends["large"].hits.Load()) +} + +func TestRoutedConnectionPassthroughRoutesToSmallestTier(t *testing.T) { + services := testTokenRouterServices() + setupRoutedConnectionTestGlobals(t, services) + backends := map[string]*recordingBackend{ + "small": newRecordingBackend("small"), + "large": newRecordingBackend("large"), + } + t.Cleanup(func() { + backends["small"].close() + backends["large"].close() + }) + router := buildContextRouter(services) + clientConnection, serverConnection := clientServerConnectionPair(t) + startRoutedHandler(t, serverConnection, router, connectToRecordingBackend(backends)) + + // Not HTTP: cannot be framed, must be blindly forwarded to the smallest tier + _, err := clientConnection.Write([]byte("PROXY custom handshake\r\n")) + assert.NoError(t, err) + statusCode, body := readResponseFromConnection(t, clientConnection) + + // The HTTP backend answers the malformed request with 400 without invoking + // the handler, proving the bytes were forwarded to the small service (the + // large one is never hit; a malformed request never reaches a handler, so + // the hit counters cannot be used here). + assert.Equal(t, http.StatusBadRequest, statusCode) + assert.NotEmpty(t, body) + assert.Equal(t, int64(0), backends["large"].hits.Load()) +} + +func TestRoutedConnectionBytesModeSwitchesOnByteCount(t *testing.T) { + services := []ServiceConfig{ + {Name: "large", ContextSizeBytes: uintPtr(1000)}, + {Name: "small", ContextSizeBytes: uintPtr(10)}, + } + setupRoutedConnectionTestGlobals(t, services) + backends := map[string]*recordingBackend{ + "small": newRecordingBackend("small"), + "large": newRecordingBackend("large"), + } + t.Cleanup(func() { + backends["small"].close() + backends["large"].close() + }) + router := buildContextRouter(services) + clientConnection, serverConnection := clientServerConnectionPair(t) + startRoutedHandler(t, serverConnection, router, connectToRecordingBackend(backends)) + + writeRequestToConnection(t, clientConnection, `{"model":"m","messages":[{"role":"user","content":"short"}]}`) + assert.Equal(t, "served-by:small", readResponseFromConnectionForBody(t, clientConnection)) + + // 50 bytes of content > 10 byte limit of the small tier + writeRequestToConnection(t, clientConnection, `{"model":"m","messages":[{"role":"user","content":"`+strings.Repeat("x", 50)+`"}]}`) + assert.Equal(t, "served-by:large", readResponseFromConnectionForBody(t, clientConnection)) +} + +func TestRoutedConnectionClientDisconnectDuringSwitchAborts(t *testing.T) { + services := testTokenRouterServices() + setupRoutedConnectionTestGlobals(t, services) + backends := map[string]*recordingBackend{ + "small": newRecordingBackend("small"), + } + t.Cleanup(func() { backends["small"].close() }) + + connect := func(serviceConfig ServiceConfig, clientDisconnected <-chan struct{}) net.Conn { + if serviceConfig.Name == "small" { + return connectToRecordingBackend(backends)(serviceConfig, clientDisconnected) + } + // Simulate a slow-to-start large service: it only "finishes starting" + // when the client goes away, then fails. + <-clientDisconnected + return nil + } + + router := buildContextRouter(services) + clientConnection, serverConnection := clientServerConnectionPair(t) + startRoutedHandler(t, serverConnection, router, connect) + + writeRequestToConnection(t, clientConnection, `{"model":"m","messages":[{"role":"user","content":"`+smallRequestContent+`"}]}`) + assert.Equal(t, "served-by:small", readResponseFromConnectionForBody(t, clientConnection)) + + // Oversized request triggers a switch; client gives up while waiting. + writeRequestToConnection(t, clientConnection, `{"model":"m","messages":[{"role":"user","content":"`+largeRequestContent()+`"}]}`) + _ = clientConnection.Close() + + waitForConnectionStats(t, ServiceConnectionStats{}, 10*time.Second, "small", "large") +} + +func TestRoutedConnectionInitialRequestTimeoutFallsBackToSmallestTier(t *testing.T) { + previousTimeout := contextRoutingInitialRequestTimeout + contextRoutingInitialRequestTimeout = 150 * time.Millisecond + t.Cleanup(func() { contextRoutingInitialRequestTimeout = previousTimeout }) + + services := testTokenRouterServices() + setupRoutedConnectionTestGlobals(t, services) + backends := map[string]*recordingBackend{ + "small": newRecordingBackend("small"), + "large": newRecordingBackend("large"), + } + t.Cleanup(func() { + backends["small"].close() + backends["large"].close() + }) + router := buildContextRouter(services) + clientConnection, serverConnection := clientServerConnectionPair(t) + startRoutedHandler(t, serverConnection, router, connectToRecordingBackend(backends)) + + // Connect but send nothing until after the initial routing timeout, then + // send a normal request: it must be forwarded blind to the smallest tier. + time.Sleep(300 * time.Millisecond) + writeRequestToConnection(t, clientConnection, `{"model":"m","messages":[{"role":"user","content":"hi"}]}`) + statusCode, body := readResponseFromConnection(t, clientConnection) + + assert.Equal(t, http.StatusOK, statusCode) + assert.Equal(t, "served-by:small", body) + assert.Equal(t, int64(1), backends["small"].hits.Load()) + assert.Equal(t, int64(0), backends["large"].hits.Load()) +} + +func TestRouteInitialRequestUsesLargestOfInitialBurst(t *testing.T) { + services := testTokenRouterServices() + setupRoutedConnectionTestGlobals(t, services) + router := buildContextRouter(services) + + // A client that pipelines several requests in one burst (legal, though rare, + // in HTTP/1.1): the routing decision must account for every request already + // completed in the burst, so the connection starts on the tier that fits + // the largest of them instead of switching immediately after. + smallRequest := testChatRequest("", `{"model":"m","messages":[{"role":"user","content":"`+smallRequestContent+`"}]}`) + largeRequest := testChatRequest("", `{"model":"m","messages":[{"role":"user","content":"`+largeRequestContent()+`"}]}`) + burst := append(append([]byte{}, smallRequest...), largeRequest...) + clientChunks := make(chan []byte, 1) + clientChunks <- burst + + decision, decided := routeInitialRequest(router, clientChunks, make(chan struct{})) + + assert.True(t, decided) + assert.Equal(t, 1, decision.tierIndex, "the 30-token request must pick the large tier") + assert.Len(t, decision.requests, 2, "both pipelined requests must be forwarded") + assert.Equal(t, uint64(31), countRequestUnits(decision.requests[1], router)) +} + +func TestRoutedConnectionContentLengthIsUsedForUnits(t *testing.T) { + services := testTokenRouterServices() + setupRoutedConnectionTestGlobals(t, services) + backends := map[string]*recordingBackend{ + "small": newRecordingBackend("small"), + "large": newRecordingBackend("large"), + } + t.Cleanup(func() { + backends["small"].close() + backends["large"].close() + }) + router := buildContextRouter(services) + clientConnection, serverConnection := clientServerConnectionPair(t) + startRoutedHandler(t, serverConnection, router, connectToRecordingBackend(backends)) + + // A chunked request whose content exceeds the small tier + chunkedRequest := "POST /v1/chat/completions HTTP/1.1\r\n" + + "Host: localhost\r\n" + + "Transfer-Encoding: chunked\r\n\r\n" + + strconv.FormatInt(int64(len(largeRequestContent())), 16) + "\r\n" + largeRequestContent() + "\r\n" + + "0\r\n\r\n" + _, err := clientConnection.Write([]byte(chunkedRequest)) + assert.NoError(t, err) + + assert.Equal(t, "served-by:large", readResponseFromConnectionForBody(t, clientConnection)) +} diff --git a/context_routing_e2e_test.go b/context_routing_e2e_test.go new file mode 100644 index 00000000..c636eb70 --- /dev/null +++ b/context_routing_e2e_test.go @@ -0,0 +1,225 @@ +package main + +import ( + "bufio" + "fmt" + "net" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +// TestContextBasedRouting verifies the context-based service selection feature +// end to end: several services share one listen port and requests are routed +// to the service with the smallest context size that fits the request, with a +// mid-connection switch to a larger service once the context grows past the +// current service's limit. Both token-based and raw-byte-based context sizes +// are covered. +func TestContextBasedRouting(t *testing.T) { + t.Parallel() + testName := t.Name() + + config := Config{ + Services: []ServiceConfig{ + { + Name: "tokens-small", + ListenPort: "2150", + ProxyTargetHost: "localhost", + ProxyTargetPort: "12150", + Command: "./test-server/test-server", + Args: "-openai-api-port 12150 -openai-api-keep-alive", + ContextSize: uintPtr(5), + Tokenizer: "qwen3.8", + }, + { + Name: "tokens-large", + ListenPort: "2150", + ProxyTargetHost: "localhost", + ProxyTargetPort: "12151", + Command: "./test-server/test-server", + Args: "-openai-api-port 12151 -openai-api-keep-alive", + ContextSize: uintPtr(5000), + Tokenizer: "qwen3.8", + }, + { + Name: "bytes-small", + ListenPort: "2151", + ProxyTargetHost: "localhost", + ProxyTargetPort: "12152", + Command: "./test-server/test-server", + Args: "-openai-api-port 12152 -openai-api-keep-alive", + ContextSizeBytes: uintPtr(10), + }, + { + Name: "bytes-large", + ListenPort: "2151", + ProxyTargetHost: "localhost", + ProxyTargetPort: "12153", + Command: "./test-server/test-server", + Args: "-openai-api-port 12153 -openai-api-keep-alive", + ContextSizeBytes: uintPtr(5000), + }, + }, + ManagementApi: ManagementApi{ListenPort: "4150"}, + } + StandardizeConfigNamesAndPaths(&config, testName) + configFilePath := createTempConfig(t, config) + + waitChannel := make(chan error, 1) + cmd, err := startLargeModelProxy("context-based-routing", configFilePath, "", waitChannel) + if err != nil { + t.Fatalf("could not start application: %v", err) + } + defer func() { + if stopErr := stopApplication(cmd, waitChannel); stopErr != nil { + t.Errorf("failed to stop application: %v", stopErr) + } + for _, address := range []string{ + "localhost:2150", "localhost:2151", "localhost:4150", + "localhost:12150", "localhost:12151", "localhost:12152", "localhost:12153", + } { + if portErr := checkPortClosed(address); portErr != nil { + t.Errorf("port %s is still open after application exit: %v", address, portErr) + } + } + }() + + managementApiAddress := "localhost:4150" + tokenRoutingAddress := "localhost:2150" + + // The first, small request must start and be served by the small service. + smallContent := "hi" + connection, err := net.Dial("tcp", tokenRoutingAddress) + if err != nil { + t.Fatalf("failed to connect to routed port: %v", err) + } + defer func() { _ = connection.Close() }() + + response := sendChatRequestOnConnection(t, connection, smallContent, 30*time.Second) + assert.Contains(t, response, `"role":"assistant"`, "expected a chat completion response") + waitForServiceState(t, managementApiAddress, testName+"_tokens-small", ServiceStateRunning, 30*time.Second) + status := getStatusFromManagementAPI(t, managementApiAddress) + verifyServiceStatus(t, status, testName+"_tokens-large", ServiceStateStopped, 0, 0, map[string]int{}) + + // A request that no longer fits the small service's context must be served + // by the large service on the same client connection. + largeContent := strings.Repeat("word ", 30) + response = sendChatRequestOnConnection(t, connection, largeContent, 60*time.Second) + assert.Contains(t, response, `"role":"assistant"`, "expected a chat completion response after switching tiers") + waitForServiceState(t, managementApiAddress, testName+"_tokens-large", ServiceStateRunning, 30*time.Second) + + // Byte-based routing: 5 bytes fit the 10-byte tier, 50 bytes do not. + bytesRoutingAddress := "localhost:2151" + bytesConnection, err := net.Dial("tcp", bytesRoutingAddress) + if err != nil { + t.Fatalf("failed to connect to byte-routed port: %v", err) + } + defer func() { _ = bytesConnection.Close() }() + + response = sendChatRequestOnConnection(t, bytesConnection, "short", 30*time.Second) + assert.Contains(t, response, `"role":"assistant"`) + waitForServiceState(t, managementApiAddress, testName+"_bytes-small", ServiceStateRunning, 30*time.Second) + + response = sendChatRequestOnConnection(t, bytesConnection, strings.Repeat("x", 50), 60*time.Second) + assert.Contains(t, response, `"role":"assistant"`) + waitForServiceState(t, managementApiAddress, testName+"_bytes-large", ServiceStateRunning, 30*time.Second) +} + +// sendChatRequestOnConnection writes one chat completion request on a raw TCP +// connection (which the proxy keeps open thanks to keep-alive) and returns the +// response body. The timeout must cover on-demand service startup, which can +// take a while for a real model. +func sendChatRequestOnConnection(t *testing.T, connection net.Conn, content string, timeout time.Duration) string { + t.Helper() + body := fmt.Sprintf(`{"model":"m","messages":[{"role":"user","content":%q}]}`, content) + request := fmt.Sprintf( + "POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\nContent-Type: application/json\r\nContent-Length: %d\r\n\r\n%s", + len(body), body) + if _, err := connection.Write([]byte(request)); err != nil { + t.Fatalf("failed to write request: %v", err) + } + _ = connection.SetReadDeadline(time.Now().Add(timeout)) + response, err := readResponse(connection) + if err != nil { + t.Fatalf("failed to read response: %v", err) + } + return response +} + +func readResponse(connection net.Conn) (string, error) { + reader := bufio.NewReader(connection) + statusLine, err := reader.ReadString('\n') + if err != nil { + return "", fmt.Errorf("failed to read status line: %w", err) + } + if !strings.Contains(statusLine, "200") { + return "", fmt.Errorf("unexpected status line: %s", strings.TrimSpace(statusLine)) + } + contentLength := -1 + chunked := false + for { + headerLine, err := reader.ReadString('\n') + if err != nil { + return "", fmt.Errorf("failed to read header line: %w", err) + } + if headerLine == "\r\n" || headerLine == "\n" { + break + } + lower := strings.ToLower(headerLine) + if strings.HasPrefix(lower, "content-length:") { + _, scanErr := fmt.Sscanf(strings.TrimSpace(headerLine[len("content-length:"):]), "%d", &contentLength) + if scanErr != nil { + return "", fmt.Errorf("failed to parse content-length %q: %w", headerLine, scanErr) + } + } + if strings.HasPrefix(lower, "transfer-encoding:") && strings.Contains(lower, "chunked") { + chunked = true + } + } + if chunked { + var body strings.Builder + for { + sizeLine, err := reader.ReadString('\n') + if err != nil { + return "", fmt.Errorf("failed to read chunk size: %w", err) + } + var chunkSize int + _, scanErr := fmt.Sscanf(strings.TrimSpace(sizeLine), "%x", &chunkSize) + if scanErr != nil { + return "", fmt.Errorf("failed to parse chunk size %q: %w", sizeLine, scanErr) + } + if chunkSize == 0 { + _, _ = reader.ReadString('\n') // trailing CRLF after last chunk + return body.String(), nil + } + chunk := make([]byte, chunkSize) + if _, err := readFull(reader, chunk); err != nil { + return "", fmt.Errorf("failed to read chunk: %w", err) + } + body.Write(chunk) + _, _ = reader.ReadString('\n') // CRLF after chunk data + } + } + if contentLength >= 0 { + body := make([]byte, contentLength) + if _, err := readFull(reader, body); err != nil { + return "", fmt.Errorf("failed to read body: %w", err) + } + return string(body), nil + } + return "", fmt.Errorf("response has neither content-length nor chunked encoding") +} + +func readFull(reader *bufio.Reader, buffer []byte) (int, error) { + total := 0 + for total < len(buffer) { + n, err := reader.Read(buffer[total:]) + total += n + if err != nil { + return total, err + } + } + return total, nil +} diff --git a/http_request_splitter.go b/http_request_splitter.go new file mode 100644 index 00000000..72cb8b23 --- /dev/null +++ b/http_request_splitter.go @@ -0,0 +1,301 @@ +package main + +import ( + "bytes" + "strconv" + "strings" +) + +// httpRequestSplitter incrementally frames complete HTTP requests out of a raw +// TCP byte stream. It is used by context-based routing to know where one +// request ends and the next begins, so the request's size can be measured +// before any of its bytes are committed to a backend service. +// +// Framing rules: +// - Requests with Content-Length complete when the full body has arrived. +// - Requests with Transfer-Encoding: chunked complete after the terminal +// chunk (and optional trailers). +// - Bodyless methods (GET, HEAD, ...) complete at the end of the headers. +// - Anything whose framing cannot be determined (methods with a body but no +// length, Expect: 100-continue, connection upgrades, non-HTTP traffic) +// switches the splitter into passthrough: all remaining bytes are handed +// back as-is and no further framing is attempted. +// +// Byte slices returned by Write and FlushIncomplete alias internal buffers and +// are only valid until the next call on the splitter. +type httpRequestSplitter struct { + buffer []byte + consumed int // offset in buffer where the in-flight request starts + state splitterFramingState + passthrough bool + bodyStart int // offset where the current request's body starts + bodyRemain int64 // content-length or current chunk bytes still expected + cursor int // chunked parsing position +} + +type splitterFramingState int + +const ( + framingHeaders splitterFramingState = iota + framingBody + framingChunkSize + framingChunkData + framingTrailers +) + +type httpRequestSplitResult struct { + completedRequests [][]byte + enteredPassthrough bool + // throughBytes is set when enteredPassthrough: all received bytes that are + // not part of a completed request, to be forwarded as-is. + throughBytes []byte +} + +// Methods likely to be seen on LLM service ports. A stream that does not start +// with one of these (as a prefix) is treated as non-HTTP traffic. +var httpMethods = []string{ + "OPTIONS", "GET", "HEAD", "POST", "PUT", "DELETE", "TRACE", "CONNECT", "PATCH", + "PROPFIND", "PROPPATCH", "MKCOL", "COPY", "MOVE", "LOCK", "UNLOCK", "REPORT", + "MKCALENDAR", "ACL", "SEARCH", +} + +func newHttpRequestSplitter() *httpRequestSplitter { + return &httpRequestSplitter{state: framingHeaders} +} + +func (s *httpRequestSplitter) Write(chunk []byte) httpRequestSplitResult { + if s.passthrough { + return httpRequestSplitResult{enteredPassthrough: true, throughBytes: chunk} + } + + if s.consumed > 0 { + // Compact completed requests out of the buffer. Callers consume + // completed requests before calling Write again (documented contract), + // so mutating the consumed region is safe and keeps memory bounded on + // long-lived keep-alive connections. + s.buffer = append(s.buffer[:0], s.buffer[s.consumed:]...) + s.bodyStart -= s.consumed + s.cursor -= s.consumed + s.consumed = 0 + } + s.buffer = append(s.buffer, chunk...) + + var result httpRequestSplitResult + if s.parse(&result) { + s.passthrough = true + result.enteredPassthrough = true + result.throughBytes = s.buffer[s.consumed:] + s.buffer = nil + } + return result +} + +// FlushIncomplete returns the bytes of the currently in-flight (not yet fully +// received) request, for forwarding when no more bytes can arrive. +func (s *httpRequestSplitter) FlushIncomplete() []byte { + if s.passthrough { + return nil + } + return s.buffer[s.consumed:] +} + +// parse consumes as many complete requests from the buffer as possible, +// appending them to result.completedRequests. It returns true when framing +// must be abandoned (passthrough). +func (s *httpRequestSplitter) parse(result *httpRequestSplitResult) bool { + for { + if s.consumed == len(s.buffer) { + return false // need more bytes + } + switch s.state { + case framingHeaders: + if !startsWithHttpMethodPrefix(s.buffer[s.consumed:]) { + return true // not HTTP traffic + } + headerEnd, terminatorLength := findHeaderTerminator(s.buffer[s.consumed:]) + if headerEnd < 0 { + return false + } + headerEnd += s.consumed + headersEndAbsolute := headerEnd + terminatorLength + method, headers, ok := parseHttpHead(s.buffer[s.consumed:headersEndAbsolute]) + if !ok { + return true + } + if hasHeaderValueToken(headers, "expect", "100-continue") { + // The client waits for a 100 Continue that only the backend can + // send, so the body cannot be counted here. + return true + } + if method == "CONNECT" { + return true + } + if hasHeaderValueToken(headers, "connection", "upgrade") || headers["upgrade"] != "" { + result.completedRequests = append(result.completedRequests, s.buffer[s.consumed:headersEndAbsolute]) + s.consumed = headersEndAbsolute + return true // the upgraded stream after the headers cannot be framed + } + if hasHeaderValueToken(headers, "transfer-encoding", "chunked") { + s.bodyStart = headersEndAbsolute + s.cursor = headersEndAbsolute + s.state = framingChunkSize + continue + } + if contentLength, present := headers["content-length"]; present { + length, err := strconv.ParseInt(contentLength, 10, 64) + if err != nil || length < 0 { + return true + } + s.bodyStart = headersEndAbsolute + s.bodyRemain = length + s.state = framingBody + continue + } + switch method { + case "GET", "HEAD", "DELETE", "OPTIONS", "TRACE": + result.completedRequests = append(result.completedRequests, s.buffer[s.consumed:headersEndAbsolute]) + s.consumed = headersEndAbsolute + continue // possibly pipelined requests follow + default: + // A request that usually carries a body but declares no length: + // its end cannot be determined without consuming it. + return true + } + case framingBody: + if int64(len(s.buffer)-s.bodyStart) >= s.bodyRemain { + bodyEnd := s.bodyStart + int(s.bodyRemain) + result.completedRequests = append(result.completedRequests, s.buffer[s.consumed:bodyEnd]) + s.consumed = bodyEnd + s.state = framingHeaders + continue + } + return false + case framingChunkSize: + line, next, ok := readLine(s.buffer, s.cursor) + if !ok { + return false + } + chunkSize, err := strconv.ParseInt(strings.TrimSpace(strings.SplitN(line, ";", 2)[0]), 16, 63) + if err != nil || chunkSize < 0 { + return true + } + s.cursor = next + if chunkSize == 0 { + s.state = framingTrailers + continue + } + s.bodyRemain = chunkSize + s.state = framingChunkData + case framingChunkData: + // chunk bytes plus the trailing CRLF + if int64(len(s.buffer)-s.cursor) >= s.bodyRemain+2 { + s.cursor += int(s.bodyRemain) + 2 + s.state = framingChunkSize + continue + } + return false + case framingTrailers: + // Trailer section: header lines until an empty line + for { + line, next, ok := readLine(s.buffer, s.cursor) + if !ok { + return false + } + s.cursor = next + if line == "" { + result.completedRequests = append(result.completedRequests, s.buffer[s.consumed:s.cursor]) + s.consumed = s.cursor + s.state = framingHeaders + break + } + } + } + } +} + +// startsWithHttpMethodPrefix reports whether data is still consistent with the +// beginning of an HTTP request (a known method followed by a space). +func startsWithHttpMethodPrefix(data []byte) bool { + for _, method := range httpMethods { + prefix := []byte(method + " ") + if len(data) <= len(prefix) && bytes.HasPrefix(prefix, data) { + return true // data is a prefix of "METHOD " + } + if bytes.HasPrefix(data, prefix) { + return true // data starts with "METHOD " + } + } + return false +} + +// findHeaderTerminator finds the earliest end-of-headers marker (CRLFCRLF or, +// leniently like net/textproto, LFLF) in data, returning the offset where the +// header block ends and the length of the terminator. +func findHeaderTerminator(data []byte) (int, int) { + crlf := bytes.Index(data, []byte("\r\n\r\n")) + lflf := bytes.Index(data, []byte("\n\n")) + switch { + case crlf >= 0 && (lflf < 0 || crlf <= lflf): + return crlf, 4 + case lflf >= 0: + return lflf, 2 + default: + return -1, 0 + } +} + +// parseHttpHead parses the request line and headers of one HTTP head block +// (headers already stripped of the terminator). +func parseHttpHead(head []byte) (string, map[string]string, bool) { + lines := strings.Split(string(head), "\n") + requestLineParts := strings.Fields(strings.TrimSuffix(lines[0], "\r")) + if len(requestLineParts) < 3 || !strings.HasPrefix(requestLineParts[len(requestLineParts)-1], "HTTP/") { + return "", nil, false + } + method := requestLineParts[0] + headers := make(map[string]string) + for _, line := range lines[1:] { + line = strings.TrimSuffix(line, "\r") + if line == "" { + continue + } + colonIndex := strings.IndexByte(line, ':') + if colonIndex < 0 { + return "", nil, false + } + key := strings.ToLower(strings.TrimSpace(line[:colonIndex])) + value := strings.TrimSpace(line[colonIndex+1:]) + if existing, present := headers[key]; present { + headers[key] = existing + ", " + value + } else { + headers[key] = value + } + } + return method, headers, true +} + +func hasHeaderValueToken(headers map[string]string, name string, expectedToken string) bool { + value, present := headers[name] + if !present { + return false + } + for _, token := range strings.Split(value, ",") { + if strings.EqualFold(strings.TrimSpace(token), expectedToken) { + return true + } + } + return false +} + +// readLine reads one line (terminated by \n, with an optional preceding \r +// stripped) starting at pos. ok is false when no complete line is buffered yet. +func readLine(data []byte, pos int) (line string, next int, ok bool) { + newlineIndex := bytes.IndexByte(data[pos:], '\n') + if newlineIndex < 0 { + return "", 0, false + } + lineEnd := pos + newlineIndex + line = string(data[pos:lineEnd]) + line = strings.TrimSuffix(line, "\r") + return line, lineEnd + 1, true +} diff --git a/http_request_splitter_test.go b/http_request_splitter_test.go new file mode 100644 index 00000000..08978457 --- /dev/null +++ b/http_request_splitter_test.go @@ -0,0 +1,289 @@ +package main + +import ( + "bytes" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" +) + +const testChatRequestBody = `{"model":"m","messages":[{"role":"user","content":"hi"}]}` + +func testChatRequest(contentLengthExtra string, body string) []byte { + return []byte("POST /v1/chat/completions HTTP/1.1\r\n" + + "Host: localhost\r\n" + + "Content-Type: application/json\r\n" + + contentLengthExtra + + "Content-Length: " + strconv.Itoa(len(body)) + "\r\n\r\n" + + body) +} + +func writeInChunks(t *testing.T, splitter *httpRequestSplitter, data []byte, chunkSize int) []httpRequestSplitResult { + t.Helper() + results := make([]httpRequestSplitResult, 0) + for offset := 0; offset < len(data); offset += chunkSize { + end := offset + chunkSize + if end > len(data) { + end = len(data) + } + results = append(results, splitter.Write(data[offset:end])) + } + return results +} + +func completedRequestBytes(results []httpRequestSplitResult) [][]byte { + var completed [][]byte + for _, result := range results { + completed = append(completed, result.completedRequests...) + } + return completed +} + +func TestSplitterCompleteRequestInOneWrite(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + request := testChatRequest("", testChatRequestBody) + result := splitter.Write(request) + + assert.False(t, result.enteredPassthrough) + assert.Len(t, result.completedRequests, 1) + assert.True(t, bytes.Equal(request, result.completedRequests[0])) +} + +func TestSplitterRequestSplitIntoTinyChunks(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + request := testChatRequest("", testChatRequestBody) + + results := writeInChunks(t, splitter, request, 1) + + var completedBeforeLast int + for _, result := range results[:len(results)-1] { + assert.False(t, result.enteredPassthrough) + completedBeforeLast += len(result.completedRequests) + } + assert.Equal(t, 0, completedBeforeLast, "request must not complete before the last byte of the body arrives") + assert.Len(t, results[len(results)-1].completedRequests, 1) + assert.True(t, bytes.Equal(request, completedRequestBytes(results)[0])) +} + +func TestSplitterRequestSplitAtBodyBoundary(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + request := testChatRequest("", testChatRequestBody) + headerEnd := bytes.Index(request, []byte("\r\n\r\n")) + 4 + + result1 := splitter.Write(request[:headerEnd]) + assert.Len(t, result1.completedRequests, 0) + + result2 := splitter.Write(request[headerEnd:]) + assert.Len(t, result2.completedRequests, 1) + assert.True(t, bytes.Equal(request, result2.completedRequests[0])) +} + +func TestSplitterPipelinedRequests(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + request1 := testChatRequest("", "first") + request2 := testChatRequest("", "second") + + result := splitter.Write(append(append([]byte{}, request1...), request2...)) + + assert.False(t, result.enteredPassthrough) + completed := result.completedRequests + assert.Len(t, completed, 2) + assert.True(t, bytes.Equal(request1, completed[0])) + assert.True(t, bytes.Equal(request2, completed[1])) +} + +func TestSplitterGetRequestCompletesAtHeaderEnd(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + request := []byte("GET /v1/models HTTP/1.1\r\nHost: localhost\r\n\r\n") + + result := splitter.Write(request) + + assert.False(t, result.enteredPassthrough) + assert.Len(t, result.completedRequests, 1) + assert.True(t, bytes.Equal(request, result.completedRequests[0])) +} + +func TestSplitterPostWithoutContentLengthEntersPassthrough(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + request := []byte("POST /v1/completions HTTP/1.1\r\nHost: localhost\r\n\r\nbody-until-close") + + result := splitter.Write(request) + + assert.True(t, result.enteredPassthrough) + assert.Empty(t, result.completedRequests) + assert.Equal(t, string(request), string(result.throughBytes), + "everything must be handed back for blind forwarding when body length is undeterminable") +} + +func TestSplitterChunkedRequest(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + chunkedRequest := []byte("POST /v1/chat/completions HTTP/1.1\r\n" + + "Host: localhost\r\n" + + "Transfer-Encoding: chunked\r\n\r\n" + + "5\r\nhello\r\n" + + "6\r\n world\r\n" + + "0\r\n\r\n") + + result := splitter.Write(chunkedRequest) + + assert.False(t, result.enteredPassthrough) + assert.Len(t, result.completedRequests, 1) + assert.True(t, bytes.Equal(chunkedRequest, result.completedRequests[0])) +} + +func TestSplitterChunkedRequestSplitAcrossWrites(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + chunkedRequest := []byte("POST /v1/chat/completions HTTP/1.1\r\n" + + "Host: localhost\r\n" + + "Transfer-Encoding: chunked\r\n\r\n" + + "5\r\nhello\r\n" + + "6\r\n world\r\n" + + "0\r\n\r\n") + + results := writeInChunks(t, splitter, chunkedRequest, 7) + + var completedCount int + for _, result := range results { + completedCount += len(result.completedRequests) + assert.False(t, result.enteredPassthrough) + } + assert.Equal(t, 1, completedCount) +} + +func TestSplitterChunkedRequestWithTrailers(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + chunkedRequest := []byte("POST /x HTTP/1.1\r\nHost: localhost\r\nTransfer-Encoding: chunked\r\n\r\n" + + "3\r\nabc\r\n" + + "0\r\n" + + "X-Trailer: v\r\n\r\n") + + result := splitter.Write(chunkedRequest) + + assert.False(t, result.enteredPassthrough) + assert.Len(t, result.completedRequests, 1) +} + +func TestSplitterGarbageEntersPassthrough(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + + result := splitter.Write([]byte("not an http request at all")) + + assert.True(t, result.enteredPassthrough) + assert.Empty(t, result.completedRequests) + assert.Equal(t, "not an http request at all", string(result.throughBytes)) +} + +func TestSplitterGarbageAfterValidRequest(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + request := testChatRequest("", "ok") + garbage := []byte("garbage that is not a request line\r\n\r\n") + + result := splitter.Write(append(append([]byte{}, request...), garbage...)) + + assert.Len(t, result.completedRequests, 1, "the framed request must still be emitted") + assert.True(t, result.enteredPassthrough) + assert.Equal(t, string(garbage), string(result.throughBytes)) +} + +func TestSplitterUpgradeRequestCompletesHeadersAndPassesThroughRest(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + headers := []byte("GET /ws HTTP/1.1\r\nHost: localhost\r\nConnection: Upgrade\r\nUpgrade: websocket\r\n\r\n") + + result := splitter.Write(append(append([]byte{}, headers...), []byte("websocket payload bytes")...)) + + assert.Len(t, result.completedRequests, 1, "header part must be emitted as a completed request") + assert.True(t, bytes.Equal(headers, result.completedRequests[0])) + assert.True(t, result.enteredPassthrough) + assert.Equal(t, "websocket payload bytes", string(result.throughBytes)) +} + +func TestSplitterExpect100ContinueEntersPassthrough(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + request := []byte("POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\nExpect: 100-continue\r\nContent-Length: 5\r\n\r\nhello") + + result := splitter.Write(request) + + // The client waits for a 100 Continue that only the backend can send, so the + // body cannot be counted here: fall back to blind forwarding. + assert.True(t, result.enteredPassthrough) + assert.Empty(t, result.completedRequests) +} + +func TestSplitterConnectMethodEntersPassthrough(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + + result := splitter.Write([]byte("CONNECT example.com:443 HTTP/1.1\r\nHost: example.com:443\r\n\r\n")) + + assert.True(t, result.enteredPassthrough) +} + +func TestSplitterPartialHeadersProduceNothing(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + + result := splitter.Write([]byte("POST /v1/chat/completions HTTP/1.1\r\nHost: localh")) + + assert.False(t, result.enteredPassthrough) + assert.Empty(t, result.completedRequests) + assert.Equal(t, "POST /v1/chat/completions HTTP/1.1\r\nHost: localh", string(splitter.FlushIncomplete())) +} + +func TestSplitterPassthroughForwardsSubsequentWrites(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + + first := splitter.Write([]byte("garbage")) + assert.True(t, first.enteredPassthrough) + + second := splitter.Write([]byte(" more bytes")) + assert.True(t, second.enteredPassthrough) + assert.Empty(t, second.completedRequests) + assert.Equal(t, " more bytes", string(second.throughBytes)) +} + +func TestSplitterContentLengthZero(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + request := []byte("POST /v1/models HTTP/1.1\r\nHost: localhost\r\nContent-Length: 0\r\n\r\n") + + result := splitter.Write(request) + + assert.False(t, result.enteredPassthrough) + assert.Len(t, result.completedRequests, 1) +} + +func TestSplitterHeaderCaseInsensitivity(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + request := []byte("POST /x HTTP/1.1\r\nHost: localhost\r\ncOnTeNt-LeNgTh: 5\r\n\r\nhello") + + result := splitter.Write(request) + + assert.Len(t, result.completedRequests, 1) +} + +func TestSplitterFlushIncompleteAfterCompleteRequest(t *testing.T) { + t.Parallel() + splitter := newHttpRequestSplitter() + request := testChatRequest("", "body") + partial := []byte("GET /v1/mod") + + result := splitter.Write(append(append([]byte{}, request...), partial...)) + assert.Len(t, result.completedRequests, 1) + assert.Equal(t, "GET /v1/mod", string(splitter.FlushIncomplete())) +} diff --git a/main.go b/main.go index 3d159fba..cceda717 100644 --- a/main.go +++ b/main.go @@ -173,9 +173,27 @@ func main() { ) } } + // Group services by listen port: a port with a single service behaves as a + // plain transparent proxy, while a port shared by several services is + // served by a context router that picks the smallest context size that + // fits each request. + servicesByListenPort := make(map[string][]ServiceConfig) + listenPortOrder := make([]string, 0) for _, service := range config.Services { - if service.ListenPort != "" { - go startProxy(service) + if service.ListenPort == "" { + continue + } + if _, seen := servicesByListenPort[service.ListenPort]; !seen { + listenPortOrder = append(listenPortOrder, service.ListenPort) + } + servicesByListenPort[service.ListenPort] = append(servicesByListenPort[service.ListenPort], service) + } + for _, listenPort := range listenPortOrder { + services := servicesByListenPort[listenPort] + if len(services) == 1 { + go startProxy(services[0]) + } else { + go startContextRoutedProxy(listenPort, buildContextRouter(services)) } } if config.OpenAiApi.ListenPort != "" { diff --git a/main_test.go b/main_test.go index 195fb626..11221879 100644 --- a/main_test.go +++ b/main_test.go @@ -11,6 +11,7 @@ import ( "net/http" "net/url" "os" + "os/exec" "strconv" "strings" "syscall" @@ -2535,10 +2536,10 @@ func TestProcessExitDuringShutdown(t *testing.T) { // forever calling stopService on it — visible in the logs as a tight, microsecond- // spaced repetition of: // -// Failed to send SIGTERM to -: no such process -// Stopping service to free resources for -// Sending SIGTERM to service process group: - -// ... +// Failed to send SIGTERM to -: no such process +// Stopping service to free resources for +// Sending SIGTERM to service process group: - +// ... // // and the requesting service never starts (its client connection hangs). // @@ -2558,6 +2559,13 @@ func TestProcessExitDuringShutdown(t *testing.T) { func TestEvictionOfAlreadyDeadProcessDoesNotLoop(t *testing.T) { t.Parallel() + // The proxy log is opened in append mode and survives between runs, while + // the loop-count assertion below must only consider this run's lines + // (same pattern as the resource-check-command SetupFunc). + if err := os.Remove("test-logs/test_eviction-already-dead-process.log"); err != nil && !os.IsNotExist(err) { + t.Fatalf("Failed to remove stale proxy log: %v", err) + } + // Hook file: monitorProcess blocks here after reaping the exited service-one // process, keeping service-one in runningServices with a dead process. hookDir := t.TempDir() @@ -2567,9 +2575,9 @@ func TestEvictionOfAlreadyDeadProcessDoesNotLoop(t *testing.T) { } const ( - managementApiAddress = "localhost:2129" - holderProxyAddress = "localhost:2130" - holderTargetPort = "12310" + managementApiAddress = "localhost:2129" + holderProxyAddress = "localhost:2130" + holderTargetPort = "12310" requesterProxyAddress = "localhost:2131" requesterTargetPort = "12311" testCaseName = "eviction-already-dead-process" @@ -2586,21 +2594,21 @@ func TestEvictionOfAlreadyDeadProcessDoesNotLoop(t *testing.T) { ManagementApi: ManagementApi{ListenPort: "2129"}, Services: []ServiceConfig{ { - Name: "holder", - ListenPort: "2130", - ProxyTargetHost: "localhost", - ProxyTargetPort: holderTargetPort, - Command: "./test-server/test-server", - Args: "-p " + holderTargetPort + " -exit-after-duration 800ms", + Name: "holder", + ListenPort: "2130", + ProxyTargetHost: "localhost", + ProxyTargetPort: holderTargetPort, + Command: "./test-server/test-server", + Args: "-p " + holderTargetPort + " -exit-after-duration 800ms", ResourceRequirements: map[string]int{"CPU": 1}, }, { - Name: "requester", - ListenPort: "2131", - ProxyTargetHost: "localhost", - ProxyTargetPort: requesterTargetPort, - Command: "./test-server/test-server", - Args: "-p " + requesterTargetPort, + Name: "requester", + ListenPort: "2131", + ProxyTargetHost: "localhost", + ProxyTargetPort: requesterTargetPort, + Command: "./test-server/test-server", + Args: "-p " + requesterTargetPort, ResourceRequirements: map[string]int{"CPU": 1}, }, }, @@ -2696,7 +2704,7 @@ func TestEvictionOfAlreadyDeadProcessDoesNotLoop(t *testing.T) { } stopCount := strings.Count(logText, "Stopping service to free resources") if stopCount > 5 { - t.Errorf("Expected the eviction to run a handful of times at most, but " + + t.Errorf("Expected the eviction to run a handful of times at most, but "+ "\"Stopping service to free resources\" appeared %d times in the log — "+ "this is the endless loop from issue #119", stopCount) } @@ -2704,6 +2712,163 @@ func TestEvictionOfAlreadyDeadProcessDoesNotLoop(t *testing.T) { } } +// TestServiceWaitingForResourcesNotStartedOnInterrupt is a regression test for the +// shutdown contract fixed in b67f4b0: a service that is starved for resources +// (parked in reserveResources waiting for a unit held by another service) must +// NOT be started once an interrupt signal is received, even though interrupting +// stops the holder and would otherwise free the resource. The proxy must simply +// close the still-unstarted service's client connection and exit. +// +// Without the interrupt exit points, the waiter could be re-evaluated while the +// proxy is tearing down (e.g. once the holder releases its resource) and proceed +// to spawn its backend mid-shutdown. This test asserts the observable contract: +// the waiter's backend process is never spawned (no "Starting"/"Service started" +// log line for it), its client connection is closed, and the proxy shuts down +// promptly. +func TestServiceWaitingForResourcesNotStartedOnInterrupt(t *testing.T) { + t.Parallel() + + const managementApiAddress = "localhost:2210" + const holderProxyAddress = "localhost:2211" + const waiterProxyAddress = "localhost:2212" + const testName = "resource-waiter-not-started-on-interrupt" + const holderServiceName = testName + "_holder" + const waiterServiceName = testName + "_waiter" + + // The waiter's max-wait is far longer than the test, so the only thing that + // can unblock it is the resource becoming free — never a self-imposed timeout. + // The holder's idle timeout is likewise far longer than the test, so it keeps + // holding TestResource (with its proxied connection open, making it + // non-evictable) until the proxy is interrupted. + maxWaitSeconds := uint(60) + holderIdleTimeoutSeconds := uint(300) + cfg := Config{ + MaxTimeToWaitForServiceToCloseConnectionBeforeGivingUpSeconds: &maxWaitSeconds, + ResourcesAvailable: map[string]ResourceAvailable{"TestResource": {Amount: 1}}, + ManagementApi: ManagementApi{ListenPort: "2210"}, + Services: []ServiceConfig{ + { + Name: "holder", + ListenPort: "2211", + ProxyTargetHost: "localhost", + ProxyTargetPort: "12210", + Command: "./test-server/test-server", + Args: "-p 12210 -sleep-after-writing-pid-duration 60s", + ShutDownAfterInactivitySeconds: holderIdleTimeoutSeconds, + ResourceRequirements: map[string]int{"TestResource": 1}, + }, + { + Name: "waiter", + ListenPort: "2212", + ProxyTargetHost: "localhost", + ProxyTargetPort: "12211", + Command: "./test-server/test-server", + Args: "-p 12211", + ResourceRequirements: map[string]int{"TestResource": 1}, + }, + }, + } + StandardizeConfigNamesAndPaths(&cfg, testName) + configFilePath := createTempConfig(t, cfg) + + // Start from a clean proxy log so the post-shutdown read only reflects this + // run (the proxy opens it with O_APPEND, so removing it here yields a fresh + // file that captures exactly this run's "Starting"/"Service started" lines). + proxyLogPath := fmt.Sprintf("test-logs/test_%s.log", testName) + _ = os.Remove(proxyLogPath) + + waitChannel := make(chan error, 1) + cmd, err := startLargeModelProxy(testName, configFilePath, "", waitChannel) + if err != nil { + t.Fatalf("could not start application: %v", err) + } + // Defensive: if the regression fires and the waiter's backend is spawned + // mid-shutdown, it is not in the proxy's stop-loop snapshot and would be + // orphaned on os.Exit. Kill any such leftover so it cannot outlive the test. + defer func() { + _ = exec.Command("pkill", "-f", "test-server/test-server -p 12211").Run() + }() + defer func() { + // The proxy exits via the explicit SIGINT below; this is a safety net in + // case the test bails out before reaching it. + if cmd.ProcessState == nil { + _ = cmd.Process.Signal(syscall.SIGINT) + select { + case <-waitChannel: + case <-time.After(15 * time.Second): + _ = cmd.Process.Kill() + } + } + for _, address := range []string{holderProxyAddress, waiterProxyAddress, managementApiAddress, "localhost:12210", "localhost:12211"} { + if err := checkPortClosed(address); err != nil { + t.Errorf("port %s is still open after application exit: %v", address, err) + } + } + }() + + // 1. Start the holder and keep its proxied connection open so it holds + // TestResource and is non-evictable. + holderConn, err := net.DialTimeout("tcp", holderProxyAddress, 3*time.Second) + if err != nil { + t.Fatalf("failed to connect to holder: %v", err) + } + defer func() { _ = holderConn.Close() }() + readPidFromOpenConnection(t, holderConn) + statusResponse := getStatusFromManagementAPI(t, managementApiAddress) + verifyServiceStatus(t, statusResponse, holderServiceName, ServiceStateRunning, 0, 1, map[string]int{"TestResource": 1}) + + // 2. A connection to the waiter must block: TestResource is held by the + // non-evictable holder, so the waiter parks in waiting_for_resources. + waiterConn, err := net.DialTimeout("tcp", waiterProxyAddress, 3*time.Second) + if err != nil { + t.Fatalf("failed to connect to waiter: %v", err) + } + defer func() { _ = waiterConn.Close() }() + waitForServiceState(t, managementApiAddress, waiterServiceName, ServiceStateWaitingForResources, 3*time.Second) + + // 3. Interrupt the proxy while the waiter is still starved. + shutdownStart := time.Now() + if err := cmd.Process.Signal(syscall.SIGINT); err != nil { + t.Fatalf("Failed to send SIGINT to proxy: %v", err) + } + + // 4. The proxy must shut down promptly, and the waiter's still-unstarted + // client connection must be closed (the proxy must not keep it open or + // start its backend). Both happen when the proxy exits; bound the waits so + // a hang fails fast. + select { + case <-waitChannel: + case <-time.After(10 * time.Second): + t.Fatalf("proxy did not shut down within 10s of SIGINT while a service was waiting for resources") + } + shutdownDuration := time.Since(shutdownStart) + t.Logf("Shutdown completed in %v", shutdownDuration) + if shutdownDuration > 8*time.Second { + t.Errorf("shutdown took %v, expected prompt (< 8s): a resource-starved service must not be started (and thus delay) shutdown", shutdownDuration) + } + // The connection is closed as part of the proxy exiting; by now it must be + // gone. Re-check (rather than relying on the exit alone) so a regression that + // leaves it open is caught explicitly. + assertRemoteClosedWithin(t, waiterConn, 3*time.Second) + + // 5. The core contract: the waiter's backend must NEVER have been started. + // runServiceCommand logs "[] Starting \"...\"" the instant resources + // are reserved and the process is spawned, and "Service started with pid" + // once fully up — neither line may appear for the waiter. + logContents, readErr := os.ReadFile(proxyLogPath) + if readErr != nil { + t.Fatalf("failed to read proxy log %s: %v", proxyLogPath, readErr) + } + logString := string(logContents) + startMarker := fmt.Sprintf("[%s] Starting \"", waiterServiceName) + if strings.Contains(logString, startMarker) { + t.Errorf("the resource-starved waiter's backend was spawned during shutdown (log contains %q); the proxy must not start a service that had not yet started when interrupted", startMarker) + } + if strings.Contains(logString, fmt.Sprintf("[%s] Service started with pid", waiterServiceName)) { + t.Errorf("the resource-starved waiter's backend fully started during shutdown; the proxy must not start a service that had not yet started when interrupted") + } +} + // TestWaitingConnectionsDecrementedOnServiceStartFailure is a regression test for // a leak where a connection that triggered a service start was counted as // "waiting" but the counter was never decremented when the service failed to diff --git a/test-configs/client-close-full.jsonc b/test-configs/client-close-full.jsonc new file mode 100644 index 00000000..8d98a28f --- /dev/null +++ b/test-configs/client-close-full.jsonc @@ -0,0 +1,34 @@ +{ + "ManagementApi": { + "ListenPort": "2029" //Does not exist in the test + }, + "ResourcesAvailable": { + "VRAM": { + "Amount": 1 + } + }, + "Services": [ + { + "Name": "test-server-close-full-service-one", + "ListenPort": "2030", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12030", + "Command": "./test-server/test-server", + "Args": "-p 12030 -sleep-after-writing-pid-duration 10s", + "ResourceRequirements": { + "VRAM": 1 + } + }, + { + "Name": "test-server-close-full-service-two", + "ListenPort": "2031", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12031", + "Command": "./test-server/test-server", + "Args": "-p 12031", + "ResourceRequirements": { + "VRAM": 1 + } + } + ] +} diff --git a/test-configs/healthcheck-stuck-timeout.jsonc b/test-configs/healthcheck-stuck-timeout.jsonc new file mode 100644 index 00000000..96b2e251 --- /dev/null +++ b/test-configs/healthcheck-stuck-timeout.jsonc @@ -0,0 +1,26 @@ +{ + "ResourcesAvailable": { + "CPU": { + "Amount": 1 + } + }, + "ManagementApi": { + "ListenPort": "2065" + }, + "Services": [ + { + "Name": "test-server-stuck-timeout", + "ListenPort": "2064", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12064", + "Command": "./test-server/test-server", + "Args": "-p 12064 -startup-duration 24h", + "HealthcheckCommand": "false", + "HealthcheckIntervalMilliseconds": 200, + "StartupTimeoutMilliseconds": 2000, + "ResourceRequirements": { + "CPU": 1 + } + } + ] +} diff --git a/test-configs/healthcheck-stuck.jsonc b/test-configs/healthcheck-stuck.jsonc new file mode 100644 index 00000000..947c1fe6 --- /dev/null +++ b/test-configs/healthcheck-stuck.jsonc @@ -0,0 +1,17 @@ +{ + "ManagementApi": { + "ListenPort": "2029", //Does not exist in the test + }, + "Services": [ + { + "Name": "test-server-stuck-service", + "ListenPort": "2005", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12005", + "Command": "./test-server/test-server", + "Args": "-p 12005 -healthcheck-port 2015", + "HealthcheckCommand": "false", + "HealthcheckIntervalMilliseconds": 200 + } + ] +} diff --git a/test-configs/idle-timeout.jsonc b/test-configs/idle-timeout.jsonc new file mode 100644 index 00000000..242be221 --- /dev/null +++ b/test-configs/idle-timeout.jsonc @@ -0,0 +1,13 @@ +{ + "ShutDownAfterInactivitySeconds": 3, + "Services": [ + { + "Name": "TestService", + "ListenPort": "2007", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12007", + "Command": "./test-server/test-server", + "Args": "-p 12007" + } + ] +} diff --git a/test-configs/invalid-template.json b/test-configs/invalid-template.json new file mode 100644 index 00000000..53c9cb5f --- /dev/null +++ b/test-configs/invalid-template.json @@ -0,0 +1,33 @@ +{ + "ResourcesAvailable": { + "CPU": 1 + }, + "DefaultServiceUrl": "http://localhost/{{.POR}}", + "ManagementApi": { + "ListenPort": "2035" + }, + "Services": [ + { + "Name": "self-dying-process", + "ListenPort": "2036", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12036", + "Command": "./test-server/test-server", + "Args": "-p 12036 -exit-after-duration 1s --sleep-after-writing-pid-duration 3s", + "ResourceRequirements": { + "CPU": 1 + } + }, + { + "Name": "not-dying-process", + "ListenPort": "2037", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12037", + "Command": "./test-server/test-server", + "Args": "-p 12037 --sleep-after-writing-pid-duration 3s", + "ResourceRequirements": { + "CPU": 1 + } + } + ] +} \ No newline at end of file diff --git a/test-configs/log-output.jsonc b/test-configs/log-output.jsonc new file mode 100644 index 00000000..7ed22712 --- /dev/null +++ b/test-configs/log-output.jsonc @@ -0,0 +1,35 @@ +{ + "Services": [ + { + "Name": "service1", + "ListenPort": "2055", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12055", + "Command": "./test-server/test-server", + "Args": "-p 12055 --plain-output" + }, + { + "Name": "Service TWO2️⃣ Два", + "ListenPort": "2056", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12056", + "Command": "./test-server/test-server", + "Args": "-p 12056 --plain-output --log-to-stdout" + }, + { + "Name": "{Service 3}", + "ListenPort": "2059", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12059", //nothing is actually listening there + "Command": "./test-server/output-test.sh" + }, + { + "Name": "[Service 4]", + "ListenPort": "2060", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12060", //nothing is actually listening there + "Command": "./test-server/output-test.sh", + "Args": "-stderr" + } + ] +} diff --git a/test-configs/multiple-connections-white-waiting-for-resources.jsonc b/test-configs/multiple-connections-white-waiting-for-resources.jsonc new file mode 100644 index 00000000..e2f1c1e5 --- /dev/null +++ b/test-configs/multiple-connections-white-waiting-for-resources.jsonc @@ -0,0 +1,35 @@ +{ + "ResourcesAvailable": { + "TestResource": { + "Amount": 1 + } + }, + "LogLevel": "Debug", + "ManagementApi": { + "ListenPort": "2091" + }, + "Services": [ + { + "Name": "ServiceOne", + "ListenPort": "2087", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12087", + "Command": "./test-server/test-server", + "Args": "-p 12087 -healthcheck-port 2089 -sleep-before-listening 3s", + "ResourceRequirements": { + "TestResource": 1 + } + }, + { + "Name": "ServiceTwo", + "ListenPort": "2088", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12088", + "Command": "./test-server/test-server", + "Args": "-p 12088 -healthcheck-port 2090 -sleep-before-listening 2s -request-processing-duration 3s", + "ResourceRequirements": { + "TestResource": 1 + } + } + ] +} diff --git a/test-configs/no-resource-requirements.json b/test-configs/no-resource-requirements.json new file mode 100644 index 00000000..5ba26627 --- /dev/null +++ b/test-configs/no-resource-requirements.json @@ -0,0 +1,26 @@ +{ + "ResourcesAvailable": { + "VRAM": 20 + }, + "Services": [ + { + "Name": "service1", + "ListenPort": "2032", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12032", + "Command": "./test-server/test-server", + "Args": "-p 12032", + "ResourceRequirements": { + "VRAM": 20 + } + }, + { + "Name": "service2", + "ListenPort": "2033", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12033", + "Command": "./test-server/test-server", + "Args": "-p 12033" + } + ] +} \ No newline at end of file diff --git a/test-configs/openai-api-models-by-id.jsonc b/test-configs/openai-api-models-by-id.jsonc new file mode 100644 index 00000000..19a3af76 --- /dev/null +++ b/test-configs/openai-api-models-by-id.jsonc @@ -0,0 +1,42 @@ +{ + "OpenAiApi": { + "ListenPort": "2071" + }, + "ShutDownAfterInactivitySeconds": 3, + "Services": [ + { + "Name": "openai-api-1", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12072", + "Command": "./test-server/test-server", + "Args": "-p 12072", + "OpenAiApi": true + }, + { + "Name": "openai-api-2", + "ListenPort": "2073", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "120723", + "Command": "./test-server/test-server", + "Args": "-p 12073", + "OpenAiApi": true, + "OpenAiApiModels": [ "fizz", "buzz", "$-_.+!*'(),проверка"]}, + { + "Name": "non-llm-1", + "ListenPort": "2074", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12074", + "Command": "./test-server/test-server", + "Args": "-p 12074", + "OpenAiApi": false + }, + { + "Name": "$-_.+!*'(),проверка-2/", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12075", + "Command": "./test-server/test-server", + "Args": "-p 12075", + "OpenAiApi": true + } + ] +} diff --git a/test-configs/resource-check-command.jsonc b/test-configs/resource-check-command.jsonc new file mode 100644 index 00000000..55042786 --- /dev/null +++ b/test-configs/resource-check-command.jsonc @@ -0,0 +1,36 @@ +{ + "ResourcesAvailable": { + "TestResource": { + "CheckCommand": "read -r original_integer < test-logs/resource-check-command.counter.txt; incremented_integer=$((original_integer + 1)); printf '%d\n' \"$incremented_integer\" | tee test-logs/resource-check-command.counter.txt", + "CheckIntervalMilliseconds": 1000 + } + }, + "LogLevel": "Debug", + "ManagementApi": { + "ListenPort": "2076" + }, + "Services": [ + { + "Name": "service-one", + "ListenPort": "2077", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12077", + "Command": "./test-server/test-server", + "Args": "-p 12077 -healthcheck-port 2080 -sleep-before-listening 10s", + "ResourceRequirements": { + "TestResource": 4 + } + }, + { + "Name": "service-two", + "ListenPort": "2079", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12079", + "Command": "./test-server/test-server", + "Args": "-p 12079 -healthcheck-port 2081", + "ResourceRequirements": { + "TestResource": 5 + } + } + ] +} diff --git a/test-configs/self-dying.json b/test-configs/self-dying.json new file mode 100644 index 00000000..a6164aea --- /dev/null +++ b/test-configs/self-dying.json @@ -0,0 +1,32 @@ +{ + "ResourcesAvailable": { + "CPU": 1 + }, + "ManagementApi": { + "ListenPort": "2035" + }, + "Services": [ + { + "Name": "self-dying-process", + "ListenPort": "2036", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12036", + "Command": "./test-server/test-server", + "Args": "-p 12036 -exit-after-duration 1s --sleep-after-writing-pid-duration 3s", + "ResourceRequirements": { + "CPU": 1 + } + }, + { + "Name": "not-dying-process", + "ListenPort": "2037", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12037", + "Command": "./test-server/test-server", + "Args": "-p 12037 --sleep-after-writing-pid-duration 3s", + "ResourceRequirements": { + "CPU": 1 + } + } + ] +} diff --git a/test-configs/should-not-use-an-outdated-resource.jsonc b/test-configs/should-not-use-an-outdated-resource.jsonc new file mode 100644 index 00000000..c250b592 --- /dev/null +++ b/test-configs/should-not-use-an-outdated-resource.jsonc @@ -0,0 +1,37 @@ +{ + "ResourcesAvailable": { + "TestResource": { + "CheckCommand": "cat test-logs/should-not-use-an-outdated-resource-check-result.resource-amount.txt", + "CheckIntervalMilliseconds": 60000, + "Amount": 2 + } + }, + "LogLevel": "Debug", + "ManagementApi": { + "ListenPort": "2086" + }, + "Services": [ + { + "Name": "ServiceOne", + "ListenPort": "2082", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12077", + "Command": "sh", + "Args": "-c \"echo '11' > test-logs/should-not-use-an-outdated-resource-check-result.resource-amount.txt &&sleep 4 && echo '0' > test-logs/should-not-use-an-outdated-resource-check-result.resource-amount.txt &&./test-server/test-server -p 12077 -healthcheck-port 2084 -exit-after-duration 2s && echo '12' > test-logs/should-not-use-an-outdated-resource-check-result.resource-amount.txt\"", + "ResourceRequirements": { + "TestResource": 10 + } + }, + { + "Name": "ServiceTwo", + "ListenPort": "2083", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12079", + "Command": "./test-server/test-server", + "Args": "-p 12079 -healthcheck-port 2085", + "ResourceRequirements": { + "TestResource": 10 + } + } + ] +} diff --git a/test-server/main.go b/test-server/main.go index 9b0dc5e2..aa69a020 100644 --- a/test-server/main.go +++ b/test-server/main.go @@ -29,6 +29,7 @@ func main() { durationToSleepBeforeListeningForHealthCheck := flag.Duration("sleep-before-listening-for-healthcheck", 0, "How much time to sleep before listening for healthcheck starts, such as \"300ms\", \"-1.5h\" or \"2h45m\". Valid time units are \"ns\", \"us\" (or \"µs\"), \"ms\", \"s\", \"m\", \"h\". ") exitAfterDuration := flag.Duration("exit-after-duration", time.Duration(1<<63-1), "How much time to exit after the program start, such as \"300ms\", \"1.5h\" or \"2h45m\". Valid time units are \"ns\", \"us\" (or \"µs\"), \"ms\", \"s\", \"m\", \"h\". ") OpenAiApiPort := flag.String("openai-api-port", "", "OpenAI API port to listen on. If not specified, OpenAI API is disabled") + openAiApiKeepAlive := flag.Bool("openai-api-keep-alive", false, "Enable HTTP keep-alive on the OpenAI API server") procPort := flag.String("procinfo-port", "", "Port to expose process information") plainOutput := flag.Bool("plain-output", false, "Do not add timestamps to log output") logToStdout := flag.Bool("log-to-stdout", false, "Send logs to stdout instead of stderr") @@ -55,7 +56,7 @@ func main() { go healthCheckListen(healthCheckApiPort, durationToSleepBeforeListeningForHealthCheck) } if *OpenAiApiPort != "" { - go OpenAiApiListen(OpenAiApiPort) + go OpenAiApiListen(OpenAiApiPort, openAiApiKeepAlive) } if *procPort != "" { go procListen(*procPort) @@ -322,7 +323,7 @@ type ChatCompletionChoice struct { FinishReason *string `json:"finish_reason,omitempty"` } -func OpenAiApiListen(port *string) { +func OpenAiApiListen(port *string, keepAlive *bool) { mux := http.NewServeMux() mux.HandleFunc("/v1/completions", handleCompletions) mux.HandleFunc("/v1/chat/completions", handleChatCompletions) @@ -331,7 +332,7 @@ func OpenAiApiListen(port *string) { Addr: ":" + *port, Handler: mux, } - server.SetKeepAlivesEnabled(false) + server.SetKeepAlivesEnabled(*keepAlive) log.Printf("OpenAI API server listening on :%s", *port) if err := server.ListenAndServe(); err != nil { log.Fatalf("Could not start OpenAI API server: %s\n", err.Error()) diff --git a/tokenizer.go b/tokenizer.go new file mode 100644 index 00000000..80ec4555 --- /dev/null +++ b/tokenizer.go @@ -0,0 +1,209 @@ +package main + +import ( + "math" + "sort" + "sync" + "unicode" + "unicode/utf8" +) + +// TokenCounter estimates how many tokens a given text occupies for a specific +// model family's tokenizer. +// +// Counters in this package are heuristics: they do not load real BPE/SentencePiece +// vocabularies, but approximate a model family's tokenization using per-model +// character-class ratios. For context-tier routing (choosing between e.g. a 4k +// and a 32k context instance of the same model) this approximation is sufficient +// as long as operators leave headroom when configuring ContextSize values. The +// run-based algorithm below tends to slightly overestimate, which is the safe +// direction: it switches to a larger context service a bit early rather than +// sending an oversized context to a service that cannot fit it. +type TokenCounter func(text string) int + +var ( + tokenCounterMutex sync.RWMutex + tokenCounters = map[string]TokenCounter{} +) + +// RegisterTokenCounter makes a token counter available by name for use in the +// Tokenizer configuration field. Adding support for a new model family is a +// single call to this function (usually from an init()). +func RegisterTokenCounter(name string, counter TokenCounter) { + tokenCounterMutex.Lock() + defer tokenCounterMutex.Unlock() + tokenCounters[name] = counter +} + +// GetTokenCounter looks up a registered token counter by name. +func GetTokenCounter(name string) (TokenCounter, bool) { + tokenCounterMutex.RLock() + defer tokenCounterMutex.RUnlock() + counter, found := tokenCounters[name] + return counter, found +} + +// RegisteredTokenCounterNames returns the sorted list of all registered +// token counter names. +func RegisteredTokenCounterNames() []string { + tokenCounterMutex.RLock() + defer tokenCounterMutex.RUnlock() + names := make([]string, 0, len(tokenCounters)) + for name := range tokenCounters { + names = append(names, name) + } + sort.Strings(names) + return names +} + +// heuristicTokenizerConfig holds the per-model-family approximation parameters. +// They describe how many units of each character class (runes, or bytes for +// non-ASCII) a tokenizer of that family produces per token on average. +type heuristicTokenizerConfig struct { + lettersPerToken float64 + digitsPerToken float64 + whitespacePerToken float64 + punctuationPerToken float64 + nonAsciiBytesPerToken float64 +} + +func init() { + // Qwen3-family BPE tokenizer (e.g. Qwen3-8B): efficient on both English and CJK + qwen := newHeuristicTokenCounter(heuristicTokenizerConfig{ + lettersPerToken: 5.0, + digitsPerToken: 3.0, + whitespacePerToken: 4.0, + punctuationPerToken: 3.0, + nonAsciiBytesPerToken: 3.0, + }) + RegisterTokenCounter("qwen3.8", qwen) + RegisterTokenCounter("qwen3", qwen) + + // Gemma-family SentencePiece tokenizer: byte-fallback for CJK makes + // non-ASCII text far more expensive than on Qwen + gemma := newHeuristicTokenCounter(heuristicTokenizerConfig{ + lettersPerToken: 5.5, + digitsPerToken: 3.0, + whitespacePerToken: 4.0, + punctuationPerToken: 3.0, + nonAsciiBytesPerToken: 1.5, + }) + RegisterTokenCounter("gemma4", gemma) + RegisterTokenCounter("gemma3", gemma) +} + +// newHeuristicTokenCounter builds a TokenCounter from per-family parameters. +// +// The text is split into maximal runs of one character class (letters, digits, +// whitespace, ASCII punctuation, everything else measured in UTF-8 bytes). +// Tokenizers never merge across script boundaries, so each run costs at least +// one token, and long runs are split proportionally to the family's ratio. A +// single inter-word space is merged into the following word run, mirroring how +// BPE/SentencePiece encode " word" as a single token. +func newHeuristicTokenCounter(config heuristicTokenizerConfig) TokenCounter { + return func(text string) int { + var tokens float64 + var pendingMergedSpace int + + for _, characterClass := range classifyRunes(text) { + mergedSpace := 0 + if characterClass.class == runeClassLetter || characterClass.class == runeClassDigit { + // A single space immediately before a word is counted together + // with the word (the " word" merge), as tokenizers do. + mergedSpace = pendingMergedSpace + } + pendingMergedSpace = 0 + + switch characterClass.class { + case runeClassLetter: + tokens += math.Ceil(float64(characterClass.runes+mergedSpace) / config.lettersPerToken) + case runeClassDigit: + tokens += math.Ceil(float64(characterClass.runes+mergedSpace) / config.digitsPerToken) + case runeClassWhitespace: + if characterClass.singleSpace { + // not consumed by a merge (e.g. trailing space): count it alone + pendingMergedSpace = characterClass.runes + continue + } + tokens += math.Ceil(float64(characterClass.runes) / config.whitespacePerToken) + case runeClassPunctuation: + tokens += math.Ceil(float64(characterClass.runes) / config.punctuationPerToken) + case runeClassOther: + tokens += math.Ceil(float64(characterClass.bytes) / config.nonAsciiBytesPerToken) + } + } + // A trailing single space that never merged still occupies roughly a + // fraction of a token; charge it so that trailing whitespace is not free. + if pendingMergedSpace > 0 { + tokens += math.Ceil(float64(pendingMergedSpace) / config.whitespacePerToken) + } + return int(tokens) + } +} + +type runeClass int + +const ( + runeClassLetter runeClass = iota + runeClassDigit + runeClassWhitespace + runeClassPunctuation + runeClassOther +) + +type runeRun struct { + class runeClass + runes int + bytes int // only used for runeClassOther + // singleSpace is true when the run is exactly one plain space character + singleSpace bool +} + +func classifyRunes(text string) []runeRun { + var runs []runeRun + appendRun := func(class runeClass, runes int, bytes int, singleSpace bool) { + runs = append(runs, runeRun{class: class, runes: runes, bytes: bytes, singleSpace: singleSpace}) + } + + var currentClass runeClass + var currentRunes int + var currentBytes int + var lastRune rune + flush := func() { + if currentRunes == 0 { + return + } + // Only a plain U+0020 space participates in the " word" merge; other + // whitespace (tabs, newlines) always forms its own run. + singleSpace := currentClass == runeClassWhitespace && currentRunes == 1 && currentBytes == 1 && lastRune == ' ' + appendRun(currentClass, currentRunes, currentBytes, singleSpace) + currentRunes = 0 + currentBytes = 0 + } + + for _, character := range text { + var class runeClass + switch { + case unicode.IsDigit(character): + class = runeClassDigit + case unicode.IsSpace(character): + class = runeClassWhitespace + case unicode.IsLetter(character) && character < utf8.RuneSelf: + class = runeClassLetter + case (unicode.IsPunct(character) || unicode.IsSymbol(character)) && character < utf8.RuneSelf: + class = runeClassPunctuation + default: + class = runeClassOther + } + + if class != currentClass { + flush() + currentClass = class + } + currentRunes++ + currentBytes += utf8.RuneLen(character) + lastRune = character + } + flush() + return runs +} diff --git a/tokenizer_test.go b/tokenizer_test.go new file mode 100644 index 00000000..40245dce --- /dev/null +++ b/tokenizer_test.go @@ -0,0 +1,126 @@ +package main + +import ( + "math" + "testing" + + "github.com/stretchr/testify/assert" +) + +// Exact token counts below are hand-computed from the documented heuristic +// tokenizer parameters (see tokenizer.go). They are intentionally exact so any +// accidental change to the counting algorithm that would silently alter routing +// decisions breaks a test. +func TestQwen38TokenCounting(t *testing.T) { + t.Parallel() + counter, found := GetTokenCounter("qwen3.8") + assert.True(t, found, "qwen3.8 tokenizer should be registered") + + assert.Equal(t, 0, counter(""), "empty text") + // "hello" -> ceil(5/5.0)=1, " world" (space merged) -> ceil(6/5.0)=2 + assert.Equal(t, 3, counter("hello world")) + // 4 CJK runes = 12 bytes -> ceil(12/3.0)=4 + assert.Equal(t, 4, counter("你好世界")) + // "abc"=1, double space (not merged)=1, "123"=ceil(3/3)=1 + assert.Equal(t, 3, counter("abc 123")) + // "Count"=1, ":"=1, " 42" (merged space + digits)=ceil(3/3)=1, "!"=1 + assert.Equal(t, 4, counter("Count: 42!")) + // single long word: ceil(20/5.0) + assert.Equal(t, 4, counter("internationalization")) + // leading space merges into the word: ceil(6/5.0) + assert.Equal(t, 2, counter(" hello")) + // newline is whitespace that cannot merge: "a"=1 + "\n"=1 + "b"=1 + assert.Equal(t, 3, counter("a\nb")) +} + +func TestGemma4TokenCounting(t *testing.T) { + t.Parallel() + counter, found := GetTokenCounter("gemma4") + assert.True(t, found, "gemma4 tokenizer should be registered") + + assert.Equal(t, 0, counter("")) + // "hello"=ceil(5/5.5)=1, " world"=ceil(6/5.5)=2 + assert.Equal(t, 3, counter("hello world")) + // Gemma falls back to byte-level tokens for CJK: ceil(12/1.5)=8 + assert.Equal(t, 8, counter("你好世界")) + // letters: ceil(20/5.5)=4 + assert.Equal(t, 4, counter("internationalization")) +} + +func TestTokenCounterAliases(t *testing.T) { + t.Parallel() + qwen38, found38 := GetTokenCounter("qwen3.8") + qwen3, found3 := GetTokenCounter("qwen3") + assert.True(t, found38) + assert.True(t, found3) + // Aliases must resolve to the same counting behavior + assert.Equal(t, qwen38("hello world"), qwen3("hello world")) + + gemma4, foundGemma4 := GetTokenCounter("gemma4") + gemma3, foundGemma3 := GetTokenCounter("gemma3") + assert.True(t, foundGemma4) + assert.True(t, foundGemma3) + assert.Equal(t, gemma4("hello world"), gemma3("hello world")) +} + +func TestGetTokenCounterUnknownName(t *testing.T) { + t.Parallel() + _, found := GetTokenCounter("does-not-exist") + assert.False(t, found) +} + +func TestRegisteredTokenCounterNames(t *testing.T) { + t.Parallel() + names := RegisteredTokenCounterNames() + assert.Contains(t, names, "qwen3.8") + assert.Contains(t, names, "gemma4") +} + +// Adding support for a new model must be a single registration call. +func TestRegisterCustomTokenCounter(t *testing.T) { + t.Parallel() + RegisterTokenCounter("test-model", func(text string) int { + return len(text) + }) + counter, found := GetTokenCounter("test-model") + assert.True(t, found, "newly registered counter should be retrievable") + assert.Equal(t, 5, counter("hello")) +} + +func TestTokenCounterMonotonicallyIncreases(t *testing.T) { + t.Parallel() + counter, _ := GetTokenCounter("qwen3.8") + previousCount := 0 + text := "" + for i := 0; i < 30; i++ { + text += "word " + count := counter(text) + assert.GreaterOrEqual(t, count, previousCount, "adding text must never decrease token count (text=%q)", text) + previousCount = count + } +} + +func TestHeuristicTokenCounterNeverReturnsNegativeOrNaN(t *testing.T) { + t.Parallel() + counter := newHeuristicTokenCounter(heuristicTokenizerConfig{ + lettersPerToken: 5.0, + digitsPerToken: 3.0, + whitespacePerToken: 4.0, + punctuationPerToken: 3.0, + nonAsciiBytesPerToken: 3.0, + }) + for _, text := range []string{"", " ", "\x00\x01", "😀😀", "a b\tc\rd"} { + count := counter(text) + assert.False(t, math.IsNaN(float64(count))) + assert.GreaterOrEqual(t, count, 0) + } +} + +// The same CJK text must cost more tokens on the byte-fallback tokenizer. +func TestGemmaCountsMoreTokensForCJKThanQwen(t *testing.T) { + t.Parallel() + qwen, _ := GetTokenCounter("qwen3.8") + gemma, _ := GetTokenCounter("gemma4") + cjk := "这是一段中文文本" + assert.Greater(t, gemma(cjk), qwen(cjk)) +} From 592132d4f0ff9667176651291370b10486f3d0bb Mon Sep 17 00:00:00 2001 From: Konstantin Pereiaslov Date: Thu, 3 Sep 2026 02:45:21 +0000 Subject: [PATCH 2/2] feat: automatic listen port selection with SQLite persistence A service can now set ListenPort to "auto" and get a free port picked automatically. The last port selected for each service is persisted in a small SQLite database (AutoListenPortDatabasePath, default auto-listen-ports.db, keyed by service name), and on startup the proxy first tries to reuse the service's last used port, so clients keep working across restarts. When that port is no longer free, a new one is requested from the OS, avoiding every port already used in the same configuration (including OpenAI API and management ports). Chosen ports are bound immediately at selection and the live listener is handed to the proxy, so a selected port can never be lost to another process between allocation and use. Resolved ports are written back into the config, so the startup logs, the management API listen_port field and ServiceUrl {{.PORT}} templates all report the real port. Each "auto" service gets a port of its own; context-routing groups still require explicit numeric ports. Persistence uses the pure-Go modernc.org/sqlite driver, keeping cross-compilation cgo-free. --- AGENTS.md | 3 + README.md | 36 ++++ auto_listen_port_e2e_test.go | 153 ++++++++++++++++ config.go | 17 ++ config_test.go | 134 ++++++++++++++ connection.go | 10 +- go.mod | 14 +- go.sum | 43 +++++ main.go | 32 +++- port_allocation.go | 176 ++++++++++++++++++ port_allocation_test.go | 343 +++++++++++++++++++++++++++++++++++ 11 files changed, 958 insertions(+), 3 deletions(-) create mode 100644 auto_listen_port_e2e_test.go create mode 100644 port_allocation.go create mode 100644 port_allocation_test.go diff --git a/AGENTS.md b/AGENTS.md index 899b0af1..3016e3e4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -69,6 +69,7 @@ Client → large-model-proxy → [Service Process] | `service_process.go` | Service process spawning/stopping, output logging, and process-exit monitoring | | `resources.go` | Resource reservation, LRU eviction, and resource release logic | | `openai_api.go` | Unified OpenAI API server and request routing to backends by model name | +| `port_allocation.go` | Automatic listen port selection with SQLite persistence of last used ports | | `context_router.go` | Context-based routing: shared-port service tiers, request size measurement, per-connection tier switching | | `http_request_splitter.go` | Incremental HTTP request framing from a raw TCP stream (used by context routing) | | `tokenizer.go` | Named token counters (heuristic, per model family) used to measure request context sizes | @@ -86,6 +87,8 @@ Client → large-model-proxy → [Service Process] | `config_test.go` | Configuration parsing and validation tests | | `management_api_test.go` | Management API endpoint tests | | `monitor_resources_test.go` | Resource monitoring tests | +| `port_allocation_test.go` | Listen port store, allocator and config resolution tests | +| `auto_listen_port_e2e_test.go` | End-to-end automatic port selection test (reuse and reassignment) | | `context_router_test.go` | Context routing tier selection and request unit counting tests | | `context_routing_connection_test.go` | In-process tests for the routed connection handler (routing, switching, passthrough) | | `context_routing_e2e_test.go` | End-to-end context routing test through the real proxy binary | diff --git a/README.md b/README.md index 6794b49e..ac90cc1d 100644 --- a/README.md +++ b/README.md @@ -287,6 +287,42 @@ smallest tier. clients do not pipeline requests (every real HTTP client sends the next request only after reading the previous response). +## Automatic listen port selection + +A service can ask the proxy to pick its listen port automatically: + +```jsonc +{ + "Name": "Qwen3-8B", + "ListenPort": "auto", // any free port, remembered across restarts + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "18085", + "Command": "llama-server", + "Args": "-m Qwen3-8B.gguf --port 18085", +} +``` + +At startup the proxy first tries the port this service used last time, so +clients keep working across restarts. If that port is no longer free (or the +service has no history yet), a free port is requested from the OS; ports +already used elsewhere in the same configuration are never picked. The chosen +port is then persisted in a small SQLite database +(`AutoListenPortDatabasePath`, default `auto-listen-ports.db` in the working +directory, keyed by service name), and the port is bound immediately when +selected so it cannot be lost to another process before the proxy starts +listening on it. + +Ways to discover the assigned port: + +- the startup log line `[ServiceName] Automatically selected listen port N` + (and the usual `Listening on port N` line), +- the `listen_port` field of the service in the Management API `/status` + response (also used for the `{{.PORT}}` variable in `ServiceUrl` templates), +- the SQLite database itself. + +Each service with `ListenPort: "auto"` gets a port of its own. Services that +share one port for context-based routing must use explicit numeric ports. + ## Management API The management API is a simple HTTP API that allows you to get the status of the proxy and the services it is proxying. diff --git a/auto_listen_port_e2e_test.go b/auto_listen_port_e2e_test.go new file mode 100644 index 00000000..9b716f55 --- /dev/null +++ b/auto_listen_port_e2e_test.go @@ -0,0 +1,153 @@ +package main + +import ( + "fmt" + "net" + "os" + "os/exec" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestAutoListenPortSelection verifies automatic listen port selection end to +// end: a service configured with ListenPort "auto" gets a free port, the +// choice is persisted to the SQLite database, a restart reuses the same port, +// and a restart while the old port is occupied picks (and persists) a new one. +func TestAutoListenPortSelection(t *testing.T) { + t.Parallel() + + repoDir, err := os.Getwd() + require.NoError(t, err) + workDir := t.TempDir() + databasePath := filepath.Join(workDir, defaultAutoListenPortDatabasePath) + + testConfig := Config{ + Services: []ServiceConfig{ + { + Name: "auto-service", + ListenPort: autoListenPort, + ProxyTargetHost: "localhost", + ProxyTargetPort: "12160", + Command: "./test-server/test-server", + Args: "-p 12160", + Workdir: repoDir, + }, + }, + ManagementApi: ManagementApi{ListenPort: "4160"}, + } + StandardizeConfigNamesAndPaths(&testConfig, t.Name()) + serviceName := testConfig.Services[0].Name + testConfig.Services[0].LogFilePath = filepath.Join(workDir, "auto-service.log") + configFilePath := createTempConfig(t, testConfig) + + startProxyInstance := func() (*stopHandle, error) { + waitChannel := make(chan error, 1) + cmd, err := startLargeModelProxy("auto-listen-port", configFilePath, workDir, waitChannel) + if err != nil { + return nil, err + } + return &stopHandle{cmd: cmd, waitChannel: waitChannel}, nil + } + + waitForAssignedPort := func() int { + t.Helper() + deadline := time.Now().Add(15 * time.Second) + for { + store, err := openListenPortStore(databasePath) + if err == nil { + port, found, getErr := store.getLastUsedPort(serviceName) + _ = store.close() + if getErr == nil && found { + return port + } + } + if time.Now().After(deadline) { + t.Fatalf("service %s never got a port persisted in %s", serviceName, databasePath) + } + time.Sleep(50 * time.Millisecond) + } + } + + verifyServiceReachableThroughPort := func(port int) { + t.Helper() + connection, err := net.DialTimeout("tcp", fmt.Sprintf("localhost:%d", port), 10*time.Second) + if err != nil { + t.Fatalf("failed to connect to automatically selected port %d: %v", port, err) + } + defer func() { _ = connection.Close() }() + _ = connection.SetReadDeadline(time.Now().Add(15 * time.Second)) + pid := readPidFromOpenConnection(t, connection) + assert.True(t, isProcessRunning(pid)) + } + + // First run: a port is selected and persisted + proxy, err := startProxyInstance() + require.NoError(t, err) + firstPort := waitForAssignedPort() + verifyServiceReachableThroughPort(firstPort) + status := getStatusFromManagementAPI(t, "localhost:4160") + serviceStatus := findServiceInStatusResponse(status, serviceName) + require.NotNil(t, serviceStatus) + assert.Equal(t, fmt.Sprintf("%d", firstPort), serviceStatus.ListenPort, + "the management API must report the resolved port") + require.NoError(t, proxy.stop()) + waitForPortClosed(t, "localhost:4160", 15*time.Second) + + // Second run with the same database: the same port must be reused + proxy, err = startProxyInstance() + require.NoError(t, err) + secondPort := waitForAssignedPort() + assert.Equal(t, firstPort, secondPort, "a free persisted port must be reused on restart") + verifyServiceReachableThroughPort(secondPort) + require.NoError(t, proxy.stop()) + waitForPortClosed(t, "localhost:4160", 15*time.Second) + + // Third run with the persisted port occupied by someone else: a new port + // must be selected and persisted instead + blocker, err := net.Listen("tcp", fmt.Sprintf(":%d", firstPort)) + require.NoError(t, err) + proxy, err = startProxyInstance() + require.NoError(t, err) + thirdPort := waitForAssignedPort() + assert.NotEqual(t, firstPort, thirdPort, "an occupied persisted port must not be reused") + verifyServiceReachableThroughPort(thirdPort) + require.NoError(t, proxy.stop()) + _ = blocker.Close() + waitForPortClosed(t, "localhost:4160", 15*time.Second) + + // The database must reflect the latest choice + store, err := openListenPortStore(databasePath) + require.NoError(t, err) + port, found, getErr := store.getLastUsedPort(serviceName) + require.NoError(t, getErr) + require.NoError(t, store.close()) + assert.True(t, found) + assert.Equal(t, thirdPort, port) +} + +type stopHandle struct { + cmd *exec.Cmd + waitChannel chan error +} + +func (h *stopHandle) stop() error { + return stopApplication(h.cmd, h.waitChannel) +} + +func waitForPortClosed(t *testing.T, address string, timeout time.Duration) { + t.Helper() + deadline := time.Now().Add(timeout) + for { + if err := checkPortClosed(address); err == nil { + return + } + if time.Now().After(deadline) { + t.Fatalf("port %s did not close within %s", address, timeout) + } + time.Sleep(50 * time.Millisecond) + } +} diff --git a/config.go b/config.go index e5f9d0f0..a623ec9f 100644 --- a/config.go +++ b/config.go @@ -111,6 +111,11 @@ type Config struct { ResourcesAvailable map[string]ResourceAvailable `json:"ResourcesAvailable"` OpenAiApi OpenAiApi ManagementApi ManagementApi + + // AutoListenPortDatabasePath is the SQLite database where ports selected + // for services with ListenPort "auto" are persisted, so services keep + // their port across restarts whenever possible. + AutoListenPortDatabasePath string `json:"AutoListenPortDatabasePath"` } type ServiceConfig struct { @@ -287,6 +292,9 @@ func loadConfigFromReader(r io.Reader) (Config, error) { if config.LogLevel == "" { config.LogLevel = LogLevelNormal } + if config.AutoListenPortDatabasePath == "" { + config.AutoListenPortDatabasePath = defaultAutoListenPortDatabasePath + } err = validateConfig(config) if err != nil { @@ -338,6 +346,12 @@ func validateConfig(cfg Config) error { portOrder := make([]string, 0) for _, svc := range cfg.Services { if svc.ListenPort != "" { + // "auto" services each get a port of their own (resolved at startup, + // persisted in the listen port database), so they never form + // shared-port groups. + if svc.ListenPort == autoListenPort { + continue + } if _, seen := servicesByPort[svc.ListenPort]; !seen { portOrder = append(portOrder, svc.ListenPort) } @@ -375,6 +389,9 @@ func validateConfig(cfg Config) error { fmt.Sprintf("service %s does not specify ListenPort", nameOrIndex)) continue } + if svc.ListenPort == autoListenPort { + continue + } portVal, err := strconv.Atoi(svc.ListenPort) if err != nil || portVal <= 0 || portVal > 65535 { issues = append(issues, diff --git a/config_test.go b/config_test.go index dccdfc6b..a4615a5c 100644 --- a/config_test.go +++ b/config_test.go @@ -1327,3 +1327,137 @@ func TestSingleServiceWithContextSizeIsValid(t *testing.T) { t.Fatalf("did not expect an error but got: %v", err) } } + +// --- automatic listen port selection --- + +func TestListenPortAutoIsValid(t *testing.T) { + t.Parallel() + cfg, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "svc", + "ListenPort": "auto", + "Command": "/bin/echo" + } + ] + }`) + if err != nil { + t.Fatalf("did not expect an error but got: %v", err) + } + assert.Equal(t, "auto", cfg.Services[0].ListenPort) +} + +func TestMultipleAutoListenPortsAreAllowed(t *testing.T) { + t.Parallel() + // Two "auto" services each get their own port; they must not be treated + // as a shared-port group (which would require context sizes). + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "one", + "ListenPort": "auto", + "Command": "/bin/echo" + }, + { + "Name": "two", + "ListenPort": "auto", + "Command": "/bin/echo" + } + ] + }`) + if err != nil { + t.Fatalf("did not expect an error but got: %v", err) + } +} + +func TestAutoListenPortWithStaticPortsIsValid(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "OpenAiApi": {"ListenPort": "7070"}, + "ManagementApi": {"ListenPort": "7071"}, + "Services": [ + { + "Name": "auto-service", + "ListenPort": "auto", + "Command": "/bin/echo" + }, + { + "Name": "static-service", + "ListenPort": "8080", + "Command": "/bin/echo" + } + ] + }`) + if err != nil { + t.Fatalf("did not expect an error but got: %v", err) + } +} + +func TestAutoListenPortDatabasePathIsParsed(t *testing.T) { + t.Parallel() + cfg, err := loadConfigFromString(t, `{ + "AutoListenPortDatabasePath": "/var/lib/large-model-proxy/ports.db", + "Services": [ + { + "Name": "svc", + "ListenPort": "auto", + "Command": "/bin/echo" + } + ] + }`) + if err != nil { + t.Fatalf("did not expect an error but got: %v", err) + } + assert.Equal(t, "/var/lib/large-model-proxy/ports.db", cfg.AutoListenPortDatabasePath) +} + +func TestAutoListenPortDatabasePathDefaults(t *testing.T) { + t.Parallel() + cfg, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "svc", + "ListenPort": "8080", + "Command": "/bin/echo" + } + ] + }`) + if err != nil { + t.Fatalf("did not expect an error but got: %v", err) + } + assert.Equal(t, defaultAutoListenPortDatabasePath, cfg.AutoListenPortDatabasePath) +} + +func TestAutoListenPortStillRequiresCommand(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "svc", + "ListenPort": "auto" + } + ] + }`) + checkExpectedErrorMessages(t, err, []string{"has no Command specified"}) +} + +func TestAutoListenPortOpenAiApiOnlyServiceStillAllowedWithoutPort(t *testing.T) { + t.Parallel() + _, err := loadConfigFromString(t, `{ + "Services": [ + { + "Name": "svc", + "OpenAiApi": true, + "Command": "/bin/echo" + }, + { + "Name": "auto-service", + "ListenPort": "auto", + "Command": "/bin/echo" + } + ] + }`) + if err != nil { + t.Fatalf("did not expect an error but got: %v", err) + } +} diff --git a/connection.go b/connection.go index f585c08c..198ece37 100644 --- a/connection.go +++ b/connection.go @@ -12,10 +12,18 @@ import ( func startProxy(serviceConfig ServiceConfig) { listener, err := net.Listen("tcp", ":"+serviceConfig.ListenPort) - log.Printf("[%s] Listening on port %s", serviceConfig.Name, serviceConfig.ListenPort) if err != nil { log.Fatalf("[%s] Fatal error: cannot listen on port %s: %v", serviceConfig.Name, serviceConfig.ListenPort, err) } + startProxyWithListener(listener, serviceConfig) +} + +// startProxyWithListener serves a service on an already bound listener. It is +// used for automatically selected ports, where the listener created during +// port allocation is handed over so the port can never be grabbed by another +// process between allocation and use. +func startProxyWithListener(listener net.Listener, serviceConfig ServiceConfig) { + log.Printf("[%s] Listening on port %s", serviceConfig.Name, serviceConfig.ListenPort) defer func(listener net.Listener) { _ = listener.Close() }(listener) diff --git a/go.mod b/go.mod index c2be81c8..87f61c88 100644 --- a/go.mod +++ b/go.mod @@ -7,6 +7,18 @@ require github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 require ( github.com/stretchr/testify v1.12.1 github.com/tidwall/jsonc v0.3.3 + modernc.org/sqlite v1.34.5 ) -require go.yaml.in/yaml/v3 v3.0.5 // indirect +require ( + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/ncruces/go-strftime v0.1.9 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + go.yaml.in/yaml/v3 v3.0.5 // indirect + golang.org/x/sys v0.22.0 // indirect + modernc.org/libc v1.55.3 // indirect + modernc.org/mathutil v1.6.0 // indirect + modernc.org/memory v1.8.0 // indirect +) diff --git a/go.sum b/go.sum index 593cf393..791ace15 100644 --- a/go.sum +++ b/go.sum @@ -1,8 +1,51 @@ +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/google/pprof v0.0.0-20240409012703-83162a5b38cd h1:gbpYu9NMq8jhDVbvlGkMFWCjLFlqqEZjEmObmhUy6Vo= +github.com/google/pprof v0.0.0-20240409012703-83162a5b38cd/go.mod h1:kf6iHlnVGwgKolg33glAes7Yg/8iWP8ukqeldJSO7jw= github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4= github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4= +github.com/ncruces/go-strftime v0.1.9/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= github.com/tidwall/jsonc v0.3.3 h1:RVQqL3xFfDkKKXIDsrBiVQiEpBtxoKbmMXONb2H/y2w= github.com/tidwall/jsonc v0.3.3/go.mod h1:dw+3CIxqHi+t8eFSpzzMlcVYxKp08UP5CD8/uSFCyJE= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= +golang.org/x/mod v0.16.0 h1:QX4fJ0Rr5cPQCF7O9lh9Se4pmwfwskqZfq5moyldzic= +golang.org/x/mod v0.16.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.22.0 h1:RI27ohtqKCnwULzJLqkv897zojh5/DwS/ENaMzUOaWI= +golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/tools v0.19.0 h1:tfGCXNR1OsFG+sVdLAitlpjAvD/I6dHDKnYrpEZUHkw= +golang.org/x/tools v0.19.0/go.mod h1:qoJWxmGSIBmAeriMx19ogtrEPrGtDbPK634QFIcLAhc= +modernc.org/cc/v4 v4.21.4 h1:3Be/Rdo1fpr8GrQ7IVw9OHtplU4gWbb+wNgeoBMmGLQ= +modernc.org/cc/v4 v4.21.4/go.mod h1:HM7VJTZbUCR3rV8EYBi9wxnJ0ZBRiGE5OeGXNA0IsLQ= +modernc.org/ccgo/v4 v4.19.2 h1:lwQZgvboKD0jBwdaeVCTouxhxAyN6iawF3STraAal8Y= +modernc.org/ccgo/v4 v4.19.2/go.mod h1:ysS3mxiMV38XGRTTcgo0DQTeTmAO4oCmJl1nX9VFI3s= +modernc.org/fileutil v1.3.0 h1:gQ5SIzK3H9kdfai/5x41oQiKValumqNTDXMvKo62HvE= +modernc.org/fileutil v1.3.0/go.mod h1:XatxS8fZi3pS8/hKG2GH/ArUogfxjpEKs3Ku3aK4JyQ= +modernc.org/gc/v2 v2.4.1 h1:9cNzOqPyMJBvrUipmynX0ZohMhcxPtMccYgGOJdOiBw= +modernc.org/gc/v2 v2.4.1/go.mod h1:wzN5dK1AzVGoH6XOzc3YZ+ey/jPgYHLuVckd62P0GYU= +modernc.org/libc v1.55.3 h1:AzcW1mhlPNrRtjS5sS+eW2ISCgSOLLNyFzRh/V3Qj/U= +modernc.org/libc v1.55.3/go.mod h1:qFXepLhz+JjFThQ4kzwzOjA/y/artDeg+pcYnY+Q83w= +modernc.org/mathutil v1.6.0 h1:fRe9+AmYlaej+64JsEEhoWuAYBkOtQiMEU7n/XgfYi4= +modernc.org/mathutil v1.6.0/go.mod h1:Ui5Q9q1TR2gFm0AQRqQUaBWFLAhQpCwNcuhBOSedWPo= +modernc.org/memory v1.8.0 h1:IqGTL6eFMaDZZhEWwcREgeMXYwmW83LYW8cROZYkg+E= +modernc.org/memory v1.8.0/go.mod h1:XPZ936zp5OMKGWPqbD3JShgd/ZoQ7899TUuQqxY+peU= +modernc.org/opt v0.1.3 h1:3XOZf2yznlhC+ibLltsDGzABUGVx8J6pnFMS3E4dcq4= +modernc.org/opt v0.1.3/go.mod h1:WdSiB5evDcignE70guQKxYUl14mgWtbClRi5wmkkTX0= +modernc.org/sortutil v1.2.0 h1:jQiD3PfS2REGJNzNCMMaLSp/wdMNieTbKX920Cqdgqc= +modernc.org/sortutil v1.2.0/go.mod h1:TKU2s7kJMf1AE84OoiGppNHJwvB753OYfNl2WRb++Ss= +modernc.org/sqlite v1.34.5 h1:Bb6SR13/fjp15jt70CL4f18JIN7p7dnMExd+UFnF15g= +modernc.org/sqlite v1.34.5/go.mod h1:YLuNmX9NKs8wRNK2ko1LW1NGYcc9FkBO69JOt1AR9JE= +modernc.org/strutil v1.2.0 h1:agBi9dp1I+eOnxXeiZawM8F4LawKv4NzGWSaLfyeNZA= +modernc.org/strutil v1.2.0/go.mod h1:/mdcBmfOibveCTBxUl5B5l6W+TTH1FXPLHZE6bTosX0= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/main.go b/main.go index cceda717..71b41d14 100644 --- a/main.go +++ b/main.go @@ -5,6 +5,7 @@ import ( "fmt" "log" "maps" + "net" "os" "os/exec" "os/signal" @@ -173,6 +174,31 @@ func main() { ) } } + // Resolve services with ListenPort "auto" to concrete ports before any + // listener starts: the resolved ports end up in config so the management + // API, service URLs and logs all report them, and the pre-bound listeners + // are handed to the proxies so a chosen port can never be lost to another + // process between allocation and use. + autoPortListeners := map[string]net.Listener{} + hasAutoListenPorts := false + for _, service := range config.Services { + if service.ListenPort == autoListenPort { + hasAutoListenPorts = true + break + } + } + if hasAutoListenPorts { + listenPortStore, err := openListenPortStore(config.AutoListenPortDatabasePath) + if err != nil { + log.Fatalf("Error opening listen port database: %v", err) + } + defer func() { _ = listenPortStore.close() }() + autoPortListeners, err = resolveAutoListenPorts(&config, listenPortStore) + if err != nil { + log.Fatalf("Error selecting listen ports automatically: %v", err) + } + } + // Group services by listen port: a port with a single service behaves as a // plain transparent proxy, while a port shared by several services is // served by a context router that picks the smallest context size that @@ -191,7 +217,11 @@ func main() { for _, listenPort := range listenPortOrder { services := servicesByListenPort[listenPort] if len(services) == 1 { - go startProxy(services[0]) + if listener, found := autoPortListeners[services[0].Name]; found { + go startProxyWithListener(listener, services[0]) + } else { + go startProxy(services[0]) + } } else { go startContextRoutedProxy(listenPort, buildContextRouter(services)) } diff --git a/port_allocation.go b/port_allocation.go new file mode 100644 index 00000000..58e55931 --- /dev/null +++ b/port_allocation.go @@ -0,0 +1,176 @@ +package main + +import ( + "database/sql" + "fmt" + "log" + "net" + "strconv" + + _ "modernc.org/sqlite" +) + +// autoListenPort is the magic ListenPort value that asks the proxy to pick a +// free port automatically. Every service using it gets a port of its own; +// context-routing groups (several services sharing one port) require explicit +// numeric ports. +const autoListenPort = "auto" + +// defaultAutoListenPortDatabasePath is where chosen ports are persisted when +// AutoListenPortDatabasePath is not configured. +const defaultAutoListenPortDatabasePath = "auto-listen-ports.db" + +// listenPortStore persists the last automatically selected listen port of +// every service in a small SQLite database, so restarts keep serving services +// on stable ports whenever possible. +type listenPortStore struct { + db *sql.DB +} + +func openListenPortStore(databasePath string) (*listenPortStore, error) { + db, err := sql.Open("sqlite", databasePath) + if err != nil { + return nil, fmt.Errorf("failed to open listen port database %s: %w", databasePath, err) + } + // SQLite handles one writer at a time; a single connection avoids + // "database is locked" errors without needing WAL tuning for what is + // a handful of startup writes. + db.SetMaxOpenConns(1) + _, err = db.Exec( + "CREATE TABLE IF NOT EXISTS service_listen_ports (" + + "service_name TEXT PRIMARY KEY," + + "port INTEGER NOT NULL," + + "updated_at TEXT NOT NULL DEFAULT (datetime('now')))", + ) + if err != nil { + _ = db.Close() + return nil, fmt.Errorf("failed to initialize listen port database %s: %w", databasePath, err) + } + return &listenPortStore{db: db}, nil +} + +// getLastUsedPort returns the port this service was last assigned, if any. +func (s *listenPortStore) getLastUsedPort(serviceName string) (int, bool, error) { + var port int + err := s.db.QueryRow( + "SELECT port FROM service_listen_ports WHERE service_name = ?", serviceName, + ).Scan(&port) + if err == sql.ErrNoRows { + return 0, false, nil + } + if err != nil { + return 0, false, fmt.Errorf("failed to look up last used port for service %s: %w", serviceName, err) + } + return port, true, nil +} + +// saveUsedPort records the port assigned to a service, replacing any previous +// value. +func (s *listenPortStore) saveUsedPort(serviceName string, port int) error { + _, err := s.db.Exec( + "INSERT INTO service_listen_ports (service_name, port, updated_at) VALUES (?, ?, datetime('now')) "+ + "ON CONFLICT(service_name) DO UPDATE SET port = excluded.port, updated_at = excluded.updated_at", + serviceName, port, + ) + if err != nil { + return fmt.Errorf("failed to save port %d for service %s: %w", port, serviceName, err) + } + return nil +} + +func (s *listenPortStore) close() error { + return s.db.Close() +} + +// allocateListenPort binds a listen port for one service. If the service used +// a port before and that port is still free, it is reused so clients keep +// working across restarts; otherwise a fresh port is requested from the OS. +// isPortTaken is consulted for every candidate and lets the caller keep +// candidates away from ports that are already spoken for (statically +// configured ports, ports assigned moments ago, ...). The returned listener +// already holds the port, closing the race between allocation and use. +func allocateListenPort(lastUsedPort int, hasLastUsedPort bool, isPortTaken func(port int) bool) (net.Listener, int, error) { + if hasLastUsedPort && lastUsedPort > 0 { + listener, err := net.Listen("tcp", ":"+strconv.Itoa(lastUsedPort)) + if err == nil && !isPortTaken(lastUsedPort) { + return listener, lastUsedPort, nil + } + if listener != nil { + _ = listener.Close() + } + log.Printf("Last used listen port %d is not available, selecting a new one", lastUsedPort) + } + + // Ask the OS for a free port. The OS does not know about ports this proxy + // will only bind later (statically configured ones), so candidates + // reported as taken are released and retried. + for attempt := 0; attempt < 100; attempt++ { + listener, err := net.Listen("tcp", ":0") + if err != nil { + return nil, 0, fmt.Errorf("failed to request a free port from the OS: %w", err) + } + port := listener.Addr().(*net.TCPAddr).Port + if !isPortTaken(port) { + return listener, port, nil + } + _ = listener.Close() + } + return nil, 0, fmt.Errorf("could not find a free listen port that does not conflict with the configuration after 100 attempts") +} + +func closeListeners(listeners map[string]net.Listener) { + for _, listener := range listeners { + _ = listener.Close() + } +} + +// resolveAutoListenPorts replaces the "auto" ListenPort of every service with +// a concrete port number, mutating config in place, and returns the already +// bound listener for each of those services. Ports are persisted through the +// store so that a service keeps its port across restarts whenever possible. +func resolveAutoListenPorts(config *Config, store *listenPortStore) (map[string]net.Listener, error) { + listeners := make(map[string]net.Listener) + + takenPorts := make(map[int]bool) + for _, port := range []string{config.OpenAiApi.ListenPort, config.ManagementApi.ListenPort} { + if portNumber, err := strconv.Atoi(port); err == nil { + takenPorts[portNumber] = true + } + } + + for serviceIndex := range config.Services { + if config.Services[serviceIndex].ListenPort != autoListenPort { + if portNumber, err := strconv.Atoi(config.Services[serviceIndex].ListenPort); err == nil { + takenPorts[portNumber] = true + } + continue + } + if store == nil { + return nil, fmt.Errorf("service %s uses ListenPort \"auto\" but no listen port database is available; configure AutoListenPortDatabasePath", config.Services[serviceIndex].Name) + } + + serviceName := config.Services[serviceIndex].Name + lastUsedPort, hasLastUsedPort, err := store.getLastUsedPort(serviceName) + if err != nil { + return nil, err + } + + isPortTaken := func(port int) bool { return takenPorts[port] } + listener, port, err := allocateListenPort(lastUsedPort, hasLastUsedPort, isPortTaken) + if err != nil { + closeListeners(listeners) + return nil, fmt.Errorf("failed to allocate a listen port for service %s: %w", serviceName, err) + } + + if err := store.saveUsedPort(serviceName, port); err != nil { + _ = listener.Close() + closeListeners(listeners) + return nil, err + } + takenPorts[port] = true + listeners[serviceName] = listener + config.Services[serviceIndex].ListenPort = strconv.Itoa(port) + log.Printf("[%s] Automatically selected listen port %d", serviceName, port) + } + return listeners, nil +} diff --git a/port_allocation_test.go b/port_allocation_test.go new file mode 100644 index 00000000..41363d24 --- /dev/null +++ b/port_allocation_test.go @@ -0,0 +1,343 @@ +package main + +import ( + "net" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newTestListenPortStore(t *testing.T) *listenPortStore { + t.Helper() + store, err := openListenPortStore(t.TempDir() + "/listen-ports.db") + if err != nil { + t.Fatalf("failed to open listen port store: %v", err) + } + t.Cleanup(func() { _ = store.close() }) + return store +} + +// findFreePort reserves a port from the OS and releases it again. +func findFreePort(t *testing.T) int { + t.Helper() + listener, err := net.Listen("tcp", ":0") + if err != nil { + t.Fatalf("failed to reserve a port: %v", err) + } + defer func() { _ = listener.Close() }() + return listener.Addr().(*net.TCPAddr).Port +} + +// --- sqlite persistence --- + +func TestListenPortStoreRoundTripsPort(t *testing.T) { + t.Parallel() + store := newTestListenPortStore(t) + port := findFreePort(t) + + require.NoError(t, store.saveUsedPort("service-a", port)) + + lastUsed, found, err := store.getLastUsedPort("service-a") + require.NoError(t, err) + assert.True(t, found) + assert.Equal(t, port, lastUsed) +} + +func TestListenPortStoreSurvivesReopen(t *testing.T) { + t.Parallel() + databasePath := t.TempDir() + "/listen-ports.db" + store, err := openListenPortStore(databasePath) + require.NoError(t, err) + port := findFreePort(t) + require.NoError(t, store.saveUsedPort("service-a", port)) + require.NoError(t, store.close()) + + reopened, err := openListenPortStore(databasePath) + require.NoError(t, err) + defer func() { _ = reopened.close() }() + + lastUsed, found, err := reopened.getLastUsedPort("service-a") + require.NoError(t, err) + assert.True(t, found) + assert.Equal(t, port, lastUsed) +} + +func TestListenPortStoreUnknownServiceHasNoPort(t *testing.T) { + t.Parallel() + store := newTestListenPortStore(t) + + lastUsed, found, err := store.getLastUsedPort("never-seen") + require.NoError(t, err) + assert.False(t, found) + assert.Equal(t, 0, lastUsed) +} + +func TestListenPortStoreOverwritesPreviousPort(t *testing.T) { + t.Parallel() + store := newTestListenPortStore(t) + + require.NoError(t, store.saveUsedPort("service-a", 45678)) + require.NoError(t, store.saveUsedPort("service-a", 45679)) + + lastUsed, found, err := store.getLastUsedPort("service-a") + require.NoError(t, err) + assert.True(t, found) + assert.Equal(t, 45679, lastUsed, "the most recently saved port must win") +} + +func TestListenPortStoresServicesIndependently(t *testing.T) { + t.Parallel() + store := newTestListenPortStore(t) + + require.NoError(t, store.saveUsedPort("service-a", 45678)) + require.NoError(t, store.saveUsedPort("service-b", 45679)) + + portA, foundA, err := store.getLastUsedPort("service-a") + require.NoError(t, err) + portB, foundB, err := store.getLastUsedPort("service-b") + require.NoError(t, err) + assert.True(t, foundA) + assert.True(t, foundB) + assert.Equal(t, 45678, portA) + assert.Equal(t, 45679, portB) +} + +func TestListenPortStoreCreatesDatabaseFile(t *testing.T) { + t.Parallel() + databasePath := t.TempDir() + "/listen-ports.db" + store, err := openListenPortStore(databasePath) + require.NoError(t, err) + require.NoError(t, store.saveUsedPort("service-a", 45678)) + require.NoError(t, store.close()) + + assert.FileExists(t, databasePath) +} + +// --- allocation --- + +func TestAllocateListenPortPrefersLastUsedPort(t *testing.T) { + t.Parallel() + // A freed port can be grabbed by unrelated concurrent listeners between + // releasing it and the allocator binding it, so give the scenario a few + // attempts: at least one must observe the reuse. + for attempt := 0; attempt < 20; attempt++ { + lastUsedPort := findFreePort(t) + + listener, allocatedPort, err := allocateListenPort(lastUsedPort, true, func(int) bool { return false }) + require.NoError(t, err) + portWasReused := allocatedPort == lastUsedPort + require.NoError(t, listener.Close()) + if portWasReused { + return + } + } + t.Fatal("allocator never reused a free last-used port") +} + +func TestAllocateListenPortAvoidsOccupiedLastUsedPort(t *testing.T) { + t.Parallel() + // Keep the last-used port occupied so it cannot be reused + blocker, err := net.Listen("tcp", ":0") + require.NoError(t, err) + defer func() { _ = blocker.Close() }() + occupiedPort := blocker.Addr().(*net.TCPAddr).Port + + listener, allocatedPort, err := allocateListenPort(occupiedPort, true, func(int) bool { return false }) + require.NoError(t, err) + defer func() { _ = listener.Close() }() + + assert.NotEqual(t, occupiedPort, allocatedPort) + assert.Greater(t, allocatedPort, 0) +} + +func TestAllocateListenPortAvoidsTakenCandidates(t *testing.T) { + t.Parallel() + var candidates []int + var firstCandidate int + // Reject the first candidate the OS offers: the allocator must close it + // and ask again instead of returning a port marked as taken. + isPortTaken := func(port int) bool { + candidates = append(candidates, port) + if len(candidates) == 1 { + firstCandidate = port + return true + } + return port == firstCandidate + } + + listener, allocatedPort, err := allocateListenPort(0, false, isPortTaken) + require.NoError(t, err) + defer func() { _ = listener.Close() }() + + assert.NotEmpty(t, candidates, "the taken-port check must be consulted") + assert.NotEqual(t, firstCandidate, allocatedPort, "a port reported as taken must never be returned") +} + +func TestAllocateListenPortReturnsHeldListener(t *testing.T) { + t.Parallel() + listener, allocatedPort, err := allocateListenPort(0, false, func(int) bool { return false }) + require.NoError(t, err) + defer func() { _ = listener.Close() }() + + // The returned listener must already hold the port, so a second bind must fail + second, secondErr := net.Listen("tcp", ":"+strconv.Itoa(allocatedPort)) + if secondErr == nil { + _ = second.Close() + t.Fatalf("expected port %d to be held by the returned listener", allocatedPort) + } +} + +// --- config resolution --- + +func TestResolveAutoListenPortsAssignsNumericPort(t *testing.T) { + t.Parallel() + store := newTestListenPortStore(t) + cfg := &Config{ + Services: []ServiceConfig{ + {Name: "auto-service", ListenPort: autoListenPort, Command: "/bin/echo"}, + {Name: "static-service", ListenPort: "8099", Command: "/bin/echo"}, + }, + } + + listeners, err := resolveAutoListenPorts(cfg, store) + require.NoError(t, err) + defer closeListeners(listeners) + + assert.NotEqual(t, autoListenPort, cfg.Services[0].ListenPort) + portNumber, convertErr := strconv.Atoi(cfg.Services[0].ListenPort) + require.NoError(t, convertErr, "the resolved listen port must be numeric") + assert.Greater(t, portNumber, 0) + assert.Equal(t, "8099", cfg.Services[1].ListenPort, "static ports must not be touched") + + listener, found := listeners["auto-service"] + require.True(t, found, "a bound listener must be returned for the auto service") + assert.Equal(t, portNumber, listener.Addr().(*net.TCPAddr).Port) + + savedPort, found, err := store.getLastUsedPort("auto-service") + require.NoError(t, err) + assert.True(t, found, "the allocated port must be persisted") + assert.Equal(t, portNumber, savedPort) +} + +func TestResolveAutoListenPortsAvoidsConfiguredPorts(t *testing.T) { + t.Parallel() + store := newTestListenPortStore(t) + cfg := &Config{ + OpenAiApi: OpenAiApi{ListenPort: "8100"}, + ManagementApi: ManagementApi{ListenPort: "8101"}, + Services: []ServiceConfig{ + {Name: "auto-service", ListenPort: autoListenPort, Command: "/bin/echo"}, + {Name: "static-service", ListenPort: "8099", Command: "/bin/echo"}, + }, + } + + listeners, err := resolveAutoListenPorts(cfg, store) + require.NoError(t, err) + defer closeListeners(listeners) + + resolvedPort := cfg.Services[0].ListenPort + assert.NotEqual(t, "8099", resolvedPort) + assert.NotEqual(t, "8100", resolvedPort) + assert.NotEqual(t, "8101", resolvedPort) +} + +func TestResolveAutoListenPortsReusesPersistedPort(t *testing.T) { + t.Parallel() + // Retry to tolerate unrelated listeners grabbing the seeded port between + // releasing it and the proxy binding it. + for attempt := 0; attempt < 20; attempt++ { + store := newTestListenPortStore(t) + persistedPort := findFreePort(t) + require.NoError(t, store.saveUsedPort("auto-service", persistedPort)) + cfg := &Config{ + Services: []ServiceConfig{ + {Name: "auto-service", ListenPort: autoListenPort, Command: "/bin/echo"}, + }, + } + + listeners, err := resolveAutoListenPorts(cfg, store) + require.NoError(t, err) + resolvedPort := cfg.Services[0].ListenPort + resolvedListenerPort := listeners["auto-service"].Addr().(*net.TCPAddr).Port + closeListeners(listeners) + + if resolvedPort == strconv.Itoa(persistedPort) && resolvedListenerPort == persistedPort { + return + } + } + t.Fatal("resolver never reused a free persisted port") +} + +func TestResolveAutoListenPortsReplacesUnavailablePersistedPort(t *testing.T) { + t.Parallel() + store := newTestListenPortStore(t) + blocker, err := net.Listen("tcp", ":0") + require.NoError(t, err) + defer func() { _ = blocker.Close() }() + unavailablePort := blocker.Addr().(*net.TCPAddr).Port + require.NoError(t, store.saveUsedPort("auto-service", unavailablePort)) + cfg := &Config{ + Services: []ServiceConfig{ + {Name: "auto-service", ListenPort: autoListenPort, Command: "/bin/echo"}, + }, + } + + listeners, err := resolveAutoListenPorts(cfg, store) + require.NoError(t, err) + defer closeListeners(listeners) + + resolvedPortNumber, convertErr := strconv.Atoi(cfg.Services[0].ListenPort) + require.NoError(t, convertErr) + assert.NotEqual(t, unavailablePort, resolvedPortNumber) + + savedPort, found, err := store.getLastUsedPort("auto-service") + require.NoError(t, err) + assert.True(t, found) + assert.Equal(t, resolvedPortNumber, savedPort, "the store must record the replacement port") +} + +func TestResolveAutoListenPortsAssignsDistinctPorts(t *testing.T) { + t.Parallel() + store := newTestListenPortStore(t) + cfg := &Config{ + Services: []ServiceConfig{ + {Name: "auto-one", ListenPort: autoListenPort, Command: "/bin/echo"}, + {Name: "auto-two", ListenPort: autoListenPort, Command: "/bin/echo"}, + }, + } + + listeners, err := resolveAutoListenPorts(cfg, store) + require.NoError(t, err) + defer closeListeners(listeners) + + assert.NotEqual(t, cfg.Services[0].ListenPort, cfg.Services[1].ListenPort, "two auto services must never share a port") + assert.Len(t, listeners, 2) +} + +func TestResolveAutoListenPortsWithoutAutoPorts(t *testing.T) { + t.Parallel() + cfg := &Config{ + Services: []ServiceConfig{ + {Name: "static-service", ListenPort: "8099", Command: "/bin/echo"}, + }, + } + + listeners, err := resolveAutoListenPorts(cfg, nil) + require.NoError(t, err) + assert.Empty(t, listeners) + assert.Equal(t, "8099", cfg.Services[0].ListenPort) +} + +func TestResolveAutoListenPortsWithoutStoreFails(t *testing.T) { + t.Parallel() + cfg := &Config{ + Services: []ServiceConfig{ + {Name: "auto-service", ListenPort: autoListenPort, Command: "/bin/echo"}, + }, + } + + _, err := resolveAutoListenPorts(cfg, nil) + assert.Error(t, err, "auto ports without a persistence store must be a configuration error") +}