diff --git a/.gitignore b/.gitignore index 6f3fb5a..4735390 100644 --- a/.gitignore +++ b/.gitignore @@ -42,4 +42,9 @@ tmp/ build/ commands.txt -gateway/.claude/ \ No newline at end of file +gateway/.claude/ + +# Local planning & Claude config — never commit +CONTEXT.md +.claudeignore +CLAUDE.md \ No newline at end of file diff --git a/gateway/cmd/iq/main.go b/gateway/cmd/iq/main.go index 8a61afc..e8912fe 100644 --- a/gateway/cmd/iq/main.go +++ b/gateway/cmd/iq/main.go @@ -6,6 +6,9 @@ package main import ( "context" + "crypto/rand" + "encoding/hex" + "errors" "fmt" "net" "net/http" @@ -21,35 +24,94 @@ import ( var version = "dev" +func generateToken() string { + b := make([]byte, 16) + if _, err := rand.Read(b); err != nil { + return fmt.Sprintf("iq-fallback-%d", time.Now().UnixNano()) + } + return hex.EncodeToString(b) +} + func main() { go checkForUpdate() - port, err := getFreePort() - if err != nil { - fmt.Fprintf(os.Stderr, "iq: failed to find free port: %v\n", err) - os.Exit(1) + if len(os.Args) < 2 { + // iq alone → run claude (backward compat, no subcommand) + runClaude([]string{}, false) + return + } + + switch os.Args[1] { + case "claude": + devMode, claudeArgs := parseDevFlag(os.Args[2:]) + runClaude(claudeArgs, devMode) + case "gemini": + fmt.Println(" iq gemini — coming soon") + case "codex": + fmt.Println(" iq codex — coming soon") + case "help", "--help", "-h": + printHelp() + default: + // Unknown subcommand → pass everything to claude (backward compat). + // Ensures `iq somefile.go` and `iq --resume` still work as before. + runClaude(os.Args[1:], false) + } +} + +func parseDevFlag(args []string) (bool, []string) { + dev := false + filtered := make([]string, 0, len(args)) + for _, a := range args { + if a == "--dev" { + dev = true + } else { + filtered = append(filtered, a) + } } + return dev, filtered +} + +func runClaude(args []string, devMode bool) { + if devMode { + os.Setenv("IQ_DEV_MODE", "1") + fmt.Fprintln(os.Stderr, " [iq] dev mode — relaxed guards, full telemetry") + } + + // Generate a per-invocation session ID so the circuit breaker scopes + // similarity checks within this iq session, not across sessions. + os.Setenv("IQ_SESSION_ID", generateToken()) - // Redirect gateway stderr to a log file so nothing leaks into the terminal. - logPath := filepath.Join(os.Getenv("HOME"), ".indexqube", "gateway.log") - os.MkdirAll(filepath.Dir(logPath), 0700) //nolint:errcheck - if lf, err := os.OpenFile(logPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644); err == nil { - syscall.Dup2(int(lf.Fd()), int(os.Stderr.Fd())) //nolint:errcheck + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + fmt.Fprintf(os.Stderr, "iq: failed to reserve local port: %v\n", err) + os.Exit(1) } + port := ln.Addr().(*net.TCPAddr).Port ctx, cancel := context.WithCancel(context.Background()) defer cancel() - os.Setenv("INDEXQUBE_LOG_LEVEL", "error") - go startProxy(ctx, port) + os.Setenv("INDEXQUBE_LOG_LEVEL", "off") + go startProxy(ctx, ln) if !waitForProxy(port) { fmt.Fprintf(os.Stderr, "iq: proxy failed to start within 2s\n") os.Exit(1) } - args := append([]string{"claude"}, os.Args[1:]...) - cmd := exec.Command(args[0], args[1:]...) + claudePath, err := exec.LookPath("claude") + if err != nil { + fmt.Fprintf(os.Stderr, "iq: could not find 'claude' in PATH: %v\n", err) + os.Exit(1) + } + + if !filepath.IsAbs(claudePath) { + fmt.Fprintf(os.Stderr, "iq: 'claude' path is not absolute: %s\n", claudePath) + os.Exit(1) + } + + cmdArgs := append([]string{claudePath}, args...) + cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...) cmd.Env = append(os.Environ(), fmt.Sprintf("ANTHROPIC_BASE_URL=http://127.0.0.1:%d", port), ) @@ -59,6 +121,7 @@ func main() { sigs := make(chan os.Signal, 1) signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM) + defer signal.Stop(sigs) go func() { sig := <-sigs if cmd.Process != nil { @@ -66,14 +129,42 @@ func main() { } }() - cmd.Run() //nolint:errcheck + if err := cmd.Run(); err != nil { + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { + os.Exit(exitErr.ExitCode()) + } + fmt.Fprintf(os.Stderr, "iq: failed to run claude: %v\n", err) + os.Exit(1) + } +} + +func printHelp() { + fmt.Print(` + iq — IndexQube CLI + + USAGE + iq Start Claude Code (default) + iq claude Start Claude Code via IndexQube + iq claude --dev Start Claude Code, dev mode (relaxed guards) + iq gemini Gemini via IndexQube (coming soon) + iq codex Codex via IndexQube (coming soon) + iq help Show this help + + DEV MODE + --dev disables velocity guard, enables verbose logging. + Use when building IndexQube with IndexQube. + +`) } const telemetryEndpoint = "https://dev-api.indexqube.com" -func startProxy(ctx context.Context, port int) { - os.Setenv("INDEXQUBE_BIND_ADDR", fmt.Sprintf("127.0.0.1:%d", port)) - os.Setenv("INDEXQUBE_DEV_TOKEN", "iq-local") +func startProxy(ctx context.Context, ln net.Listener) { + os.Setenv("INDEXQUBE_BIND_ADDR", ln.Addr().String()) + if os.Getenv("INDEXQUBE_DEV_TOKEN") == "" { + os.Setenv("INDEXQUBE_DEV_TOKEN", generateToken()) + } os.Setenv("INDEXQUBE_MODE", "optimize") os.Setenv("INDEXQUBE_ENABLE_BLOCK_OPTIMIZER", "true") // Route telemetry through the deployed gateway so Supabase credentials @@ -85,16 +176,7 @@ func startProxy(ctx context.Context, port int) { if os.Getenv("INDEXQUBE_LOG_LEVEL") == "" { os.Setenv("INDEXQUBE_LOG_LEVEL", "warn") } - server.Run(ctx) //nolint:errcheck -} - -func getFreePort() (int, error) { - ln, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - return 0, err - } - defer ln.Close() - return ln.Addr().(*net.TCPAddr).Port, nil + server.RunWithPublicListener(ctx, ln) //nolint:errcheck } func waitForProxy(port int) bool { diff --git a/gateway/cmd/iq/update.go b/gateway/cmd/iq/update.go index cd3f093..e23ddc5 100644 --- a/gateway/cmd/iq/update.go +++ b/gateway/cmd/iq/update.go @@ -31,7 +31,7 @@ func checkForUpdate() { fmt.Fprintf(os.Stderr, "\niq: new version available %s → %s\n"+ - " Run: curl -sSL https://raw.githubusercontent.com/Revanth14/indexqube/main/install.sh | bash\n\n", + " Download the latest release from: https://github.com/Revanth14/indexqube/releases/latest\n\n", version, release.TagName, ) } diff --git a/gateway/deployments/gateway.env.example b/gateway/deployments/gateway.env.example index 5bf90b5..c23162f 100644 --- a/gateway/deployments/gateway.env.example +++ b/gateway/deployments/gateway.env.example @@ -33,6 +33,14 @@ INDEXQUBE_ENABLE_BLOCK_OPTIMIZER=false # Code retry bursts from hammering the provider with huge repeated payloads. INDEXQUBE_RATE_LIMIT_COOLDOWN=30s +# --- Local runaway loop guard (circuit breaker for repeated large similar requests) --- +# IQ_ALLOW_RUNAWAY=1 bypasses this guard for the current session. +# INDEXQUBE_CIRCUIT_BREAKER_ENABLED=true +# INDEXQUBE_CIRCUIT_WINDOW_SECONDS=60 +# INDEXQUBE_CIRCUIT_MAX_SIMILAR=15 +# INDEXQUBE_CIRCUIT_MIN_TOKENS=50000 +# INDEXQUBE_CIRCUIT_RETRY_AFTER_SECONDS=60 + # --- AWS Bedrock backend (replaces direct Anthropic API calls for /v1/messages) --- # Set INDEXQUBE_BEDROCK_ENABLED=true to route /v1/messages through Bedrock instead. # Credentials come from the standard AWS SDK chain: env vars (AWS_ACCESS_KEY_ID / diff --git a/gateway/go.mod b/gateway/go.mod index 6867709..d02c5c3 100644 --- a/gateway/go.mod +++ b/gateway/go.mod @@ -32,6 +32,7 @@ require ( github.com/beorn7/perks v1.0.1 // indirect github.com/cenkalti/backoff/v5 v5.0.3 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/go-logr/logr v1.4.3 // indirect github.com/go-logr/stdr v1.2.2 // indirect @@ -41,10 +42,13 @@ require ( github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/pgx/v5 v5.9.2 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect github.com/prometheus/client_model v0.6.2 // indirect github.com/prometheus/common v0.66.1 // indirect github.com/prometheus/procfs v0.16.1 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.43.0 // indirect go.opentelemetry.io/otel/metric v1.43.0 // indirect @@ -58,4 +62,8 @@ require ( google.golang.org/genproto/googleapis/rpc v0.0.0-20260401024825-9d38bb4040a9 // indirect google.golang.org/grpc v1.80.0 // indirect google.golang.org/protobuf v1.36.11 // indirect + modernc.org/libc v1.72.3 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect + modernc.org/sqlite v1.50.1 // indirect ) diff --git a/gateway/go.sum b/gateway/go.sum index af43ea3..84ccd1e 100644 --- a/gateway/go.sum +++ b/gateway/go.sum @@ -41,6 +41,8 @@ github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XL github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +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/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= @@ -72,8 +74,12 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= +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/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= +github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o= @@ -84,6 +90,8 @@ github.com/prometheus/common v0.66.1 h1:h5E0h5/Y8niHc5DlaLlWLArTQI7tMrsfQjHV+d9Z github.com/prometheus/common v0.66.1/go.mod h1:gcaUsgf3KfRSwHY4dIMXLPV0K/Wg1oZ8+SbZk/HH/dA= github.com/prometheus/procfs v0.16.1 h1:hZ15bTNuirocR6u0JZ6BAHHmwS1p8B4P6MRqxtzMyRg= github.com/prometheus/procfs v0.16.1/go.mod h1:teAbpZRB1iIAJYREa1LsoWUXykVXA1KlTmWl8x/U+Is= +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/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= @@ -119,6 +127,7 @@ golang.org/x/net v0.52.0 h1:He/TN1l0e4mmR3QqHMT2Xab3Aj3L9qjbhRm78/6jrW0= golang.org/x/net v0.52.0/go.mod h1:R1MAz7uMZxVMualyPXb+VaqGSa3LIaUqk0eEt3w36Sw= golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.35.0 h1:JOVx6vVDFokkpaq1AEptVzLTpDe9KGpj5tR4/X+ybL8= @@ -139,3 +148,11 @@ gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EV gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +modernc.org/libc v1.72.3 h1:ZnDF4tXn4NBXFutMMQC4vtbTFSXhhKzR73fv0beZEAU= +modernc.org/libc v1.72.3/go.mod h1:dn0dZNnnn1clLyvRxLxYExxiKRZIRENOfqQ8XEeg4Qs= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/sqlite v1.50.1 h1:l+cQvn0sd0zJJtfygGHuQJ5AjlrwXmWPw4KP3ZMwr9w= +modernc.org/sqlite v1.50.1/go.mod h1:tcNzv5p84E0skkmJn038y+hWJbLQXQqEnQfeh5r2JLM= diff --git a/gateway/internal/chunker/chunker.go b/gateway/internal/chunker/chunker.go new file mode 100644 index 0000000..8501094 --- /dev/null +++ b/gateway/internal/chunker/chunker.go @@ -0,0 +1,223 @@ +// Package chunker implements Rabin-Karp content-defined chunking (CDC). +// +// Content-defined chunking splits a byte stream at positions determined by +// the content itself rather than by fixed offsets. This means that inserting +// or removing bytes near the beginning of the stream shifts the boundary of +// the affected chunk but leaves most downstream chunks unchanged — the +// "sliding-window" property that makes CDC ideal for deduplication and +// semantic caching of LLM prompts. +// +// # Algorithm +// +// A polynomial rolling hash is maintained over a sliding window of W bytes: +// +// H[l,r) = data[l]·B^(W-1) + data[l+1]·B^(W-2) + … + data[r-1]·B^0 +// +// When the window advances by one position (removing data[l], adding data[r]): +// +// H[l+1,r+1) = H[l,r) · B − data[l] · B^W + data[r] +// +// A chunk boundary is declared at byte i when H & Mask == 0. +// MinSize and MaxSize clamp the search range so that no chunk is +// pathologically small or large regardless of the data distribution. +package chunker + +// Default tuning constants. Mask gives ~4 KB average chunk size. +const ( + defaultBase uint64 = 257 + defaultWindowSize int = 64 + defaultMinSize int = 512 + defaultMaxSize int = 8 * 1024 // 8 KiB hard cap + defaultMask uint64 = (1 << 12) - 1 // boundary when lower 12 bits are zero → ~4 KiB avg +) + +// Config holds the CDC tuning parameters. +// All size values are in bytes. +type Config struct { + // WindowSize is the number of bytes covered by the rolling hash at any + // given position. Larger windows reduce false positives from short + // repeated sequences. Typical range: 32–128. + WindowSize int + + // MinSize is the minimum chunk length. Boundaries are not considered + // before this many bytes have been consumed. The very last chunk may + // be shorter if the remaining data is exhausted. + MinSize int + + // MaxSize is the hard upper bound on chunk length. A boundary is forced + // here even if the hash condition is never satisfied. + MaxSize int + + // Mask is the bit-mask used to detect boundaries: a cut is made when + // (hash & Mask) == 0. The average chunk size is approximately + // (MaxSize + MinSize) / 2 when Mask ≈ TargetSize - 1. + // Use a value of the form (1<= cfg.MaxSize { + panic("chunker: MinSize must be less than MaxSize") + } + if cfg.Base == 0 { + panic("chunker: Base must be non-zero") + } + + // Precompute Base^WindowSize (mod 2^64). + // This is the coefficient we subtract when the oldest byte leaves the window. + pop := uint64(1) + for i := 0; i < cfg.WindowSize; i++ { + pop *= cfg.Base + } + + return &Chunker{cfg: cfg, pop: pop} +} + +// Default returns a Chunker with DefaultConfig(). +func Default() *Chunker { + return New(DefaultConfig()) +} + +// Split partitions data into content-defined chunks. +// +// Properties guaranteed by this implementation: +// - Completeness: the concatenation of all Chunk.Data equals data. +// - Offset correctness: chunks[i].Offset + len(chunks[i].Data) == chunks[i+1].Offset. +// - Min/Max bounds: len(chunk.Data) <= MaxSize for every chunk; only the +// final chunk may be shorter than MinSize (when remaining bytes < MinSize). +// - Determinism: identical inputs always produce identical chunk boundaries. +// +// Returns nil for empty input. +func (c *Chunker) Split(data []byte) []Chunk { + if len(data) == 0 { + return nil + } + // Pre-allocate enough capacity for the worst case (every chunk == MinSize). + // This guarantees at most 1 allocation regardless of the actual chunk count, + // at the cost of slightly over-allocating when chunks are large. + estimated := len(data)/c.cfg.MinSize + 1 + return c.SplitInto(data, make([]Chunk, 0, estimated)) +} + +// SplitInto is the zero-allocation variant of Split for hot-path callers. +// It reuses dst as the backing slice, appending to dst[:0] each call, so a +// caller that pre-allocates dst once (e.g. via sync.Pool) incurs 0 heap +// allocations per split regardless of how many chunks are produced — as long +// as dst has sufficient capacity. +// +// var pool = sync.Pool{New: func() any { return make([]chunker.Chunk, 0, 32) }} +// +// dst := pool.Get().([]chunker.Chunk) +// dst = c.SplitInto(data, dst[:0]) +// // … use dst … +// pool.Put(dst) +func (c *Chunker) SplitInto(data []byte, dst []Chunk) []Chunk { + if len(data) == 0 { + return dst + } + dst = dst[:0] + offset := 0 + for offset < len(data) { + end := c.nextBoundary(data, offset) + dst = append(dst, Chunk{ + Data: data[offset:end], + Offset: offset, + }) + offset = end + } + return dst +} + +// nextBoundary returns the exclusive end index of the next chunk +// starting at data[start]. +func (c *Chunker) nextBoundary(data []byte, start int) int { + n := len(data) + + // Hard boundaries derived from the min/max config. + minEnd := start + c.cfg.MinSize + maxEnd := start + c.cfg.MaxSize + + maxEnd = min(maxEnd, n) + // If fewer bytes remain than MinSize, emit the rest as a single short chunk. + if minEnd >= n { + return n + } + + // ── Phase 1: fill the initial rolling-hash window ────────────────────── + // + // The window should cover the W bytes ending just before minEnd, i.e., + // [minEnd-W, minEnd). If that range extends before `start` (when + // MinSize < WindowSize), clamp it to `start`. + hash := uint64(0) + winStart := max(minEnd-c.cfg.WindowSize, start) + for i := winStart; i < minEnd; i++ { + hash = hash*c.cfg.Base + uint64(data[i]) + } + + // ── Phase 2: slide the window, looking for a boundary ───────────────── + // + // At each step we advance the window by one byte: + // • incoming byte: data[i] + // • outgoing byte: data[i-WindowSize] (if it is within the chunk) + // + // Update rule: H_new = H_old · B − data[out] · B^W + data[in] + for i := minEnd; i < maxEnd; i++ { + out := i - c.cfg.WindowSize + if out >= start { + // Full window: evict the outgoing byte using the precomputed pop. + hash = hash*c.cfg.Base + uint64(data[i]) - c.pop*uint64(data[out]) + } else { + // Window is not yet fully formed (only when WindowSize > MinSize). + hash = hash*c.cfg.Base + uint64(data[i]) + } + + if hash&c.cfg.Mask == 0 { + return i + 1 + } + } + + // No boundary found before MaxSize — force a cut at the hard cap. + return maxEnd +} diff --git a/gateway/internal/chunker/chunker_bench_test.go b/gateway/internal/chunker/chunker_bench_test.go new file mode 100644 index 0000000..c326efb --- /dev/null +++ b/gateway/internal/chunker/chunker_bench_test.go @@ -0,0 +1,141 @@ +package chunker_test + +import ( + "bytes" + "testing" + + "github.com/Revanth14/indexqube/gateway/internal/chunker" +) + +// pseudoRandData generates n bytes of deterministic pseudo-random data using +// an xorshift64 PRNG. Determinism makes benchmark results reproducible across +// runs; the data is varied enough that the Rabin-Karp hash finds real boundaries. +func pseudoRandData(n int) []byte { + data := make([]byte, n) + x := uint64(0x123456789abcdef0) + for i := range data { + x ^= x << 13 + x ^= x >> 7 + x ^= x << 17 + data[i] = byte(x) + } + return data +} + +// ── Split (1 alloc: the result slice header) ────────────────────────────────── + +func BenchmarkSplit_4KB(b *testing.B) { + benchSplit(b, 4*1024) +} + +func BenchmarkSplit_64KB(b *testing.B) { + benchSplit(b, 64*1024) +} + +func BenchmarkSplit_1MB(b *testing.B) { + benchSplit(b, 1024*1024) +} + +func BenchmarkSplit_8MB(b *testing.B) { + benchSplit(b, 8*1024*1024) +} + +// BenchmarkSplit_WorstCase uses a run of identical bytes, which is the +// degenerate case for content-defined chunking: the hash never hits a boundary +// naturally so every chunk is forced at MaxSize. +func BenchmarkSplit_WorstCase_1MB(b *testing.B) { + data := bytes.Repeat([]byte{0x42}, 1024*1024) + c := chunker.Default() + b.SetBytes(int64(len(data))) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = c.Split(data) + } +} + +func benchSplit(b *testing.B, size int) { + b.Helper() + data := pseudoRandData(size) + c := chunker.Default() + b.SetBytes(int64(size)) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = c.Split(data) + } +} + +// ── SplitInto — zero allocations on the hot path ───────────────────────────── + +func BenchmarkSplitInto_4KB(b *testing.B) { + benchSplitInto(b, 4*1024) +} + +func BenchmarkSplitInto_64KB(b *testing.B) { + benchSplitInto(b, 64*1024) +} + +func BenchmarkSplitInto_1MB(b *testing.B) { + benchSplitInto(b, 1024*1024) +} + +func BenchmarkSplitInto_8MB(b *testing.B) { + benchSplitInto(b, 8*1024*1024) +} + +func benchSplitInto(b *testing.B, size int) { + b.Helper() + data := pseudoRandData(size) + c := chunker.Default() + // Pre-warm: run once before ResetTimer so dst grows to its steady-state + // capacity and subsequent iterations are truly 0-alloc. + dst := make([]chunker.Chunk, 0, 16) + dst = c.SplitInto(data, dst[:0]) + + b.SetBytes(int64(size)) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + dst = c.SplitInto(data, dst[:0]) + } +} + +// ── AllocsPerRun assertions (machine-verifiable 0-alloc guarantee) ──────────── + +// TestSplitIntoZeroAllocs asserts that SplitInto with a pre-warmed dst slice +// incurs zero heap allocations per call. +func TestSplitIntoZeroAllocs(t *testing.T) { + c := chunker.Default() + data := pseudoRandData(256 * 1024) + + // Warm up: let Split determine the typical chunk count so dst has capacity. + initial := c.Split(data) + dst := make([]chunker.Chunk, 0, len(initial)+4) + + allocs := testing.AllocsPerRun(20, func() { + dst = c.SplitInto(data, dst[:0]) + }) + if allocs > 0 { + t.Errorf("SplitInto: %.0f allocs/op; want 0 (pre-warmed dst must not grow)", allocs) + } + // Sanity: we still produced the right number of chunks. + if len(dst) != len(initial) { + t.Errorf("SplitInto produced %d chunks; Split produced %d", len(dst), len(initial)) + } +} + +// TestSplitOneAllocPerCall asserts that Split itself makes exactly 1 allocation +// per call (the result-slice backing array), regardless of input size. +func TestSplitOneAllocPerCall(t *testing.T) { + c := chunker.Default() + for _, size := range []int{4 * 1024, 64 * 1024, 1024 * 1024} { + data := pseudoRandData(size) + allocs := testing.AllocsPerRun(10, func() { + _ = c.Split(data) + }) + if allocs > 1 { + t.Errorf("Split(%d bytes): %.0f allocs/op; want ≤ 1", size, allocs) + } + } +} diff --git a/gateway/internal/chunker/chunker_test.go b/gateway/internal/chunker/chunker_test.go new file mode 100644 index 0000000..d14b93d --- /dev/null +++ b/gateway/internal/chunker/chunker_test.go @@ -0,0 +1,312 @@ +package chunker_test + +import ( + "bytes" + "crypto/rand" + "testing" + + "github.com/Revanth14/indexqube/gateway/internal/chunker" +) + +// randBytes returns n cryptographically random bytes. +func randBytes(t *testing.T, n int) []byte { + t.Helper() + b := make([]byte, n) + if _, err := rand.Read(b); err != nil { + t.Fatalf("rand.Read: %v", err) + } + return b +} + +// repeatByte returns a slice of n copies of b. +func repeatByte(b byte, n int) []byte { + return bytes.Repeat([]byte{b}, n) +} + +// ── Construction ───────────────────────────────────────────────────────────── + +func TestDefaultConfigValues(t *testing.T) { + cfg := chunker.DefaultConfig() + if cfg.WindowSize <= 0 { + t.Errorf("WindowSize = %d; want > 0", cfg.WindowSize) + } + if cfg.MinSize <= 0 { + t.Errorf("MinSize = %d; want > 0", cfg.MinSize) + } + if cfg.MaxSize <= cfg.MinSize { + t.Errorf("MaxSize (%d) <= MinSize (%d)", cfg.MaxSize, cfg.MinSize) + } + if cfg.Mask == 0 { + t.Error("Mask = 0; want non-zero") + } + if cfg.Base == 0 { + t.Error("Base = 0; want non-zero") + } +} + +func TestNewPanicsOnBadConfig(t *testing.T) { + cases := []struct { + name string + cfg chunker.Config + }{ + {"zero WindowSize", chunker.Config{WindowSize: 0, MinSize: 512, MaxSize: 4096, Base: 257, Mask: 0xfff}}, + {"zero MinSize", chunker.Config{WindowSize: 64, MinSize: 0, MaxSize: 4096, Base: 257, Mask: 0xfff}}, + {"zero MaxSize", chunker.Config{WindowSize: 64, MinSize: 512, MaxSize: 0, Base: 257, Mask: 0xfff}}, + {"MinSize == MaxSize", chunker.Config{WindowSize: 64, MinSize: 512, MaxSize: 512, Base: 257, Mask: 0xfff}}, + {"MinSize > MaxSize", chunker.Config{WindowSize: 64, MinSize: 4096, MaxSize: 512, Base: 257, Mask: 0xfff}}, + {"zero Base", chunker.Config{WindowSize: 64, MinSize: 512, MaxSize: 4096, Base: 0, Mask: 0xfff}}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + defer func() { + if r := recover(); r == nil { + t.Errorf("New(%+v) did not panic; expected a panic", tc.cfg) + } + }() + chunker.New(tc.cfg) + }) + } +} + +// ── Edge cases ─────────────────────────────────────────────────────────────── + +func TestSplitEmpty(t *testing.T) { + c := chunker.Default() + if got := c.Split(nil); got != nil { + t.Errorf("Split(nil) = %v; want nil", got) + } + if got := c.Split([]byte{}); got != nil { + t.Errorf("Split([]byte{}) = %v; want nil", got) + } +} + +func TestSplitSingleByteInputs(t *testing.T) { + c := chunker.Default() + for _, size := range []int{1, 10, 100} { + data := randBytes(t, size) + chunks := c.Split(data) + if len(chunks) != 1 { + t.Errorf("len(Split(%d bytes)) = %d; want 1", size, len(chunks)) + } + } +} + +// Input exactly equal to MinSize must produce a single chunk. +func TestSplitExactlyMinSize(t *testing.T) { + cfg := smallConfig() + c := chunker.New(cfg) + data := randBytes(t, cfg.MinSize) + chunks := c.Split(data) + if len(chunks) != 1 { + t.Errorf("got %d chunks for input == MinSize; want 1", len(chunks)) + } +} + +// Input exactly equal to MaxSize may or may not split, but must never panic +// and must always reconstitute correctly. +func TestSplitExactlyMaxSize(t *testing.T) { + cfg := smallConfig() + c := chunker.New(cfg) + data := randBytes(t, cfg.MaxSize) + chunks := c.Split(data) + assertReconstitutes(t, data, chunks) +} + +// ── Core properties ────────────────────────────────────────────────────────── + +// TestSplitReconstitutes verifies that concatenating all chunks produces the +// original input byte-for-byte. +func TestSplitReconstitutes(t *testing.T) { + c := chunker.Default() + sizes := []int{1, 512, 1024, 4*1024, 32*1024, 256*1024} + for _, sz := range sizes { + data := randBytes(t, sz) + chunks := c.Split(data) + assertReconstitutes(t, data, chunks) + } +} + +// TestSplitDeterministic verifies that identical inputs always produce +// identical chunk boundaries. +func TestSplitDeterministic(t *testing.T) { + c := chunker.Default() + data := randBytes(t, 64*1024) + + first := c.Split(data) + for i := 0; i < 5; i++ { + repeat := c.Split(data) + if len(repeat) != len(first) { + t.Fatalf("run %d: got %d chunks; want %d", i, len(repeat), len(first)) + } + for j, ch := range repeat { + if ch.Offset != first[j].Offset || len(ch.Data) != len(first[j].Data) { + t.Fatalf("run %d chunk %d: offset=%d len=%d; want offset=%d len=%d", + i, j, ch.Offset, len(ch.Data), first[j].Offset, len(first[j].Data)) + } + } + } +} + +// TestSplitMinSizeRespected verifies no chunk (except possibly the last) +// is shorter than MinSize. +func TestSplitMinSizeRespected(t *testing.T) { + cfg := smallConfig() + c := chunker.New(cfg) + data := randBytes(t, 128*1024) + chunks := c.Split(data) + + for i, ch := range chunks[:len(chunks)-1] { + if len(ch.Data) < cfg.MinSize { + t.Errorf("chunk %d: len=%d < MinSize=%d", i, len(ch.Data), cfg.MinSize) + } + } +} + +// TestSplitMaxSizeRespected verifies no chunk exceeds MaxSize. +func TestSplitMaxSizeRespected(t *testing.T) { + cfg := smallConfig() + c := chunker.New(cfg) + // Use a run of identical bytes — worst case for hash collisions. + data := repeatByte(0xAB, 512*1024) + chunks := c.Split(data) + + for i, ch := range chunks { + if len(ch.Data) > cfg.MaxSize { + t.Errorf("chunk %d: len=%d > MaxSize=%d", i, len(ch.Data), cfg.MaxSize) + } + } +} + +// TestSplitOffsetsConsistent verifies that Chunk.Offset and len(Chunk.Data) +// are consistent: chunks tile the input with no gaps or overlaps. +func TestSplitOffsetsConsistent(t *testing.T) { + c := chunker.Default() + data := randBytes(t, 64*1024) + chunks := c.Split(data) + + cursor := 0 + for i, ch := range chunks { + if ch.Offset != cursor { + t.Fatalf("chunk %d: Offset=%d; want %d", i, ch.Offset, cursor) + } + cursor += len(ch.Data) + } + if cursor != len(data) { + t.Errorf("total bytes covered = %d; want %d", cursor, len(data)) + } +} + +// TestSplitDataIsSubslice verifies that chunk.Data points into the original +// slice (zero-copy guarantee). +func TestSplitDataIsSubslice(t *testing.T) { + c := chunker.Default() + data := randBytes(t, 32*1024) + chunks := c.Split(data) + + for i, ch := range chunks { + // If Data is a true sub-slice, its first byte's address must fall + // within the original backing array. + if len(ch.Data) == 0 { + continue + } + expected := data[ch.Offset : ch.Offset+len(ch.Data)] + if !bytes.Equal(ch.Data, expected) { + t.Errorf("chunk %d data mismatch", i) + } + } +} + +// ── Content-defined property ────────────────────────────────────────────────── + +// TestSplitContentDefined is the key property of CDC: prepending N bytes to +// the input shifts chunk boundaries only near the insertion point; after a +// few chunks the boundaries resynchronise and the tail chunks are identical. +// +// We verify that at least half the chunks from the original split are +// reproduced somewhere in the modified split, confirming partial stability. +func TestSplitContentDefined(t *testing.T) { + cfg := smallConfig() + c := chunker.New(cfg) + + // Large input so there are many downstream chunks to check. + original := randBytes(t, 512*1024) + prefix := randBytes(t, cfg.MinSize/2) // a chunk-boundary-crossing insertion + + modified := append(prefix, original...) + + origChunks := c.Split(original) + modChunks := c.Split(modified) + + if len(origChunks) < 4 { + t.Skip("too few chunks to test content-defined property meaningfully") + } + + // Build a set of (offset-normalised) chunk fingerprints from modChunks. + // We normalise by hashing the Data bytes, not the Offset. + modSet := make(map[string]struct{}, len(modChunks)) + for _, ch := range modChunks { + modSet[string(ch.Data)] = struct{}{} + } + + // Count how many of the original tail chunks appear verbatim in modChunks. + // Skip the first two chunks (most likely to straddle the insertion point). + matched := 0 + tail := origChunks[2:] + for _, ch := range tail { + if _, ok := modSet[string(ch.Data)]; ok { + matched++ + } + } + + // Expect at least 50% of tail chunks to reappear unchanged. + threshold := len(tail) / 2 + if matched < threshold { + t.Errorf("content-defined property: %d/%d tail chunks matched; want ≥%d", + matched, len(tail), threshold) + } +} + +// ── Identical-byte stress test ──────────────────────────────────────────────── + +// TestSplitRepeatedBytes verifies correct behaviour when every byte is the +// same — a degenerate case for hash functions. +func TestSplitRepeatedBytes(t *testing.T) { + cfg := smallConfig() + c := chunker.New(cfg) + data := repeatByte(0xFF, 256*1024) + chunks := c.Split(data) + + assertReconstitutes(t, data, chunks) + for i, ch := range chunks { + if len(ch.Data) > cfg.MaxSize { + t.Errorf("chunk %d: len=%d > MaxSize=%d", i, len(ch.Data), cfg.MaxSize) + } + } +} + +// ── Helpers ─────────────────────────────────────────────────────────────────── + +// smallConfig returns a Config with small MinSize/MaxSize so tests run fast +// and produce multiple chunks without needing megabytes of data. +func smallConfig() chunker.Config { + return chunker.Config{ + WindowSize: 16, + MinSize: 64, + MaxSize: 512, + Mask: (1 << 7) - 1, // ~128 byte average chunk + Base: 257, + } +} + +// assertReconstitutes checks that concatenating chunks byte-for-byte reproduces data. +func assertReconstitutes(t *testing.T, data []byte, chunks []chunker.Chunk) { + t.Helper() + var buf bytes.Buffer + for _, ch := range chunks { + buf.Write(ch.Data) + } + if !bytes.Equal(buf.Bytes(), data) { + t.Errorf("reconstituted %d bytes; original was %d bytes", buf.Len(), len(data)) + } +} diff --git a/gateway/internal/guard/budget.go b/gateway/internal/guard/budget.go new file mode 100644 index 0000000..8cbd816 --- /dev/null +++ b/gateway/internal/guard/budget.go @@ -0,0 +1,62 @@ +package guard + +import ( + "net/http" + "os" + "sync" +) + +type BudgetGuard struct { + mu sync.Mutex + cfg Config + spentUSD map[string]float64 +} + +func NewBudgetGuard(cfg Config) *BudgetGuard { + return &BudgetGuard{ + cfg: cfg, + spentUSD: make(map[string]float64), + } +} + +func (bg *BudgetGuard) Check(sig RequestSignal) Decision { + if os.Getenv(envDevMode) == "1" || os.Getenv(envAllowOverBudget) == "1" || os.Getenv(envAllowRunaway) == "1" { + return Decision{Allow: true, Reason: "override"} + } + + key := sig.SessionKey + if key == "" { + key = "no-session" + } + + bg.mu.Lock() + spent := bg.spentUSD[key] + projected := spent + sig.EstimatedCostUSD + bg.spentUSD[key] = projected + bg.mu.Unlock() + + budget := bg.cfg.SessionBudgetUSD + decision := Decision{ + BudgetUSD: budget, + SpentUSD: spent, + ProjectedUSD: projected, + } + + if projected >= budget*bg.cfg.BudgetHardRatio { + decision.Allow = false + decision.StatusCode = http.StatusTooManyRequests + decision.Reason = "budget_exceeded" + return decision + } + + if projected >= budget*bg.cfg.BudgetWarnRatio { + decision.Allow = true + decision.Warn = true + decision.Reason = "budget_warning" + return decision + } + + decision.Allow = true + decision.Reason = "allowed" + return decision +} diff --git a/gateway/internal/guard/circuit.go b/gateway/internal/guard/circuit.go new file mode 100644 index 0000000..97fd44f --- /dev/null +++ b/gateway/internal/guard/circuit.go @@ -0,0 +1,140 @@ +package guard + +import ( + "net/http" + "os" + "strings" + "sync" + "time" +) + +const maxCircuitKeys = 10_000 + +type CircuitBreaker struct { + mu sync.Mutex + cfg Config + events map[string][]event +} + +type event struct { + At time.Time + Tokens int +} + +func NewCircuitBreaker(cfg Config) *CircuitBreaker { + return &CircuitBreaker{ + cfg: cfg, + events: make(map[string][]event), + } +} + +func (cb *CircuitBreaker) Check(sig RequestSignal) (decision Decision) { + defer func() { + if recover() != nil { + decision = Decision{Allow: true, Reason: "guard_error"} + } + }() + if cb == nil { + return Decision{Allow: true, Reason: "disabled"} + } + if !cb.cfg.CircuitEnabled { + return Decision{Allow: true, Reason: "disabled"} + } + if os.Getenv(envAllowRunaway) == "1" { + return Decision{Allow: true, Reason: "override"} + } + if sig.AttemptedTokens < cb.cfg.CircuitMinAttemptedTokens { + return Decision{Allow: true, Reason: "below_min_tokens"} + } + + now := sig.Now + if now.IsZero() { + now = time.Now() + } + key := makeKey(sig) + cutoff := now.Add(-cb.cfg.CircuitWindow) + + cb.mu.Lock() + defer cb.mu.Unlock() + + events := pruneBefore(cb.events[key], cutoff) + count := len(events) + 1 + if count > cb.cfg.CircuitMaxSimilarRequests { + cb.events[key] = events + return Decision{ + Allow: false, + StatusCode: http.StatusTooManyRequests, + Reason: "similar_large_requests", + RetryAfter: cb.cfg.CircuitRetryAfter, + Count: count, + Remaining: 0, + } + } + + events = append(events, event{At: now, Tokens: sig.AttemptedTokens}) + cb.events[key] = events + if len(cb.events) > maxCircuitKeys { + cb.pruneMap(cutoff) + } + + remaining := cb.cfg.CircuitMaxSimilarRequests - count + return Decision{ + Allow: true, + Warn: remaining <= 3, + Reason: "allowed", + Count: count, + Remaining: remaining, + } +} + +func (cb *CircuitBreaker) pruneMap(cutoff time.Time) { + for key, events := range cb.events { + events = pruneBefore(events, cutoff) + if len(events) == 0 { + delete(cb.events, key) + if len(cb.events) <= maxCircuitKeys { + return + } + continue + } + cb.events[key] = events + } +} + +func pruneBefore(events []event, cutoff time.Time) []event { + if len(events) == 0 { + return events + } + idx := 0 + for idx < len(events) && events[idx].At.Before(cutoff) { + idx++ + } + if idx == 0 { + return events + } + if idx >= len(events) { + return nil + } + return events[idx:] +} + +func makeKey(sig RequestSignal) string { + parts := []string{ + safePart(sig.MachineID), + safePart(sig.SessionKey), + safePart(sig.Route), + safePart(sig.Model), + safePart(sig.Fingerprint), + } + if sig.SessionKey == "" { + parts[1] = "no-session" + } + return strings.Join(parts, "|") +} + +func safePart(s string) string { + if strings.TrimSpace(s) == "" { + return "_" + } + return strings.TrimSpace(s) +} diff --git a/gateway/internal/guard/circuit_test.go b/gateway/internal/guard/circuit_test.go new file mode 100644 index 0000000..a19721a --- /dev/null +++ b/gateway/internal/guard/circuit_test.go @@ -0,0 +1,151 @@ +package guard + +import ( + "testing" + "time" +) + +func TestAllowsBelowMinTokens(t *testing.T) { + t.Parallel() + cb := NewCircuitBreaker(Config{ + CircuitEnabled: true, + CircuitWindow: time.Minute, + CircuitMaxSimilarRequests: 2, + CircuitMinAttemptedTokens: 100, + CircuitRetryAfter: time.Minute, + }) + + got := cb.Check(RequestSignal{ + MachineID: "m", + SessionKey: "s", + Route: "/v1/messages", + Model: "claude", + Fingerprint: "fp", + AttemptedTokens: 99, + Now: time.Unix(1, 0), + }) + if !got.Allow || got.Reason != "below_min_tokens" { + t.Fatalf("decision=%+v, want allow below_min_tokens", got) + } +} + +func TestBlocksAfterMaxSimilarRequests(t *testing.T) { + t.Parallel() + cb := NewCircuitBreaker(Config{ + CircuitEnabled: true, + CircuitWindow: 10 * time.Second, + CircuitMaxSimilarRequests: 2, + CircuitMinAttemptedTokens: 100, + CircuitRetryAfter: 2 * time.Second, + }) + base := time.Unix(10, 0) + sig := RequestSignal{ + MachineID: "m", + SessionKey: "s", + Route: "/v1/messages", + Model: "claude", + Fingerprint: "fp", + AttemptedTokens: 1000, + } + + first := cb.Check(withNow(sig, base)) + second := cb.Check(withNow(sig, base.Add(1*time.Second))) + third := cb.Check(withNow(sig, base.Add(2*time.Second))) + if !first.Allow || !second.Allow { + t.Fatalf("first=%+v second=%+v, expected allowed", first, second) + } + if third.Allow || third.StatusCode != 429 || third.Reason != "similar_large_requests" { + t.Fatalf("third=%+v, expected blocked 429 similar_large_requests", third) + } +} + +func TestWindowExpiryAllowsAgain(t *testing.T) { + t.Parallel() + cb := NewCircuitBreaker(Config{ + CircuitEnabled: true, + CircuitWindow: 1 * time.Second, + CircuitMaxSimilarRequests: 1, + CircuitMinAttemptedTokens: 100, + CircuitRetryAfter: 2 * time.Second, + }) + sig := RequestSignal{ + MachineID: "m", + SessionKey: "s", + Route: "/v1/messages", + Model: "claude", + Fingerprint: "fp", + AttemptedTokens: 1000, + } + + first := cb.Check(withNow(sig, time.Unix(1, 0))) + blocked := cb.Check(withNow(sig, time.Unix(1, 500*int64(time.Millisecond)))) + again := cb.Check(withNow(sig, time.Unix(8, 0))) + if !first.Allow { + t.Fatalf("first=%+v, expected allowed", first) + } + if blocked.Allow { + t.Fatalf("blocked=%+v, expected blocked", blocked) + } + if !again.Allow { + t.Fatalf("again=%+v, expected allowed after window expiry", again) + } +} + +func TestDifferentFingerprintsDoNotShareLimit(t *testing.T) { + t.Parallel() + cb := NewCircuitBreaker(Config{ + CircuitEnabled: true, + CircuitWindow: time.Minute, + CircuitMaxSimilarRequests: 1, + CircuitMinAttemptedTokens: 100, + CircuitRetryAfter: time.Minute, + }) + base := RequestSignal{ + MachineID: "m", + SessionKey: "s", + Route: "/v1/messages", + Model: "claude", + AttemptedTokens: 1000, + } + + _ = cb.Check(withNow(withFingerprint(base, "fp-1"), time.Unix(1, 0))) + secondKey := cb.Check(withNow(withFingerprint(base, "fp-2"), time.Unix(2, 0))) + if !secondKey.Allow { + t.Fatalf("second key decision=%+v, expected allowed", secondKey) + } +} + +func TestOverrideAllowsRunaway(t *testing.T) { + t.Setenv(envAllowRunaway, "1") + cb := NewCircuitBreaker(Config{ + CircuitEnabled: true, + CircuitWindow: time.Minute, + CircuitMaxSimilarRequests: 1, + CircuitMinAttemptedTokens: 100, + CircuitRetryAfter: time.Minute, + }) + sig := RequestSignal{ + MachineID: "m", + SessionKey: "s", + Route: "/v1/messages", + Model: "claude", + Fingerprint: "fp", + AttemptedTokens: 1000, + Now: time.Unix(1, 0), + } + + got := cb.Check(sig) + if !got.Allow || got.Reason != "override" { + t.Fatalf("decision=%+v, want allow override", got) + } +} + +func withNow(sig RequestSignal, now time.Time) RequestSignal { + sig.Now = now + return sig +} + +func withFingerprint(sig RequestSignal, fp string) RequestSignal { + sig.Fingerprint = fp + return sig +} diff --git a/gateway/internal/guard/config.go b/gateway/internal/guard/config.go new file mode 100644 index 0000000..b5e2fe7 --- /dev/null +++ b/gateway/internal/guard/config.go @@ -0,0 +1,130 @@ +package guard + +import ( + "os" + "strconv" + "time" +) + +const ( + envCircuitEnabled = "INDEXQUBE_CIRCUIT_BREAKER_ENABLED" + envWindowSeconds = "INDEXQUBE_CIRCUIT_WINDOW_SECONDS" + envMaxSimilar = "INDEXQUBE_CIRCUIT_MAX_SIMILAR" + envMinTokens = "INDEXQUBE_CIRCUIT_MIN_TOKENS" + envRetryAfterSeconds = "INDEXQUBE_CIRCUIT_RETRY_AFTER_SECONDS" + envAllowRunaway = "IQ_ALLOW_RUNAWAY" + + // Budget Config + envSessionBudgetUSD = "INDEXQUBE_SESSION_BUDGET_USD" + envBudgetWarnRatio = "INDEXQUBE_BUDGET_WARN_RATIO" + envBudgetHardRatio = "INDEXQUBE_BUDGET_HARD_RATIO" + envAllowOverBudget = "IQ_ALLOW_OVER_BUDGET" + + // Velocity Config + envSpendVelocityEnabled = "INDEXQUBE_SPEND_VELOCITY_ENABLED" + envSpendVelocityWindowSecs = "INDEXQUBE_SPEND_VELOCITY_WINDOW_SECONDS" + envSpendVelocityWarnUSD = "INDEXQUBE_SPEND_VELOCITY_WARN_USD" + envSpendVelocityBlockUSD = "INDEXQUBE_SPEND_VELOCITY_BLOCK_USD" + + // Dev mode: set by `iq claude --dev` — disables velocity guard entirely. + envDevMode = "IQ_DEV_MODE" +) + +type Config struct { + CircuitEnabled bool + CircuitWindow time.Duration + CircuitMaxSimilarRequests int + CircuitMinAttemptedTokens int + CircuitRetryAfter time.Duration + + SessionBudgetUSD float64 + BudgetWarnRatio float64 + BudgetHardRatio float64 + + SpendVelocityEnabled bool + SpendVelocityWindow time.Duration + SpendVelocityWarnUSD float64 + SpendVelocityBlockUSD float64 +} + +func DefaultConfig() Config { + return Config{ + CircuitEnabled: true, + CircuitWindow: 60 * time.Second, + CircuitMaxSimilarRequests: 15, + CircuitMinAttemptedTokens: 50_000, + CircuitRetryAfter: 60 * time.Second, + + SessionBudgetUSD: 1.00, + BudgetWarnRatio: 0.80, + BudgetHardRatio: 1.00, + + SpendVelocityEnabled: true, + SpendVelocityWindow: 120 * time.Second, + SpendVelocityWarnUSD: 0.75, + SpendVelocityBlockUSD: 2.00, + } +} + +func FromEnv() Config { + cfg := DefaultConfig() + + // Circuit Breaker + cfg.CircuitEnabled = envBool(envCircuitEnabled, cfg.CircuitEnabled) + cfg.CircuitWindow = time.Duration(envInt(envWindowSeconds, int(cfg.CircuitWindow.Seconds()))) * time.Second + cfg.CircuitMaxSimilarRequests = envInt(envMaxSimilar, cfg.CircuitMaxSimilarRequests) + cfg.CircuitMinAttemptedTokens = envInt(envMinTokens, cfg.CircuitMinAttemptedTokens) + cfg.CircuitRetryAfter = time.Duration(envInt(envRetryAfterSeconds, int(cfg.CircuitRetryAfter.Seconds()))) * time.Second + + // Budget + cfg.SessionBudgetUSD = envFloat(envSessionBudgetUSD, cfg.SessionBudgetUSD) + cfg.BudgetWarnRatio = envFloat(envBudgetWarnRatio, cfg.BudgetWarnRatio) + cfg.BudgetHardRatio = envFloat(envBudgetHardRatio, cfg.BudgetHardRatio) + + // Velocity + cfg.SpendVelocityEnabled = envBool(envSpendVelocityEnabled, cfg.SpendVelocityEnabled) + cfg.SpendVelocityWindow = time.Duration(envInt(envSpendVelocityWindowSecs, int(cfg.SpendVelocityWindow.Seconds()))) * time.Second + cfg.SpendVelocityWarnUSD = envFloat(envSpendVelocityWarnUSD, cfg.SpendVelocityWarnUSD) + cfg.SpendVelocityBlockUSD = envFloat(envSpendVelocityBlockUSD, cfg.SpendVelocityBlockUSD) + + return cfg +} + +func envInt(key string, fallback int) int { + v := os.Getenv(key) + if v == "" { + return fallback + } + n, err := strconv.Atoi(v) + if err != nil || n <= 0 { + return fallback + } + return n +} + +func envFloat(key string, fallback float64) float64 { + v := os.Getenv(key) + if v == "" { + return fallback + } + n, err := strconv.ParseFloat(v, 64) + if err != nil || n <= 0 { + return fallback + } + return n +} + +func envBool(key string, fallback bool) bool { + v := os.Getenv(key) + if v == "" { + return fallback + } + switch v { + case "1", "true", "TRUE", "True", "yes", "y": + return true + case "0", "false", "FALSE", "False", "no", "n": + return false + default: + return fallback + } +} diff --git a/gateway/internal/guard/fingerprint.go b/gateway/internal/guard/fingerprint.go new file mode 100644 index 0000000..a10692c --- /dev/null +++ b/gateway/internal/guard/fingerprint.go @@ -0,0 +1,93 @@ +package guard + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "strings" +) + +type FingerprintInput struct { + Route string + Model string + MessageCount int + ContentBlockCount int + ToolResultCount int + AttemptedTokens int + BlocksAnalyzed int + BlocksPruned int + LatestUserText string + SystemText string +} + +func BuildFingerprint(input FingerprintInput) string { + normalized := strings.Join([]string{ + fmt.Sprintf("route=%s", normalize(input.Route)), + fmt.Sprintf("model=%s", normalize(input.Model)), + fmt.Sprintf("messages=%s", countBucket(input.MessageCount)), + fmt.Sprintf("content_blocks=%s", countBucket(input.ContentBlockCount)), + fmt.Sprintf("tool_results=%s", countBucket(input.ToolResultCount)), + fmt.Sprintf("tokens=%s", tokenBucket(input.AttemptedTokens)), + fmt.Sprintf("blocks_analyzed=%s", countBucket(input.BlocksAnalyzed)), + fmt.Sprintf("blocks_pruned=%s", countBucket(input.BlocksPruned)), + fmt.Sprintf("latest=%s", hashPrefix(input.LatestUserText)), + fmt.Sprintf("system=%s", hashPrefix(input.SystemText)), + }, "\n") + sum := sha256.Sum256([]byte(normalized)) + return hex.EncodeToString(sum[:]) +} + +func normalize(s string) string { + v := strings.TrimSpace(strings.ToLower(s)) + if v == "" { + return "unknown" + } + return v +} + +func hashPrefix(s string) string { + trimmed := strings.TrimSpace(s) + if trimmed == "" { + return "empty" + } + sum := sha256.Sum256([]byte(trimmed)) + return hex.EncodeToString(sum[:])[:12] +} + +func countBucket(n int) string { + switch { + case n <= 0: + return "0" + case n == 1: + return "1" + case n <= 3: + return "2-3" + case n <= 10: + return "4-10" + case n <= 25: + return "11-25" + case n <= 50: + return "26-50" + case n <= 100: + return "51-100" + default: + return "100+" + } +} + +func tokenBucket(n int) string { + switch { + case n < 10_000: + return "0-10k" + case n < 25_000: + return "10k-25k" + case n < 50_000: + return "25k-50k" + case n < 100_000: + return "50k-100k" + case n < 200_000: + return "100k-200k" + default: + return "200k+" + } +} diff --git a/gateway/internal/guard/fingerprint_test.go b/gateway/internal/guard/fingerprint_test.go new file mode 100644 index 0000000..a94ce76 --- /dev/null +++ b/gateway/internal/guard/fingerprint_test.go @@ -0,0 +1,65 @@ +package guard + +import "testing" + +func TestBuildFingerprintStableForEquivalentShape(t *testing.T) { + t.Parallel() + a := BuildFingerprint(FingerprintInput{ + Route: "/v1/messages", + Model: "claude-sonnet-4-6", + MessageCount: 17, + ContentBlockCount: 31, + ToolResultCount: 5, + AttemptedTokens: 75_000, + BlocksAnalyzed: 120, + BlocksPruned: 102, + LatestUserText: "Summarize this file", + SystemText: "You are an assistant", + }) + b := BuildFingerprint(FingerprintInput{ + Route: "/v1/messages", + Model: "claude-sonnet-4-6", + MessageCount: 18, // same bucket + ContentBlockCount: 33, // same bucket + ToolResultCount: 6, // same bucket + AttemptedTokens: 80_000, + BlocksAnalyzed: 131, + BlocksPruned: 111, + LatestUserText: "Summarize this file", + SystemText: "You are an assistant", + }) + if a != b { + t.Fatalf("fingerprints differ for equivalent bucketed shape: %q != %q", a, b) + } +} + +func TestBuildFingerprintDiffersForDifferentShape(t *testing.T) { + t.Parallel() + a := BuildFingerprint(FingerprintInput{ + Route: "/v1/messages", + Model: "claude-sonnet-4-6", + MessageCount: 10, + ContentBlockCount: 10, + ToolResultCount: 0, + AttemptedTokens: 30_000, + BlocksAnalyzed: 10, + BlocksPruned: 2, + LatestUserText: "A", + SystemText: "S", + }) + b := BuildFingerprint(FingerprintInput{ + Route: "/v1/messages", + Model: "claude-sonnet-4-6", + MessageCount: 70, + ContentBlockCount: 70, + ToolResultCount: 20, + AttemptedTokens: 210_000, + BlocksAnalyzed: 120, + BlocksPruned: 90, + LatestUserText: "B", + SystemText: "T", + }) + if a == b { + t.Fatalf("fingerprints unexpectedly match: %q", a) + } +} diff --git a/gateway/internal/guard/manager.go b/gateway/internal/guard/manager.go new file mode 100644 index 0000000..321e085 --- /dev/null +++ b/gateway/internal/guard/manager.go @@ -0,0 +1,135 @@ +package guard + +import ( + "net/http" + "sync" + "time" +) + +// recentBlock records a blocked request so retries can be detected. +type recentBlock struct { + tokenCount int64 + blockedAt time.Time +} + +type Manager struct { + Circuit *CircuitBreaker + Budget *BudgetGuard + Velocity *VelocityGuard + + recentBlocksMu sync.Mutex + recentBlocks []recentBlock +} + +func NewManager(cfg Config) *Manager { + return &Manager{ + Circuit: NewCircuitBreaker(cfg), + Budget: NewBudgetGuard(cfg), + Velocity: NewVelocityGuard(cfg), + recentBlocks: make([]recentBlock, 0, 8), + } +} + +func (m *Manager) Check(sig RequestSignal) Decision { + // Retry storm check must run before any guard so retries of a blocked + // request are never counted toward budget or velocity metrics. + if m.isRetryStorm(int64(sig.AttemptedTokens)) { + return Decision{ + Allow: false, + StatusCode: http.StatusTooManyRequests, + Reason: "retry_storm", + } + } + + // Priority: + // 1. budget exceeded → block + // 2. velocity exceeded → block + // 3. circuit breaker → block + // 4. warnings → allow with warning + + bd := m.Budget.Check(sig) + if !bd.Allow { + m.recordBlock(int64(sig.AttemptedTokens)) + return bd + } + + vd := m.Velocity.Check(sig) + if !vd.Allow { + m.recordBlock(int64(sig.AttemptedTokens)) + vd.BudgetUSD = bd.BudgetUSD + vd.SpentUSD = bd.SpentUSD + vd.ProjectedUSD = bd.ProjectedUSD + return vd + } + + cd := m.Circuit.Check(sig) + if !cd.Allow { + m.recordBlock(int64(sig.AttemptedTokens)) + cd.BudgetUSD = bd.BudgetUSD + cd.SpentUSD = bd.SpentUSD + cd.ProjectedUSD = bd.ProjectedUSD + return cd + } + + // If any guard allowed with a warning, surface it. + if bd.Warn { + return bd + } + if vd.Warn { + vd.BudgetUSD = bd.BudgetUSD + vd.SpentUSD = bd.SpentUSD + vd.ProjectedUSD = bd.ProjectedUSD + return vd + } + if cd.Warn { + cd.BudgetUSD = bd.BudgetUSD + cd.SpentUSD = bd.SpentUSD + cd.ProjectedUSD = bd.ProjectedUSD + return cd + } + + return bd +} + +func (m *Manager) recordBlock(tokenCount int64) { + m.recentBlocksMu.Lock() + defer m.recentBlocksMu.Unlock() + m.recentBlocks = append(m.recentBlocks, recentBlock{ + tokenCount: tokenCount, + blockedAt: time.Now(), + }) +} + +// isRetryStorm returns true when the incoming request looks like a retry of a +// recently blocked request: same token count (±10%) within a 2-second window. +// It also prunes expired entries from the window as a side effect. +func (m *Manager) isRetryStorm(tokenCount int64) bool { + m.recentBlocksMu.Lock() + defer m.recentBlocksMu.Unlock() + + now := time.Now() + live := m.recentBlocks[:0] + storm := false + + for _, rb := range m.recentBlocks { + if now.Sub(rb.blockedAt) >= 2*time.Second { + continue // expired, drop it + } + live = append(live, rb) + if rb.tokenCount > 0 { + diff := float64(abs64(tokenCount-rb.tokenCount)) / float64(rb.tokenCount) + if diff < 0.10 { + storm = true + } + } + } + m.recentBlocks = live + return storm +} + +func abs64(n int64) int64 { + if n < 0 { + return -n + } + return n +} diff --git a/gateway/internal/guard/types.go b/gateway/internal/guard/types.go new file mode 100644 index 0000000..ee7f004 --- /dev/null +++ b/gateway/internal/guard/types.go @@ -0,0 +1,38 @@ +package guard + +import "time" + +type RequestSignal struct { + MachineID string + SessionKey string + Route string + Model string + Fingerprint string + + AttemptedTokens int + SentTokens int + TokensSaved int + ReductionPct float64 + + BlocksAnalyzed int + BlocksPruned int + + Now time.Time + + EstimatedCostUSD float64 +} + +type Decision struct { + Allow bool + Warn bool + StatusCode int + Reason string + RetryAfter time.Duration + Count int + Remaining int + + // Add fields for Budget + BudgetUSD float64 + SpentUSD float64 + ProjectedUSD float64 +} diff --git a/gateway/internal/guard/velocity.go b/gateway/internal/guard/velocity.go new file mode 100644 index 0000000..cd3f448 --- /dev/null +++ b/gateway/internal/guard/velocity.go @@ -0,0 +1,103 @@ +package guard + +import ( + "net/http" + "os" + "sync" + "time" +) + +type SpendEvent struct { + At time.Time + EstimatedUSD float64 + AttemptedTokens int + SentTokens int +} + +type VelocityGuard struct { + mu sync.Mutex + cfg Config + events map[string][]SpendEvent +} + +func NewVelocityGuard(cfg Config) *VelocityGuard { + return &VelocityGuard{ + cfg: cfg, + events: make(map[string][]SpendEvent), + } +} + +func (vg *VelocityGuard) Check(sig RequestSignal) Decision { + if !vg.cfg.SpendVelocityEnabled { + return Decision{Allow: true, Reason: "disabled"} + } + if os.Getenv(envDevMode) == "1" { + return Decision{Allow: true, Reason: "dev_mode"} + } + if os.Getenv(envAllowOverBudget) == "1" || os.Getenv(envAllowRunaway) == "1" { + return Decision{Allow: true, Reason: "override"} + } + + key := sig.SessionKey + if key == "" { + key = "no-session" + } + now := sig.Now + if now.IsZero() { + now = time.Now() + } + cutoff := now.Add(-vg.cfg.SpendVelocityWindow) + + vg.mu.Lock() + defer vg.mu.Unlock() + + events := pruneSpendEvents(vg.events[key], cutoff) + + events = append(events, SpendEvent{ + At: now, + EstimatedUSD: sig.EstimatedCostUSD, + AttemptedTokens: sig.AttemptedTokens, + SentTokens: sig.SentTokens, + }) + vg.events[key] = events + + var sumUSD float64 + for _, e := range events { + sumUSD += e.EstimatedUSD + } + + if sumUSD >= vg.cfg.SpendVelocityBlockUSD { + return Decision{ + Allow: false, + StatusCode: http.StatusTooManyRequests, + Reason: "velocity_exceeded", + } + } + + if sumUSD >= vg.cfg.SpendVelocityWarnUSD { + return Decision{ + Allow: true, + Warn: true, + Reason: "velocity_warning", + } + } + + return Decision{Allow: true, Reason: "allowed"} +} + +func pruneSpendEvents(events []SpendEvent, cutoff time.Time) []SpendEvent { + if len(events) == 0 { + return events + } + idx := 0 + for idx < len(events) && events[idx].At.Before(cutoff) { + idx++ + } + if idx == 0 { + return events + } + if idx >= len(events) { + return nil + } + return events[idx:] +} diff --git a/gateway/internal/hdr/histogram.go b/gateway/internal/hdr/histogram.go new file mode 100644 index 0000000..8970767 --- /dev/null +++ b/gateway/internal/hdr/histogram.go @@ -0,0 +1,534 @@ +// Package hdr implements a High Dynamic Range (HDR) Histogram from scratch. +// +// # Problem +// +// Tracking LLM request latencies is hard with ordinary fixed-bucket histograms: +// a cached hit takes ~1 ms while a streaming generation can take 5 minutes — +// seven orders of magnitude. A histogram with enough buckets to cover that range +// at fine granularity wastes megabytes of memory; one with coarse buckets loses +// the tail-latency signal you care most about. +// +// # Solution +// +// The HDR Histogram (Gil Tene, Azul Systems — http://hdrhistogram.org) maintains +// configurable significant-digit precision across the entire dynamic range. +// Recording and querying are O(1) and O(buckets) respectively, and memory usage +// is O(log(max/min) × precision) — typically 25–200 KiB. +// +// # Data structure +// +// Values are stored in a flat counts array that conceptually forms a 2-D grid: +// +// Bucket 0 │ sub-bucket [0 … S) covers values [0 … S) — resolution 1 +// Bucket 1 │ sub-bucket [S/2 … S) covers values [S … 2S) — resolution 2 +// Bucket 2 │ sub-bucket [S/2 … S) covers values [2S … 4S) — resolution 4 +// … +// +// where S = subBucketCount (a power of two derived from SignificantFigures). +// Within each bucket the sub-bucket index encodes the top significant digits of +// the value; the bucket index encodes the magnitude. The index computation uses +// math/bits.LeadingZeros64 — a single instruction on modern CPUs. +// +// # Equivalence ranges +// +// Two values share a bucket when they differ by less than the resolution at +// their magnitude. For example, with 2 significant figures (S = 256): +// - Values 1–255 each occupy a unique bucket. Resolution = 1. +// - Values 256–511 are grouped into 128 buckets, two values each. Resolution = 2. +// - Values 512–1023 → 128 buckets, four values each. Resolution = 4. +// +// The relative error is always ≤ 1/(S/2) ≈ 0.78 % for sigfigs = 2. +package hdr + +import ( + "errors" + "fmt" + "math" + "math/bits" + "sync" +) + +// Defaults tuned for tracking LLM gateway latency in microseconds. +const ( + // DefaultLowestDiscernibleValue is 1 µs — values below this are clamped up. + DefaultLowestDiscernibleValue int64 = 1 + // DefaultHighestTrackableValue is 1 hour expressed in microseconds. + DefaultHighestTrackableValue int64 = 3_600_000_000 + // DefaultSignificantFigures gives ~0.78 % relative error throughout the range. + DefaultSignificantFigures int = 2 +) + +// Config holds the HDR Histogram construction parameters. +// All three fields must be set; use DefaultConfig() for sane starting values. +type Config struct { + // LowestDiscernibleValue is the smallest unit of resolution (≥ 1). + // Values smaller than this are recorded as this value. + LowestDiscernibleValue int64 + + // HighestTrackableValue is the maximum value tracked with full precision. + // Must be ≥ 2 × LowestDiscernibleValue. + // Values above this are clamped to this before recording. + HighestTrackableValue int64 + + // SignificantFigures controls decimal precision [1, 5]. + // 2 → ~25 KiB, 3 → ~200 KiB, 4 → ~1.6 MiB. + SignificantFigures int +} + +// DefaultConfig returns a Config suited for tracking LLM request latency +// in microseconds with 2 significant figures of precision. +func DefaultConfig() Config { + return Config{ + LowestDiscernibleValue: DefaultLowestDiscernibleValue, + HighestTrackableValue: DefaultHighestTrackableValue, + SignificantFigures: DefaultSignificantFigures, + } +} + +// Snapshot is an immutable point-in-time summary exported by Histogram.Snapshot. +// All latency values use the same unit as the recorded values (e.g., µs). +type Snapshot struct { + Count int64 + Min int64 + Max int64 + Mean float64 + P50 int64 + P90 int64 + P95 int64 + P99 int64 + P999 int64 // 99.9th percentile + OverflowCount int64 // values clamped above HighestTrackableValue +} + +// Histogram is an HDR Histogram safe for concurrent use. +// Construct with New or Default; the zero value must not be used. +type Histogram struct { + // Immutable geometry — computed once in New. + cfg Config + unitMagnitude int // floor(log2(LowestDiscernibleValue)) + subBucketHalfCountMagnitude int // controls sub-bucket count + subBucketCount int // S: total sub-buckets per power-of-2 range + subBucketHalfCount int // S/2: "new" sub-buckets added per range + subBucketMask int64 // (S-1) << unitMagnitude + bucketCount int // number of power-of-2 ranges needed + + // Mutable state — all access serialised through mu. + mu sync.Mutex + counts []int64 // flat counts array; length = (bucketCount+1)*subBucketHalfCount + totalCount int64 + overflowCount int64 // values clamped above HighestTrackableValue + minValue int64 // MaxInt64 when empty + maxValue int64 // 0 when empty +} + +// New constructs a Histogram from cfg. +// Returns an error if the configuration is invalid. +func New(cfg Config) (*Histogram, error) { + if cfg.LowestDiscernibleValue < 1 { + return nil, errors.New("hdr: LowestDiscernibleValue must be ≥ 1") + } + if cfg.HighestTrackableValue < 2*cfg.LowestDiscernibleValue { + return nil, fmt.Errorf("hdr: HighestTrackableValue (%d) must be ≥ 2×LowestDiscernibleValue (%d)", + cfg.HighestTrackableValue, cfg.LowestDiscernibleValue) + } + if cfg.SignificantFigures < 1 || cfg.SignificantFigures > 5 { + return nil, fmt.Errorf("hdr: SignificantFigures must be in [1, 5], got %d", cfg.SignificantFigures) + } + + // ── Derive sub-bucket geometry ────────────────────────────────────────── + // + // We need enough sub-buckets so that the two values nearest to any power-of-2 + // boundary differ by less than one unit of relative precision. Concretely: + // + // largestSingleUnit = 2 × 10^sigfigs + // subBucketCountMagnitude = ⌈log₂(largestSingleUnit)⌉ + // + // This ensures that within each power-of-2 bucket, the sub-bucket spacing + // is ≤ 1/10^sigfigs of the bucket's lower bound. + largestSingleUnit := 2 * int64(math.Pow10(cfg.SignificantFigures)) + subBucketCountMagnitude := int(math.Ceil(math.Log2(float64(largestSingleUnit)))) + subBucketCountMagnitude = max(subBucketCountMagnitude, 1) + + subBucketHalfCountMagnitude := subBucketCountMagnitude - 1 + subBucketCount := 1 << uint(subBucketHalfCountMagnitude+1) + subBucketHalfCount := subBucketCount / 2 + + unitMagnitude := int(math.Floor(math.Log2(float64(cfg.LowestDiscernibleValue)))) + unitMagnitude = max(unitMagnitude, 0) + subBucketMask := int64(subBucketCount-1) << uint(unitMagnitude) + + bucketCount := requiredBuckets(cfg.HighestTrackableValue, subBucketCount, unitMagnitude) + counts := make([]int64, (bucketCount+1)*subBucketHalfCount) + + return &Histogram{ + cfg: cfg, + unitMagnitude: unitMagnitude, + subBucketHalfCountMagnitude: subBucketHalfCountMagnitude, + subBucketCount: subBucketCount, + subBucketHalfCount: subBucketHalfCount, + subBucketMask: subBucketMask, + bucketCount: bucketCount, + counts: counts, + minValue: math.MaxInt64, + }, nil +} + +// Default returns a Histogram built with DefaultConfig. +// Panics on construction error (impossible with compile-time constants). +func Default() *Histogram { + h, err := New(DefaultConfig()) + if err != nil { + panic("hdr: default config is invalid: " + err.Error()) + } + return h +} + +// ── Recording ──────────────────────────────────────────────────────────────── + +// Record records a single observation of value. +// Values outside [0, HighestTrackableValue] are clamped to the nearest bound. +func (h *Histogram) Record(value int64) { h.RecordN(value, 1) } + +// RecordN records n occurrences of value in a single lock acquisition. +func (h *Histogram) RecordN(value, n int64) { + if value < 0 { + value = 0 + } + overflow := value > h.cfg.HighestTrackableValue + if overflow { + value = h.cfg.HighestTrackableValue + } + idx := h.countsIndexFor(value) + + h.mu.Lock() + h.counts[idx] += n + h.totalCount += n + if overflow { + h.overflowCount += n + } + if value < h.minValue { + h.minValue = value + } + if value > h.maxValue { + h.maxValue = value + } + h.mu.Unlock() +} + +// ── Point queries ───────────────────────────────────────────────────────────── + +// Count returns the total number of recorded values. +func (h *Histogram) Count() int64 { + h.mu.Lock() + c := h.totalCount + h.mu.Unlock() + return c +} + +// Min returns the smallest recorded value, or 0 if the histogram is empty. +func (h *Histogram) Min() int64 { + h.mu.Lock() + defer h.mu.Unlock() + if h.totalCount == 0 { + return 0 + } + return h.minValue +} + +// Max returns the largest recorded value, or 0 if the histogram is empty. +func (h *Histogram) Max() int64 { + h.mu.Lock() + defer h.mu.Unlock() + return h.maxValue +} + +// OverflowCount returns the number of recorded values that exceeded +// HighestTrackableValue and were clamped. A non-zero result is a signal to +// either raise HighestTrackableValue or investigate unexpectedly large latencies. +func (h *Histogram) OverflowCount() int64 { + h.mu.Lock() + c := h.overflowCount + h.mu.Unlock() + return c +} + +// Mean returns the arithmetic mean of all recorded values, or 0 if empty. +// The mean is computed from the midpoint of each equivalence range to avoid +// systematic bias. +func (h *Histogram) Mean() float64 { + h.mu.Lock() + defer h.mu.Unlock() + return h.mean() +} + +// mean is the lock-free inner implementation (caller must hold mu). +func (h *Histogram) mean() float64 { + if h.totalCount == 0 { + return 0 + } + total := float64(0) + for i, c := range h.counts { + if c != 0 { + total += float64(c) * float64(h.medianEquivalentValue(h.valueFromIndex(i))) + } + } + return total / float64(h.totalCount) +} + +// ValueAtPercentile returns the highest value v such that at least percentile % +// of all recorded values are ≤ v. percentile is clamped to [0, 100]. +func (h *Histogram) ValueAtPercentile(percentile float64) int64 { + percentile = clamp(percentile, 0, 100) + + h.mu.Lock() + defer h.mu.Unlock() + + if h.totalCount == 0 { + return 0 + } + threshold := countAtPercentile(h.totalCount, percentile) + running := int64(0) + for i, c := range h.counts { + running += c + if running >= threshold { + return h.highestEquivalentValue(h.valueFromIndex(i)) + } + } + return h.highestEquivalentValue(h.maxValue) +} + +// PercentileOf returns the percentile [0, 100] at which value falls; i.e., the +// percentage of recorded values that are ≤ value. +// Returns 100 for an empty histogram. +func (h *Histogram) PercentileOf(value int64) float64 { + if value < 0 { + value = 0 + } + h.mu.Lock() + defer h.mu.Unlock() + if h.totalCount == 0 { + return 100 + } + idx := h.countsIndexFor(value) + if idx >= len(h.counts) { + idx = len(h.counts) - 1 + } + running := int64(0) + for i := 0; i <= idx; i++ { + running += h.counts[i] + } + return 100.0 * float64(running) / float64(h.totalCount) +} + +// ── Bulk operations ─────────────────────────────────────────────────────────── + +// Reset clears all recorded values, returning the histogram to its empty state. +func (h *Histogram) Reset() { + h.mu.Lock() + for i := range h.counts { + h.counts[i] = 0 + } + h.totalCount = 0 + h.overflowCount = 0 + h.minValue = math.MaxInt64 + h.maxValue = 0 + h.mu.Unlock() +} + +// Merge adds all counts from other into h. +// Both histograms must have been constructed with identical Config values. +// Returns an error if the configurations differ. +func (h *Histogram) Merge(other *Histogram) error { + if h.cfg != other.cfg { + return errors.New("hdr: cannot merge histograms with different configurations") + } + // Snapshot other under its own lock to avoid holding two locks simultaneously. + other.mu.Lock() + otherCounts := make([]int64, len(other.counts)) + copy(otherCounts, other.counts) + otherTotal := other.totalCount + otherOverflow := other.overflowCount + otherMin := other.minValue + otherMax := other.maxValue + other.mu.Unlock() + + h.mu.Lock() + for i, c := range otherCounts { + h.counts[i] += c + } + h.totalCount += otherTotal + h.overflowCount += otherOverflow + if otherMin < h.minValue { + h.minValue = otherMin + } + if otherMax > h.maxValue { + h.maxValue = otherMax + } + h.mu.Unlock() + return nil +} + +// Snapshot returns an immutable summary of the histogram's current state. +// All five percentiles are computed in a single O(N) pass through the counts array. +func (h *Histogram) Snapshot() Snapshot { + h.mu.Lock() + defer h.mu.Unlock() + + if h.totalCount == 0 { + return Snapshot{} + } + + snap := Snapshot{ + Count: h.totalCount, + OverflowCount: h.overflowCount, + Min: h.minValue, + Max: h.maxValue, + Mean: h.mean(), + } + + // Percentile thresholds in ascending order — filled in one pass. + type target struct { + threshold int64 + dest *int64 + } + targets := [5]target{ + {countAtPercentile(h.totalCount, 50), &snap.P50}, + {countAtPercentile(h.totalCount, 90), &snap.P90}, + {countAtPercentile(h.totalCount, 95), &snap.P95}, + {countAtPercentile(h.totalCount, 99), &snap.P99}, + {countAtPercentile(h.totalCount, 99.9), &snap.P999}, + } + + running := int64(0) + next := 0 // index into targets + for i, c := range h.counts { + running += c + for next < len(targets) && running >= targets[next].threshold { + *targets[next].dest = h.highestEquivalentValue(h.valueFromIndex(i)) + next++ + } + if next >= len(targets) { + break + } + } + return snap +} + +// ── Index arithmetic ────────────────────────────────────────────────────────── + +// getBucketIndex returns the power-of-2 bucket that value falls into. +// The computation uses a single LeadingZeros64 call — one CPU instruction. +func (h *Histogram) getBucketIndex(value int64) int { + // OR with subBucketMask ensures we count enough bits for the minimum + // representable value in this unit magnitude. + pow2ceiling := 64 - bits.LeadingZeros64(uint64(value|h.subBucketMask)) + idx := pow2ceiling - h.unitMagnitude - (h.subBucketHalfCountMagnitude + 1) + if idx < 0 { + return 0 + } + return idx +} + +// getSubBucketIdx returns the sub-bucket index within the given bucket. +func (h *Histogram) getSubBucketIdx(value int64, bucketIdx int) int { + shift := max(bucketIdx+h.unitMagnitude, 0) + return int(value >> uint(shift)) +} + +// countsIndex converts a (bucket, sub-bucket) pair to a flat array index. +// +// Layout of the counts array: +// +// Bucket 0: indices [0, subBucketCount) +// Bucket 1: indices [subBucketHalfCount, subBucketHalfCount*2) ← new half only +// Bucket 2: indices [subBucketHalfCount*2, subBucketHalfCount*3) ← new half only +// … +func (h *Histogram) countsIndex(bucketIdx, subBucketIdx int) int { + return (bucketIdx+1)*h.subBucketHalfCount + subBucketIdx - h.subBucketHalfCount +} + +// countsIndexFor returns the flat counts-array index for value. +func (h *Histogram) countsIndexFor(value int64) int { + b := h.getBucketIndex(value) + s := h.getSubBucketIdx(value, b) + return h.countsIndex(b, s) +} + +// ── Inverse mapping: index → value ─────────────────────────────────────────── + +// valueFromIndex converts a flat counts-array index back to the lowest value +// in that bucket's equivalence range. +func (h *Histogram) valueFromIndex(idx int) int64 { + bucketIdx := (idx >> h.subBucketHalfCountMagnitude) - 1 + subBucketIdx := (idx & (h.subBucketHalfCount - 1)) + h.subBucketHalfCount + if bucketIdx < 0 { + subBucketIdx -= h.subBucketHalfCount + bucketIdx = 0 + } + return int64(subBucketIdx) << uint(bucketIdx+h.unitMagnitude) +} + +// ── Equivalence-range helpers ───────────────────────────────────────────────── + +// sizeOfEquivalentRange returns the width of the equivalence range containing value. +// All values in the range map to the same bucket and are indistinguishable. +func (h *Histogram) sizeOfEquivalentRange(value int64) int64 { + b := h.getBucketIndex(value) + shift := max(b+h.unitMagnitude, 0) + return int64(1) << uint(shift) +} + +// lowestEquivalentValue returns the smallest value that maps to the same bucket as value. +func (h *Histogram) lowestEquivalentValue(value int64) int64 { + b := h.getBucketIndex(value) + s := h.getSubBucketIdx(value, b) + return int64(s) << uint(b+h.unitMagnitude) +} + +// highestEquivalentValue returns the largest value that maps to the same bucket as value. +func (h *Histogram) highestEquivalentValue(value int64) int64 { + return h.lowestEquivalentValue(value) + h.sizeOfEquivalentRange(value) - 1 +} + +// medianEquivalentValue returns the midpoint of value's equivalence range. +// Used for mean computation to avoid systematic rounding bias. +func (h *Histogram) medianEquivalentValue(value int64) int64 { + return h.lowestEquivalentValue(value) + h.sizeOfEquivalentRange(value)/2 +} + +// ── Package-level helpers ───────────────────────────────────────────────────── + +// requiredBuckets computes how many power-of-2 ranges are needed to cover highest. +func requiredBuckets(highest int64, subBucketCount, unitMagnitude int) int { + // smallestUntrackable is the first value that falls outside bucket 0. + smallestUntrackable := int64(subBucketCount) << uint(unitMagnitude) + needed := 1 + for smallestUntrackable <= highest { + if smallestUntrackable > math.MaxInt64/2 { + return needed + 1 + } + smallestUntrackable <<= 1 + needed++ + } + return needed +} + +// countAtPercentile converts a percentile to the cumulative count threshold. +func countAtPercentile(total int64, percentile float64) int64 { + c := int64(math.Ceil((percentile / 100.0) * float64(total))) + if c < 1 { + return 1 + } + return c +} + +// clamp constrains v to [lo, hi]. +func clamp(v, lo, hi float64) float64 { + if v < lo { + return lo + } + if v > hi { + return hi + } + return v +} diff --git a/gateway/internal/hdr/histogram_bench_test.go b/gateway/internal/hdr/histogram_bench_test.go new file mode 100644 index 0000000..6d74af6 --- /dev/null +++ b/gateway/internal/hdr/histogram_bench_test.go @@ -0,0 +1,192 @@ +package hdr_test + +import ( + "testing" + + "github.com/Revanth14/indexqube/gateway/internal/hdr" +) + +// ── Record ─────────────────────────────────────────────────────────────────── + +// BenchmarkRecord measures the cost of a single Record call under a mutex. +// Expected: 0 allocs/op, ~20–60 ns/op (dominated by sync.Mutex). +func BenchmarkRecord(b *testing.B) { + h := hdr.Default() + b.ReportAllocs() + b.ResetTimer() + for i := b.N; i > 0; i-- { + h.Record(int64(i % 1_000_000)) + } +} + +// BenchmarkRecord_Parallel measures Record throughput across GOMAXPROCS goroutines. +// Contention on the mutex becomes visible here. +func BenchmarkRecord_Parallel(b *testing.B) { + h := hdr.Default() + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + i := int64(0) + for pb.Next() { + h.Record(i % 1_000_000) + i++ + } + }) +} + +// BenchmarkRecordN measures bulk recording (one lock acquisition per N values). +func BenchmarkRecordN(b *testing.B) { + h := hdr.Default() + b.ReportAllocs() + b.ResetTimer() + for i := b.N; i > 0; i-- { + h.RecordN(int64(i%1_000_000), 10) + } +} + +// ── Percentile queries ──────────────────────────────────────────────────────── + +// BenchmarkValueAtPercentile measures a single percentile query. +// Expected: 0 allocs/op (walks the counts array under the lock). +func BenchmarkValueAtPercentile(b *testing.B) { + h := filledHistogram() + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = h.ValueAtPercentile(99) + } +} + +// BenchmarkValueAtPercentile_P999 walks slightly further into the tail. +func BenchmarkValueAtPercentile_P999(b *testing.B) { + h := filledHistogram() + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = h.ValueAtPercentile(99.9) + } +} + +// ── Snapshot ───────────────────────────────────────────────────────────────── + +// BenchmarkSnapshot measures the cost of a full Snapshot (5 percentiles + mean +// computed in a single O(N) pass). Expected: 0 allocs/op. +func BenchmarkSnapshot(b *testing.B) { + h := filledHistogram() + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = h.Snapshot() + } +} + +// ── AllocsPerRun assertions ─────────────────────────────────────────────────── + +// TestRecordZeroAllocs asserts that Record never allocates on the heap. +func TestRecordZeroAllocs(t *testing.T) { + h := hdr.Default() + allocs := testing.AllocsPerRun(100, func() { + h.Record(12_345) + }) + if allocs > 0 { + t.Errorf("Record: %.0f allocs/op; want 0", allocs) + } +} + +// TestValueAtPercentileZeroAllocs asserts that ValueAtPercentile never allocates. +func TestValueAtPercentileZeroAllocs(t *testing.T) { + h := filledHistogram() + allocs := testing.AllocsPerRun(100, func() { + _ = h.ValueAtPercentile(99) + }) + if allocs > 0 { + t.Errorf("ValueAtPercentile: %.0f allocs/op; want 0", allocs) + } +} + +// TestSnapshotZeroAllocs asserts that Snapshot never allocates on the heap +// (it returns a plain struct by value; no slices or maps inside). +func TestSnapshotZeroAllocs(t *testing.T) { + h := filledHistogram() + allocs := testing.AllocsPerRun(100, func() { + _ = h.Snapshot() + }) + if allocs > 0 { + t.Errorf("Snapshot: %.0f allocs/op; want 0", allocs) + } +} + +// ── Overflow ────────────────────────────────────────────────────────────────── + +// TestOverflowCount verifies that out-of-range values are counted and exposed. +func TestOverflowCount(t *testing.T) { + cfg := hdr.Config{ + LowestDiscernibleValue: 1, + HighestTrackableValue: 1_000, + SignificantFigures: 2, + } + h, err := hdr.New(cfg) + if err != nil { + t.Fatalf("New: %v", err) + } + + h.Record(500) // within range + h.RecordN(9_999, 3) // above HighestTrackableValue — clamped, counted as overflow + + if h.OverflowCount() != 3 { + t.Errorf("OverflowCount() = %d; want 3", h.OverflowCount()) + } + // Total count still includes the clamped values. + if h.Count() != 4 { + t.Errorf("Count() = %d; want 4 (in-range + clamped)", h.Count()) + } + // Snapshot surfaces overflow too. + snap := h.Snapshot() + if snap.OverflowCount != 3 { + t.Errorf("Snapshot.OverflowCount = %d; want 3", snap.OverflowCount) + } +} + +func TestOverflowCountReset(t *testing.T) { + h := hdr.Default() + h.RecordN(hdr.DefaultHighestTrackableValue+1, 5) + if h.OverflowCount() != 5 { + t.Fatalf("pre-reset OverflowCount = %d; want 5", h.OverflowCount()) + } + h.Reset() + if h.OverflowCount() != 0 { + t.Errorf("post-reset OverflowCount = %d; want 0", h.OverflowCount()) + } +} + +func TestOverflowCountMerge(t *testing.T) { + cfg := hdr.DefaultConfig() + a, _ := hdr.New(cfg) + b, _ := hdr.New(cfg) + + a.RecordN(hdr.DefaultHighestTrackableValue+1, 2) + b.RecordN(hdr.DefaultHighestTrackableValue+1, 3) + + if err := a.Merge(b); err != nil { + t.Fatalf("Merge: %v", err) + } + if a.OverflowCount() != 5 { + t.Errorf("merged OverflowCount = %d; want 5", a.OverflowCount()) + } +} + +// ── Helper ──────────────────────────────────────────────────────────────────── + +// filledHistogram returns a histogram pre-loaded with 100 000 values spread +// across the µs–minute range, giving benchmarks a realistic counts distribution. +func filledHistogram() *hdr.Histogram { + h := hdr.Default() + // Bimodal: 90% fast (1–10 ms), 10% slow (1–5 s). + for i := int64(0); i < 90_000; i++ { + h.Record((i % 10_000) + 1_000) // 1 000–11 000 µs + } + for i := int64(0); i < 10_000; i++ { + h.Record((i % 4_000_000) + 1_000_000) // 1–5 s in µs + } + return h +} diff --git a/gateway/internal/hdr/histogram_test.go b/gateway/internal/hdr/histogram_test.go new file mode 100644 index 0000000..774cac3 --- /dev/null +++ b/gateway/internal/hdr/histogram_test.go @@ -0,0 +1,456 @@ +package hdr_test + +import ( + "math" + "sync" + "testing" + + "github.com/Revanth14/indexqube/gateway/internal/hdr" +) + +// ── Construction ────────────────────────────────────────────────────────────── + +func TestNewRejectsInvalidConfig(t *testing.T) { + cases := []struct { + name string + cfg hdr.Config + }{ + { + name: "zero LowestDiscernibleValue", + cfg: hdr.Config{LowestDiscernibleValue: 0, HighestTrackableValue: 1000, SignificantFigures: 2}, + }, + { + name: "negative LowestDiscernibleValue", + cfg: hdr.Config{LowestDiscernibleValue: -1, HighestTrackableValue: 1000, SignificantFigures: 2}, + }, + { + name: "HighestTrackableValue < 2×LowestDiscernibleValue", + cfg: hdr.Config{LowestDiscernibleValue: 100, HighestTrackableValue: 100, SignificantFigures: 2}, + }, + { + name: "SignificantFigures = 0", + cfg: hdr.Config{LowestDiscernibleValue: 1, HighestTrackableValue: 1000, SignificantFigures: 0}, + }, + { + name: "SignificantFigures = 6", + cfg: hdr.Config{LowestDiscernibleValue: 1, HighestTrackableValue: 1000, SignificantFigures: 6}, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + _, err := hdr.New(tc.cfg) + if err == nil { + t.Errorf("New(%+v): expected error, got nil", tc.cfg) + } + }) + } +} + +func TestDefaultConstructsSuccessfully(t *testing.T) { + h := hdr.Default() + if h == nil { + t.Fatal("Default() returned nil") + } + if h.Count() != 0 { + t.Errorf("new histogram count = %d; want 0", h.Count()) + } +} + +// ── Empty-histogram behaviour ───────────────────────────────────────────────── + +func TestEmptyHistogramReturnsZeros(t *testing.T) { + h := hdr.Default() + if got := h.Count(); got != 0 { + t.Errorf("Count() = %d; want 0", got) + } + if got := h.Min(); got != 0 { + t.Errorf("Min() = %d; want 0", got) + } + if got := h.Max(); got != 0 { + t.Errorf("Max() = %d; want 0", got) + } + if got := h.Mean(); got != 0 { + t.Errorf("Mean() = %f; want 0", got) + } + for _, p := range []float64{0, 50, 90, 99, 100} { + if got := h.ValueAtPercentile(p); got != 0 { + t.Errorf("ValueAtPercentile(%v) = %d; want 0", p, got) + } + } + snap := h.Snapshot() + if snap.Count != 0 || snap.P50 != 0 || snap.P99 != 0 { + t.Errorf("empty Snapshot not zero: %+v", snap) + } +} + +// ── Single-value recording ──────────────────────────────────────────────────── + +func TestRecordSingleValue(t *testing.T) { + h := hdr.Default() + h.Record(1000) // 1 ms in µs + + if h.Count() != 1 { + t.Errorf("Count() = %d; want 1", h.Count()) + } + if h.Min() != 1000 { + t.Errorf("Min() = %d; want 1000", h.Min()) + } + if h.Max() != 1000 { + t.Errorf("Max() = %d; want 1000", h.Max()) + } + // p50 and p100 should both resolve to an equivalent value near 1000. + p50 := h.ValueAtPercentile(50) + if !withinPct(p50, 1000, 1.0) { + t.Errorf("p50 = %d; want ~1000 (within 1%%)", p50) + } +} + +// ── Precision / equivalence-range ──────────────────────────────────────────── + +// TestPrecision verifies that the relative error stays within the expected +// bound for 2 significant figures throughout the dynamic range. +func TestPrecision(t *testing.T) { + h := hdr.Default() // 2 sigfigs ≈ 1% relative error + + // Sample values across many orders of magnitude. + probes := []int64{ + 1, 10, 100, 1_000, 10_000, 100_000, + 1_000_000, 10_000_000, 100_000_000, 1_000_000_000, + } + + for _, v := range probes { + h.Reset() + h.Record(v) + + // The retrieved percentile value should be within 1% of the recorded value. + got := h.ValueAtPercentile(100) + if !withinPct(got, v, 1.0) { + t.Errorf("value %d: p100 = %d, relative error %.3f%% > 1%%", + v, got, relativeErrorPct(got, v)) + } + } +} + +// ── Percentile accuracy ─────────────────────────────────────────────────────── + +// TestValueAtPercentile_UniformDistribution records values 1..1000 and checks +// that the retrieved percentiles are accurate to within 1%. +func TestValueAtPercentile_UniformDistribution(t *testing.T) { + h := hdr.Default() + const n = 1000 + for i := int64(1); i <= n; i++ { + h.Record(i) + } + + cases := []struct { + pct float64 + want int64 + }{ + {50, 500}, + {75, 750}, + {90, 900}, + {95, 950}, + {99, 990}, + } + + for _, tc := range cases { + got := h.ValueAtPercentile(tc.pct) + if !withinPct(got, tc.want, 1.0) { + t.Errorf("p%.0f = %d; want ~%d (within 1%%)", tc.pct, got, tc.want) + } + } +} + +// TestPercentileOf is the inverse of ValueAtPercentile: +// PercentileOf(ValueAtPercentile(p)) should be ≥ p. +func TestPercentileOf_InverseOfValueAtPercentile(t *testing.T) { + h := hdr.Default() + for i := int64(1); i <= 10_000; i++ { + h.Record(i) + } + + for _, p := range []float64{10, 25, 50, 75, 90, 95, 99} { + v := h.ValueAtPercentile(p) + got := h.PercentileOf(v) + if got < p-1.0 { // allow 1 percentage point tolerance + t.Errorf("PercentileOf(ValueAtPercentile(%.1f)) = %.2f; want ≥ %.1f", p, got, p) + } + } +} + +// ── Percentile clamping ─────────────────────────────────────────────────────── + +func TestValueAtPercentileClamps(t *testing.T) { + h := hdr.Default() + h.Record(500) + + // Percentiles outside [0, 100] should be clamped, not panic. + _ = h.ValueAtPercentile(-5) + _ = h.ValueAtPercentile(200) +} + +// ── Clamping of out-of-range values ────────────────────────────────────────── + +func TestRecordClampsNegative(t *testing.T) { + h := hdr.Default() + h.Record(-999) + // Clamped to 0 — the min should be 0 (or the lowest-equivalent value of 0). + if h.Min() < 0 { + t.Errorf("Min() = %d after recording -999; want ≥ 0", h.Min()) + } +} + +func TestRecordClampsAboveHighest(t *testing.T) { + cfg := hdr.Config{ + LowestDiscernibleValue: 1, + HighestTrackableValue: 1_000, + SignificantFigures: 2, + } + h, err := hdr.New(cfg) + if err != nil { + t.Fatalf("New: %v", err) + } + h.Record(999_999) // way above HighestTrackableValue + if h.Max() > 1_000 { + t.Errorf("Max() = %d; want ≤ 1000 (clamped)", h.Max()) + } +} + +// ── RecordN ─────────────────────────────────────────────────────────────────── + +func TestRecordN(t *testing.T) { + h := hdr.Default() + h.RecordN(42, 100) + if h.Count() != 100 { + t.Errorf("Count() = %d; want 100", h.Count()) + } +} + +// ── Min / Max ───────────────────────────────────────────────────────────────── + +func TestMinMax(t *testing.T) { + h := hdr.Default() + for _, v := range []int64{500, 1, 9999, 3, 10000} { + h.Record(v) + } + if h.Min() > 1 { + t.Errorf("Min() = %d; want ≤ 1", h.Min()) + } + if h.Max() < 10000 { + t.Errorf("Max() = %d; want ≥ 10000", h.Max()) + } +} + +// ── Mean ────────────────────────────────────────────────────────────────────── + +func TestMean(t *testing.T) { + h := hdr.Default() + // Record 1, 2, 3 — mean should be 2. + h.Record(1) + h.Record(2) + h.Record(3) + mean := h.Mean() + if math.Abs(mean-2.0) > 0.05 { + t.Errorf("Mean() = %.4f; want ~2.0", mean) + } +} + +// ── Reset ───────────────────────────────────────────────────────────────────── + +func TestReset(t *testing.T) { + h := hdr.Default() + for i := int64(1); i <= 1000; i++ { + h.Record(i) + } + if h.Count() == 0 { + t.Fatal("pre-reset Count() is 0; something is wrong") + } + h.Reset() + + if h.Count() != 0 { + t.Errorf("post-reset Count() = %d; want 0", h.Count()) + } + if h.Min() != 0 { + t.Errorf("post-reset Min() = %d; want 0", h.Min()) + } + if h.Max() != 0 { + t.Errorf("post-reset Max() = %d; want 0", h.Max()) + } + if p := h.ValueAtPercentile(99); p != 0 { + t.Errorf("post-reset p99 = %d; want 0", p) + } +} + +// ── Merge ───────────────────────────────────────────────────────────────────── + +func TestMerge_CombinesCounts(t *testing.T) { + cfg := hdr.DefaultConfig() + a, _ := hdr.New(cfg) + b, _ := hdr.New(cfg) + + for i := int64(1); i <= 500; i++ { + a.Record(i) + } + for i := int64(501); i <= 1000; i++ { + b.Record(i) + } + + if err := a.Merge(b); err != nil { + t.Fatalf("Merge: %v", err) + } + + if a.Count() != 1000 { + t.Errorf("merged Count() = %d; want 1000", a.Count()) + } + p50 := a.ValueAtPercentile(50) + if !withinPct(p50, 500, 1.0) { + t.Errorf("merged p50 = %d; want ~500", p50) + } +} + +func TestMerge_RejectsIncompatibleConfig(t *testing.T) { + a := hdr.Default() + b, _ := hdr.New(hdr.Config{ + LowestDiscernibleValue: 1, + HighestTrackableValue: 1_000, + SignificantFigures: 2, + }) + if err := a.Merge(b); err == nil { + t.Error("Merge with incompatible config: expected error, got nil") + } +} + +// ── Snapshot ────────────────────────────────────────────────────────────────── + +func TestSnapshot_PercentilesConsistent(t *testing.T) { + h := hdr.Default() + for i := int64(1); i <= 10_000; i++ { + h.Record(i) + } + snap := h.Snapshot() + + // Percentiles must be non-decreasing. + order := []int64{snap.P50, snap.P90, snap.P95, snap.P99, snap.P999} + for i := 1; i < len(order); i++ { + if order[i] < order[i-1] { + t.Errorf("percentile order violated: %v", order) + } + } + + // Spot-check against direct ValueAtPercentile calls. + if !withinPct(snap.P99, h.ValueAtPercentile(99), 0.01) { + t.Errorf("Snapshot.P99 (%d) differs from ValueAtPercentile(99) (%d)", + snap.P99, h.ValueAtPercentile(99)) + } + if snap.Count != h.Count() { + t.Errorf("Snapshot.Count = %d; want %d", snap.Count, h.Count()) + } +} + +func TestSnapshot_MeanMatchesMean(t *testing.T) { + h := hdr.Default() + for _, v := range []int64{100, 200, 300} { + h.Record(v) + } + snap := h.Snapshot() + if math.Abs(snap.Mean-h.Mean()) > 0.001 { + t.Errorf("Snapshot.Mean = %.4f; h.Mean() = %.4f", snap.Mean, h.Mean()) + } +} + +// ── Concurrent safety ───────────────────────────────────────────────────────── + +// TestConcurrentRecord runs N goroutines each recording M values. +// With -race this detects data races; the final Count must be exact. +func TestConcurrentRecord(t *testing.T) { + h := hdr.Default() + const goroutines = 50 + const perGoroutine = 1_000 + + var wg sync.WaitGroup + wg.Add(goroutines) + for g := 0; g < goroutines; g++ { + go func(start int64) { + defer wg.Done() + for i := start; i < start+perGoroutine; i++ { + h.Record(i % 1_000_000) + } + }(int64(g * perGoroutine)) + } + wg.Wait() + + want := int64(goroutines * perGoroutine) + if h.Count() != want { + t.Errorf("Count() = %d; want %d", h.Count(), want) + } +} + +// TestConcurrentRecordAndSnapshot checks that Snapshot can be called +// concurrently with Record without panicking or deadlocking. +func TestConcurrentRecordAndSnapshot(t *testing.T) { + h := hdr.Default() + done := make(chan struct{}) + + go func() { + defer close(done) + for i := int64(0); i < 10_000; i++ { + h.Record(i) + } + }() + + for { + select { + case <-done: + return + default: + _ = h.Snapshot() + } + } +} + +// ── Gateway-realistic scenario ──────────────────────────────────────────────── + +// TestGatewayLatencyScenario simulates recording LLM request latencies +// (cache hits vs. generation) and verifies that the bimodal distribution +// is reflected correctly in the snapshot. +func TestGatewayLatencyScenario(t *testing.T) { + h := hdr.Default() + + // 90% cache hits: ~5 ms (5_000 µs) + for i := 0; i < 900; i++ { + h.RecordN(5_000, 1) + } + // 10% generation requests: ~3 s (3_000_000 µs) + for i := 0; i < 100; i++ { + h.RecordN(3_000_000, 1) + } + + snap := h.Snapshot() + + // p50 should be near the cache-hit latency. + if !withinPct(snap.P50, 5_000, 1.0) { + t.Errorf("p50 = %d µs; want ~5000 µs (cache hit)", snap.P50) + } + // p99 should be near the generation latency. + if !withinPct(snap.P99, 3_000_000, 1.0) { + t.Errorf("p99 = %d µs; want ~3_000_000 µs (generation)", snap.P99) + } + if snap.Count != 1000 { + t.Errorf("Count = %d; want 1000", snap.Count) + } +} + +// ── Helpers ─────────────────────────────────────────────────────────────────── + +// withinPct reports whether got is within pct% of want. +func withinPct(got, want int64, pct float64) bool { + if want == 0 { + return got == 0 + } + return relativeErrorPct(got, want) <= pct +} + +func relativeErrorPct(got, want int64) float64 { + return math.Abs(float64(got)-float64(want)) / float64(want) * 100 +} diff --git a/gateway/internal/proxy/claude_messages.go b/gateway/internal/proxy/claude_messages.go index cfcb536..88fd8c4 100644 --- a/gateway/internal/proxy/claude_messages.go +++ b/gateway/internal/proxy/claude_messages.go @@ -13,10 +13,13 @@ import ( "log/slog" "net/http" "net/url" + "os" "runtime" + "strconv" "strings" "time" + "github.com/Revanth14/indexqube/gateway/internal/guard" "github.com/Revanth14/indexqube/gateway/internal/memory" "github.com/Revanth14/indexqube/gateway/internal/middleware" "github.com/Revanth14/indexqube/gateway/internal/telemetry" @@ -41,6 +44,14 @@ type anthropicMessage struct { Content json.RawMessage `json:"content"` } +type claudeRequestShape struct { + MessageCount int + ContentBlockCount int + ToolResultCount int + LatestUserText string + SystemText string +} + // claudeOptimizerStats holds per-request accounting populated by prepareClaudeBody. // All byte/token fields reflect the state after optimization (or the original if skipped). type claudeOptimizerStats struct { @@ -153,55 +164,118 @@ func (p *Proxy) handleClaudeMessages(w http.ResponseWriter, r *http.Request) { return } + // overheadStart marks the beginning of pure IndexQube processing time. + // It stops just before dispatching to upstream so ProxyOverheadMs + // reflects only the optimizer + guard work, not the model round-trip. + overheadStart := time.Now() + sessionKey := claudeSessionKey(r, cfg.DevToken) var earlyMeta struct { Model string `json:"model"` } _ = json.Unmarshal(body, &earlyMeta) - if cooldown, ok := p.claudeCooldowns.Get("anthropic", earlyMeta.Model, time.Now()); ok { - remaining := time.Until(cooldown.Until) - if value := retryAfterSeconds(remaining); value != "" { - w.Header().Set("Retry-After", value) - } - p.writeError(w, r, errorPayload{ - HTTPStatus: http.StatusTooManyRequests, - Type: "upstream_error", - Code: "provider_rate_limited", - Message: "Provider is cooling down after a recent rate limit. Retry shortly or switch model.", - }) - streamStats := claudeStreamStats{ - Status: "error", - StatusCode: http.StatusTooManyRequests, - UpstreamErrorCode: firstNonEmpty(cooldown.UpstreamCode, "provider_rate_limited"), - UpstreamErrorType: cooldown.UpstreamType, - UpstreamRequestID: cooldown.UpstreamRequestID, - RetryAfter: remaining, - CircuitOpen: true, - CircuitCooldown: remaining, - } - duration := time.Since(started) - if cfg.SessionStore != nil { - cfg.SessionStore.RecordUsage(sessionKey, memory.UsageTotals{ - Requests: 1, - TokensIn: estimateTokens(len(body)), - BytesIn: len(body), + if os.Getenv("IQ_DEV_MODE") != "1" { + if cooldown, ok := p.claudeCooldowns.Get("anthropic", earlyMeta.Model, time.Now()); ok { + remaining := time.Until(cooldown.Until) + if value := retryAfterSeconds(remaining); value != "" { + w.Header().Set("Retry-After", value) + } + p.writeError(w, r, errorPayload{ + HTTPStatus: http.StatusTooManyRequests, + Type: "upstream_error", + Code: "provider_rate_limited", + Message: "Provider is cooling down after a recent rate limit. Retry shortly or switch model.", }) + streamStats := claudeStreamStats{ + Status: "error", + StatusCode: http.StatusTooManyRequests, + UpstreamErrorCode: firstNonEmpty(cooldown.UpstreamCode, "provider_rate_limited"), + UpstreamErrorType: cooldown.UpstreamType, + UpstreamRequestID: cooldown.UpstreamRequestID, + RetryAfter: remaining, + CircuitOpen: true, + CircuitCooldown: remaining, + } + duration := time.Since(started) + if cfg.SessionStore != nil { + cfg.SessionStore.RecordUsage(sessionKey, memory.UsageTotals{ + Requests: 1, + TokensIn: estimateTokens(len(body)), + BytesIn: len(body), + }) + } + p.logClaudeRequestComplete(r.Context(), requestID, cfg.Mode, earlyMeta.Model, sessionKey, len(body), claudeOptimizerStats{ + BytesBefore: len(body), + BytesAfter: len(body), + EstimatedTokensBefore: estimateTokens(len(body)), + EstimatedTokensAfter: estimateTokens(len(body)), + }, streamStats, duration) + return } - p.logClaudeRequestComplete(r.Context(), requestID, cfg.Mode, earlyMeta.Model, sessionKey, len(body), claudeOptimizerStats{ - BytesBefore: len(body), - BytesAfter: len(body), - EstimatedTokensBefore: estimateTokens(len(body)), - EstimatedTokensAfter: estimateTokens(len(body)), - }, streamStats, duration) - return } - forwardBody, meta, optStats, err := p.prepareClaudeBody(r.Context(), cfg, sessionKey, body) + forwardBody, meta, optStats, shape, err := p.prepareClaudeBody(r.Context(), cfg, sessionKey, body) if err != nil { p.writeError(w, r, errorPayload{HTTPStatus: http.StatusBadRequest, Type: "invalid_request_error", Code: "invalid_anthropic_request", Message: err.Error()}) return } + guardDecision := guard.Decision{Allow: true, Reason: "disabled"} + if p.guardManager != nil { + fingerprint := guard.BuildFingerprint(guard.FingerprintInput{ + Route: r.URL.Path, + Model: meta.Model, + MessageCount: shape.MessageCount, + ContentBlockCount: shape.ContentBlockCount, + ToolResultCount: shape.ToolResultCount, + AttemptedTokens: optStats.EstimatedTokensBefore, + BlocksAnalyzed: optStats.BlocksSeen, + BlocksPruned: optStats.BlocksPruned, + LatestUserText: shape.LatestUserText, + SystemText: shape.SystemText, + }) + guardDecision = p.guardManager.Check(guard.RequestSignal{ + MachineID: telemetry.GetMachineID(), + SessionKey: sessionKey, + Route: r.URL.Path, + Model: meta.Model, + Fingerprint: fingerprint, + AttemptedTokens: optStats.EstimatedTokensBefore, + SentTokens: optStats.EstimatedTokensAfter, + TokensSaved: optStats.EstimatedTokensSaved, + ReductionPct: optStats.ReductionRatio * 100, + BlocksAnalyzed: optStats.BlocksSeen, + BlocksPruned: optStats.BlocksPruned, + Now: time.Now(), + EstimatedCostUSD: telemetry.EstimateCost(cfg.Mode, meta.Model, optStats.EstimatedTokensBefore), + }) + if !guardDecision.Allow { + p.writeGuardResponse(w, r, guardDecision) + duration := time.Since(started) + if cfg.SessionStore != nil { + cfg.SessionStore.RecordUsage(sessionKey, memory.UsageTotals{ + Requests: 1, + TokensIn: estimateTokens(len(body)), + BytesIn: len(body), + }) + } + streamStats := claudeStreamStats{ + Status: "error", + StatusCode: http.StatusTooManyRequests, + UpstreamErrorCode: "indexqube_guard_blocked", + RetryAfter: guardDecision.RetryAfter, + } + p.logClaudeRequestComplete(r.Context(), requestID, cfg.Mode, meta.Model, sessionKey, len(body), optStats, streamStats, duration) + p.emitClaudeUsageEvent(r, sessionKey, meta.Model, optStats, duration, 0, guardDecision, time.Since(overheadStart).Milliseconds()) + return + } + if guardDecision.Warn { + w.Header().Set("X-IndexQube-Guard-Warning", guardDecision.Reason) + w.Header().Set("X-IndexQube-Guard-Remaining", strconv.Itoa(guardDecision.Remaining)) + } + } + + overheadMs := time.Since(overheadStart).Milliseconds() streamStats := p.forwardClaudeMessages(w, r, cfg, forwardBody) if streamStats.StatusCode == http.StatusTooManyRequests { entry := p.claudeCooldowns.Open("anthropic", meta.Model, http.StatusTooManyRequests, streamStats.upstreamMeta(), cfg.RateLimitCooldown, time.Now()) @@ -219,23 +293,7 @@ func (p *Proxy) handleClaudeMessages(w http.ResponseWriter, r *http.Request) { }) } p.logClaudeRequestComplete(r.Context(), requestID, cfg.Mode, meta.Model, sessionKey, len(body), optStats, streamStats, duration) - if p.usageTracker != nil { - p.usageTracker.Track(telemetry.UsageEvent{ - MachineID: telemetry.GetMachineID(), - OsArch: runtime.GOOS + "/" + runtime.GOARCH, - IqVersion: Version, - CliAgent: r.Header.Get("User-Agent"), - ModelTarget: meta.Model, - InputTokensAttempted: optStats.EstimatedTokensBefore, - InputTokensSent: optStats.EstimatedTokensAfter, - TokensSaved: optStats.EstimatedTokensSaved, - ReductionRatio: optStats.ReductionRatio * 100, - BlocksAnalyzed: optStats.BlocksSeen, - BlocksPruned: optStats.BlocksPruned, - TotalLatencyMs: int(duration.Milliseconds()), - UpstreamStatus: streamStats.StatusCode, - }) - } + p.emitClaudeUsageEvent(r, sessionKey, meta.Model, optStats, duration, streamStats.StatusCode, guardDecision, overheadMs) } func (p *Proxy) handleClaudeCountTokens(w http.ResponseWriter, r *http.Request) { @@ -309,7 +367,11 @@ func (p *Proxy) claudeDefaults() ClaudeMessagesConfig { cfg.HTTPClient = http.DefaultClient } if cfg.RateLimitCooldown <= 0 { - cfg.RateLimitCooldown = 30 * time.Second + if os.Getenv("IQ_DEV_MODE") == "1" { + cfg.RateLimitCooldown = 1 * time.Second + } else { + cfg.RateLimitCooldown = 30 * time.Second + } } // Apply optimizer defaults when not explicitly configured. if cfg.Optimizer.MinSpanBytes <= 0 { @@ -338,9 +400,9 @@ func (c ClaudeMessagesConfig) validate() error { } func validClaudeDevToken(r *http.Request, want string) bool { - auth := strings.TrimSpace(r.Header.Get("Authorization")) - token := strings.TrimSpace(strings.TrimPrefix(auth, "Bearer ")) - return token != "" + auth := strings.TrimSpace(r.Header.Get("Authorization")) + token := strings.TrimSpace(strings.TrimPrefix(auth, "Bearer ")) + return token != "" } func claudeSessionKey(r *http.Request, fallback string) string { @@ -350,23 +412,30 @@ func claudeSessionKey(r *http.Request, fallback string) string { auth := strings.TrimSpace(r.Header.Get("Authorization")) if auth != "" { sum := sha256.Sum256([]byte(auth)) - return hex.EncodeToString(sum[:8]) + key := hex.EncodeToString(sum[:8]) + // Suffix with the per-invocation session ID so the circuit breaker + // scopes similar-request counts to this iq session, not across sessions. + if sid := os.Getenv("IQ_SESSION_ID"); sid != "" { + return key + "-" + sid[:8] + } + return key } sum := sha256.Sum256([]byte(fallback)) return hex.EncodeToString(sum[:8]) } -func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, sessionKey string, body []byte) ([]byte, anthropicMessagesRequest, claudeOptimizerStats, error) { +func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, sessionKey string, body []byte) ([]byte, anthropicMessagesRequest, claudeOptimizerStats, claudeRequestShape, error) { if err := ctx.Err(); err != nil { - return nil, anthropicMessagesRequest{}, claudeOptimizerStats{}, err + return nil, anthropicMessagesRequest{}, claudeOptimizerStats{}, claudeRequestShape{}, err } var root map[string]any if err := json.Unmarshal(body, &root); err != nil { - return nil, anthropicMessagesRequest{}, claudeOptimizerStats{}, fmt.Errorf("parse anthropic messages body: %w", err) + return nil, anthropicMessagesRequest{}, claudeOptimizerStats{}, claudeRequestShape{}, fmt.Errorf("parse anthropic messages body: %w", err) } req := anthropicMessagesRequest{} req.Model, _ = root["model"].(string) req.Stream, _ = root["stream"].(bool) + shape := extractClaudeRequestShape(root) stats := claudeOptimizerStats{ BytesBefore: len(body), BytesAfter: len(body), @@ -374,7 +443,7 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, EstimatedTokensAfter: estimateTokens(len(body)), } if cfg.SessionStore == nil { - return body, req, stats, nil + return body, req, stats, shape, nil } minSpanBytes := cfg.Optimizer.MinSpanBytes @@ -396,7 +465,7 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, }) } } - return body, req, stats, nil + return body, req, stats, shape, nil } // optimize / dry_run: full span accounting and class-aware pruning. @@ -462,14 +531,14 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, } if len(prunableSpans) == 0 { - return body, req, stats, nil + return body, req, stats, shape, nil } // Rewrite on a fresh parse so the original stays intact for dry_run. var rewriteRoot map[string]any if err := json.Unmarshal(body, &rewriteRoot); err != nil { p.logger.WarnContext(ctx, "claude optimize re-parse failed; forwarding original", slog.Any("err", err)) - return body, req, stats, nil + return body, req, stats, shape, nil } pruneCount := 0 @@ -491,7 +560,7 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, optimized, err := json.Marshal(rewriteRoot) if err != nil { p.logger.WarnContext(ctx, "claude optimize marshal failed; forwarding original", slog.Any("err", err)) - return body, req, stats, nil + return body, req, stats, shape, nil } stats.BlocksPruned = pruneCount @@ -503,10 +572,10 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, } if cfg.Mode == "optimize" { - return optimized, req, stats, nil + return optimized, req, stats, shape, nil } // dry_run: report accurate stats but forward the original body. - return body, req, stats, nil + return body, req, stats, shape, nil } // isEligibleSpanClass returns true if the span class is eligible for pruning @@ -745,10 +814,10 @@ func bedrockClientFor(ctx context.Context, cfg BedrockConfig) (*bedrockruntime.C // (without regional prefix). Package-level to avoid per-request allocation. var bedrockKnownModels = map[string]string{ // Claude 3 series - "claude-3-sonnet-20240229": "anthropic.claude-3-sonnet-20240229-v1:0", - "claude-3-haiku-20240307": "anthropic.claude-3-haiku-20240307-v1:0", - "claude-3-5-haiku-20241022": "anthropic.claude-3-5-haiku-20241022-v1:0", - "claude-3-5-haiku-latest": "anthropic.claude-3-5-haiku-20241022-v1:0", + "claude-3-sonnet-20240229": "anthropic.claude-3-sonnet-20240229-v1:0", + "claude-3-haiku-20240307": "anthropic.claude-3-haiku-20240307-v1:0", + "claude-3-5-haiku-20241022": "anthropic.claude-3-5-haiku-20241022-v1:0", + "claude-3-5-haiku-latest": "anthropic.claude-3-5-haiku-20241022-v1:0", // Claude 4 series (IDs from Bedrock ListFoundationModels) "claude-sonnet-4-20250514": "anthropic.claude-sonnet-4-20250514-v1:0", "claude-haiku-4-5-20251001": "anthropic.claude-haiku-4-5-20251001-v1:0", @@ -1025,6 +1094,222 @@ func anthropicUsageOutputTokens(payload string) int { return ev.Usage.OutputTokens } +func (p *Proxy) writeGuardResponse(w http.ResponseWriter, r *http.Request, d guard.Decision) { + status := d.StatusCode + if status == 0 { + status = http.StatusTooManyRequests + } + + w.Header().Set("Content-Type", "application/json") + + var errType, msg, override string + + switch d.Reason { + case "budget_exceeded": + w.Header().Set("X-IndexQube-Guard", "budget") + errType = "indexqube_budget_exceeded" + msg = "IndexQube stopped this session because it reached the configured budget." + override = "Set IQ_ALLOW_OVER_BUDGET=1 to continue." + w.WriteHeader(status) + if err := json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "type": errType, + "message": msg, + "budget_usd": d.BudgetUSD, + "estimated_spend_usd": d.ProjectedUSD, + "override": override, + }, + }); err != nil { + p.logger.ErrorContext(r.Context(), "failed to encode guard response", slog.Any("err", err)) + } + return + case "velocity_exceeded": + w.Header().Set("X-IndexQube-Guard", "velocity") + errType = "indexqube_velocity_exceeded" + msg = "IndexQube stopped this request because spend velocity exceeded the maximum allowed." + override = "Set IQ_ALLOW_OVER_BUDGET=1 or IQ_ALLOW_RUNAWAY=1 to continue." + w.WriteHeader(status) + if err := json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "type": errType, + "message": msg, + "override": override, + }, + }); err != nil { + p.logger.ErrorContext(r.Context(), "failed to encode guard response", slog.Any("err", err)) + } + return + default: + // Default to circuit breaker + retryAfter := int(d.RetryAfter.Seconds()) + if retryAfter <= 0 { + retryAfter = 60 + } + w.Header().Set("Retry-After", strconv.Itoa(retryAfter)) + w.Header().Set("X-IndexQube-Guard", "circuit-breaker") + errType = "indexqube_circuit_breaker" + msg = "IndexQube stopped a likely runaway agent loop. This session sent too many similar large requests in a short period." + override = "Set IQ_ALLOW_RUNAWAY=1 to disable this protection for the current session." + w.WriteHeader(status) + if err := json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "type": errType, + "message": msg, + "reason": firstNonEmpty(d.Reason, "similar_large_requests"), + "retry_after_seconds": retryAfter, + "override": override, + }, + }); err != nil { + p.logger.ErrorContext(r.Context(), "failed to encode guard response", slog.Any("err", err)) + } + } +} + +func (p *Proxy) emitClaudeUsageEvent(r *http.Request, sessionKey, model string, optStats claudeOptimizerStats, duration time.Duration, upstreamStatus int, d guard.Decision, overheadMs int64) { + if p.usageTracker != nil { + event := telemetry.UsageEvent{ + MachineID: telemetry.GetMachineID(), + OsArch: runtime.GOOS + "/" + runtime.GOARCH, + IqVersion: Version, + CliAgent: r.Header.Get("User-Agent"), + ModelTarget: model, + InputTokensAttempted: optStats.EstimatedTokensBefore, + InputTokensSent: optStats.EstimatedTokensAfter, + TokensSaved: optStats.EstimatedTokensSaved, + ReductionRatio: optStats.ReductionRatio * 100, + BlocksAnalyzed: optStats.BlocksSeen, + BlocksPruned: optStats.BlocksPruned, + SkipReasons: guardSkipReasons(d), + TotalLatencyMs: int(duration.Milliseconds()), + ProxyOverheadMs: float64(overheadMs), + UpstreamStatus: upstreamStatus, + } + if !d.Allow { + event.InputTokensSent = 0 + event.TokensSaved = 0 + event.ReductionRatio = 0 + event.UpstreamStatus = http.StatusTooManyRequests + } + p.usageTracker.Track(event) + } + + if p.sessionTracker != nil { + p.sessionTracker.Record(sessionKey, telemetry.RequestOutcome{ + TokensAttempted: optStats.EstimatedTokensBefore, + TokensSent: optStats.EstimatedTokensAfter, + TokensSaved: optStats.EstimatedTokensSaved, + DollarsSaved: telemetry.EstimateCostSaved("optimize", model, optStats.EstimatedTokensSaved), + GuardReason: d.Reason, + Killed: !d.Allow, + Warned: d.Warn, + }) + } +} + +func guardSkipReasons(d guard.Decision) map[string]int { + out := map[string]int{} + if d.Allow { + out["guard_allowed"] = 1 + if d.Warn { + out["guard_warning"] = 1 + out["guard_warning_similar_large_requests"] = 1 + } + } else { + out["guard_blocked"] = 1 + } + reason := strings.ToLower(strings.TrimSpace(d.Reason)) + if reason != "" { + out["guard_reason_"+strings.ReplaceAll(reason, " ", "_")] = 1 + } + return out +} + +func extractClaudeRequestShape(root map[string]any) claudeRequestShape { + var shape claudeRequestShape + msgs, _ := root["messages"].([]any) + shape.MessageCount = len(msgs) + for _, raw := range msgs { + msg, ok := raw.(map[string]any) + if !ok { + continue + } + blocks, toolResults, text := summarizeAnthropicContent(msg["content"]) + shape.ContentBlockCount += blocks + shape.ToolResultCount += toolResults + role, _ := msg["role"].(string) + if strings.EqualFold(role, "user") { + trimmed := strings.TrimSpace(text) + if trimmed != "" { + shape.LatestUserText = trimmed + } + } + } + _, _, systemText := summarizeAnthropicContent(root["system"]) + shape.SystemText = strings.TrimSpace(systemText) + return shape +} + +func summarizeAnthropicContent(content any) (blocks int, toolResults int, text string) { + switch v := content.(type) { + case string: + if strings.TrimSpace(v) == "" { + return 0, 0, "" + } + return 1, 0, v + case []any: + var sb strings.Builder + for _, item := range v { + blocks++ + m, ok := item.(map[string]any) + if !ok { + continue + } + typ, _ := m["type"].(string) + switch typ { + case "tool_result": + toolResults++ + appendText(&sb, m["content"]) + default: + appendText(&sb, m["text"]) + } + } + return blocks, toolResults, strings.TrimSpace(sb.String()) + case map[string]any: + var sb strings.Builder + blocks = 1 + if typ, _ := v["type"].(string); typ == "tool_result" { + toolResults = 1 + appendText(&sb, v["content"]) + } else { + appendText(&sb, v["text"]) + } + return blocks, toolResults, strings.TrimSpace(sb.String()) + default: + return 0, 0, "" + } +} + +func appendText(sb *strings.Builder, value any) { + switch v := value.(type) { + case string: + trimmed := strings.TrimSpace(v) + if trimmed == "" { + return + } + if sb.Len() > 0 { + sb.WriteByte('\n') + } + sb.WriteString(trimmed) + case []any: + for _, item := range v { + appendText(sb, item) + } + case map[string]any: + appendText(sb, v["text"]) + appendText(sb, v["content"]) + } +} + func (p *Proxy) logClaudeRequestComplete(ctx context.Context, requestID, mode, model, sessionKey string, bytesBefore int, opt claudeOptimizerStats, stream claudeStreamStats, dur time.Duration) { level := slog.LevelInfo if stream.Status == "error" || stream.Status == "stream_error" { @@ -1158,4 +1443,3 @@ func estimateTokens(chars int) int { } return (chars + 3) / 4 } - diff --git a/gateway/internal/proxy/handlers.go b/gateway/internal/proxy/handlers.go index acbb79f..a4388e6 100644 --- a/gateway/internal/proxy/handlers.go +++ b/gateway/internal/proxy/handlers.go @@ -574,6 +574,40 @@ func (p *Proxy) writeError(w http.ResponseWriter, r *http.Request, payload error } } +// handleAgentSessions returns a snapshot of all in-memory agent sessions and +// the recent kill log. This is the data source for the iq ui dashboard's +// "Agent Sessions" panel and the interview story around agentic observability. +func (p *Proxy) handleAgentSessions(w http.ResponseWriter, _ *http.Request) { + type response struct { + Sessions []telemetry.AgentSession `json:"sessions"` + KillLog []telemetry.KillEvent `json:"kill_log"` + TotalSessions int `json:"total_sessions"` + TotalKills int `json:"total_kills"` + TotalDollarsSaved float64 `json:"total_dollars_saved"` + } + + var resp response + if p.sessionTracker != nil { + resp.Sessions = p.sessionTracker.Snapshot() + resp.KillLog = p.sessionTracker.KillLog() + } + if resp.Sessions == nil { + resp.Sessions = []telemetry.AgentSession{} + } + if resp.KillLog == nil { + resp.KillLog = []telemetry.KillEvent{} + } + resp.TotalSessions = len(resp.Sessions) + for _, s := range resp.Sessions { + resp.TotalKills += s.KillEvents + resp.TotalDollarsSaved += s.DollarsSaved + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) +} + // handleTelemetry accepts a UsageEvent from the iq binary and forwards it to // the configured telemetry sink (Supabase). This keeps Supabase credentials // server-side only — the distributed iq binary never sees them. diff --git a/gateway/internal/proxy/proxy.go b/gateway/internal/proxy/proxy.go index 5b9982d..82649c8 100644 --- a/gateway/internal/proxy/proxy.go +++ b/gateway/internal/proxy/proxy.go @@ -13,6 +13,7 @@ import ( "time" "github.com/Revanth14/indexqube/gateway/internal/domain" + "github.com/Revanth14/indexqube/gateway/internal/guard" "github.com/Revanth14/indexqube/gateway/internal/memory" "github.com/Revanth14/indexqube/gateway/internal/telemetry" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" @@ -40,6 +41,8 @@ type Proxy struct { claude ClaudeMessagesConfig claudeCooldowns *providerCooldowns usageTracker telemetry.Sink + guardManager *guard.Manager + sessionTracker *telemetry.AgentSessionStore } // BedrockConfig routes /v1/messages to AWS Bedrock instead of Anthropic's API. @@ -136,6 +139,14 @@ func WithUsageTracker(c telemetry.Sink) Option { } } +// WithAgentSessionStore wires the in-memory agent session tracker. +// Nil disables session aggregation silently. +func WithAgentSessionStore(s *telemetry.AgentSessionStore) Option { + return func(p *Proxy) { + p.sessionTracker = s + } +} + // New returns a wired Proxy. A nil governor is a programmer error and // panics fast at boot rather than 5xx-ing every request in production. func New(gov Governor, opts ...Option) *Proxy { @@ -149,6 +160,7 @@ func New(gov Governor, opts ...Option) *Proxy { maxRequestSize: 8 << 20, // default 8 MiB optimizeTimeout: 30 * time.Second, claudeCooldowns: newProviderCooldowns(), + guardManager: guard.NewManager(guard.FromEnv()), } for _, opt := range opts { opt(p) @@ -175,6 +187,7 @@ func (p *Proxy) registerRoutes() { w.WriteHeader(http.StatusOK) }) p.mux.HandleFunc("GET /readyz", p.handleReady) + p.mux.HandleFunc("GET /v1/agent-sessions", p.handleAgentSessions) p.mux.HandleFunc("GET /v1/diagnostics", p.handleDiagnostics) p.mux.HandleFunc("GET /v1/models", p.handleModels) p.mux.HandleFunc("POST /v1/messages/count_tokens", p.handleClaudeCountTokens) diff --git a/gateway/internal/proxy/proxy_test.go b/gateway/internal/proxy/proxy_test.go index 5458c6b..7e14c01 100644 --- a/gateway/internal/proxy/proxy_test.go +++ b/gateway/internal/proxy/proxy_test.go @@ -17,6 +17,7 @@ import ( "github.com/Revanth14/indexqube/gateway/internal/domain" govpkg "github.com/Revanth14/indexqube/gateway/internal/governor" "github.com/Revanth14/indexqube/gateway/internal/memory" + "github.com/Revanth14/indexqube/gateway/internal/telemetry" ) // fakeGovernor lets tests script the governor's streaming behavior. @@ -498,6 +499,66 @@ func TestClaudeMessages_MissingAuth(t *testing.T) { } } +func TestClaudeMessages_AgentSessionsRecordWithoutUsageTracker(t *testing.T) { + t.Parallel() + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "event: message_stop\n") + _, _ = io.WriteString(w, `data: {"type":"message_stop"}`+"\n\n") + })) + t.Cleanup(upstream.Close) + + p := New(&fakeGovernor{}, + WithAgentSessionStore(telemetry.NewAgentSessionStore(time.Hour)), + WithClaudeMessages(ClaudeMessagesConfig{ + Mode: "observe", + DevToken: "iq-dev-local", + AnthropicAPIKey: "sk-ant-test", + AnthropicBaseURL: upstream.URL, + AnthropicVersion: "2023-06-01", + SessionStore: memory.NewStore(time.Hour), + }), + ) + srv := httptest.NewServer(p.Handler()) + t.Cleanup(srv.Close) + + req, err := http.NewRequest(http.MethodPost, srv.URL+"/v1/messages", strings.NewReader(`{"model":"claude-sonnet-4-6","stream":true,"messages":[{"role":"user","content":"hi"}]}`)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Authorization", "Bearer iq-dev-local") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("POST /v1/messages: %v", err) + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status=%d, want 200", resp.StatusCode) + } + + resp, err = http.Get(srv.URL + "/v1/agent-sessions") + if err != nil { + t.Fatalf("GET /v1/agent-sessions: %v", err) + } + defer resp.Body.Close() + var body struct { + TotalSessions int `json:"total_sessions"` + Sessions []struct { + RequestsTotal int `json:"requests_total"` + } `json:"sessions"` + } + if err := json.NewDecoder(resp.Body).Decode(&body); err != nil { + t.Fatalf("decode response: %v", err) + } + if body.TotalSessions != 1 || len(body.Sessions) != 1 { + t.Fatalf("sessions=%d len=%d, want one session", body.TotalSessions, len(body.Sessions)) + } + if body.Sessions[0].RequestsTotal != 1 { + t.Fatalf("requests_total=%d, want 1", body.Sessions[0].RequestsTotal) + } +} + func TestClaudeCountTokens_Passthrough(t *testing.T) { t.Parallel() var gotPath, gotKey string @@ -582,6 +643,95 @@ func TestClaudeMessages_RateLimitOpensModelCooldown(t *testing.T) { } } +func TestClaudeMessages_CircuitBreakerBlocksRepeatedRequests(t *testing.T) { + t.Setenv("INDEXQUBE_CIRCUIT_BREAKER_ENABLED", "true") + t.Setenv("INDEXQUBE_CIRCUIT_MIN_TOKENS", "1") + t.Setenv("INDEXQUBE_CIRCUIT_MAX_SIMILAR", "1") + t.Setenv("INDEXQUBE_CIRCUIT_WINDOW_SECONDS", "60") + t.Setenv("INDEXQUBE_CIRCUIT_RETRY_AFTER_SECONDS", "60") + + var calls int + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls++ + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "event: message_stop\n") + _, _ = io.WriteString(w, `data: {"type":"message_stop"}`+"\n\n") + })) + t.Cleanup(upstream.Close) + srv := newClaudeTestServer(t, upstream.URL, memory.NewStore(time.Hour), "optimize", true) + + body := `{"model":"claude-sonnet-4-6","stream":true,"messages":[{"role":"user","content":"repeat me"}]}` + for i := 0; i < 2; i++ { + req, err := http.NewRequest(http.MethodPost, srv.URL+"/v1/messages", strings.NewReader(body)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Authorization", "Bearer iq-dev-local") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("POST /v1/messages: %v", err) + } + gotBody, _ := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if i == 0 && resp.StatusCode != http.StatusOK { + t.Fatalf("first request status=%d body=%s, want 200", resp.StatusCode, gotBody) + } + if i == 1 { + if resp.StatusCode != http.StatusTooManyRequests { + t.Fatalf("second request status=%d body=%s, want 429", resp.StatusCode, gotBody) + } + if guard := resp.Header.Get("X-IndexQube-Guard"); guard != "circuit-breaker" { + t.Fatalf("X-IndexQube-Guard=%q, want circuit-breaker", guard) + } + if retry := resp.Header.Get("Retry-After"); retry != "60" { + t.Fatalf("Retry-After=%q, want 60", retry) + } + } + } + if calls != 1 { + t.Fatalf("upstream calls=%d, want 1", calls) + } +} + +func TestClaudeMessages_CircuitBreakerOverrideBypassesBlock(t *testing.T) { + t.Setenv("INDEXQUBE_CIRCUIT_BREAKER_ENABLED", "true") + t.Setenv("INDEXQUBE_CIRCUIT_MIN_TOKENS", "1") + t.Setenv("INDEXQUBE_CIRCUIT_MAX_SIMILAR", "1") + t.Setenv("INDEXQUBE_CIRCUIT_WINDOW_SECONDS", "60") + t.Setenv("IQ_ALLOW_RUNAWAY", "1") + + var calls int + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls++ + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "event: message_stop\n") + _, _ = io.WriteString(w, `data: {"type":"message_stop"}`+"\n\n") + })) + t.Cleanup(upstream.Close) + srv := newClaudeTestServer(t, upstream.URL, memory.NewStore(time.Hour), "optimize", true) + + body := `{"model":"claude-sonnet-4-6","stream":true,"messages":[{"role":"user","content":"repeat me"}]}` + for i := 0; i < 2; i++ { + req, err := http.NewRequest(http.MethodPost, srv.URL+"/v1/messages", strings.NewReader(body)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Authorization", "Bearer iq-dev-local") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("POST /v1/messages: %v", err) + } + gotBody, _ := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("request %d status=%d body=%s, want 200", i+1, resp.StatusCode, gotBody) + } + } + if calls != 2 { + t.Fatalf("upstream calls=%d, want 2", calls) + } +} + func TestClaudeMessages_OptimizePrunesRepeatedTextBlock(t *testing.T) { t.Parallel() var bodies []string diff --git a/gateway/internal/server/run.go b/gateway/internal/server/run.go index 9fd1693..a04f7ec 100644 --- a/gateway/internal/server/run.go +++ b/gateway/internal/server/run.go @@ -39,6 +39,17 @@ const ( // Run starts the gateway, reading all config from the environment, and blocks // until ctx is cancelled or a fatal startup error occurs. func Run(ctx context.Context) error { + return run(ctx, nil) +} + +// RunWithPublicListener starts the gateway with an already-bound public +// listener. This is useful for callers that need to reserve an ephemeral port +// before startup. +func RunWithPublicListener(ctx context.Context, publicListener net.Listener) error { + return run(ctx, publicListener) +} + +func run(ctx context.Context, publicListener net.Listener) error { bootLogger := slog.New(slog.NewJSONHandler(os.Stderr, nil)) cfg, err := config.Load() @@ -156,12 +167,15 @@ func Run(ctx context.Context) error { logger.Info("supabase usage telemetry enabled") } + agentSessions := telemetry.NewAgentSessionStore(0) // 4 h TTL default + p := proxy.New(gov, proxy.WithLogger(logger), proxy.WithMaxRequestSize(cfg.Governor.MaxRequestSize), proxy.WithMetrics(tp.Metrics), proxy.WithOptimizeTimeout(cfg.Governor.OptimizeTimeout), proxy.WithUsageTracker(usageTracker), + proxy.WithAgentSessionStore(agentSessions), proxy.WithClaudeMessages(proxy.ClaudeMessagesConfig{ Mode: cfg.ClaudeCode.Mode, DevToken: cfg.ClaudeCode.DevToken, @@ -199,8 +213,19 @@ func Run(ctx context.Context) error { serverErr := make(chan error, 2) go func() { - logger.Info("public server listening", slog.String("addr", publicServer.Addr)) - if err := publicServer.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + listenAddr := publicServer.Addr + if publicListener != nil { + listenAddr = publicListener.Addr().String() + } + logger.Info("public server listening", slog.String("addr", listenAddr)) + + var err error + if publicListener != nil { + err = publicServer.Serve(publicListener) + } else { + err = publicServer.ListenAndServe() + } + if err != nil && !errors.Is(err, http.ErrServerClosed) { serverErr <- fmt.Errorf("public server: %w", err) } }() @@ -212,7 +237,7 @@ func Run(ctx context.Context) error { }() janitorCtx, stopJanitor := context.WithCancel(context.Background()) - go runMemoryJanitor(janitorCtx, claudeStore, memoryJanitorInterval, logger) + go runMemoryJanitor(janitorCtx, claudeStore, agentSessions, memoryJanitorInterval, logger) // Block until ctx is cancelled or a server fails fatally. var runErr error @@ -253,7 +278,7 @@ func newUpstreamTransport() *http.Transport { } } -func runMemoryJanitor(ctx context.Context, store *memory.Store, interval time.Duration, logger *slog.Logger) { +func runMemoryJanitor(ctx context.Context, store *memory.Store, sessions *telemetry.AgentSessionStore, interval time.Duration, logger *slog.Logger) { if store == nil || interval <= 0 { return } @@ -265,6 +290,9 @@ func runMemoryJanitor(ctx context.Context, store *memory.Store, interval time.Du return case <-t.C: store.CleanupExpired() + if sessions != nil { + sessions.CleanupExpired() + } } } } diff --git a/gateway/internal/store/lsm/bench_test.go b/gateway/internal/store/lsm/bench_test.go new file mode 100644 index 0000000..2bf1459 --- /dev/null +++ b/gateway/internal/store/lsm/bench_test.go @@ -0,0 +1,304 @@ +package lsm_test + +// Benchmark suite: LSM-Tree vs SQLite B-tree +// +// Workload: 70% writes / 30% reads - mirrors IndexQube's actual token-cache ratio. +// Scales: 1K, 10K, 100K entries (1M omitted from CI; run manually with -benchtime=1x). +// +// How to read the output: +// +// ns/op - latency per operation +// MB/s - throughput (key+value payload, not wire bytes) +// write_amp - disk bytes written / logical bytes written (lower = better for LSM) +// disk_bytes - total on-disk footprint after the run +// +// Key insight this demonstrates: +// +// LSM batches random writes into sequential I/O (MemTable -> SSTable flush). +// SQLite's B-tree does in-place updates: every write is a random page seek. +// At the 70/30 write-heavy workload IndexQube sees, LSM wins on throughput. +// At read-heavy workloads the Bloom filters keep LSM competitive, but the +// B-tree's O(log N) indexed lookup begins to close the gap. + +import ( + "database/sql" + "fmt" + "math/rand/v2" + "os" + "path/filepath" + "testing" + + "github.com/Revanth14/indexqube/gateway/internal/store/lsm" + _ "modernc.org/sqlite" // pure-Go SQLite, no CGO required +) + +// Workload shape. +const ( + writeFraction = 0.70 // 70% Put, 30% Get + keyBytes = 32 // SHA-256 hash size - typical IndexQube chunk key + valBytes = 512 // representative chunk content excerpt +) + +// benchScales lists the entry counts to benchmark at. +// 1M is excluded from the default run to keep CI under 60 s; add it locally. +var benchScales = []int{1_000, 10_000, 100_000} + +// --- key / value generation --------------------------------------------------- + +// makeKey produces a deterministic fixed-width key for entry i. +// Using a Fibonacci-scrambled index spreads keys across the key space +// (avoids pathological B-tree sequential-insert behaviour that would +// +// flatter the B-tree unfairly). +func makeKey(i int) []byte { + scrambled := uint64(i) * 0x9e3779b97f4a7c15 // Knuth multiplicative hash + return []byte(fmt.Sprintf("%016x%016x", scrambled, uint64(i))) +} + +// makeVal produces a deterministic value of exactly valBytes. +func makeVal(i int) []byte { + v := make([]byte, valBytes) + for j := range v { + v[j] = byte((i*31 + j) & 0xff) + } + return v +} + +// --- helpers ------------------------------------------------------------------ + +// dirSize returns the total byte count of all regular files under dir. +// Used as a write-amplification proxy: disk_bytes / logical_bytes_written. +func dirSize(dir string) int64 { + var total int64 + _ = filepath.Walk(dir, func(_ string, fi os.FileInfo, _ error) error { + if fi != nil && !fi.IsDir() { + total += fi.Size() + } + return nil + }) + return total +} + +// reportWriteAmp computes and reports write amplification and disk footprint. +func reportWriteAmp(b *testing.B, dir string) { + b.Helper() + diskBytes := dirSize(dir) + logicalWrites := int64(float64(b.N) * writeFraction * float64(keyBytes+valBytes)) + if logicalWrites > 0 { + b.ReportMetric(float64(diskBytes)/float64(logicalWrites), "write_amp") + } + b.ReportMetric(float64(diskBytes)/1024/1024, "disk_MiB") +} + +// --- LSM benchmarks ----------------------------------------------------------- + +// BenchmarkLSM_70w30r benchmarks the custom LSM-Tree engine at the 70/30 ratio. +func BenchmarkLSM_70w30r(b *testing.B) { + for _, n := range benchScales { + n := n + b.Run(fmt.Sprintf("N=%d", n), func(b *testing.B) { + dir := b.TempDir() + eng, err := lsm.Open(dir, lsm.Options{}) + if err != nil { + b.Fatal(err) + } + defer eng.Close() + + // Pre-populate n/2 keys so reads have a realistic hit rate. + for i := range n / 2 { + if err := eng.Put(makeKey(i), makeVal(i)); err != nil { + b.Fatal(err) + } + } + + rng := rand.New(rand.NewPCG(42, 0)) + b.ResetTimer() + b.SetBytes(int64(keyBytes + valBytes)) + + for range b.N { + idx := rng.IntN(n) + if rng.Float64() < writeFraction { + if err := eng.Put(makeKey(idx), makeVal(idx)); err != nil { + b.Fatal(err) + } + } else { + if _, _, err := eng.Get(makeKey(idx)); err != nil { + b.Fatal(err) + } + } + } + + b.StopTimer() + reportWriteAmp(b, dir) + }) + } +} + +// BenchmarkLSM_ReadHeavy benchmarks a read-dominant workload (10w/90r) to show +// how Bloom filters keep LSM competitive even when reads dominate. +func BenchmarkLSM_ReadHeavy(b *testing.B) { + for _, n := range benchScales { + n := n + b.Run(fmt.Sprintf("N=%d", n), func(b *testing.B) { + dir := b.TempDir() + eng, err := lsm.Open(dir, lsm.Options{}) + if err != nil { + b.Fatal(err) + } + defer eng.Close() + + for i := range n { + if err := eng.Put(makeKey(i), makeVal(i)); err != nil { + b.Fatal(err) + } + } + + rng := rand.New(rand.NewPCG(99, 0)) + b.ResetTimer() + b.SetBytes(int64(keyBytes + valBytes)) + + for range b.N { + idx := rng.IntN(n) + if rng.Float64() < 0.10 { + eng.Put(makeKey(idx), makeVal(idx)) //nolint:errcheck + } else { + eng.Get(makeKey(idx)) //nolint:errcheck + } + } + + b.StopTimer() + reportWriteAmp(b, dir) + }) + } +} + +// --- SQLite B-tree benchmarks ------------------------------------------------- + +// openSQLiteKV opens a SQLite database with a simple KV table tuned to be a +// fair B-tree comparison: WAL mode (same durability model as LSM flush), +// WITHOUT ROWID (clustered primary-key index, minimises SQLite overhead), +// NORMAL synchronous (matching our LSM's lack of fsync per write). +func openSQLiteKV(b *testing.B, dir string) *sql.DB { + b.Helper() + db, err := sql.Open("sqlite", filepath.Join(dir, "bench.db")) + if err != nil { + b.Fatal(err) + } + db.SetMaxOpenConns(1) // WAL mode still only allows one writer + _, err = db.Exec(` + PRAGMA journal_mode = WAL; + PRAGMA synchronous = NORMAL; + PRAGMA cache_size = -8192; + CREATE TABLE IF NOT EXISTS kv (k BLOB PRIMARY KEY, v BLOB) WITHOUT ROWID; + `) + if err != nil { + b.Fatal(err) + } + return db +} + +// BenchmarkSQLite_70w30r benchmarks SQLite's B-tree at the same 70/30 workload. +// This is the apples-to-apples comparison: same key/value sizes, same hit rate. +func BenchmarkSQLite_70w30r(b *testing.B) { + for _, n := range benchScales { + n := n + b.Run(fmt.Sprintf("N=%d", n), func(b *testing.B) { + dir := b.TempDir() + db := openSQLiteKV(b, dir) + defer db.Close() + + // Pre-populate n/2 keys in a single transaction (same as LSM pre-load). + tx, err := db.Begin() + if err != nil { + b.Fatal(err) + } + ins, err := tx.Prepare(`INSERT OR REPLACE INTO kv(k,v) VALUES(?,?)`) + if err != nil { + b.Fatal(err) + } + for i := range n / 2 { + if _, err := ins.Exec(makeKey(i), makeVal(i)); err != nil { + b.Fatal(err) + } + } + ins.Close() + if err := tx.Commit(); err != nil { + b.Fatal(err) + } + + putStmt, err := db.Prepare(`INSERT OR REPLACE INTO kv(k,v) VALUES(?,?)`) + if err != nil { + b.Fatal(err) + } + getStmt, err := db.Prepare(`SELECT v FROM kv WHERE k = ?`) + if err != nil { + b.Fatal(err) + } + defer putStmt.Close() + defer getStmt.Close() + + rng := rand.New(rand.NewPCG(42, 0)) + b.ResetTimer() + b.SetBytes(int64(keyBytes + valBytes)) + + for range b.N { + idx := rng.IntN(n) + if rng.Float64() < writeFraction { + if _, err := putStmt.Exec(makeKey(idx), makeVal(idx)); err != nil { + b.Fatal(err) + } + } else { + row := getStmt.QueryRow(makeKey(idx)) + var val []byte + _ = row.Scan(&val) + } + } + + b.StopTimer() + reportWriteAmp(b, dir) + }) + } +} + +// BenchmarkSQLite_ReadHeavy is the SQLite counterpart to BenchmarkLSM_ReadHeavy. +func BenchmarkSQLite_ReadHeavy(b *testing.B) { + for _, n := range benchScales { + n := n + b.Run(fmt.Sprintf("N=%d", n), func(b *testing.B) { + dir := b.TempDir() + db := openSQLiteKV(b, dir) + defer db.Close() + + tx, _ := db.Begin() + ins, _ := tx.Prepare(`INSERT OR REPLACE INTO kv(k,v) VALUES(?,?)`) + for i := range n { + ins.Exec(makeKey(i), makeVal(i)) //nolint:errcheck + } + ins.Close() + tx.Commit() //nolint:errcheck + + putStmt, _ := db.Prepare(`INSERT OR REPLACE INTO kv(k,v) VALUES(?,?)`) + getStmt, _ := db.Prepare(`SELECT v FROM kv WHERE k = ?`) + defer putStmt.Close() + defer getStmt.Close() + + rng := rand.New(rand.NewPCG(99, 0)) + b.ResetTimer() + b.SetBytes(int64(keyBytes + valBytes)) + + for range b.N { + idx := rng.IntN(n) + if rng.Float64() < 0.10 { + putStmt.Exec(makeKey(idx), makeVal(idx)) //nolint:errcheck + } else { + row := getStmt.QueryRow(makeKey(idx)) + var val []byte + row.Scan(&val) //nolint:errcheck + } + } + + b.StopTimer() + reportWriteAmp(b, dir) + }) + } +} diff --git a/gateway/internal/store/lsm/bloom.go b/gateway/internal/store/lsm/bloom.go new file mode 100644 index 0000000..9a312b6 --- /dev/null +++ b/gateway/internal/store/lsm/bloom.go @@ -0,0 +1,130 @@ +package lsm + +import ( + "encoding/binary" + "fmt" + "hash/fnv" + "math" +) + +// Bloom is a space-efficient probabilistic membership filter. +// +// # Optimal parameters +// +// Given n expected insertions and a target false-positive rate p ∈ (0, 1): +// +// m = ⌈ −n · ln p / (ln 2)² ⌉ — bit-array size +// k = ⌈ (m/n) · ln 2 ⌉ — number of hash functions +// +// # Double hashing (Kirsch-Mitzenmacher) +// +// Instead of k independent hash functions we generate: +// +// g_i(x) = ( h1(x) + i · h2(x) ) mod m, i = 0 … k-1 +// +// using only two FNV-1a-64 evaluations per key. h2 is always forced odd so +// the probe sequence visits every bit position when m is a power of two. +type Bloom struct { + bits []uint64 // bit array, word-aligned + m uint64 // total number of bits + k int // number of hash probes per key +} + +// NewBloom constructs a Bloom filter for n expected insertions at false-positive +// probability p. Both n and p must be positive; p must be < 1. +func NewBloom(n int, p float64) *Bloom { + if n < 1 { + n = 1 + } + m := bloomM(n, p) + k := bloomK(m, n) + return &Bloom{ + bits: make([]uint64, (m+63)/64), + m: m, + k: k, + } +} + +// Add inserts key into the filter. +func (b *Bloom) Add(key []byte) { + h1, h2 := bloomHash(key) + for i := uint64(0); i < uint64(b.k); i++ { + bit := (h1 + i*h2) % b.m + b.bits[bit>>6] |= 1 << (bit & 63) + } +} + +// MayContain reports whether key is possibly in the set. +// A false return is a definitive miss. A true return may be a false positive. +func (b *Bloom) MayContain(key []byte) bool { + h1, h2 := bloomHash(key) + for i := uint64(0); i < uint64(b.k); i++ { + bit := (h1 + i*h2) % b.m + if b.bits[bit>>6]&(1<<(bit&63)) == 0 { + return false + } + } + return true +} + +// Encode serialises the filter for storage inside an SSTable. +// +// Wire format: k:4 | m:8 | bits… (little-endian, 12-byte header) +func (b *Bloom) Encode() []byte { + out := make([]byte, 12+len(b.bits)*8) + binary.LittleEndian.PutUint32(out[0:], uint32(b.k)) + binary.LittleEndian.PutUint64(out[4:], b.m) + for i, w := range b.bits { + binary.LittleEndian.PutUint64(out[12+i*8:], w) + } + return out +} + +// DecodeBloom deserialises a filter produced by Encode. +func DecodeBloom(data []byte) (*Bloom, error) { + const hdr = 12 + if len(data) < hdr { + return nil, fmt.Errorf("lsm/bloom: data too short (%d bytes)", len(data)) + } + k := int(binary.LittleEndian.Uint32(data[0:])) + m := binary.LittleEndian.Uint64(data[4:]) + words := (m + 63) / 64 + need := uint64(hdr) + words*8 + if uint64(len(data)) < need { + return nil, fmt.Errorf("lsm/bloom: truncated bit array (have %d, need %d)", len(data), need) + } + bits := make([]uint64, words) + for i := range bits { + bits[i] = binary.LittleEndian.Uint64(data[hdr+i*8:]) + } + return &Bloom{bits: bits, m: m, k: k}, nil +} + +// bloomHash returns two independent 64-bit hashes of key. +// h2 is forced odd so gcd(h2, m) = 1 when m is a power of two. +func bloomHash(key []byte) (h1, h2 uint64) { + a := fnv.New64a() + a.Write(key) + h1 = a.Sum64() + + // Mix in a fixed suffix to de-correlate h2 from h1. + b := fnv.New64a() + b.Write(key) + b.Write([]byte{0xCA, 0xFE, 0xBA, 0xBE}) + h2 = b.Sum64() | 1 + return +} + +// bloomM computes the optimal bit-array size. +func bloomM(n int, p float64) uint64 { + return uint64(math.Ceil(-float64(n) * math.Log(p) / (math.Log(2) * math.Log(2)))) +} + +// bloomK computes the optimal number of hash functions. +func bloomK(m uint64, n int) int { + k := int(math.Round(float64(m) / float64(n) * math.Log(2))) + if k < 1 { + k = 1 + } + return k +} diff --git a/gateway/internal/store/lsm/compactor.go b/gateway/internal/store/lsm/compactor.go new file mode 100644 index 0000000..9952fcc --- /dev/null +++ b/gateway/internal/store/lsm/compactor.go @@ -0,0 +1,338 @@ +package lsm + +import ( + "bytes" + "container/heap" + "fmt" + "os" +) + +// compactor runs background leveled compaction for the Engine. +// +// # Leveled compaction rules +// +// - L0 holds direct MemTable flushes; keys may overlap across files. +// Compaction is triggered when |L0| ≥ opts.MaxL0Tables. +// All L0 tables are merged into new L1 table(s). +// +// - L1+ holds non-overlapping key-range files sorted by first key. +// Compaction is triggered when ΣfileSize(Lk) > levelBudget(k). +// The oldest file in the overflowing level is merged into Lk+1. +// +// The compactor is the only writer to the SSTable file list; it holds +// the engine write-lock only for the brief atomic swap at the end. +type compactor struct { + eng *Engine +} + +// maybeCompact checks all levels and runs one round of compaction if needed. +// It is called from the background goroutine after each flush or previous compaction. +func (c *compactor) maybeCompact() { + eng := c.eng + + eng.mu.RLock() + l0count := len(eng.levels[0]) + eng.mu.RUnlock() + + if l0count >= eng.opts.MaxL0Tables { + if err := c.compactL0(); err != nil { + // Log and continue; compaction is best-effort. + _ = err + } + } + + // Check higher levels. + for lvl := 1; ; lvl++ { + eng.mu.RLock() + if lvl >= len(eng.levels) { + eng.mu.RUnlock() + break + } + tables := make([]*SSTable, len(eng.levels[lvl])) + copy(tables, eng.levels[lvl]) + eng.mu.RUnlock() + + budget := eng.levelBudget(lvl) + if totalSize(tables) <= budget { + break + } + if err := c.compactLevel(lvl, tables); err != nil { + break + } + } +} + +// compactL0 merges all L0 SSTables into new L1 SSTable(s). +func (c *compactor) compactL0() error { + eng := c.eng + + eng.mu.RLock() + l0 := make([]*SSTable, len(eng.levels[0])) + copy(l0, eng.levels[0]) + var l1 []*SSTable + if len(eng.levels) > 1 { + l1 = make([]*SSTable, len(eng.levels[1])) + copy(l1, eng.levels[1]) + } + eng.mu.RUnlock() + + // Build iterators: L0 newest-first (highest priority), then L1. + // When keys collide, the higher-priority (newer) entry wins. + iters := make([]iterWithPrio, 0, len(l0)+len(l1)) + for i, t := range l0 { + it := newSSTIter(t) + it.Next() + if it.Valid() { + iters = append(iters, iterWithPrio{it: it, prio: len(l0) - i}) + } + } + for _, t := range l1 { + it := newSSTIter(t) + it.Next() + if it.Valid() { + iters = append(iters, iterWithPrio{it: it, prio: 0}) + } + } + + newL1, err := c.mergeToLevel(1, iters) + if err != nil { + return err + } + + // Atomically replace L0 and L1 with the compaction output. + eng.mu.Lock() + old := append(l0, l1...) + eng.ensureLevel(1) + eng.levels[0] = nil + eng.levels[1] = newL1 + eng.mu.Unlock() + + // Delete old files after releasing the lock. + for _, t := range old { + t.Close() + os.Remove(t.Path()) + } + return nil +} + +// compactLevel merges the oldest SSTable in level lvl into level lvl+1. +func (c *compactor) compactLevel(lvl int, tables []*SSTable) error { + eng := c.eng + + // Pick the oldest table (index 0 = oldest by insertion order). + victim := tables[0] + + eng.mu.Lock() + eng.ensureLevel(lvl + 1) + next := make([]*SSTable, len(eng.levels[lvl+1])) + copy(next, eng.levels[lvl+1]) + eng.mu.Unlock() + + // Find overlapping tables in lvl+1. + overlapping := overlappingTables(next, victim) + + inputs := append([]*SSTable{victim}, overlapping...) + iters := make([]iterWithPrio, 0, len(inputs)) + for i, t := range inputs { + it := newSSTIter(t) + it.Next() + if it.Valid() { + iters = append(iters, iterWithPrio{it: it, prio: len(inputs) - i}) + } + } + + newTables, err := c.mergeToLevel(lvl+1, iters) + if err != nil { + return err + } + + eng.mu.Lock() + // Remove victim from lvl. + remaining := make([]*SSTable, 0, len(tables)-1) + for _, t := range tables { + if t != victim { + remaining = append(remaining, t) + } + } + eng.levels[lvl] = remaining + // Replace overlapping in lvl+1 with newTables. + eng.ensureLevel(lvl + 1) + replacement := make([]*SSTable, 0, len(next)-len(overlapping)+len(newTables)) + overlapSet := make(map[*SSTable]bool, len(overlapping)) + for _, t := range overlapping { + overlapSet[t] = true + } + for _, t := range next { + if !overlapSet[t] { + replacement = append(replacement, t) + } + } + replacement = append(replacement, newTables...) + eng.levels[lvl+1] = replacement + eng.mu.Unlock() + + for _, t := range inputs { + t.Close() + os.Remove(t.Path()) + } + return nil +} + +// mergeToLevel runs a k-way merge of iters and writes new SSTables at the +// given level. Each output file is capped at opts.MaxTableSize bytes. +// The last tombstone seen for a key at the deepest level is dropped. +func (c *compactor) mergeToLevel(lvl int, iters []iterWithPrio) ([]*SSTable, error) { + eng := c.eng + mh := make(mergeHeap, 0, len(iters)) + for _, ip := range iters { + mh = append(mh, ip) + } + heap.Init(&mh) + + var ( + out []*SSTable + builder *Builder + written int64 + lastKey []byte + ) + + flush := func() error { + if builder == nil { + return nil + } + if err := builder.Finish(); err != nil { + return err + } + t, err := openSSTable(builder.path) + if err != nil { + return err + } + out = append(out, t) + builder = nil + written = 0 + return nil + } + + for mh.Len() > 0 { + top := heap.Pop(&mh).(iterWithPrio) + + key := top.it.Key() + val := top.it.Val() + + // Skip duplicate keys — only the highest-priority entry counts. + isDup := bytes.Equal(key, lastKey) + lastKey = append(lastKey[:0], key...) + + // Advance the iterator and push it back if still valid. + top.it.Next() + if top.it.Valid() { + heap.Push(&mh, top) + } + + // Also drain any lower-priority entries for the same key. + for mh.Len() > 0 && bytes.Equal(mh[0].it.Key(), key) { + dup := heap.Pop(&mh).(iterWithPrio) + dup.it.Next() + if dup.it.Valid() { + heap.Push(&mh, dup) + } + } + + if isDup { + continue + } + // Drop tombstones at the bottom level (no older data can exist). + if val == nil && lvl >= len(eng.levels)-1 { + continue + } + + if builder == nil { + path := eng.newTablePath(lvl) + var err error + builder, err = newBuilder(path, eng.opts) + if err != nil { + return nil, fmt.Errorf("lsm: compaction builder: %w", err) + } + } + if err := builder.Add(key, val); err != nil { + return nil, err + } + written += int64(len(key) + len(val) + 6) + if written >= eng.opts.MaxTableSize { + if err := flush(); err != nil { + return nil, err + } + } + } + if err := flush(); err != nil { + return nil, err + } + return out, nil +} + +// ─── k-way merge helpers ────────────────────────────────────────────────────── + +// iterWithPrio wraps an SSTIter with a priority for tie-breaking. +// Higher priority = more recent data (wins when keys collide). +type iterWithPrio struct { + it *SSTIter + prio int +} + +// mergeHeap implements heap.Interface over iterWithPrio. +// It sorts by key ASC, then by priority DESC (higher priority = pops first). +type mergeHeap []iterWithPrio + +func (h mergeHeap) Len() int { return len(h) } +func (h mergeHeap) Less(i, j int) bool { + cmp := bytes.Compare(h[i].it.Key(), h[j].it.Key()) + if cmp != 0 { + return cmp < 0 + } + return h[i].prio > h[j].prio // higher priority first +} +func (h mergeHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } +func (h *mergeHeap) Push(x any) { *h = append(*h, x.(iterWithPrio)) } +func (h *mergeHeap) Pop() any { + old := *h + n := len(old) + x := old[n-1] + *h = old[:n-1] + return x +} + +// ─── Utility helpers ────────────────────────────────────────────────────────── + +// totalSize sums the on-disk sizes of the given SSTables. +func totalSize(tables []*SSTable) int64 { + total := int64(0) + for _, t := range tables { + if fi, err := t.f.Stat(); err == nil { + total += fi.Size() + } + } + return total +} + +// overlappingTables returns tables from next whose key range overlaps victim. +func overlappingTables(next []*SSTable, victim *SSTable) []*SSTable { + if len(victim.index) == 0 { + return nil + } + vMin := victim.index[0].firstKey + vMax := victim.index[len(victim.index)-1].firstKey + + var out []*SSTable + for _, t := range next { + if len(t.index) == 0 { + continue + } + tMin := t.index[0].firstKey + tMax := t.index[len(t.index)-1].firstKey + // Overlap if vMin ≤ tMax and tMin ≤ vMax. + if bytes.Compare(vMin, tMax) <= 0 && bytes.Compare(tMin, vMax) <= 0 { + out = append(out, t) + } + } + return out +} diff --git a/gateway/internal/store/lsm/engine.go b/gateway/internal/store/lsm/engine.go new file mode 100644 index 0000000..3cfc1bd --- /dev/null +++ b/gateway/internal/store/lsm/engine.go @@ -0,0 +1,414 @@ +// Package lsm implements a Log-Structured Merge-Tree (LSM-Tree) storage engine. +// +// # Architecture +// +// Write path: Put/Delete → MemTable (skip list, in-memory sorted map) +// → when full, freeze & flush → L0 SSTable on disk +// → when |L0| ≥ MaxL0Tables, compact → L1 SSTable(s) +// → repeat for L1 → L2 → … when size budget exceeded +// +// Read path: Get → MemTable → immutable MemTable → L0 (newest-first) +// → L1 … (Bloom filter short-circuits most disk reads) +// +// # File naming +// +// {level:02d}-{seq:010d}.sst e.g. "00-0000000001.sst" (L0), "01-0000000002.sst" (L1) +// +// # Durability +// +// This implementation does not include a Write-Ahead Log (WAL). Entries in the +// active MemTable are lost on crash. Add a WAL for production use. +// +// # Concurrency +// +// Engine.Get is safe for concurrent use. +// Engine.Put / Engine.Delete acquire a write lock for the MemTable update only. +// SSTable reads use os.File.ReadAt which is concurrent-safe on all major OSes. +package lsm + +import ( + "fmt" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "sync/atomic" +) + +// Options configures an LSM-Tree Engine. +type Options struct { + // MemTableSize is the size threshold (bytes) that triggers a MemTable flush. + // Default: 4 MiB. + MemTableSize int64 + + // MaxL0Tables is the number of L0 SSTables that triggers L0 → L1 compaction. + // Default: 4. + MaxL0Tables int + + // BlockSize is the target data block size inside each SSTable (bytes). + // Default: 4096. + BlockSize int + + // BloomFPRate is the target false-positive rate for Bloom filters. + // Default: 0.01 (1%). + BloomFPRate float64 + + // BloomExpected is the expected number of keys per SSTable for sizing the + // Bloom filter. Default: 65536. + BloomExpected int + + // MaxTableSize caps individual SSTable file sizes during compaction (bytes). + // Default: 2 MiB. + MaxTableSize int64 + + // L1Budget is the total SSTable size allowed at L1 before L1→L2 compaction. + // Each subsequent level is 10× larger. Default: 10 MiB. + L1Budget int64 +} + +func (o *Options) setDefaults() { + if o.MemTableSize == 0 { + o.MemTableSize = 4 * 1024 * 1024 + } + if o.MaxL0Tables == 0 { + o.MaxL0Tables = 4 + } + if o.BlockSize == 0 { + o.BlockSize = 4096 + } + if o.BloomFPRate == 0 { + o.BloomFPRate = 0.01 + } + if o.BloomExpected == 0 { + o.BloomExpected = 65536 + } + if o.MaxTableSize == 0 { + o.MaxTableSize = 2 * 1024 * 1024 + } + if o.L1Budget == 0 { + o.L1Budget = 10 * 1024 * 1024 + } +} + +// Engine is an LSM-Tree key-value store. +// Open it with Open; close it with Close when done. +type Engine struct { + mu sync.RWMutex + mem *MemTable // active mutable MemTable + imm *MemTable // immutable MemTable waiting to be flushed (nil if none) + + // levels[0] = L0 (may have overlapping key ranges) + // levels[k] = Lk (non-overlapping after compaction) + levels [][]*SSTable + + dir string + opts Options + seq atomic.Int64 // monotonically increasing SSTable sequence number + + compactC chan struct{} // signal background compactor to wake up + done chan struct{} + wg sync.WaitGroup +} + +// Open opens or creates the Engine at dir. +// The directory is created if it does not exist. +func Open(dir string, opts Options) (*Engine, error) { + opts.setDefaults() + if err := os.MkdirAll(dir, 0o755); err != nil { + return nil, fmt.Errorf("lsm: mkdir %s: %w", dir, err) + } + + e := &Engine{ + dir: dir, + opts: opts, + mem: newMemTable(opts.MemTableSize), + levels: make([][]*SSTable, 1), // start with L0 + compactC: make(chan struct{}, 1), + done: make(chan struct{}), + } + + if err := e.loadExistingTables(); err != nil { + return nil, err + } + + e.wg.Add(1) + go e.backgroundLoop() + return e, nil +} + +// loadExistingTables scans dir for *.sst files and opens them. +func (e *Engine) loadExistingTables() error { + entries, err := os.ReadDir(e.dir) + if err != nil { + return fmt.Errorf("lsm: read dir: %w", err) + } + + type tableInfo struct { + lvl int + seq int64 + path string + } + var tables []tableInfo + maxSeq := int64(0) + + for _, de := range entries { + name := de.Name() + if !strings.HasSuffix(name, ".sst") { + continue + } + // Expected: "{lvl:02d}-{seq:010d}.sst" + parts := strings.SplitN(strings.TrimSuffix(name, ".sst"), "-", 2) + if len(parts) != 2 { + continue + } + lvl, err1 := strconv.Atoi(parts[0]) + seq, err2 := strconv.ParseInt(parts[1], 10, 64) + if err1 != nil || err2 != nil || lvl < 0 { + continue + } + tables = append(tables, tableInfo{lvl: lvl, seq: seq, path: filepath.Join(e.dir, name)}) + if seq > maxSeq { + maxSeq = seq + } + } + e.seq.Store(maxSeq) + + // Sort within each level by sequence number (ascending = oldest first). + sort.Slice(tables, func(i, j int) bool { + if tables[i].lvl != tables[j].lvl { + return tables[i].lvl < tables[j].lvl + } + return tables[i].seq < tables[j].seq + }) + + for _, ti := range tables { + t, err := openSSTable(ti.path) + if err != nil { + return fmt.Errorf("lsm: open %s: %w", ti.path, err) + } + e.ensureLevel(ti.lvl) + e.levels[ti.lvl] = append(e.levels[ti.lvl], t) + } + return nil +} + +// Get retrieves the value for key. +// +// - (val, true, nil): key found +// - (nil, false, nil): key absent +// - (_, _, err): I/O error +// +// Tombstones (deleted keys) are invisible: (nil, false, nil) is returned. +func (e *Engine) Get(key []byte) (val []byte, found bool, err error) { + e.mu.RLock() + defer e.mu.RUnlock() + + // 1. Active MemTable. + if v, ok, tomb := e.mem.Get(key); ok { + if tomb { + return nil, false, nil + } + return v, true, nil + } + + // 2. Immutable MemTable (being flushed). + if e.imm != nil { + if v, ok, tomb := e.imm.Get(key); ok { + if tomb { + return nil, false, nil + } + return v, true, nil + } + } + + // 3. L0 — newest-first (L0 tables may have overlapping key ranges). + l0 := e.levels[0] + for i := len(l0) - 1; i >= 0; i-- { + v, ok, err := l0[i].Get(key) + if err != nil { + return nil, false, err + } + if ok { + if v == nil { + return nil, false, nil // tombstone + } + return v, true, nil + } + } + + // 4. L1 and deeper — non-overlapping key ranges; check each level linearly + // (Bloom filters make most checks free). + for lvl := 1; lvl < len(e.levels); lvl++ { + for _, t := range e.levels[lvl] { + v, ok, err := t.Get(key) + if err != nil { + return nil, false, err + } + if ok { + if v == nil { + return nil, false, nil // tombstone + } + return v, true, nil + } + } + } + return nil, false, nil +} + +// Put inserts or updates key → val. +func (e *Engine) Put(key, val []byte) error { + e.mu.Lock() + full := e.mem.Put(key, val) + e.mu.Unlock() + if full { + e.triggerFlush() + } + return nil +} + +// Delete removes key by writing a tombstone. +func (e *Engine) Delete(key []byte) error { + e.mu.Lock() + full := e.mem.Delete(key) + e.mu.Unlock() + if full { + e.triggerFlush() + } + return nil +} + +// Close flushes all pending writes and shuts down the background goroutine. +func (e *Engine) Close() error { + // Flush the active MemTable. + e.mu.Lock() + if e.mem.Count() > 0 { + e.imm = e.mem + e.mem = newMemTable(e.opts.MemTableSize) + } + e.mu.Unlock() + + if e.imm != nil { + if err := e.flushImm(); err != nil { + return err + } + } + + // Signal background loop to exit. + close(e.done) + e.wg.Wait() + + // Close all SSTable file handles. + e.mu.Lock() + defer e.mu.Unlock() + for _, level := range e.levels { + for _, t := range level { + t.Close() + } + } + return nil +} + +// ─── Internal: flush & compaction ──────────────────────────────────────────── + +// triggerFlush freezes the active MemTable and signals the background loop. +func (e *Engine) triggerFlush() { + e.mu.Lock() + if e.imm == nil { + e.imm = e.mem + e.mem = newMemTable(e.opts.MemTableSize) + } + e.mu.Unlock() + select { + case e.compactC <- struct{}{}: + default: + } +} + +// backgroundLoop handles flushes and compaction in a single goroutine. +func (e *Engine) backgroundLoop() { + defer e.wg.Done() + c := &compactor{eng: e} + for { + select { + case <-e.done: + return + case <-e.compactC: + // Flush immutable MemTable first. + e.mu.RLock() + hasImm := e.imm != nil + e.mu.RUnlock() + if hasImm { + _ = e.flushImm() + } + // Then check whether compaction is needed. + c.maybeCompact() + } + } +} + +// flushImm writes the immutable MemTable to a new L0 SSTable. +func (e *Engine) flushImm() error { + e.mu.RLock() + imm := e.imm + e.mu.RUnlock() + if imm == nil { + return nil + } + + path := e.newTablePath(0) + b, err := newBuilder(path, e.opts) + if err != nil { + return err + } + + it := imm.Iterator() + for it.Valid() { + if err := b.Add(it.Key(), it.Val()); err != nil { + return err + } + it.Next() + } + if err := b.Finish(); err != nil { + return err + } + + t, err := openSSTable(path) + if err != nil { + return err + } + + e.mu.Lock() + e.ensureLevel(0) + e.levels[0] = append(e.levels[0], t) + e.imm = nil + e.mu.Unlock() + return nil +} + +// ─── Level / path helpers ───────────────────────────────────────────────────── + +// newTablePath allocates a new unique SSTable file path for the given level. +func (e *Engine) newTablePath(lvl int) string { + seq := e.seq.Add(1) + return filepath.Join(e.dir, fmt.Sprintf("%02d-%010d.sst", lvl, seq)) +} + +// levelBudget returns the maximum total SSTable size allowed at level lvl. +// Budget grows 10× per level: L1 = L1Budget, L2 = 10×L1Budget, etc. +func (e *Engine) levelBudget(lvl int) int64 { + budget := e.opts.L1Budget + for i := 1; i < lvl; i++ { + budget *= 10 + } + return budget +} + +// ensureLevel guarantees that e.levels has at least lvl+1 slots. +// Caller must hold e.mu for writing. +func (e *Engine) ensureLevel(lvl int) { + for len(e.levels) <= lvl { + e.levels = append(e.levels, nil) + } +} diff --git a/gateway/internal/store/lsm/engine_test.go b/gateway/internal/store/lsm/engine_test.go new file mode 100644 index 0000000..9547547 --- /dev/null +++ b/gateway/internal/store/lsm/engine_test.go @@ -0,0 +1,514 @@ +package lsm_test + +import ( + "fmt" + "math/rand/v2" + "os" + "path/filepath" + "sort" + "testing" + "time" + + "github.com/Revanth14/indexqube/gateway/internal/store/lsm" +) + +// ── Bloom filter ────────────────────────────────────────────────────────────── + +func TestBloom_NoFalseNegatives(t *testing.T) { + b := lsm.NewBloom(1000, 0.01) + keys := make([][]byte, 1000) + for i := range keys { + keys[i] = []byte(fmt.Sprintf("key-%04d", i)) + b.Add(keys[i]) + } + for _, k := range keys { + if !b.MayContain(k) { + t.Errorf("false negative for key %q", k) + } + } +} + +func TestBloom_FalsePositiveRate(t *testing.T) { + const n = 10_000 + const wantFPR = 0.01 + b := lsm.NewBloom(n, wantFPR) + for i := range n { + b.Add([]byte(fmt.Sprintf("inserted-%d", i))) + } + fps := 0 + const probes = 100_000 + for i := range probes { + // Keys that were never inserted. + if b.MayContain([]byte(fmt.Sprintf("absent-%d", i))) { + fps++ + } + } + gotFPR := float64(fps) / probes + // Allow 3× the target rate as slack for the statistical nature of Bloom filters. + if gotFPR > wantFPR*3 { + t.Errorf("false positive rate %.4f exceeds 3× target %.4f", gotFPR, wantFPR) + } +} + +func TestBloom_EncodeDecode(t *testing.T) { + b := lsm.NewBloom(500, 0.01) + for i := range 500 { + b.Add([]byte(fmt.Sprintf("k%d", i))) + } + enc := b.Encode() + b2, err := lsm.DecodeBloom(enc) + if err != nil { + t.Fatalf("DecodeBloom: %v", err) + } + for i := range 500 { + if !b2.MayContain([]byte(fmt.Sprintf("k%d", i))) { + t.Errorf("decoded bloom: false negative for k%d", i) + } + } +} + +// ── MemTable ────────────────────────────────────────────────────────────────── + +func TestMemTable_PutGet(t *testing.T) { + m := lsm.NewMemTableForTest(4 * 1024 * 1024) + m.Put([]byte("hello"), []byte("world")) + v, found, tomb := m.Get([]byte("hello")) + if !found || tomb || string(v) != "world" { + t.Errorf("Get(hello) = %q, %v, %v; want world, true, false", v, found, tomb) + } +} + +func TestMemTable_UpdateInPlace(t *testing.T) { + m := lsm.NewMemTableForTest(4 * 1024 * 1024) + m.Put([]byte("k"), []byte("v1")) + m.Put([]byte("k"), []byte("v2")) + v, _, _ := m.Get([]byte("k")) + if string(v) != "v2" { + t.Errorf("after update: got %q; want v2", v) + } +} + +func TestMemTable_Delete(t *testing.T) { + m := lsm.NewMemTableForTest(4 * 1024 * 1024) + m.Put([]byte("x"), []byte("1")) + m.Delete([]byte("x")) + _, found, tomb := m.Get([]byte("x")) + if !found || !tomb { + t.Errorf("after Delete: found=%v tomb=%v; want true, true", found, tomb) + } +} + +func TestMemTable_Iterator_Order(t *testing.T) { + m := lsm.NewMemTableForTest(4 * 1024 * 1024) + keys := []string{"banana", "apple", "cherry", "apricot", "date"} + for _, k := range keys { + m.Put([]byte(k), []byte("v")) + } + sort.Strings(keys) + it := m.Iterator() + for _, want := range keys { + if !it.Valid() { + t.Fatalf("iterator ended early; want key %q", want) + } + if got := string(it.Key()); got != want { + t.Errorf("key = %q; want %q", got, want) + } + it.Next() + } + if it.Valid() { + t.Errorf("iterator has extra entries after all keys consumed") + } +} + +// ── SSTable builder + reader ────────────────────────────────────────────────── + +func TestSSTable_BuildRead(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "test.sst") + opts := lsm.DefaultOptions() + + b, err := lsm.NewBuilderForTest(path, opts) + if err != nil { + t.Fatalf("newBuilder: %v", err) + } + entries := [][2]string{ + {"apple", "1"}, {"banana", "2"}, {"cherry", "3"}, + {"date", "4"}, {"elderberry", "5"}, + } + for _, e := range entries { + b.Add([]byte(e[0]), []byte(e[1])) + } + if err := b.Finish(); err != nil { + t.Fatalf("Finish: %v", err) + } + + sst, err := lsm.OpenSSTableForTest(path) + if err != nil { + t.Fatalf("openSSTable: %v", err) + } + defer sst.Close() + + for _, e := range entries { + v, found, err := sst.Get([]byte(e[0])) + if err != nil { + t.Fatalf("Get(%q): %v", e[0], err) + } + if !found { + t.Errorf("Get(%q): not found", e[0]) + continue + } + if string(v) != e[1] { + t.Errorf("Get(%q) = %q; want %q", e[0], v, e[1]) + } + } + + // Absent key — Bloom should short-circuit. + v, found, err := sst.Get([]byte("zzz")) + if err != nil || found || v != nil { + t.Errorf("Get(zzz) = %q, %v, %v; want nil, false, nil", v, found, err) + } +} + +func TestSSTable_Tombstone(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "tomb.sst") + opts := lsm.DefaultOptions() + + b, _ := lsm.NewBuilderForTest(path, opts) + b.Add([]byte("alive"), []byte("yes")) + b.Add([]byte("dead"), nil) // tombstone + b.Finish() + + sst, _ := lsm.OpenSSTableForTest(path) + defer sst.Close() + + v, found, _ := sst.Get([]byte("alive")) + if !found || string(v) != "yes" { + t.Errorf("alive key: found=%v val=%q", found, v) + } + v, found, _ = sst.Get([]byte("dead")) + if !found || v != nil { + t.Errorf("dead key: found=%v val=%q; want found=true val=nil", found, v) + } +} + +func TestSSTable_Iterator(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "iter.sst") + opts := lsm.DefaultOptions() + + keys := []string{"a", "b", "c", "d", "e", "f", "g"} + b, _ := lsm.NewBuilderForTest(path, opts) + for _, k := range keys { + b.Add([]byte(k), []byte("v")) + } + b.Finish() + + sst, _ := lsm.OpenSSTableForTest(path) + defer sst.Close() + + it := lsm.NewSSTIterForTest(sst) + var got []string + for it.Next() { + got = append(got, string(it.Key())) + } + if len(got) != len(keys) { + t.Fatalf("iterated %d keys; want %d", len(got), len(keys)) + } + for i, k := range keys { + if got[i] != k { + t.Errorf("key[%d] = %q; want %q", i, got[i], k) + } + } +} + +// ── Engine: end-to-end ──────────────────────────────────────────────────────── + +func TestEngine_PutGet(t *testing.T) { + e := openTestEngine(t) + defer e.Close() + + if err := e.Put([]byte("foo"), []byte("bar")); err != nil { + t.Fatalf("Put: %v", err) + } + v, found, err := e.Get([]byte("foo")) + if err != nil || !found || string(v) != "bar" { + t.Errorf("Get(foo) = %q, %v, %v; want bar, true, nil", v, found, err) + } +} + +func TestEngine_Delete(t *testing.T) { + e := openTestEngine(t) + defer e.Close() + + e.Put([]byte("k"), []byte("v")) + e.Delete([]byte("k")) + _, found, _ := e.Get([]byte("k")) + if found { + t.Error("Get after Delete: found=true; want false") + } +} + +func TestEngine_Overwrite(t *testing.T) { + e := openTestEngine(t) + defer e.Close() + + e.Put([]byte("k"), []byte("v1")) + e.Put([]byte("k"), []byte("v2")) + v, _, _ := e.Get([]byte("k")) + if string(v) != "v2" { + t.Errorf("Get after overwrite = %q; want v2", v) + } +} + +func TestEngine_AbsentKey(t *testing.T) { + e := openTestEngine(t) + defer e.Close() + + _, found, err := e.Get([]byte("nosuchkey")) + if err != nil || found { + t.Errorf("Get(absent) = found=%v err=%v; want false, nil", found, err) + } +} + +func TestEngine_FlushAndRecover(t *testing.T) { + dir := t.TempDir() + opts := lsm.DefaultOptions() + opts.MemTableSize = 512 // tiny threshold to force flushes + + e, err := lsm.Open(dir, opts) + if err != nil { + t.Fatalf("Open: %v", err) + } + + const n = 200 + for i := range n { + key := []byte(fmt.Sprintf("key-%04d", i)) + val := []byte(fmt.Sprintf("val-%04d", i)) + if err := e.Put(key, val); err != nil { + t.Fatalf("Put %d: %v", i, err) + } + } + // Allow background flush to complete. + time.Sleep(50 * time.Millisecond) + if err := e.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + // Re-open and verify SSTable persistence. + e2, err := lsm.Open(dir, opts) + if err != nil { + t.Fatalf("reopen: %v", err) + } + defer e2.Close() + + // At least some keys should survive in L0 SSTables. + hits := 0 + for i := range n { + key := []byte(fmt.Sprintf("key-%04d", i)) + want := fmt.Sprintf("val-%04d", i) + v, found, err := e2.Get(key) + if err != nil { + t.Fatalf("Get key-%04d: %v", i, err) + } + if found && string(v) == want { + hits++ + } + } + if hits == 0 { + t.Error("no keys survived after engine re-open") + } + t.Logf("re-open: %d/%d keys found in SSTables", hits, n) +} + +// TestEngine_L0Compaction writes enough data to trigger L0→L1 compaction +// and verifies all keys remain readable. +func TestEngine_L0Compaction(t *testing.T) { + dir := t.TempDir() + opts := lsm.DefaultOptions() + opts.MemTableSize = 1024 // flush every ~1 KB + opts.MaxL0Tables = 2 // compact L0 after 2 tables + opts.BloomExpected = 32 + + e, err := lsm.Open(dir, opts) + if err != nil { + t.Fatalf("Open: %v", err) + } + defer e.Close() + + const n = 500 + for i := range n { + key := []byte(fmt.Sprintf("ckey-%06d", i)) + val := []byte(fmt.Sprintf("cval-%06d", i)) + if err := e.Put(key, val); err != nil { + t.Fatalf("Put %d: %v", i, err) + } + } + // Give background compactor time to run. + time.Sleep(100 * time.Millisecond) + + for i := range n { + key := []byte(fmt.Sprintf("ckey-%06d", i)) + want := fmt.Sprintf("cval-%06d", i) + v, found, err := e.Get(key) + if err != nil { + t.Fatalf("Get[%d]: %v", i, err) + } + if !found { + t.Errorf("key-%06d: not found after compaction", i) + continue + } + if string(v) != want { + t.Errorf("key-%06d: got %q; want %q", i, v, want) + } + } +} + +// TestEngine_DeleteAfterFlush verifies that a tombstone written after a key has +// been flushed to L0 correctly hides the value on Get. +func TestEngine_DeleteAfterFlush(t *testing.T) { + dir := t.TempDir() + opts := lsm.DefaultOptions() + opts.MemTableSize = 512 + + e, _ := lsm.Open(dir, opts) + defer e.Close() + + e.Put([]byte("tkey"), []byte("tval")) + time.Sleep(30 * time.Millisecond) // allow flush to L0 + e.Delete([]byte("tkey")) + + _, found, _ := e.Get([]byte("tkey")) + if found { + t.Error("key visible after Delete; tombstone did not mask flushed value") + } +} + +// ── Benchmarks ──────────────────────────────────────────────────────────────── + +// BenchmarkEngine_Put measures Put throughput with all I/O serialised. +func BenchmarkEngine_Put(b *testing.B) { + e := openBenchEngine(b) + defer e.Close() + key := make([]byte, 16) + val := make([]byte, 64) + b.SetBytes(int64(len(key) + len(val))) + b.ReportAllocs() + b.ResetTimer() + for i := b.N; i > 0; i-- { + writeUint64(key, uint64(i)) + e.Put(key, val) //nolint:errcheck + } +} + +// BenchmarkEngine_Get_MemTable measures Get on an in-memory-only dataset +// (no disk I/O, no Bloom filter cost). +func BenchmarkEngine_Get_MemTable(b *testing.B) { + e := openBenchEngine(b) + defer e.Close() + + const n = 10_000 + keys := make([][]byte, n) + for i := range keys { + keys[i] = []byte(fmt.Sprintf("bench-key-%06d", i)) + e.Put(keys[i], []byte("bench-value")) + } + b.ReportAllocs() + b.ResetTimer() + for i := b.N; i > 0; i-- { + e.Get(keys[i%n]) //nolint:errcheck + } +} + +// BenchmarkBloom_Add and BenchmarkBloom_MayContain measure the filter's hot path. +func BenchmarkBloom_Add(b *testing.B) { + bl := lsm.NewBloom(b.N+1, 0.01) + key := make([]byte, 16) + b.ReportAllocs() + b.ResetTimer() + for i := b.N; i > 0; i-- { + writeUint64(key, uint64(i)) + bl.Add(key) + } +} + +func BenchmarkBloom_MayContain(b *testing.B) { + const n = 100_000 + bl := lsm.NewBloom(n, 0.01) + keys := make([][]byte, n) + for i := range keys { + keys[i] = []byte(fmt.Sprintf("bkey-%d", i)) + bl.Add(keys[i]) + } + b.ReportAllocs() + b.ResetTimer() + for i := b.N; i > 0; i-- { + bl.MayContain(keys[i%n]) + } +} + +// BenchmarkMemTable_Put measures skip-list insertion. +func BenchmarkMemTable_Put(b *testing.B) { + m := lsm.NewMemTableForTest(1 << 30) + key := make([]byte, 16) + val := make([]byte, 64) + b.ReportAllocs() + b.ResetTimer() + for i := b.N; i > 0; i-- { + writeUint64(key, uint64(i)) + m.Put(key, val) + } +} + +// BenchmarkMemTable_Get measures skip-list point-lookup. +func BenchmarkMemTable_Get(b *testing.B) { + const n = 100_000 + m := lsm.NewMemTableForTest(1 << 30) + keys := make([][]byte, n) + for i := range keys { + keys[i] = []byte(fmt.Sprintf("k%d", rand.IntN(1<<24))) + m.Put(keys[i], []byte("v")) + } + b.ReportAllocs() + b.ResetTimer() + for i := b.N; i > 0; i-- { + m.Get(keys[i%n]) + } +} + +// ── Helpers ─────────────────────────────────────────────────────────────────── + +func openTestEngine(t *testing.T) *lsm.Engine { + t.Helper() + dir := t.TempDir() + e, err := lsm.Open(dir, lsm.DefaultOptions()) + if err != nil { + t.Fatalf("Open: %v", err) + } + return e +} + +func openBenchEngine(b *testing.B) *lsm.Engine { + b.Helper() + dir, err := os.MkdirTemp("", "lsm-bench-*") + if err != nil { + b.Fatalf("MkdirTemp: %v", err) + } + b.Cleanup(func() { os.RemoveAll(dir) }) + e, err := lsm.Open(dir, lsm.DefaultOptions()) + if err != nil { + b.Fatalf("Open: %v", err) + } + return e +} + +func writeUint64(b []byte, v uint64) { + b[0] = byte(v >> 56) + b[1] = byte(v >> 48) + b[2] = byte(v >> 40) + b[3] = byte(v >> 32) + b[4] = byte(v >> 24) + b[5] = byte(v >> 16) + b[6] = byte(v >> 8) + b[7] = byte(v) +} diff --git a/gateway/internal/store/lsm/export_test.go b/gateway/internal/store/lsm/export_test.go new file mode 100644 index 0000000..480668d --- /dev/null +++ b/gateway/internal/store/lsm/export_test.go @@ -0,0 +1,31 @@ +package lsm + +// export_test.go exposes internal constructors and types to the _test package +// without polluting the public API. + +// NewMemTableForTest creates a MemTable with the given max size. +func NewMemTableForTest(maxSize int64) *MemTable { return newMemTable(maxSize) } + +// MemTable.Put/Get/Delete/Iterator are already exported on the type. + +// NewBuilderForTest opens a Builder for testing. +func NewBuilderForTest(path string, opts Options) (*Builder, error) { + return newBuilder(path, opts) +} + +// Builder.Add and Builder.Finish are already exported on the type. + +// OpenSSTableForTest opens an SSTable for testing. +func OpenSSTableForTest(path string) (*SSTable, error) { return openSSTable(path) } + +// SSTable.Get, Close, Path are already exported on the type. + +// NewSSTIterForTest creates an SSTable iterator for testing. +func NewSSTIterForTest(sst *SSTable) *SSTIter { return newSSTIter(sst) } + +// DefaultOptions returns options with all defaults applied. +func DefaultOptions() Options { + o := Options{} + o.setDefaults() + return o +} diff --git a/gateway/internal/store/lsm/memtable.go b/gateway/internal/store/lsm/memtable.go new file mode 100644 index 0000000..b3efec8 --- /dev/null +++ b/gateway/internal/store/lsm/memtable.go @@ -0,0 +1,200 @@ +package lsm + +import ( + "bytes" + "math/rand/v2" + "sync" +) + +// skipList tuning constants. +const ( + maxSkipLevel = 12 // supports ~4 M entries at p=0.25 + skipP = 0.25 // probability of promoting to the next level + nodeOverhead = 96 // conservative per-node GC/pointer overhead in bytes +) + +// skipNode is one node in the probabilistic skip list. +type skipNode struct { + key []byte + val []byte // nil signals a tombstone (deleted key) + next []*skipNode // len(next) == node level +} + +// MemTable is a write-optimised in-memory sorted map backed by a skip list. +// +// # Complexity +// +// - Put / Delete / Get: O(log N) expected +// - Ordered iteration: O(N) via the level-0 linked list +// +// A MemTable is safe for concurrent use. +type MemTable struct { + mu sync.RWMutex + head *skipNode // sentinel; never holds a real key + level int // highest level currently in use (1-indexed) + count int // total slots including tombstones + size int64 // approximate memory usage in bytes + maxSize int64 // size threshold that triggers a flush +} + +// newMemTable returns a MemTable that reports "should flush" once its +// approximate memory usage reaches maxSize. +func newMemTable(maxSize int64) *MemTable { + return &MemTable{ + head: &skipNode{next: make([]*skipNode, maxSkipLevel)}, + level: 1, + maxSize: maxSize, + } +} + +// Put inserts or updates key → val. +// Returns true when the MemTable has reached its size threshold. +func (m *MemTable) Put(key, val []byte) bool { + m.mu.Lock() + m.set(key, val) + full := m.size >= m.maxSize + m.mu.Unlock() + return full +} + +// Delete marks key as deleted (tombstone). +// Returns true when the MemTable should be flushed. +func (m *MemTable) Delete(key []byte) bool { + m.mu.Lock() + m.set(key, nil) + full := m.size >= m.maxSize + m.mu.Unlock() + return full +} + +// Get returns (value, found, isTombstone). +// - found=false: key is not present +// - found=true, tomb=true: key was deleted +// - found=true, tomb=false: key is present with the returned value +func (m *MemTable) Get(key []byte) (val []byte, found, tombstone bool) { + m.mu.RLock() + defer m.mu.RUnlock() + + cur := m.head + for i := m.level - 1; i >= 0; i-- { + for cur.next[i] != nil && bytes.Compare(cur.next[i].key, key) < 0 { + cur = cur.next[i] + } + } + node := cur.next[0] + if node == nil || !bytes.Equal(node.key, key) { + return nil, false, false + } + if node.val == nil { + return nil, true, true // tombstone + } + return node.val, true, false +} + +// Size returns the approximate memory used by the table in bytes. +func (m *MemTable) Size() int64 { + m.mu.RLock() + s := m.size + m.mu.RUnlock() + return s +} + +// Count returns the number of entries (including tombstones). +func (m *MemTable) Count() int { + m.mu.RLock() + c := m.count + m.mu.RUnlock() + return c +} + +// Iterator returns a snapshot-like in-order iterator. +// The iterator is not protected by a lock after construction; callers must +// ensure no concurrent writes happen while iterating (the engine does this +// by only iterating immutable MemTables). +func (m *MemTable) Iterator() *MemIter { + m.mu.RLock() + first := m.head.next[0] + m.mu.RUnlock() + return &MemIter{cur: first} +} + +// set inserts or updates the entry for key. Caller must hold mu for writing. +func (m *MemTable) set(key, val []byte) { + // Walk down from the top level collecting nodes that must be updated. + update := make([]*skipNode, maxSkipLevel) + cur := m.head + for i := m.level - 1; i >= 0; i-- { + for cur.next[i] != nil && bytes.Compare(cur.next[i].key, key) < 0 { + cur = cur.next[i] + } + update[i] = cur + } + + target := cur.next[0] + if target != nil && bytes.Equal(target.key, key) { + // Update in place — adjust size delta. + oldValLen := int64(len(target.val)) + newValLen := int64(len(val)) + m.size += newValLen - oldValLen + target.val = val + if val != nil { + target.val = append([]byte(nil), val...) + } + return + } + + // New key: insert a node at a random level. + lvl := m.randomLevel() + if lvl > m.level { + for i := m.level; i < lvl; i++ { + update[i] = m.head + } + m.level = lvl + } + + n := &skipNode{ + key: append([]byte(nil), key...), + next: make([]*skipNode, lvl), + } + if val != nil { + n.val = append([]byte(nil), val...) + } + for i := 0; i < lvl; i++ { + n.next[i] = update[i].next[i] + update[i].next[i] = n + } + m.size += int64(len(key)+len(val)) + nodeOverhead + m.count++ +} + +// randomLevel generates a level in [1, maxSkipLevel] with geometric distribution. +func (m *MemTable) randomLevel() int { + level := 1 + for level < maxSkipLevel && rand.Float64() < skipP { + level++ + } + return level +} + +// ─── Iterator ───────────────────────────────────────────────────────────────── + +// MemIter is a forward-only in-order iterator over a MemTable. +// It follows the level-0 linked list — O(1) per step, O(N) total. +type MemIter struct { + cur *skipNode +} + +// Valid reports whether the iterator is positioned at a valid entry. +func (it *MemIter) Valid() bool { return it.cur != nil } + +// Key returns the current key. Only valid when Valid() is true. +func (it *MemIter) Key() []byte { return it.cur.key } + +// Val returns the current value, or nil for a tombstone. +func (it *MemIter) Val() []byte { return it.cur.val } + +// IsTombstone reports whether the current entry is a deletion marker. +func (it *MemIter) IsTombstone() bool { return it.cur != nil && it.cur.val == nil } + +// Next advances to the next entry. +func (it *MemIter) Next() { it.cur = it.cur.next[0] } diff --git a/gateway/internal/store/lsm/sstable.go b/gateway/internal/store/lsm/sstable.go new file mode 100644 index 0000000..b3d5f2d --- /dev/null +++ b/gateway/internal/store/lsm/sstable.go @@ -0,0 +1,462 @@ +package lsm + +import ( + "bufio" + "bytes" + "encoding/binary" + "errors" + "fmt" + "os" +) + +// SSTable binary format +// +// ┌──────────────────────────────────────────────────────┐ +// │ Data Region │ +// │ ┌─────────────────────────────────────────────┐ │ +// │ │ Block 0 (≤ blockSize bytes) │ │ +// │ │ entry: keyLen:2 | key | valLen:4 | val │ │ +// │ │ entry: … │ │ +// │ └─────────────────────────────────────────────┘ │ +// │ … more blocks … │ +// ├──────────────────────────────────────────────────────┤ +// │ Index Region │ +// │ blockCount:4 │ +// │ [ offset:8 | blockLen:4 | keyLen:2 | firstKey ]* │ +// ├──────────────────────────────────────────────────────┤ +// │ Bloom Filter Region │ +// │ k:4 | m:8 | bits… │ +// ├──────────────────────────────────────────────────────┤ +// │ Footer (32 bytes, fixed) │ +// │ indexOffset:8 | indexLen:4 │ +// │ bloomOffset:8 | bloomLen:4 │ +// │ magic:8 ("INDEXQUB" = 0x42555145584E4449) │ +// └──────────────────────────────────────────────────────┘ +// +// tombstoneLen (0xFFFFFFFF) in the valLen field signals a deleted key; +// no value bytes follow it. +const ( + sstMagic = uint64(0x42555145584E4449) // "INDEXQUB" little-endian + footerSize = 32 + tombstoneLen = ^uint32(0) // 0xFFFFFFFF +) + +// indexEntry records the file position and first key of one data block. +type indexEntry struct { + offset int64 + blockLen uint32 + firstKey []byte +} + +// ─── Builder ────────────────────────────────────────────────────────────────── + +// Builder writes a new SSTable to disk from an ordered stream of key-value pairs. +// Keys must be supplied in strictly ascending lexicographic order. +// +// Usage: +// +// b, _ := newBuilder(path, opts) +// b.Add(key1, val1) +// b.Add(key2, val2) // key2 > key1 +// b.Finish() +type Builder struct { + path string + f *os.File + bw *bufio.Writer + + blockBuf bytes.Buffer // current data block accumulator + blockOff int64 // file offset where blockBuf started + index []indexEntry + bloom *Bloom + opts Options +} + +func newBuilder(path string, opts Options) (*Builder, error) { + f, err := os.Create(path) + if err != nil { + return nil, err + } + return &Builder{ + path: path, + f: f, + bw: bufio.NewWriterSize(f, 64*1024), + bloom: NewBloom(max(opts.BloomExpected, 1), opts.BloomFPRate), + opts: opts, + }, nil +} + +// Add appends a key-value pair. val == nil writes a tombstone. +func (b *Builder) Add(key, val []byte) error { + if len(key) > 0xFFFF { + return fmt.Errorf("lsm: key too long (%d bytes)", len(key)) + } + + // Start a new index entry when the current block is empty. + if b.blockBuf.Len() == 0 { + b.index = append(b.index, indexEntry{ + offset: b.blockOff, + firstKey: append([]byte(nil), key...), + }) + } + + b.bloom.Add(key) + + // Entry wire format: keyLen:2 | key | valLen:4 | val + var hdr [6]byte + binary.LittleEndian.PutUint16(hdr[0:], uint16(len(key))) + if val == nil { + binary.LittleEndian.PutUint32(hdr[2:], tombstoneLen) + } else { + binary.LittleEndian.PutUint32(hdr[2:], uint32(len(val))) + } + b.blockBuf.Write(hdr[:]) + b.blockBuf.Write(key) + if val != nil { + b.blockBuf.Write(val) + } + + if b.blockBuf.Len() >= b.opts.BlockSize { + return b.flushBlock() + } + return nil +} + +func (b *Builder) flushBlock() error { + if b.blockBuf.Len() == 0 { + return nil + } + n := b.blockBuf.Len() + b.index[len(b.index)-1].blockLen = uint32(n) + if _, err := b.bw.Write(b.blockBuf.Bytes()); err != nil { + return err + } + b.blockOff += int64(n) + b.blockBuf.Reset() + return nil +} + +// Finish flushes the final block and writes the index, Bloom filter, and footer. +// The file is closed; the Builder must not be used after Finish returns. +func (b *Builder) Finish() error { + if err := b.flushBlock(); err != nil { + return err + } + if len(b.index) == 0 { + // No entries were added — write an empty but valid SSTable. + b.index = []indexEntry{} + } + + // ── Index block ─────────────────────────────────────────────────────── + indexStart := b.blockOff + var ibuf bytes.Buffer + binary.Write(&ibuf, binary.LittleEndian, uint32(len(b.index))) //nolint:errcheck + for _, e := range b.index { + binary.Write(&ibuf, binary.LittleEndian, e.offset) //nolint:errcheck + binary.Write(&ibuf, binary.LittleEndian, e.blockLen) //nolint:errcheck + binary.Write(&ibuf, binary.LittleEndian, uint16(len(e.firstKey))) //nolint:errcheck + ibuf.Write(e.firstKey) + } + indexData := ibuf.Bytes() + if _, err := b.bw.Write(indexData); err != nil { + return err + } + + // ── Bloom filter ────────────────────────────────────────────────────── + bloomStart := indexStart + int64(len(indexData)) + bloomData := b.bloom.Encode() + if _, err := b.bw.Write(bloomData); err != nil { + return err + } + + // ── Footer (32 bytes) ───────────────────────────────────────────────── + var footer [footerSize]byte + binary.LittleEndian.PutUint64(footer[0:], uint64(indexStart)) + binary.LittleEndian.PutUint32(footer[8:], uint32(len(indexData))) + binary.LittleEndian.PutUint64(footer[12:], uint64(bloomStart)) + binary.LittleEndian.PutUint32(footer[20:], uint32(len(bloomData))) + binary.LittleEndian.PutUint64(footer[24:], sstMagic) + if _, err := b.bw.Write(footer[:]); err != nil { + return err + } + + if err := b.bw.Flush(); err != nil { + return err + } + return b.f.Close() +} + +// ─── Reader ─────────────────────────────────────────────────────────────────── + +// SSTable is an open, immutable, read-only view of a single on-disk SSTable file. +// It is safe for concurrent reads via Get and Iterator. +type SSTable struct { + f *os.File + index []indexEntry + bloom *Bloom + path string +} + +// openSSTable opens the SSTable at path for reading. +func openSSTable(path string) (*SSTable, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + fi, err := f.Stat() + if err != nil { + f.Close() + return nil, err + } + if fi.Size() < footerSize { + f.Close() + return nil, fmt.Errorf("lsm: %s is too small to be a valid SSTable", path) + } + + var foot [footerSize]byte + if _, err := f.ReadAt(foot[:], fi.Size()-footerSize); err != nil { + f.Close() + return nil, fmt.Errorf("lsm: read footer of %s: %w", path, err) + } + if binary.LittleEndian.Uint64(foot[24:]) != sstMagic { + f.Close() + return nil, fmt.Errorf("lsm: %s has invalid magic", path) + } + + indexOff := int64(binary.LittleEndian.Uint64(foot[0:])) + indexLen := int(binary.LittleEndian.Uint32(foot[8:])) + bloomOff := int64(binary.LittleEndian.Uint64(foot[12:])) + bloomLen := int(binary.LittleEndian.Uint32(foot[20:])) + + indexData := make([]byte, indexLen) + if _, err := f.ReadAt(indexData, indexOff); err != nil { + f.Close() + return nil, fmt.Errorf("lsm: read index of %s: %w", path, err) + } + idx, err := decodeIndex(indexData) + if err != nil { + f.Close() + return nil, fmt.Errorf("lsm: decode index of %s: %w", path, err) + } + + bloomData := make([]byte, bloomLen) + if _, err := f.ReadAt(bloomData, bloomOff); err != nil { + f.Close() + return nil, fmt.Errorf("lsm: read bloom of %s: %w", path, err) + } + bloom, err := DecodeBloom(bloomData) + if err != nil { + f.Close() + return nil, fmt.Errorf("lsm: decode bloom of %s: %w", path, err) + } + + return &SSTable{f: f, index: idx, bloom: bloom, path: path}, nil +} + +func decodeIndex(data []byte) ([]indexEntry, error) { + if len(data) < 4 { + return nil, errors.New("index too short") + } + count := int(binary.LittleEndian.Uint32(data)) + data = data[4:] + entries := make([]indexEntry, 0, count) + for i := range count { + const fixedPart = 8 + 4 + 2 // offset + blockLen + keyLen + if len(data) < fixedPart { + return nil, fmt.Errorf("index entry %d truncated", i) + } + off := int64(binary.LittleEndian.Uint64(data[0:])) + blen := binary.LittleEndian.Uint32(data[8:]) + klen := int(binary.LittleEndian.Uint16(data[12:])) + data = data[fixedPart:] + if len(data) < klen { + return nil, fmt.Errorf("index entry %d: key truncated", i) + } + entries = append(entries, indexEntry{ + offset: off, + blockLen: blen, + firstKey: append([]byte(nil), data[:klen]...), + }) + data = data[klen:] + } + return entries, nil +} + +// Get looks up key. +// - (val, true, nil) — key found with value val +// - (nil, true, nil) — key found but is a tombstone (deleted) +// - (nil, false, nil) — key is absent +func (s *SSTable) Get(key []byte) (val []byte, found bool, err error) { + // Fast path: Bloom filter is certain the key is absent. + if !s.bloom.MayContain(key) { + return nil, false, nil + } + + // Binary search index: find the rightmost block whose firstKey ≤ key. + bi := s.findBlock(key) + if bi < 0 { + return nil, false, nil + } + + block, err := s.readBlock(bi) + if err != nil { + return nil, false, err + } + return scanBlock(block, key) +} + +// findBlock returns the index of the block that could contain key, +// or -1 if key is smaller than every firstKey. +func (s *SSTable) findBlock(key []byte) int { + lo, hi, result := 0, len(s.index)-1, -1 + for lo <= hi { + mid := (lo + hi) / 2 + if bytes.Compare(s.index[mid].firstKey, key) <= 0 { + result = mid + lo = mid + 1 + } else { + hi = mid - 1 + } + } + return result +} + +// readBlock reads raw block bytes from disk using ReadAt (concurrent-safe). +func (s *SSTable) readBlock(i int) ([]byte, error) { + e := s.index[i] + buf := make([]byte, e.blockLen) + if _, err := s.f.ReadAt(buf, e.offset); err != nil { + return nil, fmt.Errorf("lsm: read block %d of %s: %w", i, s.path, err) + } + return buf, nil +} + +// scanBlock linearly scans a raw data block for key. +func scanBlock(block, key []byte) (val []byte, found bool, err error) { + for len(block) >= 6 { + klen := int(binary.LittleEndian.Uint16(block[0:])) + vlen := binary.LittleEndian.Uint32(block[2:]) + block = block[6:] + + if len(block) < klen { + return nil, false, errors.New("lsm: corrupt block: key truncated") + } + k := block[:klen] + block = block[klen:] + + if vlen == tombstoneLen { + if bytes.Equal(k, key) { + return nil, true, nil // deleted + } + continue + } + + if int(vlen) > len(block) { + return nil, false, errors.New("lsm: corrupt block: val truncated") + } + v := block[:vlen] + block = block[vlen:] + + if bytes.Equal(k, key) { + return append([]byte(nil), v...), true, nil + } + } + return nil, false, nil +} + +// Close releases the file handle. +func (s *SSTable) Close() error { return s.f.Close() } + +// Path returns the file path of the SSTable. +func (s *SSTable) Path() string { return s.path } + +// ─── Iterator ───────────────────────────────────────────────────────────────── + +// SSTIter iterates over every entry in an SSTable in ascending key order. +// It is not safe for concurrent use. +type SSTIter struct { + sst *SSTable + blockIdx int // next block to load (-1 = not started) + block []byte // remaining bytes in the current block + key []byte + val []byte // nil = tombstone + valid bool + err error +} + +// newSSTIter creates an iterator positioned before the first entry. +// Call Next() to advance to the first entry. +func newSSTIter(sst *SSTable) *SSTIter { + return &SSTIter{sst: sst, blockIdx: -1} +} + +// Valid reports whether the iterator is positioned at a valid entry. +func (it *SSTIter) Valid() bool { return it.valid } + +// Key returns the current key. Caller must not modify the returned slice. +func (it *SSTIter) Key() []byte { return it.key } + +// Val returns the current value, or nil if the entry is a tombstone. +func (it *SSTIter) Val() []byte { return it.val } + +// IsTombstone reports whether the current entry is a deletion marker. +func (it *SSTIter) IsTombstone() bool { return it.valid && it.val == nil } + +// Err returns the first error encountered, if any. +func (it *SSTIter) Err() error { return it.err } + +// Next advances to the next entry. Returns true if a valid entry was reached. +func (it *SSTIter) Next() bool { + for { + // Parse the next entry from the current block. + if len(it.block) >= 6 { + if it.parseEntry() { + return true + } + if it.err != nil { + return false + } + } + // Load the next block. + it.blockIdx++ + if it.blockIdx >= len(it.sst.index) { + it.valid = false + return false + } + block, err := it.sst.readBlock(it.blockIdx) + if err != nil { + it.err = err + it.valid = false + return false + } + it.block = block + } +} + +func (it *SSTIter) parseEntry() bool { + klen := int(binary.LittleEndian.Uint16(it.block[0:])) + vlen := binary.LittleEndian.Uint32(it.block[2:]) + it.block = it.block[6:] + + if len(it.block) < klen { + it.err = errors.New("lsm: corrupt block: key truncated") + it.valid = false + return false + } + it.key = it.block[:klen] + it.block = it.block[klen:] + + if vlen == tombstoneLen { + it.val = nil + it.valid = true + return true + } + if int(vlen) > len(it.block) { + it.err = errors.New("lsm: corrupt block: val truncated") + it.valid = false + return false + } + it.val = it.block[:vlen] + it.block = it.block[vlen:] + it.valid = true + return true +} diff --git a/gateway/internal/telemetry/agentsession.go b/gateway/internal/telemetry/agentsession.go new file mode 100644 index 0000000..de7b845 --- /dev/null +++ b/gateway/internal/telemetry/agentsession.go @@ -0,0 +1,198 @@ +package telemetry + +import ( + "sort" + "sync" + "time" +) + +// AgentSession is a per-session aggregate of guard decisions and token spend. +// Its shape mirrors the agent_sessions schema in CONTEXT.md so it can be +// persisted to a database layer without structural changes. +type AgentSession struct { + SessionID string `json:"session_id"` + StartedAt int64 `json:"started_at"` // Unix seconds + LastSeenAt int64 `json:"last_seen_at"` + TokensAttempted int64 `json:"tokens_attempted"` + TokensSent int64 `json:"tokens_sent"` + TokensSaved int64 `json:"tokens_saved"` + RequestsTotal int `json:"requests_total"` + LoopDetected int `json:"loop_detected"` // velocity/circuit warn events + KillEvents int `json:"kill_events"` // guard blocks (HTTP 429) + KillReason string `json:"kill_reason"` // most recent block reason + DollarsSaved float64 `json:"dollars_saved"` + Status string `json:"status"` // "active" | "killed" | "ended" +} + +// KillEvent is emitted whenever a guard blocks a request (HTTP 429). +// It is the loop-detection audit trail: how many times a session was stopped, +// why, and how much spend was prevented. +type KillEvent struct { + SessionID string `json:"session_id"` + Timestamp int64 `json:"timestamp"` // Unix seconds + Reason string `json:"reason"` // "velocity_exceeded" | "budget_exceeded" | … + TokensSaved int64 `json:"tokens_saved"` // tokens not forwarded due to block + DollarsSaved float64 `json:"dollars_saved"` // cost of those tokens +} + +// RequestOutcome is the per-request signal passed to AgentSessionStore.Record. +// It is extracted from the guard decision and optimizer stats at the end of +// every request, whether allowed or blocked. +type RequestOutcome struct { + TokensAttempted int + TokensSent int + TokensSaved int + DollarsSaved float64 + GuardReason string + Killed bool // guard returned !Allow (HTTP 429 was sent) + Warned bool // guard returned Allow+Warn +} + +// AgentSessionStore tracks per-session observability in memory. +// It is safe for concurrent use and is intentionally lightweight — no +// persistence, no external dependencies. A future version can flush to the +// agent_sessions table defined in CONTEXT.md. +type AgentSessionStore struct { + mu sync.Mutex + sessions map[string]*agentSessionEntry + killLog []KillEvent + ttl time.Duration + maxKillLog int +} + +type agentSessionEntry struct { + AgentSession + lastKilledAt time.Time +} + +// NewAgentSessionStore returns a store that expires sessions after ttl of +// inactivity. A ttl of 0 defaults to 4 hours. +func NewAgentSessionStore(ttl time.Duration) *AgentSessionStore { + if ttl <= 0 { + ttl = 4 * time.Hour + } + return &AgentSessionStore{ + sessions: make(map[string]*agentSessionEntry), + ttl: ttl, + maxKillLog: 500, + } +} + +// Record updates the session aggregate with the outcome of one request. +// sessionKey is the per-invocation key produced by claudeSessionKey. +func (s *AgentSessionStore) Record(sessionKey string, out RequestOutcome) { + if sessionKey == "" { + return + } + now := time.Now() + nowUnix := now.Unix() + + s.mu.Lock() + defer s.mu.Unlock() + + entry := s.getOrCreateLocked(sessionKey, now) + + entry.LastSeenAt = nowUnix + entry.RequestsTotal++ + entry.TokensAttempted += int64(out.TokensAttempted) + entry.TokensSent += int64(out.TokensSent) + entry.TokensSaved += int64(out.TokensSaved) + entry.DollarsSaved += out.DollarsSaved + + if out.Killed { + entry.KillEvents++ + entry.KillReason = out.GuardReason + entry.lastKilledAt = now + entry.Status = "killed" + + kill := KillEvent{ + SessionID: sessionKey, + Timestamp: nowUnix, + Reason: out.GuardReason, + TokensSaved: int64(out.TokensAttempted - out.TokensSent), + DollarsSaved: out.DollarsSaved, + } + s.killLog = append(s.killLog, kill) + if len(s.killLog) > s.maxKillLog { + // Drop oldest entries when the log fills up. + s.killLog = s.killLog[len(s.killLog)-s.maxKillLog:] + } + } else if out.Warned { + entry.LoopDetected++ + if entry.Status != "killed" { + entry.Status = "active" + } + } else { + if entry.Status != "killed" { + entry.Status = "active" + } + } +} + +// Snapshot returns a point-in-time copy of all sessions, sorted by last-seen +// descending (most recently active first). Expired sessions are excluded. +func (s *AgentSessionStore) Snapshot() []AgentSession { + s.mu.Lock() + defer s.mu.Unlock() + + now := time.Now() + out := make([]AgentSession, 0, len(s.sessions)) + for _, e := range s.sessions { + if now.Unix()-e.LastSeenAt > int64(s.ttl.Seconds()) { + continue + } + snap := e.AgentSession + // Promote "killed" sessions that have had no activity for >30s to "ended". + if snap.Status == "killed" && now.Sub(e.lastKilledAt) > 30*time.Second { + snap.Status = "ended" + } + out = append(out, snap) + } + sort.Slice(out, func(i, j int) bool { + return out[i].LastSeenAt > out[j].LastSeenAt + }) + return out +} + +// KillLog returns the recent kill events, newest first. +func (s *AgentSessionStore) KillLog() []KillEvent { + s.mu.Lock() + defer s.mu.Unlock() + + out := make([]KillEvent, len(s.killLog)) + copy(out, s.killLog) + // Reverse: newest first. + for i, j := 0, len(out)-1; i < j; i, j = i+1, j-1 { + out[i], out[j] = out[j], out[i] + } + return out +} + +// CleanupExpired removes sessions that have been inactive for longer than ttl. +// Call periodically from a background goroutine. +func (s *AgentSessionStore) CleanupExpired() { + now := time.Now() + s.mu.Lock() + defer s.mu.Unlock() + for key, e := range s.sessions { + if now.Unix()-e.LastSeenAt > int64(s.ttl.Seconds()) { + delete(s.sessions, key) + } + } +} + +func (s *AgentSessionStore) getOrCreateLocked(key string, now time.Time) *agentSessionEntry { + if e, ok := s.sessions[key]; ok { + return e + } + e := &agentSessionEntry{ + AgentSession: AgentSession{ + SessionID: key, + StartedAt: now.Unix(), + LastSeenAt: now.Unix(), + Status: "active", + }, + } + s.sessions[key] = e + return e +} diff --git a/gateway/internal/telemetry/finops.go b/gateway/internal/telemetry/finops.go index 06b1c5c..037f107 100644 --- a/gateway/internal/telemetry/finops.go +++ b/gateway/internal/telemetry/finops.go @@ -1,9 +1,8 @@ package telemetry -// EstimateCostSaved calculates the saved USD amount based on token counts. -// Prices are estimated per 1M input tokens as of early 2024. -func EstimateCostSaved(provider, model string, tokensSaved int) float64 { - if tokensSaved <= 0 { +// EstimateCost calculates the estimated USD amount based on input token counts. +func EstimateCost(provider, model string, tokens int) float64 { + if tokens <= 0 { return 0 } @@ -26,5 +25,11 @@ func EstimateCostSaved(provider, model string, tokensSaved int) float64 { pricePerMillion = 3.0 } - return (float64(tokensSaved) / 1000000.0) * pricePerMillion + return (float64(tokens) / 1000000.0) * pricePerMillion +} + +// EstimateCostSaved calculates the saved USD amount based on token counts. +// Prices are estimated per 1M input tokens as of early 2024. +func EstimateCostSaved(provider, model string, tokensSaved int) float64 { + return EstimateCost(provider, model, tokensSaved) } diff --git a/gateway/internal/telemetry/telemetry.go b/gateway/internal/telemetry/telemetry.go index 9ae378e..aaa9ee5 100644 --- a/gateway/internal/telemetry/telemetry.go +++ b/gateway/internal/telemetry/telemetry.go @@ -211,6 +211,8 @@ func parseLogLevel(s string) slog.Level { return slog.LevelWarn case "error", "ERROR": return slog.LevelError + case "off", "OFF", "none", "NONE": + return slog.Level(100) // above all standard levels — effectively silent default: return slog.LevelInfo } diff --git a/k8s/deployment.yaml b/k8s/deployment.yaml new file mode 100644 index 0000000..f950313 --- /dev/null +++ b/k8s/deployment.yaml @@ -0,0 +1,168 @@ +# IndexQube Gateway - Kubernetes manifests +# +# Includes: +# Deployment - 2 replicas, liveness + readiness probes, resource limits, +# graceful SIGTERM drain (15 s, handled in Go signal handler) +# Service - ClusterIP on 8080 (proxy) and 9100 (metrics/admin) +# +# Deploy: +# kubectl apply -f k8s/deployment.yaml +# +# Customise: +# - Set image.tag to a pinned SHA or semver in CI +# - Tune replicas and resource limits for your node class +# - Set ANTHROPIC_BASE_URL to a real upstream if not using the default +# --- + +apiVersion: apps/v1 +kind: Deployment +metadata: + name: indexqube-gateway + labels: + app: indexqube-gateway +spec: + replicas: 2 + selector: + matchLabels: + app: indexqube-gateway + # Rolling update: bring up one new pod before taking down an old one. + # maxUnavailable=0 ensures the Service always has a healthy pod during deploys. + strategy: + type: RollingUpdate + rollingUpdate: + maxSurge: 1 + maxUnavailable: 0 + template: + metadata: + labels: + app: indexqube-gateway + annotations: + # Prometheus scrape annotations - works with kube-prometheus and most + # Prometheus Operator setups without a separate ServiceMonitor. + prometheus.io/scrape: "true" + prometheus.io/port: "9100" + prometheus.io/path: "/metrics" + spec: + # Allow up to 45 s for graceful drain (15 s server shutdown plus + # kube-proxy/readiness propagation buffer). Must exceed the + # shutdownGrace constant in internal/server/run.go. + terminationGracePeriodSeconds: 45 + + containers: + - name: gateway + image: indexqube/gateway:latest + imagePullPolicy: IfNotPresent + + ports: + - name: http + containerPort: 8080 + protocol: TCP + - name: admin + containerPort: 9100 + protocol: TCP + + env: + - name: PORT + value: "8080" + # Expose admin server on all interfaces so the readiness probe and + # Prometheus can reach /healthz, /readyz, and /metrics on port 9100. + - name: ADMIN_BIND_ADDR + value: "0.0.0.0" + - name: ADMIN_PORT + value: "9100" + # Structured JSON logging - required for log aggregators. + - name: INDEXQUBE_LOG_LEVEL + value: "info" + # Embed the pod name in logs for per-replica tracing. + - name: INDEXQUBE_INSTANCE + valueFrom: + fieldRef: + fieldPath: metadata.name + + resources: + requests: + cpu: "100m" + memory: "64Mi" + limits: + cpu: "500m" + memory: "256Mi" + + # Liveness probe - "is the process alive?" + # Hits /healthz on the public port (8080). Returns 200 unconditionally + # as long as the Go runtime is healthy. A failed liveness probe causes + # k8s to restart the container. + livenessProbe: + httpGet: + path: /healthz + port: http + initialDelaySeconds: 5 + periodSeconds: 10 + timeoutSeconds: 2 + failureThreshold: 3 + + # Readiness probe - "is the gateway ready to serve traffic?" + # Hits /readyz on the public port. Returns 503 during startup or + # after SIGTERM (server.Shutdown drains in-flight SSE streams). + # A failed readiness probe removes the pod from the Service endpoints + # - k8s stops routing new requests to it without restarting it. + readinessProbe: + httpGet: + path: /readyz + port: http + initialDelaySeconds: 2 + periodSeconds: 5 + timeoutSeconds: 2 + failureThreshold: 2 + + # No preStop hook - rationale: + # The standard pattern is `preStop: exec: sleep 5` to give kube-proxy + # time to drain iptables before SIGTERM hits the process. However, + # distroless/static has no shell or sleep binary. Two paths forward: + # (a) Switch Dockerfile final stage to distroless/base-debian12 + # (adds busybox ~3 MiB) and add: lifecycle.preStop.exec: ["/bin/sleep","5"] + # (b) Keep distroless/static - acceptable because: + # - terminationGracePeriodSeconds=45 gives kube-proxy ample time + # - the Go server's 15 s shutdownGrace (server/run.go) drains + # in-flight SSE streams before exit + # - readiness probe fails on SIGTERM -> k8s stops routing new + # requests within one probe period (5 s) anyway + # Current choice: (b). Switch to (a) if you observe 5xx spikes during + # rolling deploys in a high-QPS environment. + + securityContext: + # Matches the UID baked into distroless/static-debian12:nonroot. + runAsUser: 65532 + runAsGroup: 65532 + runAsNonRoot: true + readOnlyRootFilesystem: true + allowPrivilegeEscalation: false + capabilities: + drop: ["ALL"] + + # Pod-level security: no privilege escalation at the pod level either. + securityContext: + runAsNonRoot: true + seccompProfile: + type: RuntimeDefault + +--- + +apiVersion: v1 +kind: Service +metadata: + name: indexqube-gateway + labels: + app: indexqube-gateway +spec: + type: ClusterIP + selector: + app: indexqube-gateway + ports: + - name: http + port: 80 + targetPort: http + protocol: TCP + - name: admin + port: 9100 + targetPort: admin + protocol: TCP