diff --git a/BUGS b/BUGS new file mode 100644 index 000000000..f6a39f9c7 --- /dev/null +++ b/BUGS @@ -0,0 +1,3 @@ +TODO: + -if there is a sigterm in between dequeue and execution before the renqueue the job is lost + - releaseForNextPoll is just a release, useless diff --git a/TODO b/TODO new file mode 100644 index 000000000..671574abf --- /dev/null +++ b/TODO @@ -0,0 +1,220 @@ +# TODO + +## internal/util/semaphore: double-release callback is not a safety mechanism + +The `onDoubleRelease` callback in `semaphore.New` (internal/util/semaphore/semaphore.go) is +misleadingly named and does not actually protect concurrency limits. Because the semaphore is +backed by a buffered channel, `Release()` fires the callback only when the channel is already +empty — meaning all tokens have been returned and no goroutines are running. At that point there +is nothing to corrupt, so the callback provides no real protection. + +Two cleaner alternatives: + +1. Have `Acquire` return an opaque token that `Release` requires as an argument. Without the + token the caller cannot invoke `Release`, making double-release structurally impossible. + Example shape: + + type Token struct{ _ struct{} } // unexported field prevents external construction + + func (s *semaphore) Acquire(ctx context.Context) (Token, error) + func (s *semaphore) Release(t Token) + +2. Remove `onDoubleRelease` entirely and make `Release` a no-op when the channel is already + empty, which is the correct and safe behavior for this channel-based implementation. + +The wiring in `worker.go` (the `makeGuard` / `stopAccepting` pattern) would need to be removed +or replaced with a different shutdown trigger if option 2 is chosen. + +## internal/processor/worker: crash recovery relies on the filesystem, not the DB + +Recovery from crashes (internal/processor/worker/recovery.go) works by scanning leftover job +directories on `emptyDir` at container startup. This means it only handles container-level +restarts within the same pod. Pod deletion, node eviction, or any scenario where `emptyDir` is +destroyed means the filesystem-based recovery cannot find those jobs. + +The jobs themselves are not unrecoverable orphans: the orphan reconciler in batch-gc +(internal/gc/reconciler/reconciler.go) periodically cross-references DB, queue, and in-flight +hash to detect orphans and triage them — re-enqueueing `validating` jobs and CAS-transitioning +`in_progress`/`finalizing` jobs to `failed`. What is unrecoverable is the **partial output +data** (output.jsonl/error.jsonl content) that existed only on the destroyed `emptyDir` and had +not yet been uploaded to shared storage. + +The root cause is that intermediate job state (partial output/error files, request counts) is +written to local disk rather than to the DB. A more resilient design would checkpoint progress +into the DB so recovery can preserve partial results from any pod, not just the one that +crashed. + +## internal/processor/worker: releaseForNextPoll is a pointless wrapper + +`releaseForNextPoll` (worker.go) is a one-line wrapper around `release()`, which is itself a +one-line wrapper around `p.tokens.Release()`. It adds no behaviour and the name implies +semantics that don't exist — it does not schedule or influence the next poll in any way. +Replace all call sites with `p.release()` and remove `releaseForNextPoll`. + +## internal/processor/worker: pollCtx is a useless context derivation + +In runPollingLoop (worker.go), `pollCtx` is derived from `pollingCtx` solely to carry a logger +with `jobId` attached. However, every function that receives `pollCtx` and logs the job ID does +so explicitly as a key-value argument (e.g. poller.go:90), so the logger value in the context +is never actually used. `pollCtx` and `pollLogger` can be removed — just use `pollingCtx` +directly and pass `pollLogger` explicitly where needed, or drop it entirely. + +## internal/processor/worker: job pickup should be transactional, not eventually consistent + +In runPollingLoop (worker.go), there is a window between dequeueOne (removing the job from the +Redis queue) and InFlightSet (recording it in the Redis in-flight hash) where a crash leaves +the job invisible — it is neither in the queue nor in the in-flight table. More broadly, even +when InFlightSet succeeds, the job is only recoverable after a full GC reconciler cycle (~60 +minutes by default). During that window the job sits idle, consuming SLO budget while no +processor works on it. + +The current design is eventually consistent: the orphan reconciler in batch-gc periodically +cross-references DB, queue, and in-flight hash to detect and triage orphans. This works +correctly but has two drawbacks: + +1. **Recovery latency**: A crashed job waits up to one full reconciler interval before being + detected and re-enqueued or failed. For latency-sensitive batches this delay can exhaust the + SLO, turning a recoverable job into an expired one. + +2. **Distributed bookkeeping**: Job ownership is spread across three data stores (PostgreSQL + status, Redis queue, Redis in-flight hash) with no transactional guarantee tying them + together. Every crash boundary between these stores is a potential inconsistency that the + reconciler must paper over after the fact. + +A more robust design would make job pickup transactional: atomically dequeue the job and +transition its DB status (e.g. `validating` → `in_progress`) in a single operation. With the +DB record as the sole source of truth, any processor — including the one that restarts after a +crash — can detect incomplete jobs at startup by querying for `in_progress` jobs that have no +active processor, and resume or fail them immediately without waiting for a GC cycle. This +eliminates both the dequeue-to-in-flight gap and the dependency on a periodic reconciler for +timely recovery. + +## internal/processor/worker: concurrency controller is coupled to AIMDController concrete type + +The signal-classification switch in executor.go (~L477-496) is guarded by a nil-check on +`*semaphore.AIMDController` and calls its concrete methods directly. The `endpointLimit` struct +(worker.go:45) stores `aimd *semaphore.AIMDController` as a concrete pointer, so swapping to a +different congestion-control algorithm (e.g. CUBIC, PID, or static) requires editing both the +executor goroutine and the wiring in `Run()`. + +`AIMDController` itself is already well-decoupled — it takes a `setFn func(int)` callback and +knows nothing about semaphores. The missing piece is an interface at the consumer side: + + type ConcurrencyController interface { + RecordSuccess() + RecordRateLimit(reason string) + Limit() int + } + +With this interface: + +- `endpointLimit.aimd` becomes `endpointLimit.controller ConcurrencyController` (nil = no + adaptive control). +- The executor switch calls the interface methods unchanged. +- `worker.Run()` wires in `AIMDController` (or any future algorithm) based on config — the + construction code in worker.go:159-173 stays roughly the same, just returning the interface. +- New algorithms can be added without touching executor.go. + +## internal/processor/worker: worker token conflates job orchestration with capacity reservation + +The worker token semaphore (`p.tokens`, sized by `NumWorkers`) is acquired before dequeuing a +job and held for the **entire job lifecycle** — input download, all inference requests, output +upload, status finalization (`job_runner.go:82`: `defer p.release()`). This means a job that is +mostly idle (waiting on I/O, waiting on slow endpoints, or — in a future async-inference model +— waiting for results from llm-d-async) still occupies a worker slot the whole time, blocking +other jobs from starting even when the system has spare capacity. + +The token conflates two concerns: + +1. **Job orchestration** — tracking active jobs, managing their state machine, bounding memory + for per-job artifacts (plan files, output buffers, local disk usage on `emptyDir`). +2. **Capacity reservation** — limiting load on inference endpoints and the processor itself. + +For synchronous inference these roughly align: a running job is actively consuming endpoint +capacity. For asynchronous inference (e.g. llm-d-async) they diverge completely: a job that has +submitted its requests and is polling for results consumes near-zero inference capacity but +still holds its worker slot. + +Increasing `NumWorkers` is a workaround, not a fix — it raises the ceiling but the coupling +remains. A cleaner design would separate the two: + +- **Job orchestration** becomes unbounded (or bounded by a much larger limit tied to memory / + disk, not inference throughput). Any number of jobs can be in the "waiting for results" phase. +- **Capacity reservation** is handled solely by the global and per-endpoint semaphores, which + already exist and already do the right thing at the request level. + +This separation would matter most when integrating with llm-d-async, where the job lifecycle +shifts from "dispatch requests synchronously" to "submit requests, then wait." Without it, +`NumWorkers` becomes a throughput bottleneck unrelated to actual system load. + +## internal/processor/worker: plan files are written to disk unnecessarily + +Plan files are per-model binary index files produced by `preProcessJob`. Each entry is 16 bytes +(`offset int64, length uint32, prefixHash uint32`) pointing into the local `input.jsonl` copy. +They serve two purposes: + +1. **Per-model partitioning** — the input JSONL can mix requests for different models. The + preprocessor splits them into separate plan files so the executor can process each model + independently, dispatching requests to the correct inference endpoint. +2. **Prefix-hash sorting** — within each model, entries are sorted by FNV-32a hash of the + system prompt. This groups requests with the same system prompt together during dispatch, + improving KV-cache hit rates on the inference gateway. + +The preprocessor accumulates plan entries in memory (`planAccumulator`), sorts them by prefix +hash, writes them to binary files on `emptyDir` (`planner.go:152-195`), and the executor reads +them back from disk during execution — all within the same process and container lifecycle. + +The plan files are never used for crash recovery. `recovery.go` only cares about `output.jsonl` +and `error.jsonl` (partial results to upload). When a crashed job in `validating` or +`in_progress` (without partial output) is recovered, it is re-enqueued and the next worker +re-runs `preProcessJob` from scratch — re-downloading from S3 and regenerating everything. + +The entries are small: 16 bytes each, so 50k requests (the OpenAI max) is ~800 KB per model. +They are already fully materialized in memory before the disk write. The disk round-trip (write +in `Finalize`, read back in `executeJob`) adds complexity and I/O for no benefit. The entries +should stay in memory and be passed directly from `preProcessJob` to `executeJob` as a +`map[string][]planEntry`, eliminating the plan file write/read cycle entirely. + +## internal/processor/worker: local input.jsonl copy can be replaced with S3 ranged gets + +The preprocessor streams the input file from S3 and writes a full copy to local disk +(`preprocessor.go:156-158`) so the executor can later do random-access reads by byte offset +(each plan entry stores `(offset, length)` into this local file). With 20 workers at the +OpenAI max of 200 MB per file, this is up to 4 GB of `emptyDir` consumed by input copies +alone, plus the output/error files growing during execution. + +The AWS SDK already supports ranged gets — `s3.GetObjectInput` has a `Range` field that +accepts standard HTTP byte ranges (`bytes=offset-offset+length-1`). All S3-compatible stores +(AWS, MinIO, Ceph) support this. The current `BatchFilesClient` interface just doesn't expose +it (`Retrieve` always fetches the full object). + +Adding a ranged-read method to the interface (or using the S3 client directly in the executor) +would let the executor fetch each request line on demand from shared storage, eliminating the +local input copy entirely. The per-request latency overhead (~5-50ms for an S3 GET) is +negligible compared to inference latency (hundreds of ms to seconds). Adjacent plan entries +could also be batched into larger range reads to amortize further. + +Combined with the plan-files-in-memory change above, this would remove the `emptyDir` +dependency for input data entirely, leaving only `output.jsonl` and `error.jsonl` on local +disk (which could themselves be streamed to shared storage via S3 multipart upload as a +follow-up). + +## charts/batch-gateway: default emptyDir budget is tight for 20 workers + +The default Helm values set `numWorkers: 20` and `workDirVolume.sizeLimit: 10Gi` +(values.yaml:237,307). Each worker copies its job's input file to local disk (up to 200 MB per +the OpenAI spec) and accumulates output/error JSONL during execution. In the worst case, 20 +concurrent jobs with max-size inputs consume ~4 GB for input copies alone, leaving ~6 GB for +output files, error files, and plan data across all 20 jobs. Output files can grow large — +inference responses (especially from chat completions) are typically much bigger than the +corresponding request lines — so 6 GB shared across 20 jobs is not generous. + +If the `emptyDir` fills up, the kubelet evicts the pod, losing all in-progress work. The +`medium: ""` default (values.yaml:238) means the volume uses node disk rather than tmpfs, so +it doesn't consume RAM, but it competes with other pods for the node's ephemeral storage +budget. + +This is a symptom of the broader filesystem dependency discussed in the entries above. Using +S3 ranged gets for input and streaming output to shared storage would eliminate the disk +pressure entirely. Short of that, the `sizeLimit` should either be increased or documented as +requiring tuning relative to `numWorkers` and expected job sizes. diff --git a/charts/batch-gateway/templates/_helpers.tpl b/charts/batch-gateway/templates/_helpers.tpl index 0127bacf0..675a2eadc 100644 --- a/charts/batch-gateway/templates/_helpers.tpl +++ b/charts/batch-gateway/templates/_helpers.tpl @@ -166,6 +166,13 @@ app.kubernetes.io/instance: {{ .Release.Name }} app.kubernetes.io/component: processor {{- end }} +{{/* +Processor selector labels as a flat comma-separated string for label selectors. +*/}} +{{- define "batch-gateway.processor.selectorLabelsFlat" -}} +app.kubernetes.io/name={{ include "batch-gateway.name" . }}-processor,app.kubernetes.io/instance={{ .Release.Name }},app.kubernetes.io/component=processor +{{- end }} + {{/* Processor service account name */}} diff --git a/charts/batch-gateway/templates/gc-configmap.yaml b/charts/batch-gateway/templates/gc-configmap.yaml index 1c3684818..059bd59c6 100644 --- a/charts/batch-gateway/templates/gc-configmap.yaml +++ b/charts/batch-gateway/templates/gc-configmap.yaml @@ -45,4 +45,6 @@ data: reconciler: enabled: {{ .Values.gc.config.reconciler.enabled }} interval: {{ .Values.gc.config.reconciler.interval | quote }} + processor_label_selector: {{ include "batch-gateway.processor.selectorLabelsFlat" . | quote }} + processor_statefulset: {{ include "batch-gateway.processor.fullname" . | quote }} {{- end }} diff --git a/charts/batch-gateway/templates/gc-role.yaml b/charts/batch-gateway/templates/gc-role.yaml new file mode 100644 index 000000000..74951bf58 --- /dev/null +++ b/charts/batch-gateway/templates/gc-role.yaml @@ -0,0 +1,15 @@ +{{- if and .Values.gc.enabled .Values.gc.config.reconciler.enabled -}} +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: {{ include "batch-gateway.gc.fullname" . }}-pod-watcher + labels: + {{- include "batch-gateway.gc.labels" . | nindent 4 }} +rules: +- apiGroups: [""] + resources: ["pods"] + verbs: ["get", "list", "watch"] +- apiGroups: ["apps"] + resources: ["statefulsets"] + verbs: ["get", "list", "watch"] +{{- end }} diff --git a/charts/batch-gateway/templates/gc-rolebinding.yaml b/charts/batch-gateway/templates/gc-rolebinding.yaml new file mode 100644 index 000000000..42e11d700 --- /dev/null +++ b/charts/batch-gateway/templates/gc-rolebinding.yaml @@ -0,0 +1,16 @@ +{{- if and .Values.gc.enabled .Values.gc.config.reconciler.enabled -}} +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: {{ include "batch-gateway.gc.fullname" . }}-pod-watcher + labels: + {{- include "batch-gateway.gc.labels" . | nindent 4 }} +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: {{ include "batch-gateway.gc.fullname" . }}-pod-watcher +subjects: +- kind: ServiceAccount + name: {{ include "batch-gateway.gc.serviceAccountName" . }} + namespace: {{ .Release.Namespace }} +{{- end }} diff --git a/charts/batch-gateway/templates/processor-configmap.yaml b/charts/batch-gateway/templates/processor-configmap.yaml index 6410d7ea8..c7a8ee845 100644 --- a/charts/batch-gateway/templates/processor-configmap.yaml +++ b/charts/batch-gateway/templates/processor-configmap.yaml @@ -25,7 +25,6 @@ data: terminate_on_observability_failure: {{ .Values.processor.config.terminateOnObservabilityFailure }} shutdown_timeout: {{ .Values.processor.config.shutdownTimeout | quote }} - heartbeat_interval: {{ .Values.processor.config.heartbeatInterval | quote }} work_dir: {{ .Values.processor.config.workDir | quote }} db_client: diff --git a/charts/batch-gateway/templates/processor-headless-service.yaml b/charts/batch-gateway/templates/processor-headless-service.yaml new file mode 100644 index 000000000..db237f6cc --- /dev/null +++ b/charts/batch-gateway/templates/processor-headless-service.yaml @@ -0,0 +1,12 @@ +{{- if .Values.processor.enabled -}} +apiVersion: v1 +kind: Service +metadata: + name: {{ include "batch-gateway.processor.fullname" . }} + labels: + {{- include "batch-gateway.processor.labels" . | nindent 4 }} +spec: + clusterIP: None + selector: + {{- include "batch-gateway.processor.selectorLabels" . | nindent 4 }} +{{- end }} diff --git a/charts/batch-gateway/templates/processor-deployment.yaml b/charts/batch-gateway/templates/processor-statefulset.yaml similarity index 97% rename from charts/batch-gateway/templates/processor-deployment.yaml rename to charts/batch-gateway/templates/processor-statefulset.yaml index 0b078e1a5..c040dc9d8 100644 --- a/charts/batch-gateway/templates/processor-deployment.yaml +++ b/charts/batch-gateway/templates/processor-statefulset.yaml @@ -1,12 +1,14 @@ {{- if .Values.processor.enabled -}} apiVersion: apps/v1 -kind: Deployment +kind: StatefulSet metadata: name: {{ include "batch-gateway.processor.fullname" . }} labels: {{- include "batch-gateway.processor.labels" . | nindent 4 }} spec: replicas: {{ .Values.processor.replicaCount }} + serviceName: {{ include "batch-gateway.processor.fullname" . }} + podManagementPolicy: Parallel selector: matchLabels: {{- include "batch-gateway.processor.selectorLabels" . | nindent 6 }} diff --git a/charts/batch-gateway/tests/deployment_test.yaml b/charts/batch-gateway/tests/deployment_test.yaml index d4f19e87d..fa5c8674b 100644 --- a/charts/batch-gateway/tests/deployment_test.yaml +++ b/charts/batch-gateway/tests/deployment_test.yaml @@ -2,7 +2,7 @@ suite: deployments templates: - templates/apiserver-deployment.yaml - templates/apiserver-configmap.yaml - - templates/processor-deployment.yaml + - templates/processor-statefulset.yaml - templates/processor-configmap.yaml - templates/gc-deployment.yaml - templates/gc-configmap.yaml @@ -20,15 +20,15 @@ tests: value: 1 template: templates/apiserver-deployment.yaml - - it: should render processor Deployment + - it: should render processor StatefulSet asserts: - isKind: - of: Deployment - template: templates/processor-deployment.yaml + of: StatefulSet + template: templates/processor-statefulset.yaml - equal: path: spec.replicas value: 1 - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - it: should render gc Deployment asserts: @@ -72,7 +72,7 @@ tests: - equal: path: spec.replicas value: 5 - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - it: should not render apiserver Deployment when disabled set: @@ -82,13 +82,13 @@ tests: count: 0 template: templates/apiserver-deployment.yaml - - it: should not render processor Deployment when disabled + - it: should not render processor StatefulSet when disabled set: processor.enabled: false asserts: - hasDocuments: count: 0 - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - it: should not render gc Deployment when disabled set: @@ -114,14 +114,14 @@ tests: - equal: path: spec.template.spec.securityContext.runAsNonRoot value: true - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - isNull: path: spec.template.spec.securityContext.runAsUser - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - equal: path: spec.template.spec.securityContext.seccompProfile.type value: RuntimeDefault - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - equal: path: spec.template.spec.securityContext.runAsNonRoot value: true @@ -145,7 +145,7 @@ tests: - equal: path: spec.template.spec.securityContext.runAsUser value: 1000 - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - isNull: path: spec.template.spec.securityContext.runAsUser template: templates/apiserver-deployment.yaml @@ -164,7 +164,7 @@ tests: - equal: path: spec.template.spec.serviceAccountName value: RELEASE-NAME-batch-gateway-processor - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - equal: path: spec.template.spec.serviceAccountName value: RELEASE-NAME-batch-gateway-gc @@ -183,7 +183,7 @@ tests: - equal: path: spec.template.spec.serviceAccountName value: custom-processor-sa - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - equal: path: spec.template.spec.serviceAccountName value: custom-gc-sa @@ -202,7 +202,7 @@ tests: - equal: path: spec.template.spec.serviceAccountName value: default - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - equal: path: spec.template.spec.serviceAccountName value: default @@ -254,7 +254,7 @@ tests: - it: should not inject OTel env vars on processor when endpoint is empty templates: - - templates/processor-deployment.yaml + - templates/processor-statefulset.yaml asserts: - notContains: path: spec.template.spec.containers[0].env @@ -269,7 +269,7 @@ tests: - it: should inject all OTel env vars on processor when endpoint is set templates: - - templates/processor-deployment.yaml + - templates/processor-statefulset.yaml set: global.otel.endpoint: "http://jaeger:4317" asserts: @@ -300,7 +300,7 @@ tests: - matchRegex: path: spec.template.spec.containers[0].image pattern: "^.+:.+$" - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - matchRegex: path: spec.template.spec.containers[0].image pattern: "^.+:.+$" @@ -322,7 +322,7 @@ tests: - equal: path: spec.template.spec.containers[0].image value: "ghcr.io/llm-d/batch-gateway-processor@sha256:def456" - template: templates/processor-deployment.yaml + template: templates/processor-statefulset.yaml - it: should use digest when set for gc set: @@ -347,7 +347,7 @@ tests: - it: should fail when fs pvcName is empty templates: - - templates/processor-deployment.yaml + - templates/processor-statefulset.yaml set: global.fileClient.type: "fs" global.fileClient.fs.pvcName: "" diff --git a/charts/batch-gateway/values.yaml b/charts/batch-gateway/values.yaml index 40afbe770..917f21619 100644 --- a/charts/batch-gateway/values.yaml +++ b/charts/batch-gateway/values.yaml @@ -349,7 +349,6 @@ processor: count: 12 terminateOnObservabilityFailure: false shutdownTimeout: "30s" - heartbeatInterval: "5m" workDir: "/var/lib/batch-gateway/processor" # REQUIRED: Configure exactly one of globalInferenceGateway or modelGateways. # The processor will fail to start if neither is set. diff --git a/cmd/batch-gc/main.go b/cmd/batch-gc/main.go index 325d834f5..8cc5423d7 100644 --- a/cmd/batch-gc/main.go +++ b/cmd/batch-gc/main.go @@ -39,6 +39,7 @@ import ( "github.com/llm-d/llm-d-batch-gateway/internal/gc/collector" gcconfig "github.com/llm-d/llm-d-batch-gateway/internal/gc/config" gcmetrics "github.com/llm-d/llm-d-batch-gateway/internal/gc/metrics" + "github.com/llm-d/llm-d-batch-gateway/internal/gc/podwatcher" "github.com/llm-d/llm-d-batch-gateway/internal/gc/reconciler" "github.com/llm-d/llm-d-batch-gateway/internal/util/clientset" ucom "github.com/llm-d/llm-d-batch-gateway/internal/util/com" @@ -116,19 +117,26 @@ func run() error { onCycle := func(r *reconciler.Result) { gcmetrics.RecordCycleDuration(r.Duration) if !cfg.DryRun { - gcmetrics.RecordOrphansRecovered(gcmetrics.ActionCancelled, r.Cancelled) gcmetrics.RecordOrphansRecovered(gcmetrics.ActionExpired, r.Expired) gcmetrics.RecordOrphansRecovered(gcmetrics.ActionReEnqueued, r.ReEnqueued) - gcmetrics.RecordOrphansRecovered(gcmetrics.ActionFailed, r.Failed) - gcmetrics.RecordStaleCleanup(r.StaleCleanup) } gcmetrics.RecordCASConflicts(r.Conflicts) gcmetrics.RecordErrors(r.Errors) } - rec, err := reconciler.NewReconciler(clients.BatchDB, clients.Queue, clients.InFlight, cfg.Reconciler.Interval, cfg.DryRun, onCycle) + rec, err := reconciler.NewReconciler(clients.BatchDB, clients.Queue, cfg.Reconciler.Interval, cfg.DryRun, onCycle) if err != nil { return fmt.Errorf("failed to create reconciler: %w", err) } + + pw, err := podwatcher.New(cfg.Reconciler.ProcessorStatefulSet, cfg.Reconciler.ProcessorLabelSelector, func(live map[string]bool) { + rec.SetLiveProcessors(live) + rec.Trigger() + }) + if err != nil { + return fmt.Errorf("failed to create pod watcher: %w", err) + } + + g.Go(func() error { return pw.Run(gCtx) }) g.Go(func() error { return rec.RunLoop(gCtx) }) } diff --git a/cmd/batch-processor/config.yaml b/cmd/batch-processor/config.yaml index 00e105243..921969b79 100644 --- a/cmd/batch-processor/config.yaml +++ b/cmd/batch-processor/config.yaml @@ -34,11 +34,6 @@ terminate_on_observability_failure: false # Shutdown timeout for processor (should be less than k8s terminationGracePeriodSeconds) shutdown_timeout: "30s" -# Heartbeat interval for in-flight entry refresh. -# Must be shorter than the GC reconciler's interval so active jobs are not -# mistaken for orphans. Default: 5m (matches GC default of 60m). -heartbeat_interval: "5m" - # Work directory for processor work_dir: "/var/lib/batch-gateway/processor" diff --git a/examples/deploy-demo/common.sh b/examples/deploy-demo/common.sh index df49ed08b..281c2357e 100755 --- a/examples/deploy-demo/common.sh +++ b/examples/deploy-demo/common.sh @@ -117,21 +117,29 @@ wait_for_deployment() { local namespace="$2" local timeout="${3:-180s}" - step "Waiting for deployment '${deploy_name}' to be ready..." - + # Auto-detect whether the resource is a Deployment or StatefulSet. + local kind="deploy" local retries=0 local max_retries=30 - while ! kubectl get deploy "${deploy_name}" -n "${namespace}" &>/dev/null; do + while true; do + if kubectl get deploy "${deploy_name}" -n "${namespace}" &>/dev/null; then + kind="deploy" + break + elif kubectl get statefulset "${deploy_name}" -n "${namespace}" &>/dev/null; then + kind="statefulset" + break + fi retries=$((retries + 1)) if [ "$retries" -ge "$max_retries" ]; then - die "Deployment '${deploy_name}' did not become visible after $((max_retries * 2))s" + die "'${deploy_name}' did not become visible after $((max_retries * 2))s" fi - warn "Deployment not yet visible, retrying in 2s... ($retries/$max_retries)" + warn "Resource not yet visible, retrying in 2s... ($retries/$max_retries)" sleep 2 done - kubectl rollout status deploy/"${deploy_name}" -n "${namespace}" --timeout="${timeout}" - log "Deployment '${deploy_name}' is ready." + step "Waiting for ${kind}/${deploy_name} to be ready..." + kubectl rollout status "${kind}/${deploy_name}" -n "${namespace}" --timeout="${timeout}" + log "${kind}/${deploy_name} is ready." } wait_for_subscription() { @@ -686,7 +694,12 @@ do_deploy_batch_gateway_helm() { local mismatch=false for component in apiserver processor gc; do local actual_image - actual_image=$(kubectl get deploy "${BATCH_INSTANCE_NAME}-${component}" -n "${BATCH_NAMESPACE}" \ + # processor is a StatefulSet; apiserver and gc are Deployments. + local kind="deploy" + if kubectl get statefulset "${BATCH_INSTANCE_NAME}-${component}" -n "${BATCH_NAMESPACE}" &>/dev/null; then + kind="statefulset" + fi + actual_image=$(kubectl get "${kind}" "${BATCH_INSTANCE_NAME}-${component}" -n "${BATCH_NAMESPACE}" \ -o jsonpath='{.spec.template.spec.containers[0].image}') local actual_tag="${actual_image##*:}" log " ${component}: ${actual_image}" diff --git a/go.mod b/go.mod index 6559a4a4f..7aac4a5d0 100644 --- a/go.mod +++ b/go.mod @@ -28,25 +28,55 @@ require ( go.opentelemetry.io/otel/trace v1.44.0 golang.org/x/sync v0.22.0 gopkg.in/yaml.v3 v3.0.1 + k8s.io/api v0.36.2 + k8s.io/apimachinery v0.36.2 + k8s.io/client-go v0.36.2 k8s.io/klog/v2 v2.140.0 + k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2 ) require ( + github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect + github.com/emicklei/go-restful/v3 v3.13.0 // indirect github.com/felixge/httpsnoop v1.1.0 // indirect + github.com/fxamacker/cbor/v2 v2.9.0 // indirect github.com/go-logr/stdr v1.2.2 // indirect + github.com/go-openapi/jsonpointer v0.21.0 // indirect + github.com/go-openapi/jsonreference v0.20.2 // indirect + github.com/go-openapi/swag v0.23.0 // indirect + github.com/google/gnostic-models v0.7.0 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect + github.com/josharian/intern v1.0.0 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect + github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/redis/go-redis/extra/rediscmd/v9 v9.21.0 // indirect + github.com/x448/float16 v0.8.4 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.44.0 // indirect go.opentelemetry.io/otel/metric v1.44.0 // indirect go.opentelemetry.io/proto/otlp v1.10.0 // indirect + go.yaml.in/yaml/v2 v2.4.4 // indirect + go.yaml.in/yaml/v3 v3.0.4 // indirect + golang.org/x/oauth2 v0.36.0 // indirect + golang.org/x/term v0.45.0 // indirect golang.org/x/text v0.40.0 // indirect + golang.org/x/time v0.14.0 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20260630182238-925bb5da69e7 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260630182238-925bb5da69e7 // indirect google.golang.org/grpc v1.82.0 // indirect + gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect + gopkg.in/inf.v0 v0.9.1 // indirect + k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a // indirect + sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 // indirect + sigs.k8s.io/randfill v1.0.0 // indirect + sigs.k8s.io/structured-merge-diff/v6 v6.3.2 // indirect + sigs.k8s.io/yaml v1.6.0 // indirect ) require ( @@ -75,5 +105,5 @@ require ( go.uber.org/atomic v1.11.0 // indirect golang.org/x/net v0.57.0 // indirect golang.org/x/sys v0.47.0 // indirect - google.golang.org/protobuf v1.36.11 // indirect + google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect ) diff --git a/go.sum b/go.sum index 75e11ac90..43c915f9f 100644 --- a/go.sum +++ b/go.sum @@ -48,24 +48,41 @@ github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1x github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/emicklei/go-restful/v3 v3.13.0 h1:C4Bl2xDndpU6nJ4bc1jXd+uTmYPVUwkD6bFY/oTyCes= +github.com/emicklei/go-restful/v3 v3.13.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc= github.com/exaring/otelpgx v0.11.1 h1:pE79fIg/qh/Lpu00kvswFC5dKfqyJJhMJ4Y4N3w5Lj4= github.com/exaring/otelpgx v0.11.1/go.mod h1:3OojrUKhhy3lTbYIMBijP3YjMey/jo14eHAW5cXcUdk= github.com/felixge/httpsnoop v1.1.0 h1:3YtUj32ZZkqZtt3sZZsClsymw/QDuVfpNhoA31zeORc= github.com/felixge/httpsnoop v1.1.0/go.mod h1:Zqxgdd+1Rkcz8euOqdr7lqgCRJztwr5hp9vDSi5UZCE= +github.com/fxamacker/cbor/v2 v2.9.0 h1:NpKPmjDBgUfBms6tr6JZkTHtfFGcMKsw3eGcmD/sapM= +github.com/fxamacker/cbor/v2 v2.9.0/go.mod h1:vM4b+DJCtHn+zz7h3FFp/hDAI9WNWCsZj23V5ytsSxQ= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.4 h1:tG4xh9yMsRCAiodLVTxyrkzSZ9+o0L1Kg/+cPVcbP/8= github.com/go-logr/logr v1.4.4/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/go-openapi/jsonpointer v0.19.6/go.mod h1:osyAmYz/mB/C3I+WsTTSgw1ONzaLJoLCyoi6/zppojs= +github.com/go-openapi/jsonpointer v0.21.0 h1:YgdVicSA9vH5RiHs9TZW5oyafXZFc6+2Vc1rr/O9oNQ= +github.com/go-openapi/jsonpointer v0.21.0/go.mod h1:IUyH9l/+uyhIYQ/PXVA41Rexl+kOkAPDdXEYns6fzUY= +github.com/go-openapi/jsonreference v0.20.2 h1:3sVjiK66+uXK/6oQ8xgcRKcFgQ5KXa2KvnJRumpMGbE= +github.com/go-openapi/jsonreference v0.20.2/go.mod h1:Bl1zwGIM8/wsvqjsOQLJ/SH+En5Ap4rVB5KVcIDZG2k= +github.com/go-openapi/swag v0.22.3/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+VcZ0yl14= +github.com/go-openapi/swag v0.23.0 h1:vsEVJDUo2hPJ2tu0/Xc+4noaxyEffXNIs3cOULZ+GrE= +github.com/go-openapi/swag v0.23.0/go.mod h1:esZ8ITTYEsH1V2trKHjAN8Ai7xHb8RV+YSZ577vPjgQ= github.com/go-resty/resty/v2 v2.17.2 h1:FQW5oHYcIlkCNrMD2lloGScxcHJ0gkjshV3qcQAyHQk= github.com/go-resty/resty/v2 v2.17.2/go.mod h1:kCKZ3wWmwJaNc7S29BRtUhJwy7iqmn+2mLtQrOyQlVA= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/gnostic-models v0.7.0 h1:qwTtogB15McXDaNqTZdzPJRHvaVJlAl+HVQnLmJEJxo= +github.com/google/gnostic-models v0.7.0/go.mod h1:whL5G0m6dmc5cPxKc5bdKdEN3UjI7OUGxBlw57miDrQ= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF28c5ZQfqCBQ5g2xfk= @@ -78,12 +95,19 @@ github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= +github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJnhSlk= github.com/klauspost/compress v1.19.1/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE= github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= 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= @@ -92,6 +116,14 @@ github.com/llm-d/llm-d-async/api v0.9.0 h1:5kQD7UChMZtD/r8iED+uWm/ArlaIGLGCig1gt github.com/llm-d/llm-d-async/api v0.9.0/go.mod h1:hzjFDTFBJEyW9/1vrHAD2dtCVLTzqggU1GBW+7YhxlQ= github.com/llm-d/llm-d-async/producer v0.9.0 h1:PL4l0RaL0zm6Rm7rC5m+iRu+QkzY2wcxJU+fP8uxc4U= github.com/llm-d/llm-d-async/producer v0.9.0/go.mod h1:c4vtdCfWFAdSKr0pUTI2wUWkkYQfgXJCihNhMc5DAKs= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee h1:W5t00kpgFdJifH4BDsTlE89Zl93FEloxaWZfGcifgq8= +github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= 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/pashagolub/pgxmock/v4 v4.9.0 h1:itlO8nrVRnzkdMBXLs8pWUyyB2PC3Gku0WGIj/gGl7I= @@ -117,11 +149,22 @@ github.com/redis/go-redis/v9 v9.21.0 h1:FPBE4hhbAke+TLmcY3WkpbDffJEomdqPn3HYiqAt github.com/redis/go-redis/v9 v9.21.0/go.mod h1:v/M13XI1PVCDcm01VtPFOADfZtHf8YW3baQf57KlIkA= 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/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY= +github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM= +github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg= github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M= github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw= github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= @@ -152,16 +195,22 @@ go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ= go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ= +go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= +go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= golang.org/x/net v0.57.0 h1:K5+3DljvIuDG9/Jv9rvyMywYNFCQ9RSUY6OOTTkT+tE= golang.org/x/net v0.57.0/go.mod h1:KpXc8iv+r3XplLAG/f7Jsf9RPszJzdR0f58q9vGOuEU= +golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= +golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/term v0.45.0 h1:NwWyBmoJCbfTHpxrWoZ9C6/VxOf7ic219I8xZZFdrf0= +golang.org/x/term v0.45.0/go.mod h1:9aqxs0blBcrm/n0L9QW0aRVD+ktan8ssZromtqJC43w= golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= -golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE= -golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg= +golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI= +golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= google.golang.org/genproto/googleapis/api v0.0.0-20260630182238-925bb5da69e7 h1:jQ9p21COKWjP3VwuFrNRiiOTMh3mPpN45R7SLrH/HUU= @@ -170,13 +219,35 @@ google.golang.org/genproto/googleapis/rpc v0.0.0-20260630182238-925bb5da69e7 h1: google.golang.org/genproto/googleapis/rpc v0.0.0-20260630182238-925bb5da69e7/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= google.golang.org/grpc v1.82.0 h1:vguDnZUPjE26w09A63VoxZPnvPjB5Riyc0mkXPFmAIU= google.golang.org/grpc v1.82.0/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3+/ZA= -google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= -google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/evanphx/json-patch.v4 v4.13.0 h1:czT3CmqEaQ1aanPc5SdlgQrrEIb8w/wwCvWWnfEbYzo= +gopkg.in/evanphx/json-patch.v4 v4.13.0/go.mod h1:p8EYWUEYMpynmqDbY58zCKCFZw8pRWMG4EsWvDvM72M= +gopkg.in/inf.v0 v0.9.1 h1:73M5CoZyi3ZLMOyDlQh031Cx6N9NDJ2Vvfl76EDAgDc= +gopkg.in/inf.v0 v0.9.1/go.mod h1:cWUDdTG/fYaXco+Dcufb5Vnc6Gp2YChqWtbxRZE0mXw= 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= +k8s.io/api v0.36.2 h1:TF6YDLIzKfccK7cq9YpTcGX8TJmEkHVRv78DM51fRYY= +k8s.io/api v0.36.2/go.mod h1:F4LbMO4brjZYh7yFkXWhynSvtB7YauxV4c+HHkNRGNg= +k8s.io/apimachinery v0.36.2 h1:0PE/W/WNy1UX61NLbXY5TMbJ6UwLL6E6lAPkYrKFxbQ= +k8s.io/apimachinery v0.36.2/go.mod h1:fvf/HOLXq9RId0rnDIbN1OEBvHXdQbLMM8nu0LcBUf4= +k8s.io/client-go v0.36.2 h1:bfgxmFKc9CgqsgX4xKLAAdmTQlWee7Ob/HlDOrJ5TBI= +k8s.io/client-go v0.36.2/go.mod h1:1vgO4OAlfPnoLcb+Rze2GF5rAr14w8qjrYMoyXJzQj0= k8s.io/klog/v2 v2.140.0 h1:Tf+J3AH7xnUzZyVVXhTgGhEKnFqye14aadWv7bzXdzc= k8s.io/klog/v2 v2.140.0/go.mod h1:o+/RWfJ6PwpnFn7OyAG3QnO47BFsymfEfrz6XyYSSp0= +k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a h1:xCeOEAOoGYl2jnJoHkC3hkbPJgdATINPMAxaynU2Ovg= +k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a/go.mod h1:uGBT7iTA6c6MvqUvSXIaYZo9ukscABYi2btjhvgKGZ0= +k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2 h1:AZYQSJemyQB5eRxqcPky+/7EdBj0xi3g0ZcxxJ7vbWU= +k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2/go.mod h1:xDxuJ0whA3d0I4mf/C4ppKHxXynQ+fxnkmQH0vTHnuk= +sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 h1:IpInykpT6ceI+QxKBbEflcR5EXP7sU1kvOlxwZh5txg= +sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730/go.mod h1:mdzfpAEoE6DHQEN0uh9ZbOCuHbLK5wOm7dK4ctXE9Tg= +sigs.k8s.io/randfill v1.0.0 h1:JfjMILfT8A6RbawdsK2JXGBR5AQVfd+9TbzrlneTyrU= +sigs.k8s.io/randfill v1.0.0/go.mod h1:XeLlZ/jmk4i1HRopwe7/aU3H5n1zNUcX6TM94b3QxOY= +sigs.k8s.io/structured-merge-diff/v6 v6.3.2 h1:kwVWMx5yS1CrnFWA/2QHyRVJ8jM6dBA80uLmm0wJkk8= +sigs.k8s.io/structured-merge-diff/v6 v6.3.2/go.mod h1:M3W8sfWvn2HhQDIbGWj3S099YozAsymCo/wrT5ohRUE= +sigs.k8s.io/yaml v1.6.0 h1:G8fkbMSAFqgEFgh4b1wmtzDnioxFCUgTZhlbj5P9QYs= +sigs.k8s.io/yaml v1.6.0/go.mod h1:796bPqUfzR/0jLAl6XjHl3Ck7MiyVv8dbTdyT3/pMf4= diff --git a/internal/apiserver/batch/batch_handler.go b/internal/apiserver/batch/batch_handler.go index 5d1bdd218..97fa20b2b 100644 --- a/internal/apiserver/batch/batch_handler.go +++ b/internal/apiserver/batch/batch_handler.go @@ -23,7 +23,6 @@ import ( "encoding/json" "fmt" "net/http" - "strconv" "time" "go.opentelemetry.io/otel" @@ -184,9 +183,7 @@ func (c *BatchAPIHandler) CreateBatch(w http.ResponseWriter, r *http.Request) { // TODO: output_expires_after_anchor and output_expires_after_seconds are saved to database as tag. The cleanup service should delete the output file by this value // Note that the output_expires_after_anchor is the file creation time, not the time the batch is created. - tags := api.Tags{ - batch_types.TagSLO: fmt.Sprintf("%d", slo.UnixMicro()), - } + tags := api.Tags{} if batchReq.OutputExpiresAfter != nil { tags[batch_types.TagOutputExpiresAfterAnchor] = batchReq.OutputExpiresAfter.Anchor tags[batch_types.TagOutputExpiresAfterSeconds] = fmt.Sprintf("%d", batchReq.OutputExpiresAfter.Seconds) @@ -520,38 +517,37 @@ func (c *BatchAPIHandler) CancelBatch(w http.ResponseWriter, r *http.Request) { } // Try to remove from the priority queue first. - // Reconstruct the exact SLO score from the stored tag. - removedFromQueue := false - sloStr, hasSLO := item.Tags[batch_types.TagSLO] - sloMicro, parseErr := strconv.ParseInt(sloStr, 10, 64) - if hasSLO && parseErr == nil { - slo := time.UnixMicro(sloMicro).UTC() - jobPriority := &api.BatchJobPriority{ - ID: batch.ID, - SLO: slo, + // PQDelete atomically transitions the job to cancelled (with cancelled_at) + // if it is still unclaimed. Returns 1 if cancelled, 0 if already claimed. + nDeleted, err := c.clients.Queue.PQDelete(ctx, &api.BatchJobPriority{ID: batch.ID}) + if err != nil { + logger.Error(err, "failed to remove batch from queue") + common.WriteInternalServerError(w, r) + return + } + + if nDeleted > 0 { + // Job was in queue — PQDelete already transitioned it to cancelled. + // Re-read the updated state and return it. + freshItem, apiErr := c.getBatchItemFromDB(r, "cancel") + if apiErr != nil { + common.WriteAPIError(w, r, *apiErr) + return } - nDeleted, err := c.clients.Queue.PQDelete(ctx, jobPriority) - if err != nil { - logger.Error(err, "failed to remove batch from queue") + freshBatch, convErr := converter.DBItemToBatch(freshItem) + if convErr != nil { + logger.Error(convErr, "failed to convert database item to batch") common.WriteInternalServerError(w, r) return } - removedFromQueue = nDeleted > 0 - } else { - logger.Info("SLO tag missing or malformed, skipping queue removal", "key", batch_types.TagSLO, "hasSLO", hasSLO, "error", parseErr) + common.WriteJSONResponse(w, r, http.StatusOK, freshBatch) + return } - if removedFromQueue { - // Job was in queue (not yet being processed) - directly cancel it - batch.Status = openai.BatchStatusCancelled - cancelledAt := time.Now().UTC().Unix() - batch.CancelledAt = &cancelledAt - } else { - // Job is being processed - mark as cancelling and send cancel event - batch.Status = openai.BatchStatusCancelling - cancellingAt := time.Now().UTC().Unix() - batch.CancellingAt = &cancellingAt - } + // Job is being processed — mark as cancelling and send cancel event. + batch.Status = openai.BatchStatusCancelling + cancellingAt := time.Now().UTC().Unix() + batch.CancellingAt = &cancellingAt // Persist the status change *before* sending the cancel event to prevent a // write-write race between the API server and the worker. @@ -580,21 +576,19 @@ func (c *BatchAPIHandler) CancelBatch(w http.ResponseWriter, r *http.Request) { return } - // If the job is being processed, send the cancel event *after* DB update succeeds. - if !removedFromQueue { - event := []api.BatchEvent{ - { - ID: batch.ID, - Type: api.BatchEventCancel, - TTL: c.config.BatchAPI.GetBatchEventTTLSeconds(), - }, - } - _, err = c.clients.Event.ECProducerSendEvents(ctx, event) - if err != nil { - logger.Error(err, "failed to send cancel event") - common.WriteInternalServerError(w, r) - return - } + // Send the cancel event *after* DB update succeeds. + event := []api.BatchEvent{ + { + ID: batch.ID, + Type: api.BatchEventCancel, + TTL: c.config.BatchAPI.GetBatchEventTTLSeconds(), + }, + } + _, err = c.clients.Event.ECProducerSendEvents(ctx, event) + if err != nil { + logger.Error(err, "failed to send cancel event") + common.WriteInternalServerError(w, r) + return } common.WriteJSONResponse(w, r, http.StatusOK, batch) diff --git a/internal/apiserver/batch/batch_handler_test.go b/internal/apiserver/batch/batch_handler_test.go index e5326a771..722f4111f 100644 --- a/internal/apiserver/batch/batch_handler_test.go +++ b/internal/apiserver/batch/batch_handler_test.go @@ -947,10 +947,7 @@ func TestBatchHandler(t *testing.T) { }, }, } - slo := time.Now().UTC().Add(24 * time.Hour) - item, err := converter.BatchToDBItem(&batch, common.DefaultTenantID, map[string]string{ - batch_types.TagSLO: fmt.Sprintf("%d", slo.UnixMicro()), - }) + item, err := converter.BatchToDBItem(&batch, common.DefaultTenantID, nil) if err != nil { t.Fatalf("Failed to convert batch to DB item: %v", err) } @@ -1053,10 +1050,7 @@ func TestBatchHandler(t *testing.T) { Status: openai.BatchStatusInProgress, }, } - slo := time.Now().UTC().Add(24 * time.Hour) - item, err := converter.BatchToDBItem(&batch, common.DefaultTenantID, map[string]string{ - batch_types.TagSLO: fmt.Sprintf("%d", slo.UnixMicro()), - }) + item, err := converter.BatchToDBItem(&batch, common.DefaultTenantID, nil) if err != nil { t.Fatalf("Failed to convert batch to DB item: %v", err) } @@ -1093,10 +1087,7 @@ func TestBatchHandler(t *testing.T) { CancellingAt: &cancellingAt, }, } - slo := time.Now().UTC().Add(24 * time.Hour) - item, err := converter.BatchToDBItem(&batch, common.DefaultTenantID, map[string]string{ - batch_types.TagSLO: fmt.Sprintf("%d", slo.UnixMicro()), - }) + item, err := converter.BatchToDBItem(&batch, common.DefaultTenantID, nil) if err != nil { t.Fatalf("Failed to convert batch to database item: %v", err) } @@ -1161,10 +1152,7 @@ func TestBatchHandler(t *testing.T) { CancellingAt: &cancellingAt, }, } - slo := time.Now().UTC().Add(24 * time.Hour) - item, err := converter.BatchToDBItem(&batch, common.DefaultTenantID, map[string]string{ - batch_types.TagSLO: fmt.Sprintf("%d", slo.UnixMicro()), - }) + item, err := converter.BatchToDBItem(&batch, common.DefaultTenantID, nil) if err != nil { t.Fatalf("Failed to convert batch to database item: %v", err) } diff --git a/internal/apiserver/common/config.go b/internal/apiserver/common/config.go index af68edacd..ef7d1f641 100644 --- a/internal/apiserver/common/config.go +++ b/internal/apiserver/common/config.go @@ -222,7 +222,7 @@ func (c *ServerConfig) applyDefaults() { c.ObservabilityPort = "8081" } if c.DBClientCfg.Type == "" { - c.DBClientCfg.Type = sharedcfg.DBTypeRedis + c.DBClientCfg.Type = sharedcfg.DBTypePostgreSQL } if c.ReadHeaderTimeoutSeconds <= 0 { c.ReadHeaderTimeoutSeconds = DefaultReadHeaderTimeoutSeconds diff --git a/internal/database/api/batch_item.go b/internal/database/api/batch_item.go index 041b8ded9..ace9bdf4b 100644 --- a/internal/database/api/batch_item.go +++ b/internal/database/api/batch_item.go @@ -20,12 +20,33 @@ package api type BatchItem struct { BaseIndexes BaseContents + + // ProcessorID identifies the processor pod that owns this job. + // Set atomically with the status transition to in_progress during dequeue. + // Empty when the job is queued (validating) or in a terminal state. + ProcessorID string + + // Priority determines dequeue order (lower = higher priority). + // Stores SLO.UnixMicro() — jobs with earlier deadlines are dequeued first. + Priority int64 + + // Epoch is a fencing token incremented on every ownership change (dequeue, + // recovery, GC reclaim). Processor writes include WHERE epoch = N so a + // zombie whose lease was reclaimed cannot overwrite the new owner's state. + Epoch int64 + + // BumpEpoch signals that DBUpdate should atomically increment the epoch + // in addition to checking it. Set by the GC reconciler when reclaiming + // an orphan — this is an ownership change (like a Raft term bump). + BumpEpoch bool } // BatchQuery specifies parameters for retrieving batches from the database. type BatchQuery struct { BaseQuery - NonTerminal bool + NonTerminal bool + ProcessorID string + HasProcessorID bool // filter for processor_id IS NOT NULL (owned jobs) } // BatchDBClient is the typed database client for batch objects. diff --git a/internal/database/api/database.go b/internal/database/api/database.go index 53ceb3322..0ba0c5e83 100644 --- a/internal/database/api/database.go +++ b/internal/database/api/database.go @@ -104,10 +104,11 @@ var LogicalCondNames = map[LogicalCond]string{ // -- Batch jobs priority queue -- type BatchJobPriority struct { - ID string `json:"id,omitempty"` // [mandatory] ID of the batch job. - SLO time.Time `json:"slo,omitempty"` // [mandatory] The SLO value determines the priority of the job. - Data []byte `json:"data,omitempty"` // [optional] User defined data. - TTL int // [optional] TTL in seconds applied on the entire queue. If used, this should be set to a sufficiently large value to prevent premature removal of items. + ID string `json:"id,omitempty"` // [mandatory] ID of the batch job. + SLO time.Time `json:"slo,omitempty"` // [mandatory] The SLO value determines the priority of the job. + Data []byte `json:"data,omitempty"` // [optional] User defined data. + TTL int `json:"ttl,omitempty"` // [optional] TTL in seconds applied on the entire queue. If used, this should be set to a sufficiently large value to prevent premature removal of items. + Epoch int64 `json:"epoch,omitempty"` // Fencing token incremented on every ownership change. } func (bj *BatchJobPriority) IsValid() error { @@ -215,28 +216,3 @@ type BatchStatusClient interface { // StatusDelete deletes the status data for a job. StatusDelete(ctx context.Context, ID string) (nDeleted int, err error) } - -// -- In-flight job tracking -- - -// InFlightEntry records which processor owns a dequeued job and when it last -// reported liveness. -type InFlightEntry struct { - ProcessorID string `json:"pid"` - LastSeen int64 `json:"ls"` -} - -// InFlightClient tracks jobs that have been dequeued and are being processed. -type InFlightClient interface { - store.BatchClientAdmin - - // InFlightSet records or refreshes the in-flight entry for a job. - // Called after dequeue and periodically as a heartbeat. - InFlightSet(ctx context.Context, jobID, processorID string) error - - // InFlightDelete removes the in-flight entry for a job. - // Called when the job reaches a terminal state. - InFlightDelete(ctx context.Context, jobID string) error - - // InFlightGetAll returns all in-flight entries keyed by job ID. - InFlightGetAll(ctx context.Context) (map[string]*InFlightEntry, error) -} diff --git a/internal/database/mock/mock_db_client.go b/internal/database/mock/mock_db_client.go index 9e9d6fce8..6424aa24f 100644 --- a/internal/database/mock/mock_db_client.go +++ b/internal/database/mock/mock_db_client.go @@ -18,6 +18,7 @@ limitations under the License. package mock import ( + "bytes" "context" "fmt" "reflect" @@ -32,6 +33,10 @@ type MockDBClient[T any, Q any] struct { items sync.Map idGetter func(*T) string baseQueryGetter func(*Q) *api.BaseQuery + + // QueryFilter is an optional callback for query-specific filtering beyond BaseQuery. + // When set, it is called for each item during DBGet. Return true to include the item. + QueryFilter func(item *T, query *Q) bool } // NewMockDBClient creates a new mock DB client. @@ -64,12 +69,22 @@ func (m *MockDBClient[T, Q]) DBGet( bq := m.baseQueryGetter(query) var allMatches []*T + matchesAll := func(item *T) bool { + if !m.matchesFilters(*item, bq) { + return false + } + if m.QueryFilter != nil && !m.QueryFilter(item, query) { + return false + } + return true + } + // If IDs are specified, get by IDs if len(bq.IDs) > 0 { for _, id := range bq.IDs { if value, ok := m.items.Load(id); ok { if item, ok := value.(*T); ok { - if m.matchesFilters(*item, bq) { + if matchesAll(item) { allMatches = append(allMatches, item) } } @@ -79,7 +94,7 @@ func (m *MockDBClient[T, Q]) DBGet( // Collect all items, applying filters m.items.Range(func(key, value any) bool { if item, ok := value.(*T); ok { - if m.matchesFilters(*item, bq) { + if matchesAll(item) { allMatches = append(allMatches, item) } } @@ -114,9 +129,36 @@ func (m *MockDBClient[T, Q]) DBUpdate(ctx context.Context, item *T, expectedStat if id == "" { return fmt.Errorf("item has empty ID") } - if _, ok := m.items.Load(id); !ok { + existing, ok := m.items.Load(id) + if !ok { return fmt.Errorf("cannot update item with ID '%s': item doesn't exist", id) } + + if existingItem, ok := existing.(*T); ok { + val := reflect.ValueOf(*existingItem) + + // CAS: check expectedStatus matches current Status field. + if expectedStatus != nil { + statusField := val.FieldByName("Status") + if statusField.IsValid() && statusField.Kind() == reflect.Slice { + currentStatus, _ := statusField.Interface().([]byte) + if !bytes.Equal(currentStatus, expectedStatus) { + return fmt.Errorf("DBUpdate: %w", api.ErrConflict) + } + } + } + + // Epoch fencing: if the update item has Epoch > 0, check it matches. + updateVal := reflect.ValueOf(*item) + epochField := updateVal.FieldByName("Epoch") + if epochField.IsValid() && epochField.Kind() == reflect.Int64 && epochField.Int() > 0 { + existingEpoch := val.FieldByName("Epoch") + if existingEpoch.IsValid() && existingEpoch.Int() != epochField.Int() { + return fmt.Errorf("DBUpdate: %w", api.ErrConflict) + } + } + } + m.items.Store(id, item) return nil } diff --git a/internal/database/mock/mock_inflight_client.go b/internal/database/mock/mock_inflight_client.go deleted file mode 100644 index 7c8b4b44d..000000000 --- a/internal/database/mock/mock_inflight_client.go +++ /dev/null @@ -1,95 +0,0 @@ -/* -Copyright 2026 The llm-d Authors - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -package mock - -import ( - "context" - "fmt" - "sync" - "time" - - "github.com/llm-d/llm-d-batch-gateway/internal/database/api" -) - -var _ api.InFlightClient = (*MockInFlightClient)(nil) - -type MockInFlightClient struct { - mu sync.Mutex - entries map[string]*api.InFlightEntry -} - -func NewMockInFlightClient() *MockInFlightClient { - return &MockInFlightClient{ - entries: make(map[string]*api.InFlightEntry), - } -} - -func (m *MockInFlightClient) InFlightSet(_ context.Context, jobID, processorID string) error { - if jobID == "" { - return fmt.Errorf("jobID is empty") - } - if processorID == "" { - return fmt.Errorf("processorID is empty") - } - - m.mu.Lock() - defer m.mu.Unlock() - m.entries[jobID] = &api.InFlightEntry{ - ProcessorID: processorID, - LastSeen: time.Now().Unix(), - } - return nil -} - -func (m *MockInFlightClient) InFlightDelete(_ context.Context, jobID string) error { - if jobID == "" { - return fmt.Errorf("jobID is empty") - } - - m.mu.Lock() - defer m.mu.Unlock() - delete(m.entries, jobID) - return nil -} - -func (m *MockInFlightClient) InFlightGetAll(_ context.Context) (map[string]*api.InFlightEntry, error) { - m.mu.Lock() - defer m.mu.Unlock() - - result := make(map[string]*api.InFlightEntry, len(m.entries)) - for k, v := range m.entries { - copied := *v - result[k] = &copied - } - return result, nil -} - -// SetLastSeen overrides the LastSeen timestamp for a specific entry (test helper). -func (m *MockInFlightClient) SetLastSeen(jobID string, lastSeen int64) { - m.mu.Lock() - defer m.mu.Unlock() - if entry, ok := m.entries[jobID]; ok { - entry.LastSeen = lastSeen - } -} - -func (m *MockInFlightClient) Close() error { - m.mu.Lock() - defer m.mu.Unlock() - m.entries = nil - return nil -} diff --git a/internal/database/mock/mock_queue_client.go b/internal/database/mock/mock_queue_client.go index 2907507b2..6f47a8ca7 100644 --- a/internal/database/mock/mock_queue_client.go +++ b/internal/database/mock/mock_queue_client.go @@ -31,6 +31,11 @@ var _ api.BatchPriorityQueueClient = (*MockBatchPriorityQueueClient)(nil) type MockBatchPriorityQueueClient struct { mu sync.Mutex queue []*api.BatchJobPriority + + // OnDelete is called when PQDelete successfully removes a job from the + // queue. This mirrors the Postgres PQDelete behavior of atomically + // transitioning the job to cancelled in the DB. + OnDelete func(ctx context.Context, id string) error } func NewMockBatchPriorityQueueClient() *MockBatchPriorityQueueClient { @@ -109,6 +114,14 @@ func (m *MockBatchPriorityQueueClient) PQDelete(ctx context.Context, jobPriority if jp.ID == jobPriority.ID { // Remove the item m.queue = append(m.queue[:i], m.queue[i+1:]...) + if m.OnDelete != nil { + m.mu.Unlock() + err := m.OnDelete(ctx, jp.ID) + m.mu.Lock() + if err != nil { + return 0, err + } + } return 1, nil } } diff --git a/internal/database/postgresql/batch_db.go b/internal/database/postgresql/batch_db.go index f8164d063..a3c227410 100644 --- a/internal/database/postgresql/batch_db.go +++ b/internal/database/postgresql/batch_db.go @@ -44,15 +44,23 @@ func buildNonTerminalCondition() string { return colStatus + `::jsonb->>'status' NOT IN (` + strings.Join(quoted, ",") + `)` } +const ( + colProcessorID = "processor_id" + colPriority = "priority" + colEpoch = "epoch" +) + // Compile-time check: batchDescriptor implements TableDescriptor. var _ TableDescriptor = (*batchDescriptor)(nil) // batchDescriptor implements TableDescriptor for batch items. type batchDescriptor struct{} -func (batchDescriptor) TableName() string { return "batch_items" } -func (batchDescriptor) Schema() string { return batchSchemaSql } -func (batchDescriptor) ExtraColumns() []string { return nil } +func (batchDescriptor) TableName() string { return "batch_items" } +func (batchDescriptor) Schema() string { return batchSchemaSql } +func (batchDescriptor) ExtraColumns() []string { + return []string{colProcessorID, colPriority, colEpoch} +} // PostgresBatchDBClient implements api.BatchDBClient using PostgreSQL. type PostgresBatchDBClient struct { @@ -85,7 +93,11 @@ func (c *PostgresBatchDBClient) DBStore(ctx context.Context, item *api.BatchItem err = fmt.Errorf("item is nil") return } - if err = c.store(ctx, &item.BaseIndexes, &item.BaseContents, nil); err != nil { + if err = c.store(ctx, &item.BaseIndexes, &item.BaseContents, map[string]any{ + colProcessorID: item.ProcessorID, + colPriority: item.Priority, + colEpoch: item.Epoch, + }); err != nil { return } return @@ -103,18 +115,32 @@ func (c *PostgresBatchDBClient) DBGet( if query.NonTerminal { rawConditions = append(rawConditions, nonTerminalCondition) } + if query.HasProcessorID { + rawConditions = append(rawConditions, colProcessorID+" IS NOT NULL") + } - indexes, contents, _, cursor, expectMore, err := c.get( - ctx, &query.BaseQuery, includeStatic, start, limit, nil, rawConditions) + var extraFilters map[string]any + if query.ProcessorID != "" { + extraFilters = map[string]any{colProcessorID: query.ProcessorID} + } + + indexes, contents, extras, cursor, expectMore, err := c.get( + ctx, &query.BaseQuery, includeStatic, start, limit, extraFilters, rawConditions) if err != nil { return } items = make([]*api.BatchItem, len(indexes)) for i := range indexes { + processorID, _ := extras[i][colProcessorID].(string) + priority, _ := extras[i][colPriority].(int64) + epoch, _ := extras[i][colEpoch].(int64) items[i] = &api.BatchItem{ BaseIndexes: *indexes[i], BaseContents: *contents[i], + ProcessorID: processorID, + Priority: priority, + Epoch: epoch, } } @@ -126,7 +152,15 @@ func (c *PostgresBatchDBClient) DBUpdate(ctx context.Context, item *api.BatchIte err = fmt.Errorf("item is nil") return } - if err = c.update(ctx, &item.BaseIndexes, &item.BaseContents, expectedStatus); err != nil { + var epochFence map[string]any + if item.Epoch > 0 { + epochFence = map[string]any{colEpoch: item.Epoch} + } + var rawSets []string + if item.BumpEpoch { + rawSets = append(rawSets, colEpoch+" = "+colEpoch+" + 1") + } + if err = c.update(ctx, &item.BaseIndexes, &item.BaseContents, expectedStatus, epochFence, rawSets); err != nil { return } return diff --git a/internal/database/postgresql/batch_db_test.go b/internal/database/postgresql/batch_db_test.go index 0440ce904..5fa2e8370 100644 --- a/internal/database/postgresql/batch_db_test.go +++ b/internal/database/postgresql/batch_db_test.go @@ -72,7 +72,7 @@ func TestBatchStore(t *testing.T) { item := newTestBatchItem("batch-1", testTenantID) mock.ExpectExec("INSERT INTO "+testTable). - WithArgs("batch-1", testTenantID, pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). + WithArgs("batch-1", testTenantID, pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnResult(pgxmock.NewResult("INSERT", 1)) if err := client.DBStore(ctx, item); err != nil { @@ -117,8 +117,8 @@ func TestBatchGet(t *testing.T) { item := newTestBatchItem("batch-1", testTenantID) tags, _ := packTags(item.Tags) - rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus, colSpec}). - AddRow(item.ID, item.TenantID, item.Expiry, &tags, item.Status, item.Spec) + rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colProcessorID, colPriority, colEpoch, colStatus, colSpec}). + AddRow(item.ID, item.TenantID, item.Expiry, &tags, "", int64(0), int64(0), item.Status, item.Spec) // The SQL LIMIT is limit+1 because get() fetches an extra row to determine // if more results exist beyond the requested page. @@ -157,8 +157,8 @@ func TestBatchGet(t *testing.T) { item := newTestBatchItem("batch-1", testTenantID) tags, _ := packTags(item.Tags) - rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus}). - AddRow(item.ID, item.TenantID, item.Expiry, &tags, item.Status) + rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colProcessorID, colPriority, colEpoch, colStatus}). + AddRow(item.ID, item.TenantID, item.Expiry, &tags, "", int64(0), int64(0), item.Status) mock.ExpectQuery("SELECT .+ FROM "+testTable+" WHERE"). WithArgs(testTenantID, 0, 11). @@ -191,9 +191,9 @@ func TestBatchGet(t *testing.T) { tags1, _ := packTags(item1.Tags) tags2, _ := packTags(item2.Tags) - rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus, colSpec}). - AddRow(item1.ID, item1.TenantID, item1.Expiry, &tags1, item1.Status, item1.Spec). - AddRow(item2.ID, item2.TenantID, item2.Expiry, &tags2, item2.Status, item2.Spec) + rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colProcessorID, colPriority, colEpoch, colStatus, colSpec}). + AddRow(item1.ID, item1.TenantID, item1.Expiry, &tags1, "", int64(0), int64(0), item1.Status, item1.Spec). + AddRow(item2.ID, item2.TenantID, item2.Expiry, &tags2, "", int64(0), int64(0), item2.Status, item2.Spec) mock.ExpectQuery("SELECT .+ FROM "+testTable+" WHERE"). WithArgs([]string{"batch-1", "batch-2"}, 0, 11). @@ -222,8 +222,8 @@ func TestBatchGet(t *testing.T) { item.Tags = api.Tags{"env": "prod", "team": "ml"} tags, _ := packTags(item.Tags) - rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus, colSpec}). - AddRow(item.ID, item.TenantID, item.Expiry, &tags, item.Status, item.Spec) + rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colProcessorID, colPriority, colEpoch, colStatus, colSpec}). + AddRow(item.ID, item.TenantID, item.Expiry, &tags, "", int64(0), int64(0), item.Status, item.Spec) mock.ExpectQuery("SELECT .+ FROM "+testTable+" WHERE"). WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), 0, 11). @@ -254,8 +254,8 @@ func TestBatchGet(t *testing.T) { item := newTestBatchItem("batch-1", testTenantID) tags, _ := packTags(item.Tags) - rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus, colSpec}). - AddRow(item.ID, item.TenantID, item.Expiry, &tags, item.Status, item.Spec) + rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colProcessorID, colPriority, colEpoch, colStatus, colSpec}). + AddRow(item.ID, item.TenantID, item.Expiry, &tags, "", int64(0), int64(0), item.Status, item.Spec) mock.ExpectQuery("SELECT .+ FROM "+testTable+" WHERE .+NOT IN"). WithArgs(0, 11). @@ -283,8 +283,8 @@ func TestBatchGet(t *testing.T) { item := newTestBatchItem("batch-1", testTenantID) tags, _ := packTags(item.Tags) - rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus, colSpec}). - AddRow(item.ID, item.TenantID, item.Expiry, &tags, item.Status, item.Spec) + rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colProcessorID, colPriority, colEpoch, colStatus, colSpec}). + AddRow(item.ID, item.TenantID, item.Expiry, &tags, "", int64(0), int64(0), item.Status, item.Spec) mock.ExpectQuery("SELECT .+ FROM "+testTable+" WHERE .+NOT IN"). WithArgs(testTenantID, 0, 11). @@ -315,8 +315,8 @@ func TestBatchGet(t *testing.T) { item := newTestBatchItem("batch-1", testTenantID) tags, _ := packTags(item.Tags) - rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus, colSpec}). - AddRow(item.ID, item.TenantID, item.Expiry, &tags, item.Status, item.Spec) + rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colProcessorID, colPriority, colEpoch, colStatus, colSpec}). + AddRow(item.ID, item.TenantID, item.Expiry, &tags, "", int64(0), int64(0), item.Status, item.Spec) mock.ExpectQuery("SELECT .+ FROM "+testTable+" WHERE"). WithArgs(pgxmock.AnyArg(), 0, 11). diff --git a/internal/database/postgresql/batch_queue.go b/internal/database/postgresql/batch_queue.go new file mode 100644 index 000000000..d2451f3d4 --- /dev/null +++ b/internal/database/postgresql/batch_queue.go @@ -0,0 +1,219 @@ +/* +Copyright 2026 The llm-d Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package postgresql + +import ( + "context" + "fmt" + "time" + + "github.com/go-logr/logr" + + "github.com/llm-d/llm-d-batch-gateway/internal/database/api" + "github.com/llm-d/llm-d-batch-gateway/internal/util/logging" +) + +// PostgresBatchQueueClient implements api.BatchPriorityQueueClient using +// the batch_items table as a native Postgres queue. Jobs with status +// 'validating' and processor_id IS NULL are the queue. Dequeue atomically +// claims jobs using a CTE with SELECT FOR UPDATE SKIP LOCKED + UPDATE +// in a single statement. +type PostgresBatchQueueClient struct { + *pgCore + processorID string +} + +var _ api.BatchPriorityQueueClient = (*PostgresBatchQueueClient)(nil) + +func NewPostgresBatchQueueClient(ctx context.Context, config *PostgreSQLConfig, processorID string) (*PostgresBatchQueueClient, error) { + if ctx == nil { + ctx = context.Background() + } + + pgCore, err := newPgCore(ctx, config, batchDescriptor{}) + if err != nil { + return nil, err + } + + logr.FromContextOrDiscard(ctx).Info("NewPostgresBatchQueueClient: client created successfully") + return &PostgresBatchQueueClient{pgCore: pgCore, processorID: processorID}, nil +} + +func (c *PostgresBatchQueueClient) Close() error { + return c.close() +} + +// PQEnqueue makes a job available for dequeue and notifies listening processors. +// For new jobs (created via DBStore), the row already exists with status 'validating' +// and processor_id = NULL — the WHERE guard makes the UPDATE a no-op and only the +// NOTIFY fires. +// For re-enqueued jobs (recovery/GC), this resets status to 'validating' and clears +// processor_id so the job becomes visible to PQDequeue again. +// The UPDATE is guarded to only affect non-terminal jobs that are currently claimed +// by a processor, preventing accidental resurrection of completed/failed work. +func (c *PostgresBatchQueueClient) PQEnqueue(ctx context.Context, jobPriority *api.BatchJobPriority) error { + if jobPriority == nil { + return fmt.Errorf("PQEnqueue: nil job priority") + } + if jobPriority.ID == "" { + return fmt.Errorf("PQEnqueue: empty job ID") + } + _, err := c.pool.Exec(ctx, + `WITH re_enqueued AS ( + UPDATE batch_items + SET processor_id = NULL, + status = jsonb_set(status, '{status}', '"validating"'), + epoch = epoch + 1 + WHERE id = $1 + AND processor_id IS NOT NULL + AND `+nonTerminalCondition+` + ) + -- TODO: the processor polling loop (worker.go) could LISTEN on this channel + -- to wake up immediately instead of waiting for the next poll interval. + SELECT pg_notify('batch_jobs_available', '')`, + jobPriority.ID, + ) + if err != nil { + return fmt.Errorf("PQEnqueue: %w", err) + } + return nil +} + +// PQDequeue atomically claims up to maxItems unclaimed jobs. A single CTE +// statement selects validating jobs with no processor_id (ordered by priority, +// earliest SLO first), locks them with FOR UPDATE SKIP LOCKED, and sets +// processor_id in one atomic SQL statement. +// +// The timeout parameter is unused — the Postgres implementation is a +// non-blocking query. The caller (worker polling loop) controls retry cadence. +func (c *PostgresBatchQueueClient) PQDequeue(ctx context.Context, _ time.Duration, maxItems int) ([]*api.BatchJobPriority, error) { + if c.processorID == "" { + return nil, fmt.Errorf("PQDequeue: processor ID is empty, only processors can dequeue") + } + logger := logr.FromContextOrDiscard(ctx) + + rows, err := c.pool.Query(ctx, + `WITH claimed AS ( + SELECT id FROM batch_items + WHERE processor_id IS NULL + AND status IS NOT NULL + AND status::jsonb->>'status' = 'validating' + ORDER BY priority ASC + LIMIT $1 + FOR UPDATE SKIP LOCKED + ) + UPDATE batch_items + SET processor_id = $2, + epoch = epoch + 1 + FROM claimed + WHERE batch_items.id = claimed.id + RETURNING batch_items.id, batch_items.priority, batch_items.epoch`, + maxItems, c.processorID, + ) + if err != nil { + return nil, fmt.Errorf("PQDequeue: %w", err) + } + defer rows.Close() + + var result []*api.BatchJobPriority + for rows.Next() { + var id string + var priority, epoch int64 + if err := rows.Scan(&id, &priority, &epoch); err != nil { + return nil, fmt.Errorf("PQDequeue: scan: %w", err) + } + result = append(result, &api.BatchJobPriority{ + ID: id, + SLO: time.UnixMicro(priority), + Epoch: epoch, + }) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("PQDequeue: rows: %w", err) + } + + if len(result) > 0 { + logger.V(logging.DEBUG).Info("PQDequeue: claimed jobs", "count", len(result)) + } + return result, nil +} + +// PQDelete atomically removes a job from the queue by transitioning it to +// cancelled with a cancelled_at timestamp, but only if it is still unclaimed +// (validating with no processor_id). +// Returns 1 if the job was cancelled, 0 if it was already claimed by a processor. +// The FOR UPDATE SKIP LOCKED prevents races with concurrent PQDequeue calls. +func (c *PostgresBatchQueueClient) PQDelete(ctx context.Context, jobPriority *api.BatchJobPriority) (int, error) { + if jobPriority == nil { + return 0, fmt.Errorf("PQDelete: nil job priority") + } + + now := time.Now().UTC().Unix() + result, err := c.pool.Exec(ctx, + `WITH queued AS ( + SELECT id FROM batch_items + WHERE id = $1 + AND processor_id IS NULL + AND status IS NOT NULL + AND status::jsonb->>'status' = 'validating' + FOR UPDATE SKIP LOCKED + ) + UPDATE batch_items + SET status = jsonb_set( + jsonb_set(status, '{status}', '"cancelled"'), + '{cancelled_at}', to_jsonb($2::bigint) + ), + epoch = epoch + 1 + FROM queued + WHERE batch_items.id = queued.id`, + jobPriority.ID, now, + ) + if err != nil { + return 0, fmt.Errorf("PQDelete: %w", err) + } + + return int(result.RowsAffected()), nil +} + +// PQGetIDs returns the set of all job IDs currently in the queue +// (validating with no processor_id). +func (c *PostgresBatchQueueClient) PQGetIDs(ctx context.Context) (map[string]bool, error) { + rows, err := c.pool.Query(ctx, + `SELECT id FROM batch_items + WHERE processor_id IS NULL + AND status IS NOT NULL + AND status::jsonb->>'status' = 'validating'`, + ) + if err != nil { + return nil, fmt.Errorf("PQGetIDs: %w", err) + } + defer rows.Close() + + ids := make(map[string]bool) + for rows.Next() { + var id string + if err := rows.Scan(&id); err != nil { + return nil, fmt.Errorf("PQGetIDs: scan: %w", err) + } + ids[id] = true + } + + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("PQGetIDs: rows: %w", err) + } + return ids, nil +} diff --git a/internal/database/postgresql/batch_queue_test.go b/internal/database/postgresql/batch_queue_test.go new file mode 100644 index 000000000..fd0a8553c --- /dev/null +++ b/internal/database/postgresql/batch_queue_test.go @@ -0,0 +1,311 @@ +/* +Copyright 2026 The llm-d Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package postgresql + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/pashagolub/pgxmock/v4" + + "github.com/llm-d/llm-d-batch-gateway/internal/database/api" +) + +const testProcessorID = "processor-0" + +func newTestQueueClient(t *testing.T) (*PostgresBatchQueueClient, pgxmock.PgxPoolIface) { + t.Helper() + mock, err := pgxmock.NewPool() + if err != nil { + t.Fatalf("failed to create pgxmock pool: %v", err) + } + client := &PostgresBatchQueueClient{ + pgCore: &pgCore{pool: mock, desc: batchDescriptor{}}, + processorID: testProcessorID, + } + return client, mock +} + +func TestPQEnqueue(t *testing.T) { + ctx := context.Background() + + t.Run("enqueues job and notifies", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + mock.ExpectExec("WITH re_enqueued AS"). + WithArgs("batch-1"). + WillReturnResult(pgxmock.NewResult("UPDATE", 1)) + + if err := client.PQEnqueue(ctx, &api.BatchJobPriority{ID: "batch-1"}); err != nil { + t.Fatalf("expected no error, got %v", err) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet expectations: %v", err) + } + }) + + t.Run("returns error for nil job priority", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + if err := client.PQEnqueue(ctx, nil); err == nil { + t.Fatal("expected error for nil job priority") + } + }) + + t.Run("returns error for empty ID", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + if err := client.PQEnqueue(ctx, &api.BatchJobPriority{}); err == nil { + t.Fatal("expected error for empty ID") + } + }) + + t.Run("returns error on failure", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + mock.ExpectExec("WITH re_enqueued AS"). + WithArgs("batch-1"). + WillReturnError(fmt.Errorf("connection refused")) + + if err := client.PQEnqueue(ctx, &api.BatchJobPriority{ID: "batch-1"}); err == nil { + t.Fatal("expected error") + } + }) +} + +func TestPQDequeue(t *testing.T) { + ctx := context.Background() + + t.Run("claims a job", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + slo := time.Now().Add(time.Hour) + rows := pgxmock.NewRows([]string{"id", "priority", "epoch"}). + AddRow("batch-1", slo.UnixMicro(), int64(1)) + + mock.ExpectQuery("WITH claimed AS"). + WithArgs(1, testProcessorID). + WillReturnRows(rows) + + result, err := client.PQDequeue(ctx, 0, 1) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if len(result) != 1 { + t.Fatalf("expected 1 job, got %d", len(result)) + } + if result[0].ID != "batch-1" { + t.Errorf("expected ID batch-1, got %s", result[0].ID) + } + if result[0].SLO.UnixMicro() != slo.UnixMicro() { + t.Errorf("expected SLO %v, got %v", slo.UnixMicro(), result[0].SLO.UnixMicro()) + } + if result[0].Epoch != 1 { + t.Errorf("expected Epoch 1, got %d", result[0].Epoch) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet expectations: %v", err) + } + }) + + t.Run("returns nil when no jobs available", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + rows := pgxmock.NewRows([]string{"id", "priority", "epoch"}) + + mock.ExpectQuery("WITH claimed AS"). + WithArgs(1, testProcessorID). + WillReturnRows(rows) + + result, err := client.PQDequeue(ctx, 0, 1) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if len(result) != 0 { + t.Fatalf("expected empty result, got %v", result) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet expectations: %v", err) + } + }) + + t.Run("claims multiple jobs", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + slo1 := time.Now().Add(time.Hour) + slo2 := time.Now().Add(2 * time.Hour) + rows := pgxmock.NewRows([]string{"id", "priority", "epoch"}). + AddRow("batch-1", slo1.UnixMicro(), int64(1)). + AddRow("batch-2", slo2.UnixMicro(), int64(1)) + + mock.ExpectQuery("WITH claimed AS"). + WithArgs(5, testProcessorID). + WillReturnRows(rows) + + result, err := client.PQDequeue(ctx, 0, 5) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if len(result) != 2 { + t.Fatalf("expected 2 jobs, got %d", len(result)) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet expectations: %v", err) + } + }) + + t.Run("returns error on query failure", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + mock.ExpectQuery("WITH claimed AS"). + WithArgs(1, testProcessorID). + WillReturnError(fmt.Errorf("connection refused")) + + _, err := client.PQDequeue(ctx, 0, 1) + if err == nil { + t.Fatal("expected error") + } + }) +} + +func TestPQDelete(t *testing.T) { + ctx := context.Background() + + t.Run("cancels unclaimed job", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + mock.ExpectExec("WITH queued AS"). + WithArgs("batch-1", pgxmock.AnyArg()). + WillReturnResult(pgxmock.NewResult("UPDATE", 1)) + + n, err := client.PQDelete(ctx, &api.BatchJobPriority{ID: "batch-1"}) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if n != 1 { + t.Errorf("expected 1 affected row, got %d", n) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet expectations: %v", err) + } + }) + + t.Run("returns 0 for already claimed job", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + mock.ExpectExec("WITH queued AS"). + WithArgs("batch-1", pgxmock.AnyArg()). + WillReturnResult(pgxmock.NewResult("UPDATE", 0)) + + n, err := client.PQDelete(ctx, &api.BatchJobPriority{ID: "batch-1"}) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if n != 0 { + t.Errorf("expected 0 affected rows, got %d", n) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet expectations: %v", err) + } + }) + + t.Run("returns error for nil job priority", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + _, err := client.PQDelete(ctx, nil) + if err == nil { + t.Fatal("expected error for nil job priority") + } + }) +} + +func TestPQGetIDs(t *testing.T) { + ctx := context.Background() + + t.Run("returns queued job IDs", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + rows := pgxmock.NewRows([]string{"id"}). + AddRow("batch-1"). + AddRow("batch-2") + + mock.ExpectQuery("SELECT id FROM batch_items"). + WillReturnRows(rows) + + ids, err := client.PQGetIDs(ctx) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if len(ids) != 2 { + t.Fatalf("expected 2 IDs, got %d", len(ids)) + } + if !ids["batch-1"] || !ids["batch-2"] { + t.Errorf("expected batch-1 and batch-2, got %v", ids) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet expectations: %v", err) + } + }) + + t.Run("returns empty map when no jobs queued", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + rows := pgxmock.NewRows([]string{"id"}) + + mock.ExpectQuery("SELECT id FROM batch_items"). + WillReturnRows(rows) + + ids, err := client.PQGetIDs(ctx) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if len(ids) != 0 { + t.Errorf("expected empty map, got %v", ids) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet expectations: %v", err) + } + }) + + t.Run("returns error on query failure", func(t *testing.T) { + client, mock := newTestQueueClient(t) + defer mock.Close() + + mock.ExpectQuery("SELECT id FROM batch_items"). + WillReturnError(fmt.Errorf("connection refused")) + + _, err := client.PQGetIDs(ctx) + if err == nil { + t.Fatal("expected error") + } + }) +} diff --git a/internal/database/postgresql/batch_schema.sql b/internal/database/postgresql/batch_schema.sql index 5c342322a..df7e3de42 100644 --- a/internal/database/postgresql/batch_schema.sql +++ b/internal/database/postgresql/batch_schema.sql @@ -12,14 +12,29 @@ -- limitations under the License. CREATE TABLE IF NOT EXISTS batch_items ( - id TEXT PRIMARY KEY, - tenant_id TEXT NOT NULL, - expiry BIGINT, - tags JSONB, - spec JSONB, - status JSONB + id TEXT PRIMARY KEY, + tenant_id TEXT NOT NULL, + expiry BIGINT, + tags JSONB, + spec JSONB, + status JSONB, + processor_id TEXT, + priority BIGINT, + epoch BIGINT NOT NULL DEFAULT 0 ); CREATE INDEX IF NOT EXISTS idx_batch_items_tenant_id ON batch_items(tenant_id); CREATE INDEX IF NOT EXISTS idx_batch_items_expiry ON batch_items(expiry) WHERE expiry IS NOT NULL; CREATE INDEX IF NOT EXISTS idx_batch_items_tags ON batch_items USING GIN (tags) WHERE tags IS NOT NULL; + +-- Queue index: unclaimed jobs ordered by priority (SLO deadline, earliest first). +CREATE INDEX IF NOT EXISTS idx_batch_items_queue + ON batch_items (priority ASC) + WHERE processor_id IS NULL + AND status IS NOT NULL + AND status::jsonb->>'status' = 'validating'; + +-- Processor ownership index: find jobs owned by a specific processor for crash recovery. +CREATE INDEX IF NOT EXISTS idx_batch_items_processor + ON batch_items (processor_id) + WHERE processor_id IS NOT NULL; diff --git a/internal/database/postgresql/db_core.go b/internal/database/postgresql/db_core.go index ffe083ea8..b8ae083ea 100644 --- a/internal/database/postgresql/db_core.go +++ b/internal/database/postgresql/db_core.go @@ -244,7 +244,7 @@ func (c *pgCore) buildGetQuery( // scanRow scans a single row into BaseIndexes + BaseContents + extra values. // The returned map is keyed by column name for each ExtraColumns() entry. -func (c *pgCore) scanRow(rows pgx.Rows, includeStatic bool) (*api.BaseIndexes, *api.BaseContents, map[string]string, error) { +func (c *pgCore) scanRow(rows pgx.Rows, includeStatic bool) (*api.BaseIndexes, *api.BaseContents, map[string]any, error) { var ( id string tenant string @@ -256,7 +256,7 @@ func (c *pgCore) scanRow(rows pgx.Rows, includeStatic bool) (*api.BaseIndexes, * scanArgs := []any{&id, &tenant, &expiry, &tagsStr} - extraDest := make([]string, len(extraCols)) + extraDest := make([]any, len(extraCols)) for i := range extraDest { scanArgs = append(scanArgs, &extraDest[i]) } @@ -294,7 +294,7 @@ func (c *pgCore) scanRow(rows pgx.Rows, includeStatic bool) (*api.BaseIndexes, * Status: status, } - extras := make(map[string]string, len(extraCols)) + extras := make(map[string]any, len(extraCols)) for i, col := range extraCols { extras[col] = extraDest[i] } @@ -309,7 +309,7 @@ func (c *pgCore) get( ctx context.Context, bq *api.BaseQuery, includeStatic bool, start, limit int, extraFilters map[string]any, rawConditions []string, ) ( - indexes []*api.BaseIndexes, contents []*api.BaseContents, extras []map[string]string, + indexes []*api.BaseIndexes, contents []*api.BaseContents, extras []map[string]any, cursor int, expectMore bool, err error, ) { // Request one extra row beyond the limit to determine if more results exist. @@ -358,8 +358,11 @@ func (c *pgCore) get( // update updates the dynamic fields of an existing item. // When expectedStatus is non-nil, the update is conditional (CAS): it only // succeeds if the current status column matches expectedStatus exactly. -// Returns api.ErrConflict on mismatch. -func (c *pgCore) update(ctx context.Context, idx *api.BaseIndexes, contents *api.BaseContents, expectedStatus []byte) error { +// extraConditions add additional equality checks to the WHERE clause (e.g., +// epoch fencing). rawSetClauses are appended verbatim to the SET clause +// (e.g., "epoch = epoch + 1"). Returns api.ErrConflict when any condition +// prevents the update from matching. +func (c *pgCore) update(ctx context.Context, idx *api.BaseIndexes, contents *api.BaseContents, expectedStatus []byte, extraConditions map[string]any, rawSetClauses []string) error { if err := idx.Validate(); err != nil { return err } @@ -393,11 +396,23 @@ func (c *pgCore) update(ctx context.Context, idx *api.BaseIndexes, contents *api whereClause := fmt.Sprintf(colID+" = $%d", argIdx) argIdx++ + hasConditions := false if len(expectedStatus) > 0 { args = append(args, expectedStatus) whereClause += fmt.Sprintf(" AND "+colStatus+" = $%d", argIdx) + argIdx++ + hasConditions = true + } + + for col, val := range extraConditions { + args = append(args, val) + whereClause += fmt.Sprintf(" AND %s = $%d", col, argIdx) + argIdx++ + hasConditions = true } + setClauses = append(setClauses, rawSetClauses...) + sql := fmt.Sprintf( "UPDATE %s SET %s WHERE %s", c.desc.TableName(), strings.Join(setClauses, ", "), whereClause, @@ -409,7 +424,7 @@ func (c *pgCore) update(ctx context.Context, idx *api.BaseIndexes, contents *api } if result.RowsAffected() == 0 { - if len(expectedStatus) > 0 { + if hasConditions { return fmt.Errorf("DBUpdate: %w", api.ErrConflict) } return fmt.Errorf("DBUpdate: item %s not found", idx.ID) diff --git a/internal/database/postgresql/db_core_test.go b/internal/database/postgresql/db_core_test.go index 86140ad23..ab81b9f69 100644 --- a/internal/database/postgresql/db_core_test.go +++ b/internal/database/postgresql/db_core_test.go @@ -126,10 +126,11 @@ func TestCoreStore_DBFailure(t *testing.T) { contents := &api.BaseContents{Spec: []byte(`{}`), Status: []byte(`{}`)} mock.ExpectExec("INSERT INTO"). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnError(fmt.Errorf("connection refused")) - if err := core.store(context.Background(), idx, contents, nil); err == nil { + extras := map[string]any{colProcessorID: "", colPriority: int64(0), colEpoch: int64(0)} + if err := core.store(context.Background(), idx, contents, extras); err == nil { t.Fatal("expected error on DB failure") } } @@ -142,10 +143,11 @@ func TestCoreStore_NilTags(t *testing.T) { contents := &api.BaseContents{Spec: []byte(`{}`), Status: []byte(`{}`)} mock.ExpectExec("INSERT INTO"). - WithArgs("id-1", "t1", pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). + WithArgs("id-1", "t1", pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnResult(pgxmock.NewResult("INSERT", 1)) - if err := core.store(context.Background(), idx, contents, nil); err != nil { + extras := map[string]any{colProcessorID: "", colPriority: int64(0), colEpoch: int64(0)} + if err := core.store(context.Background(), idx, contents, extras); err != nil { t.Fatalf("expected no error, got %v", err) } if err := mock.ExpectationsWereMet(); err != nil { @@ -194,8 +196,8 @@ func TestCoreGet_Expired(t *testing.T) { defer mock.Close() tags := `{"purpose":"batch"}` - rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus, colSpec}). - AddRow("id-1", "t1", int64(100), &tags, []byte(`{}`), []byte(`{}`)) + rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colProcessorID, colPriority, colEpoch, colStatus, colSpec}). + AddRow("id-1", "t1", int64(100), &tags, "", int64(0), int64(0), []byte(`{}`), []byte(`{}`)) mock.ExpectQuery("SELECT .+ FROM batch_items WHERE"). WithArgs(0, 11). @@ -221,10 +223,10 @@ func TestCoreGet_Pagination(t *testing.T) { tags := `{"k":"v"}` // Return limit+1 rows (3) to indicate more results exist. - rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus}). - AddRow("id-1", "t1", int64(0), &tags, []byte(`{}`)). - AddRow("id-2", "t1", int64(0), &tags, []byte(`{}`)). - AddRow("id-3", "t1", int64(0), &tags, []byte(`{}`)) + rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colProcessorID, colPriority, colEpoch, colStatus}). + AddRow("id-1", "t1", int64(0), &tags, "", int64(0), int64(0), []byte(`{}`)). + AddRow("id-2", "t1", int64(0), &tags, "", int64(0), int64(0), []byte(`{}`)). + AddRow("id-3", "t1", int64(0), &tags, "", int64(0), int64(0), []byte(`{}`)) // get() requests limit+1 rows from the DB. mock.ExpectQuery("SELECT .+ FROM batch_items WHERE"). @@ -256,9 +258,9 @@ func TestCoreGet_Pagination(t *testing.T) { tags := `{"k":"v"}` // Return exactly limit rows (2) — no extra row means no more results. - rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus}). - AddRow("id-1", "t1", int64(0), &tags, []byte(`{}`)). - AddRow("id-2", "t1", int64(0), &tags, []byte(`{}`)) + rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colProcessorID, colPriority, colEpoch, colStatus}). + AddRow("id-1", "t1", int64(0), &tags, "", int64(0), int64(0), []byte(`{}`)). + AddRow("id-2", "t1", int64(0), &tags, "", int64(0), int64(0), []byte(`{}`)) mock.ExpectQuery("SELECT .+ FROM batch_items WHERE"). WithArgs("t1", 0, 3). @@ -293,7 +295,7 @@ func TestCoreUpdate_ValidationError(t *testing.T) { idx := &api.BaseIndexes{ID: "", TenantID: "t1"} contents := &api.BaseContents{Status: []byte(`{}`)} - if err := core.update(context.Background(), idx, contents, nil); err == nil { + if err := core.update(context.Background(), idx, contents, nil, nil, nil); err == nil { t.Fatal("expected validation error for empty ID") } } @@ -309,7 +311,7 @@ func TestCoreUpdate_NonExistentID(t *testing.T) { WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), "missing"). WillReturnResult(pgxmock.NewResult("UPDATE", 0)) - if err := core.update(context.Background(), idx, contents, nil); err == nil { + if err := core.update(context.Background(), idx, contents, nil, nil, nil); err == nil { t.Fatal("expected error for non-existent ID") } if err := mock.ExpectationsWereMet(); err != nil { @@ -324,7 +326,7 @@ func TestCoreUpdate_NoFieldsToUpdate(t *testing.T) { idx := &api.BaseIndexes{ID: "id-1", TenantID: "t1", Tags: nil} contents := &api.BaseContents{} - if err := core.update(context.Background(), idx, contents, nil); err != nil { + if err := core.update(context.Background(), idx, contents, nil, nil, nil); err != nil { t.Fatalf("expected nil for no-op update, got %v", err) } } @@ -340,7 +342,7 @@ func TestCoreUpdate_DBFailure(t *testing.T) { WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), "id-1"). WillReturnError(fmt.Errorf("connection refused")) - if err := core.update(context.Background(), idx, contents, nil); err == nil { + if err := core.update(context.Background(), idx, contents, nil, nil, nil); err == nil { t.Fatal("expected error on DB failure") } if err := mock.ExpectationsWereMet(); err != nil { diff --git a/internal/database/postgresql/file_db.go b/internal/database/postgresql/file_db.go index 206874574..018c19bfa 100644 --- a/internal/database/postgresql/file_db.go +++ b/internal/database/postgresql/file_db.go @@ -106,7 +106,7 @@ func (c *PostgresFileDBClient) DBGet( items[i] = &api.FileItem{ BaseIndexes: *indexes[i], BaseContents: *contents[i], - Purpose: extras[i][colPurpose], + Purpose: extras[i][colPurpose].(string), } } @@ -118,7 +118,7 @@ func (c *PostgresFileDBClient) DBUpdate(ctx context.Context, item *api.FileItem, err = fmt.Errorf("item is nil") return } - if err = c.update(ctx, &item.BaseIndexes, &item.BaseContents, expectedStatus); err != nil { + if err = c.update(ctx, &item.BaseIndexes, &item.BaseContents, expectedStatus, nil, nil); err != nil { return } return diff --git a/internal/database/redis/redis_db.go b/internal/database/redis/redis_db.go deleted file mode 100644 index 6df2de536..000000000 --- a/internal/database/redis/redis_db.go +++ /dev/null @@ -1,613 +0,0 @@ -/* -Copyright 2026 The llm-d Authors - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -// This file provides a redis database client implementation. - -package redis - -import ( - "context" - _ "embed" - "encoding/json" - "fmt" - "strconv" - "time" - - "github.com/go-logr/logr" - db_api "github.com/llm-d/llm-d-batch-gateway/internal/database/api" - goredis "github.com/redis/go-redis/v9" -) - -func (c *BatchDBClientRedis) DBStore(ctx context.Context, item *db_api.BatchItem) (err error) { - - if ctx == nil { - ctx = context.Background() - } - if err = item.Validate(); err != nil { - return - } - return c.dbStore(ctx, &item.BaseIndexes, &item.BaseContents, - itemTypeBatch, nil) -} - -func (c *FileDBClientRedis) DBStore(ctx context.Context, item *db_api.FileItem) (err error) { - - if ctx == nil { - ctx = context.Background() - } - if err = item.Validate(); err != nil { - return - } - return c.dbStore(ctx, &item.BaseIndexes, &item.BaseContents, - itemTypeFile, []any{item.Purpose}) -} - -func (c *DSClientRedis) dbStore(ctx context.Context, - indexes *db_api.BaseIndexes, contents *db_api.BaseContents, - itemType string, extraFields []any) (err error) { - - if ctx == nil { - ctx = context.Background() - } - logger := logr.FromContextOrDiscard(ctx).WithValues("ID", indexes.ID) - - ptags, err := packTags(indexes.Tags) - if err != nil { - return err - } - args := []any{itemType, versionV1, indexes.ID, indexes.TenantID, - indexes.Expiry, ptags, contents.Status, contents.Spec, ttlSecDefault} - if len(extraFields) > 0 { - args = append(args, extraFields...) - } - - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - res, err := redisScriptStore.Run(cctx, c.redisClient, - []string{getKeyForStore(indexes.ID, itemType)}, args...).Text() - if err != nil { - return err - } - if len(res) > 0 { - err = fmt.Errorf("%s", res) - return - } - - logger.Info("DBStore: succeeded") - return nil -} - -func getUpdateFields(status []byte, tags db_api.Tags) ( - fields []any, updateStatus, updateTags bool, err error) { - - fields = make([]any, 0, 2) - if len(status) > 0 { - fields = append(fields, fieldNameStatus, status) - updateStatus = true - } - if len(tags) > 0 { - var ptags string - ptags, err = packTags(tags) - if err != nil { - return - } - fields = append(fields, fieldNameTags, ptags) - updateTags = true - } - return -} - -func (c *BatchDBClientRedis) DBUpdate(ctx context.Context, item *db_api.BatchItem, expectedStatus []byte) (err error) { - - if ctx == nil { - ctx = context.Background() - } - if err = item.Validate(); err != nil { - return - } - return c.dbUpdate(ctx, &item.BaseIndexes, &item.BaseContents, itemTypeBatch, "DBUpdate[Batch]", expectedStatus) -} - -func (c *FileDBClientRedis) DBUpdate(ctx context.Context, item *db_api.FileItem, expectedStatus []byte) (err error) { - - if ctx == nil { - ctx = context.Background() - } - if err = item.Validate(); err != nil { - return - } - return c.dbUpdate(ctx, &item.BaseIndexes, &item.BaseContents, itemTypeFile, "DBUpdate[File]", expectedStatus) -} - -func (c *DSClientRedis) dbUpdate(ctx context.Context, - indexes *db_api.BaseIndexes, contents *db_api.BaseContents, - itemType, logPref string, expectedStatus []byte) (err error) { - - if ctx == nil { - ctx = context.Background() - } - logger := logr.FromContextOrDiscard(ctx).WithValues("ID", indexes.ID) - - fields, updatedStatus, updatedTags, err := getUpdateFields(contents.Status, indexes.Tags) - if err != nil { - return err - } - if len(fields) == 0 { - logger.Info(logPref + ": nothing to update") - return nil - } - - key := getKeyForStore(indexes.ID, itemType) - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - - if len(expectedStatus) > 0 { - args := append([]any{expectedStatus}, fields...) - result, err := redisScriptUpdateCAS.Run(cctx, c.redisClient, - []string{key}, args...).Text() - if err != nil { - return err - } - switch result { - case "OK": - // success - case "CONFLICT", "NOT_FOUND": - return fmt.Errorf("DBUpdate: %w", db_api.ErrConflict) - default: - return fmt.Errorf("DBUpdate: unexpected CAS script result %q", result) - } - } else { - err = c.redisClient.HSet(cctx, key, fields...).Err() - if err != nil { - return - } - } - - logger.Info(logPref+": succeeded", "updatedStatus", updatedStatus, "updatedTags", updatedTags) - return nil -} - -func (c *BatchDBClientRedis) DBDelete(ctx context.Context, IDs []string) ( - deletedIDs []string, err error) { - return c.dBDelete(ctx, IDs, itemTypeBatch, "DBDelete[Batch]") -} - -func (c *FileDBClientRedis) DBDelete(ctx context.Context, IDs []string) ( - deletedIDs []string, err error) { - return c.dBDelete(ctx, IDs, itemTypeFile, "DBDelete[File]") -} - -func (c *DSClientRedis) dBDelete(ctx context.Context, IDs []string, itemType, logPref string) ( - deletedIDs []string, err error) { - - if ctx == nil { - ctx = context.Background() - } - logger := logr.FromContextOrDiscard(ctx) - - // Delete the items. - resMap := make(map[string]*goredis.IntCmd) - var cmds []goredis.Cmder - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - cmds, err = c.redisClient.Pipelined(cctx, func(pipe goredis.Pipeliner) error { - for _, id := range IDs { - res := pipe.Del(cctx, getKeyForStore(id, itemType)) - resMap[id] = res - } - return nil - }) - if err != nil { - return - } - for _, cmd := range cmds { - if cmd.Err() != nil && cmd.Err() != goredis.Nil { - err = cmd.Err() - break - } - } - deletedIDs = make([]string, 0, len(resMap)) - for id, res := range resMap { - if res != nil && res.Err() == nil && res.Val() > 0 { - deletedIDs = append(deletedIDs, id) - } - } - - logger.Info(logPref+": succeeded", "nItems", len(deletedIDs), "IDs", deletedIDs) - return -} - -func (c *DSClientRedis) dbGet( - ctx context.Context, itemType, logPref string, start, limit int, includeStatic bool, - IDs []string, tagSelectors db_api.Tags, tagsLogicalCond db_api.LogicalCond, - expired bool, tenantID, purpose string, nonTerminal bool) (res []any, err error) { - - if ctx == nil { - ctx = context.Background() - } - includeSpec := strconv.FormatBool(includeStatic) - - if len(IDs) > 0 { - - keys := make([]string, 0, len(IDs)) - for _, id := range IDs { - keys = append(keys, getKeyForStore(id, itemType)) - } - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - res, err = redisScriptGetByIDs.Run(cctx, c.redisClient, - keys, tenantID, includeSpec).Slice() - if err != nil { - return - } - - } else if len(tagSelectors) > 0 { - - cond, found := db_api.LogicalCondNames[tagsLogicalCond] - if !found { - err = fmt.Errorf("invalid logical condition value: %d", tagsLogicalCond) - return - } - ctags := convertTags(tagSelectors) - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - res, err = redisScriptGetByTags.Run(cctx, c.redisClient, - ctags, cond, getKeyPatternForStore(itemType), start, limit, tenantID, includeSpec).Slice() - if err != nil { - return - } - - } else if expired { - - curTimestamp := time.Now().Unix() - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - res, err = redisScriptGetByExpiry.Run(cctx, c.redisClient, - []string{}, curTimestamp, getKeyPatternForStore(itemType), - start, limit, tenantID, includeSpec).Slice() - if err != nil { - return - } - - } else if len(purpose) > 0 { - - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - res, err = redisScriptGetByPurpose.Run(cctx, c.redisClient, - []string{}, purpose, getKeyPatternForStore(itemType), - start, limit, tenantID, includeSpec).Slice() - if err != nil { - return - } - - } else if nonTerminal { - - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - res, err = redisScriptGetByNonTerminal.Run(cctx, c.redisClient, - []string{}, getKeyPatternForStore(itemType), - start, limit, tenantID, includeSpec).Slice() - if err != nil { - return - } - - } else if len(tenantID) > 0 { - - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - res, err = redisScriptGetByTenant.Run(cctx, c.redisClient, - []string{}, tenantID, getKeyPatternForStore(itemType), - start, limit, includeSpec).Slice() - if err != nil { - return - } - - } - - return -} - -func (c *BatchDBClientRedis) DBGet( - ctx context.Context, query *db_api.BatchQuery, - includeStatic bool, start, limit int) ( - items []*db_api.BatchItem, cursor int, expectMore bool, err error) { - - if ctx == nil { - ctx = context.Background() - } - logger := logr.FromContextOrDiscard(ctx) - if query == nil { - logger.Info("DBGet[Batch]: empty query") - return - } - - var res []any - res, err = c.dbGet(ctx, itemTypeBatch, "DBGet[Batch]", start, limit, includeStatic, - query.IDs, query.TagSelectors, query.TagsLogicalCond, query.Expired, query.TenantID, "", query.NonTerminal) - if err != nil { - return - } - if res != nil { - cursor, expectMore, items, err = processGetScriptResultBatch(res) - if err != nil { - return - } - } - - logger.Info("DBGet[Batch]: succeeded", "nItems", len(items)) - return -} - -func (c *FileDBClientRedis) DBGet( - ctx context.Context, query *db_api.FileQuery, - includeStatic bool, start, limit int) ( - items []*db_api.FileItem, cursor int, expectMore bool, err error) { - - if ctx == nil { - ctx = context.Background() - } - logger := logr.FromContextOrDiscard(ctx) - if query == nil { - logger.Info("DBGet[File]: empty query") - return - } - - var res []any - res, err = c.dbGet(ctx, itemTypeFile, "DBGet[File]", start, limit, includeStatic, - query.IDs, query.TagSelectors, query.TagsLogicalCond, query.Expired, query.TenantID, query.Purpose, false) - if err != nil { - return - } - if res != nil { - cursor, expectMore, items, err = processGetScriptResultFile(res) - if err != nil { - return - } - } - - logger.Info("DBGet[File]: succeeded", "nItems", len(items)) - return -} - -func processGetScriptResultBatch(res []any) ( - cursor int, expectMore bool, items []*db_api.BatchItem, err error) { - - if len(res) != 2 { - err = fmt.Errorf("unexpected result from script") - return - } - resItems, ok := res[1].([]any) - if !ok { - err = fmt.Errorf("unexpected result type from script: %T", res[1]) - return - } - resCursor, ok := res[0].(int64) - if !ok { - err = fmt.Errorf("unexpected result type from script: %T", res[0]) - return - } - items = make([]*db_api.BatchItem, 0, len(resItems)) - for i, resItem := range resItems { - fields, ok := resItem.([]any) - if !ok { - return 0, false, nil, fmt.Errorf("unexpected item type at index %d: %T", i, resItem) - } - item, err := batchItemFromHget(fields) - if err != nil { - return 0, false, nil, err - } - if item != nil { - items = append(items, item) - } - } - cursor = int(resCursor) - expectMore = (cursor != 0) - - return -} - -func processGetScriptResultFile(res []any) ( - cursor int, expectMore bool, items []*db_api.FileItem, err error) { - - if len(res) != 2 { - err = fmt.Errorf("unexpected result from script") - return - } - resItems, ok := res[1].([]any) - if !ok { - err = fmt.Errorf("unexpected result type from script: %T", res[1]) - return - } - resCursor, ok := res[0].(int64) - if !ok { - err = fmt.Errorf("unexpected result type from script: %T", res[0]) - return - } - items = make([]*db_api.FileItem, 0, len(resItems)) - for i, resItem := range resItems { - fields, ok := resItem.([]any) - if !ok { - return 0, false, nil, fmt.Errorf("unexpected item type at index %d: %T", i, resItem) - } - item, err := fileItemFromHget(fields) - if err != nil { - return 0, false, nil, err - } - if item != nil { - items = append(items, item) - } - } - cursor = int(resCursor) - expectMore = (cursor != 0) - - return -} - -func getKeyForStore(key, itemType string) string { - return storeKeysPrefix + itemType + ":" + key -} - -func getKeyPatternForStore(itemType string) string { - return storeKeysPrefix + itemType + ":*" -} - -func packTags(tags map[string]string) (string, error) { - if len(tags) == 0 { - return "", nil - } - json, err := json.Marshal(tags) - if err != nil { - return "", err - } - return string(json), nil -} - -func unpackTags(tagsPacked string) (map[string]string, error) { - if len(tagsPacked) == 0 { - return nil, nil - } - var tags map[string]string - err := json.Unmarshal([]byte(tagsPacked), &tags) - if err != nil { - return nil, err - } - return tags, nil -} - -func convertTags(tags map[string]string) (ctags []string) { - if len(tags) > 0 { - ctags = make([]string, 0, len(tags)) - for key, val := range tags { - ctags = append(ctags, fmt.Sprintf("\"%s\":\"%s\"", key, val)) - } - } - return -} - -func batchItemFromHget(vals []any) (item *db_api.BatchItem, err error) { - - ID, tenantID, expiry, tags, _, status, spec, err := itemFromHget(vals) - if err != nil { - return nil, err - } - - item = &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: ID, - TenantID: tenantID, - Expiry: expiry, - Tags: tags, - }, - BaseContents: db_api.BaseContents{ - Spec: spec, - Status: status, - }, - } - - return -} - -func fileItemFromHget(vals []any) (item *db_api.FileItem, err error) { - - ID, tenantID, expiry, tags, purpose, status, spec, err := itemFromHget(vals) - if err != nil { - return nil, err - } - - item = &db_api.FileItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: ID, - TenantID: tenantID, - Expiry: expiry, - Tags: tags, - }, - Purpose: purpose, - BaseContents: db_api.BaseContents{ - Spec: spec, - Status: status, - }, - } - - return -} - -func itemFromHget(vals []any) ( - ID, tenantID string, expiry int64, tags db_api.Tags, - purpose string, status, spec []byte, err error) { - - if len(vals)%2 != 0 { - err = fmt.Errorf("unexpected result contents from HGETALL (odd length): %v", vals) - return - } - - // HGETALL returns a flat array: [field1, value1, field2, value2, ...]. - // Build a map from the flat array. - hash := make(map[string]string) - for i := 0; i < len(vals); i += 2 { - fieldName, ok := vals[i].(string) - if !ok { - err = fmt.Errorf("invalid field name at index %d: %v", i, vals[i]) - return - } - fieldValue, ok := vals[i+1].(string) - if !ok { - err = fmt.Errorf("invalid field value at index %d: %v", i+1, vals[i+1]) - return - } - hash[fieldName] = fieldValue - } - - // Extract ID (required). - ID = hash["ID"] - if len(ID) == 0 { - err = fmt.Errorf("missing or invalid id field") - return - } - - // Extract tenantID. - tenantID = hash["tenantID"] - - // Extract expiry. - if expiryStr := hash["expiry"]; len(expiryStr) > 0 { - expiry, err = strconv.ParseInt(expiryStr, 10, 64) - if err != nil { - err = fmt.Errorf("invalid expiry field %q: %w", expiryStr, err) - return - } - } - - // Extract tags. - tagsStr := hash["tags"] - tags, err = unpackTags(tagsStr) - if err != nil { - return - } - - // Extract purpose. - purpose = hash["purpose"] - - // Extract status. - if statusStr := hash["status"]; len(statusStr) > 0 { - status = []byte(statusStr) - } - - // Extract spec. - if specStr := hash["spec"]; len(specStr) > 0 { - spec = []byte(specStr) - } - - return -} diff --git a/internal/database/redis/redis_ds_client.go b/internal/database/redis/redis_ds_client.go index 24f01f7a2..209228866 100644 --- a/internal/database/redis/redis_ds_client.go +++ b/internal/database/redis/redis_ds_client.go @@ -48,7 +48,6 @@ const ( eventKeysPrefix = keysPrefix + "event:" statusKeysPrefix = keysPrefix + "status:" priorityQueueKeyName = queueKeysPrefix + "priority" - inFlightKeyName = keysPrefix + "inflight" eventChanSize = 100 eventReadCount = 4 eventReadTimeout = 20 * time.Second @@ -61,51 +60,13 @@ const ( ) var ( - //go:embed redis_common.lua - commonLua string - - //go:embed redis_store.lua - storeLua string - redisScriptStore = goredis.NewScript(storeLua) - - //go:embed redis_get_by_ids.lua - getByIDsLua string - redisScriptGetByIDs = goredis.NewScript(commonLua + "\n" + getByIDsLua) - - //go:embed redis_get_by_tags.lua - getByTagsLua string - redisScriptGetByTags = goredis.NewScript(commonLua + "\n" + getByTagsLua) - - //go:embed redis_get_by_expiry.lua - getByExpiryLua string - redisScriptGetByExpiry = goredis.NewScript(commonLua + "\n" + getByExpiryLua) - - //go:embed redis_get_by_purpose.lua - getByPurposeLua string - redisScriptGetByPurpose = goredis.NewScript(commonLua + "\n" + getByPurposeLua) - - //go:embed redis_get_by_tenant.lua - getByTenantLua string - redisScriptGetByTenant = goredis.NewScript(commonLua + "\n" + getByTenantLua) - - //go:embed redis_get_by_non_terminal.lua - getByNonTerminalLua string - redisScriptGetByNonTerminal = goredis.NewScript(commonLua + "\n" + getByNonTerminalLua) - - //go:embed redis_update_cas.lua - updateCASLua string - redisScriptUpdateCAS = goredis.NewScript(updateCASLua) - //go:embed redis_pq_get_ids.lua pqGetIDsLua string redisScriptPQGetIDs = goredis.NewScript(pqGetIDsLua) - _ db_api.BatchDBClient = (*BatchDBClientRedis)(nil) - _ db_api.FileDBClient = (*FileDBClientRedis)(nil) _ db_api.BatchPriorityQueueClient = (*ExchangeDBClientRedis)(nil) _ db_api.BatchEventChannelClient = (*ExchangeDBClientRedis)(nil) _ db_api.BatchStatusClient = (*ExchangeDBClientRedis)(nil) - _ db_api.InFlightClient = (*ExchangeDBClientRedis)(nil) ) type DSClientRedis struct { @@ -117,50 +78,10 @@ type DSClientRedis struct { onceClose *sync.Once } -type BatchDBClientRedis struct { - *DSClientRedis -} - -type FileDBClientRedis struct { - *DSClientRedis -} - type ExchangeDBClientRedis struct { *DSClientRedis } -// NewBatchDBClientRedis returns a new redis based batch db client. -// Provide either an already created baseRedisClient (that can be shared between multiple higher level redis based clients), -// or a conf and opTimeout for creating a new base redis client dedicated to this higher level client. -func NewBatchDBClientRedis(ctx context.Context, baseRedisClient *DSClientRedis, conf *uredis.RedisClientConfig, opTimeout time.Duration) ( - redisClient *BatchDBClientRedis, err error) { - - if baseRedisClient == nil { - baseRedisClient, err = NewDSClientRedis(ctx, conf, opTimeout) - if err != nil { - return nil, err - } - } - redisClient = &BatchDBClientRedis{DSClientRedis: baseRedisClient} - return -} - -// NewFilesDBClientRedis returns a new redis based file db client. -// Provide either an already created baseRedisClient (that can be shared between multiple higher level redis based clients), -// or a conf and opTimeout for creating a new base redis client dedicated to this higher level client. -func NewFileDBClientRedis(ctx context.Context, baseRedisClient *DSClientRedis, conf *uredis.RedisClientConfig, opTimeout time.Duration) ( - redisClient *FileDBClientRedis, err error) { - - if baseRedisClient == nil { - baseRedisClient, err = NewDSClientRedis(ctx, conf, opTimeout) - if err != nil { - return nil, err - } - } - redisClient = &FileDBClientRedis{DSClientRedis: baseRedisClient} - return -} - // NewExchangeDBClientRedis returns a new redis based exchange db client. // Provide either an already created baseRedisClient (that can be shared between multiple higher level redis based clients), // or a conf and opTimeout for creating a new base redis client dedicated to this higher level client. diff --git a/internal/database/redis/redis_ds_client_test.go b/internal/database/redis/redis_ds_client_test.go index c97eba259..346f2c897 100644 --- a/internal/database/redis/redis_ds_client_test.go +++ b/internal/database/redis/redis_ds_client_test.go @@ -21,11 +21,8 @@ package redis_test import ( "bytes" "context" - "errors" "fmt" - "maps" "os" - "sync" "testing" "time" @@ -34,14 +31,13 @@ import ( "github.com/alicebob/miniredis/v2" db_api "github.com/llm-d/llm-d-batch-gateway/internal/database/api" dbredis "github.com/llm-d/llm-d-batch-gateway/internal/database/redis" - ucom "github.com/llm-d/llm-d-batch-gateway/internal/util/com" uredis "github.com/llm-d/llm-d-batch-gateway/internal/util/redis" utls "github.com/llm-d/llm-d-batch-gateway/internal/util/tls" goredis "github.com/redis/go-redis/v9" ) func setupRedisDSClients(t *testing.T, redisUrl, redisCaCert string) ( - *dbredis.DSClientRedis, *dbredis.BatchDBClientRedis, *dbredis.FileDBClientRedis, *dbredis.ExchangeDBClientRedis) { + *dbredis.DSClientRedis, *dbredis.ExchangeDBClientRedis) { t.Helper() cfg := &uredis.RedisClientConfig{ Url: redisUrl, @@ -58,34 +54,18 @@ func setupRedisDSClients(t *testing.T, redisUrl, redisCaCert string) ( if err != nil { t.Fatalf("Failed to create base redis client: %v", err) } - batchClient, err := dbredis.NewBatchDBClientRedis(ctx, baseClient, nil, 0) - if err != nil { - t.Fatalf("Failed to create batch redis client: %v", err) - } - fileClient, err := dbredis.NewFileDBClientRedis(ctx, baseClient, nil, 0) - if err != nil { - t.Fatalf("Failed to create file redis client: %v", err) - } exchClient, err := dbredis.NewExchangeDBClientRedis(ctx, baseClient, nil, 0) if err != nil { t.Fatalf("Failed to create exchange redis client: %v", err) } - return baseClient, batchClient, fileClient, exchClient + return baseClient, exchClient } func TestRedisDSClient(t *testing.T) { redisUrl := os.Getenv("TEST_REDIS_URL") redisCaCert := os.Getenv("TEST_REDIS_CACERT_PATH") - var ( - minirds *miniredis.Miniredis - tagKey1 = "key-tag-1" - tagKey2 = "key-tag-2" - tagKey3 = "key-tag-3" - tagVal1 = "val-tag-1" - tagVal2 = "val-tag-2" - tagVal3 = "val-tag-3" - ) + var minirds *miniredis.Miniredis // Start miniredis if no external redis URL is provided. if redisUrl == "" { @@ -101,579 +81,21 @@ func TestRedisDSClient(t *testing.T) { t.Run("Create clients", func(t *testing.T) { t.Parallel() - baseClient, batchClient, fileClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) - t.Logf("Memory address of the clients: base=%p batch=%p file=%p exchange=%p", - baseClient, batchClient, fileClient, exchClient) - if baseClient == nil || batchClient == nil || fileClient == nil || exchClient == nil { + if baseClient == nil || exchClient == nil { t.Fatalf("Expected redis clients to be non-nil") } }) - t.Run("Batch db operations", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - // Store. - nBatches := 20 - nBatchesRmv := 10 - var wg sync.WaitGroup - batches, batchesRmv := make(map[string]*db_api.BatchItem), make(map[string]*db_api.BatchItem) - var batchesIDs, batchesAllIDs []string - for i := 0; i < nBatchesRmv; i++ { - batchID := uuid.New().String() - batch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: batchID, - TenantID: "Tnt2", - Expiry: time.Now().Add(time.Second).Unix(), - Tags: map[string]string{tagKey1: tagVal1, tagKey2: tagVal2}, - }, - BaseContents: db_api.BaseContents{ - Spec: []byte("spec"), - Status: []byte("status"), - }, - } - batchesRmv[batchID] = batch - batchesAllIDs = append(batchesAllIDs, batchID) - wg.Add(1) - go func() { - defer wg.Done() - err := batchClient.DBStore(context.Background(), batch) - if err != nil { - t.Errorf("Failed to store item: %v", err) - } - }() - } - wg.Wait() - for i := 0; i < nBatches; i++ { - batchID := uuid.New().String() - batch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: batchID, - TenantID: "Tnt1", - Expiry: time.Now().Add(time.Hour).Unix(), - Tags: map[string]string{tagKey1: tagVal1, tagKey3: tagVal3}, - }, - BaseContents: db_api.BaseContents{ - Spec: []byte("spec"), - Status: []byte("status"), - }, - } - batches[batchID] = batch - batchesIDs = append(batchesIDs, batchID) - batchesAllIDs = append(batchesAllIDs, batchID) - wg.Add(1) - go func() { - defer wg.Done() - err := batchClient.DBStore(context.Background(), batch) - if err != nil { - t.Errorf("Failed to store item: %v", err) - } - }() - } - wg.Wait() - time.Sleep(3 * time.Second) // To pass the expiry time of the short expiry items. - - // Get expired (filter by tenant to avoid interference from parallel tests). - expectMore := true - nRet, cursor := 0, 0 - for expectMore { - resItems, cur, expectM, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - Expired: true, - TenantID: "Tnt2", - }, - }, true, cursor, nBatchesRmv*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - sameMembersBatch(t, resItems, batchesRmv) - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nBatchesRmv { - t.Fatalf("Invalid number of items %d != %d", nRet, nBatchesRmv) - } - - // Get by IDs. - expectMore = true - nRet, cursor = 0, 0 - for expectMore { - resItems, cur, expectM, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: batchesIDs, - }, - }, true, cursor, nBatches*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - sameMembersBatch(t, resItems, batches) - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nBatches { - t.Fatalf("Invalid number of items %d != %d", nRet, nBatches) - } - - // Get by tenant. - expectMore = true - nRet, cursor = 0, 0 - for expectMore { - resItems, cur, expectM, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - TenantID: "Tnt2", - }, - }, true, cursor, nBatchesRmv*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - sameMembersBatch(t, resItems, batchesRmv) - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nBatchesRmv { - t.Fatalf("Invalid number of items %d != %d", nRet, nBatchesRmv) - } - - // Get by tags. - expectMore = true - nRet, cursor = 0, 0 - for expectMore { - resItems, cur, expectM, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - TagSelectors: db_api.Tags{tagKey1: tagVal1, tagKey3: tagVal3}, - TagsLogicalCond: db_api.LogicalCondAnd, - }, - }, true, cursor, nBatches*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - sameMembersBatch(t, resItems, batches) - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nBatches { - t.Fatalf("Invalid number of items %d != %d", nRet, nBatches) - } - expectMore = true - nRet, cursor = 0, 0 - for expectMore { - resItems, cur, expectM, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - TagSelectors: db_api.Tags{tagKey1: tagVal1, tagKey3: tagVal3}, - TagsLogicalCond: db_api.LogicalCondOr, - }, - }, true, cursor, nBatches*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nBatches+nBatchesRmv { - t.Fatalf("Invalid number of items %d != %d", nRet, nBatches+nBatchesRmv) - } - - // Get by tags and tenant. - expectMore = true - nRet, cursor = 0, 0 - for expectMore { - resItems, cur, expectM, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - TenantID: "Tnt1", - TagSelectors: db_api.Tags{tagKey1: tagVal1, tagKey3: tagVal3}, - TagsLogicalCond: db_api.LogicalCondOr, - }, - }, true, cursor, nBatches*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - sameMembersBatch(t, resItems, batches) - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nBatches { - t.Fatalf("Invalid number of items %d != %d", nRet, nBatches) - } - - // Update. - updId := batchesIDs[0] - updBatch := batches[updId] - updBatch.Status = []byte("statusUpdated") - err := batchClient.DBUpdate(context.Background(), updBatch, nil) - if err != nil { - t.Fatalf("Failed to update item: %v", err) - } - resItems, _, expectM, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: []string{updId}, - }, - }, true, 0, 1) - if err != nil { - t.Fatalf("Failed to get item: %v", err) - } - if expectM { - t.Fatalf("Invalid expect more") - } - if len(resItems) != 1 { - t.Fatalf("Invalid number of returned items") - } - isEqualBatchItem(t, updBatch, resItems[0]) - - // Delete. - deletedIDs, err := batchClient.DBDelete(context.Background(), batchesAllIDs) - if err != nil { - t.Fatalf("Failed to delete items: %v", err) - } - if deletedIDs == nil || len(deletedIDs) != len(batchesAllIDs) { - t.Fatalf("Failed to delete items: %d", len(deletedIDs)) - } - if !ucom.SameMembersInStrSlice(deletedIDs, batchesAllIDs) { - t.Fatalf("Deletion IDs mismatch: %v != %v", deletedIDs, batchesAllIDs) - } - - }) - - t.Run("Batch NonTerminal query", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - ctx := context.Background() - - // Store items with different statuses. - statuses := map[string]string{ - "nt-validating": `{"status":"validating"}`, - "nt-inprogress": `{"status":"in_progress"}`, - "nt-completed": `{"status":"completed"}`, - "nt-failed": `{"status":"failed"}`, - "nt-expired": `{"status":"expired"}`, - "nt-cancelled": `{"status":"cancelled"}`, - } - var allIDs []string - for id, status := range statuses { - item := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: id, - TenantID: "tenant-nt", - Expiry: time.Now().Add(time.Hour).Unix(), - Tags: map[string]string{"test": "nonterminal"}, - }, - BaseContents: db_api.BaseContents{ - Spec: []byte(`{"endpoint":"/v1/chat/completions"}`), - Status: []byte(status), - }, - } - if err := batchClient.DBStore(ctx, item); err != nil { - t.Fatalf("failed to store item %s: %v", id, err) - } - allIDs = append(allIDs, id) - } - - // Query non-terminal items filtered by tenant (required for test isolation - // since parallel subtests share the same miniredis instance). - items, _, _, err := batchClient.DBGet(ctx, - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{TenantID: "tenant-nt"}, - NonTerminal: true, - }, false, 0, 100) - if err != nil { - t.Fatalf("NonTerminal query failed: %v", err) - } - - // Should return only the 2 non-terminal items (validating, in_progress). - if len(items) != 2 { - var gotIDs []string - for _, item := range items { - gotIDs = append(gotIDs, item.ID) - } - t.Fatalf("expected 2 non-terminal items, got %d: %v", len(items), gotIDs) - } - - gotIDs := make(map[string]bool) - for _, item := range items { - gotIDs[item.ID] = true - } - if !gotIDs["nt-validating"] || !gotIDs["nt-inprogress"] { - t.Errorf("expected validating and in_progress, got %v", gotIDs) - } - - // Cleanup. - _, err = batchClient.DBDelete(ctx, allIDs) - if err != nil { - t.Fatalf("failed to delete items: %v", err) - } - }) - - t.Run("File db operations", func(t *testing.T) { - t.Parallel() - baseClient, _, fileClient, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - // Store. - nFiles := 20 - nFilesRmv := 10 - var wg sync.WaitGroup - files, filesRmv := make(map[string]*db_api.FileItem), make(map[string]*db_api.FileItem) - var filesIDs, filesAllIDs []string - for i := 0; i < nFilesRmv; i++ { - fileID := uuid.New().String() - file := &db_api.FileItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: fileID, - TenantID: "Tnt2", - Expiry: time.Now().Add(time.Second).Unix(), - Tags: map[string]string{tagKey1: tagVal1, tagKey2: tagVal2}, - }, - Purpose: "file", - BaseContents: db_api.BaseContents{ - Spec: []byte("spec"), - Status: []byte("status"), - }, - } - filesRmv[fileID] = file - filesAllIDs = append(filesAllIDs, fileID) - wg.Add(1) - go func() { - defer wg.Done() - err := fileClient.DBStore(context.Background(), file) - if err != nil { - t.Errorf("Failed to store item: %v", err) - } - }() - } - wg.Wait() - for i := 0; i < nFiles; i++ { - fileID := uuid.New().String() - file := &db_api.FileItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: fileID, - TenantID: "Tnt1", - Expiry: time.Now().Add(time.Hour).Unix(), - Tags: map[string]string{tagKey1: tagVal1, tagKey3: tagVal3}, - }, - BaseContents: db_api.BaseContents{ - Spec: []byte("spec"), - Status: []byte("status"), - }, - } - files[fileID] = file - filesIDs = append(filesIDs, fileID) - filesAllIDs = append(filesAllIDs, fileID) - wg.Add(1) - go func() { - defer wg.Done() - err := fileClient.DBStore(context.Background(), file) - if err != nil { - t.Errorf("Failed to store item: %v", err) - } - }() - } - wg.Wait() - time.Sleep(3 * time.Second) // To pass the expiry time of the short expiry items. - - // Get expired (filter by tenant to avoid interference from parallel tests). - expectMore := true - nRet, cursor := 0, 0 - for expectMore { - resItems, cur, expectM, err := fileClient.DBGet(context.Background(), - &db_api.FileQuery{ - BaseQuery: db_api.BaseQuery{ - Expired: true, - TenantID: "Tnt2", - }, - }, true, cursor, nFilesRmv*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - sameMembersFile(t, resItems, filesRmv) - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nFilesRmv { - t.Fatalf("Invalid number of items %d != %d", nRet, nFilesRmv) - } - - // Get by IDs. - expectMore = true - nRet, cursor = 0, 0 - for expectMore { - resItems, cur, expectM, err := fileClient.DBGet(context.Background(), - &db_api.FileQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: filesIDs, - }, - }, true, cursor, nFiles*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - sameMembersFile(t, resItems, files) - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nFiles { - t.Fatalf("Invalid number of items %d != %d", nRet, nFiles) - } - - // Get by tenant. - expectMore = true - nRet, cursor = 0, 0 - for expectMore { - resItems, cur, expectM, err := fileClient.DBGet(context.Background(), - &db_api.FileQuery{ - BaseQuery: db_api.BaseQuery{ - TenantID: "Tnt2", - }, - }, true, cursor, nFilesRmv*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - sameMembersFile(t, resItems, filesRmv) - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nFilesRmv { - t.Fatalf("Invalid number of items %d != %d", nRet, nFilesRmv) - } - - // Get by tags. - expectMore = true - nRet, cursor = 0, 0 - for expectMore { - resItems, cur, expectM, err := fileClient.DBGet(context.Background(), - &db_api.FileQuery{ - BaseQuery: db_api.BaseQuery{ - TagSelectors: db_api.Tags{tagKey1: tagVal1, tagKey3: tagVal3}, - TagsLogicalCond: db_api.LogicalCondAnd, - }, - }, true, cursor, nFiles*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - sameMembersFile(t, resItems, files) - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nFiles { - t.Fatalf("Invalid number of items %d != %d", nRet, nFiles) - } - expectMore = true - nRet, cursor = 0, 0 - for expectMore { - resItems, cur, expectM, err := fileClient.DBGet(context.Background(), - &db_api.FileQuery{ - BaseQuery: db_api.BaseQuery{ - TagSelectors: db_api.Tags{tagKey1: tagVal1, tagKey3: tagVal3}, - TagsLogicalCond: db_api.LogicalCondOr, - }, - }, true, cursor, nFiles*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nFiles+nFilesRmv { - t.Fatalf("Invalid number of items %d != %d", nRet, nFiles+nFilesRmv) - } - - // Get by tags and tenant. - expectMore = true - nRet, cursor = 0, 0 - for expectMore { - resItems, cur, expectM, err := fileClient.DBGet(context.Background(), - &db_api.FileQuery{ - BaseQuery: db_api.BaseQuery{ - TenantID: "Tnt1", - TagSelectors: db_api.Tags{tagKey1: tagVal1, tagKey3: tagVal3}, - TagsLogicalCond: db_api.LogicalCondOr, - }, - }, true, cursor, nFiles*2) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - sameMembersFile(t, resItems, files) - nRet += len(resItems) - expectMore = expectM - cursor = cur - } - if nRet != nFiles { - t.Fatalf("Invalid number of items %d != %d", nRet, nFiles) - } - - // Update. - updId := filesIDs[0] - updFile := files[updId] - updFile.Status = []byte("statusUpdated") - err := fileClient.DBUpdate(context.Background(), updFile, nil) - if err != nil { - t.Fatalf("Failed to update item: %v", err) - } - resItems, _, expectM, err := fileClient.DBGet(context.Background(), - &db_api.FileQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: []string{updId}, - }, - }, true, 0, 1) - if err != nil { - t.Fatalf("Failed to get item: %v", err) - } - if expectM { - t.Fatalf("Invalid expect more") - } - if len(resItems) != 1 { - t.Fatalf("Invalid number of returned items") - } - isEqualFileItem(t, updFile, resItems[0]) - - // Delete. - deletedIDs, err := fileClient.DBDelete(context.Background(), filesAllIDs) - if err != nil { - t.Fatalf("Failed to delete items: %v", err) - } - if deletedIDs == nil || len(deletedIDs) != len(filesAllIDs) { - t.Fatalf("Failed to delete items: %d", len(deletedIDs)) - } - if !ucom.SameMembersInStrSlice(deletedIDs, filesAllIDs) { - t.Fatalf("Deletion IDs mismatch: %v != %v", deletedIDs, filesAllIDs) - } - - }) - t.Run("Event exchange operations", func(t *testing.T) { t.Parallel() if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -733,7 +155,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -833,7 +255,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -875,7 +297,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -917,7 +339,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -982,7 +404,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -1003,7 +425,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -1052,7 +474,7 @@ func TestRedisDSClient(t *testing.T) { t.Run("Status exchange operations", func(t *testing.T) { t.Parallel() - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -1111,7 +533,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -1176,7 +598,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -1258,7 +680,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -1361,7 +783,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -1416,7 +838,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -1452,1044 +874,11 @@ func TestRedisDSClient(t *testing.T) { } }) - t.Run("includeStatic parameter - Batch", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - // Store batch with spec. - batchID := uuid.New().String() - spec := []byte("important spec data") - batch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: batchID, - TenantID: "IncludeStaticBatchTnt", - Expiry: time.Now().Add(time.Hour).Unix(), - Tags: map[string]string{tagKey1: tagVal1}, - }, - BaseContents: db_api.BaseContents{ - Spec: spec, - Status: []byte("status"), - }, - } - err := batchClient.DBStore(context.Background(), batch) - if err != nil { - t.Fatalf("Failed to store batch: %v", err) - } - - // Get with includeStatic=true. - resItems, _, _, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: []string{batchID}, - }, - }, true, 0, 10) - if err != nil { - t.Fatalf("Failed to get batch: %v", err) - } - if len(resItems) != 1 { - t.Fatalf("Expected 1 item, got %d", len(resItems)) - } - if !bytes.Equal(resItems[0].Spec, spec) { - t.Fatalf("Spec should be included when includeStatic=true") - } - - // Get with includeStatic=false. - resItems, _, _, err = batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: []string{batchID}, - }, - }, false, 0, 10) - if err != nil { - t.Fatalf("Failed to get batch: %v", err) - } - if len(resItems) != 1 { - t.Fatalf("Expected 1 item, got %d", len(resItems)) - } - if len(resItems[0].Spec) != 0 { - t.Fatalf("Spec should be excluded when includeStatic=false, got: %v", resItems[0].Spec) - } - // Status should still be present. - if len(resItems[0].Status) == 0 { - t.Fatalf("Status should still be present") - } - - // Cleanup. - _, _ = batchClient.DBDelete(context.Background(), []string{batchID}) - }) - - t.Run("includeStatic parameter - File", func(t *testing.T) { - t.Parallel() - baseClient, _, fileClient, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - // Store file with spec. - fileID := uuid.New().String() - spec := []byte("important spec data") - file := &db_api.FileItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: fileID, - TenantID: "Tnt1", - Expiry: time.Now().Add(time.Hour).Unix(), - Tags: map[string]string{tagKey1: tagVal1}, - }, - Purpose: "test", - BaseContents: db_api.BaseContents{ - Spec: spec, - Status: []byte("status"), - }, - } - err := fileClient.DBStore(context.Background(), file) - if err != nil { - t.Fatalf("Failed to store file: %v", err) - } - - // Get with includeStatic=true. - resItems, _, _, err := fileClient.DBGet(context.Background(), - &db_api.FileQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: []string{fileID}, - }, - }, true, 0, 10) - if err != nil { - t.Fatalf("Failed to get file: %v", err) - } - if len(resItems) != 1 { - t.Fatalf("Expected 1 item, got %d", len(resItems)) - } - if !bytes.Equal(resItems[0].Spec, spec) { - t.Fatalf("Spec should be included when includeStatic=true") - } - - // Get with includeStatic=false. - resItems, _, _, err = fileClient.DBGet(context.Background(), - &db_api.FileQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: []string{fileID}, - }, - }, false, 0, 10) - if err != nil { - t.Fatalf("Failed to get file: %v", err) - } - if len(resItems) != 1 { - t.Fatalf("Expected 1 item, got %d", len(resItems)) - } - if len(resItems[0].Spec) != 0 { - t.Fatalf("Spec should be excluded when includeStatic=false, got: %v", resItems[0].Spec) - } - - // Cleanup. - _, _ = fileClient.DBDelete(context.Background(), []string{fileID}) - }) - - t.Run("Negative cases - Batch", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - // Store with empty ID should fail validation. - invalidBatch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: "", - }, - BaseContents: db_api.BaseContents{ - Spec: []byte("spec"), - Status: []byte("status"), - }, - } - err := batchClient.DBStore(context.Background(), invalidBatch) - if err == nil { - t.Fatalf("Expected error when storing batch with empty ID") - } - - // Get with non-existent IDs. - resItems, _, _, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: []string{"non-existent-id-1", "non-existent-id-2"}, - }, - }, true, 0, 10) - if err != nil { - t.Fatalf("Get should not error for non-existent IDs: %v", err) - } - if len(resItems) != 0 { - t.Fatalf("Expected 0 items for non-existent IDs, got %d", len(resItems)) - } - - // Get with empty query. - resItems, _, _, err = batchClient.DBGet(context.Background(), nil, true, 0, 10) - if err != nil { - t.Fatalf("Get should handle nil query gracefully: %v", err) - } - if len(resItems) != 0 { - t.Fatalf("Expected 0 items for nil query, got %d", len(resItems)) - } - - // Get with empty IDs list. - resItems, _, expectMore, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: []string{}, - }, - }, true, 0, 10) - if err != nil { - t.Fatalf("Get should handle empty IDs list: %v", err) - } - if len(resItems) != 0 || expectMore { - t.Fatalf("Expected 0 items and no more for empty IDs") - } - - // Update non-existent item. - nonExistentBatch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: "non-existent-update-id", - TenantID: "Tnt1", - }, - BaseContents: db_api.BaseContents{ - Status: []byte("updated"), - }, - } - err = batchClient.DBUpdate(context.Background(), nonExistentBatch, nil) - if err != nil { - t.Fatalf("Update of non-existent item should not error: %v", err) - } - // Cleanup: delete the key created by the update. - _, _ = batchClient.DBDelete(context.Background(), []string{"non-existent-update-id"}) - - // Update with empty ID should fail validation. - invalidUpdate := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: "", - }, - } - err = batchClient.DBUpdate(context.Background(), invalidUpdate, nil) - if err == nil { - t.Fatalf("Expected error when updating batch with empty ID") - } - - // Delete non-existent items. - deletedIDs, err := batchClient.DBDelete(context.Background(), []string{"non-existent-1", "non-existent-2"}) - if err != nil { - t.Fatalf("Delete should not error for non-existent IDs: %v", err) - } - if len(deletedIDs) != 0 { - t.Fatalf("Expected 0 deleted IDs, got %d", len(deletedIDs)) - } - - // Delete with empty IDs list. - deletedIDs, err = batchClient.DBDelete(context.Background(), []string{}) - if err != nil { - t.Fatalf("Delete should handle empty IDs list: %v", err) - } - if len(deletedIDs) != 0 { - t.Fatalf("Expected 0 deleted IDs for empty list, got %d", len(deletedIDs)) - } - }) - - t.Run("Negative cases - File", func(t *testing.T) { - t.Parallel() - baseClient, _, fileClient, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - // Store with empty ID should fail validation. - invalidFile := &db_api.FileItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: "", - }, - Purpose: "test", - } - err := fileClient.DBStore(context.Background(), invalidFile) - if err == nil { - t.Fatalf("Expected error when storing file with empty ID") - } - - // Get by purpose with empty string. - resItems, _, _, err := fileClient.DBGet(context.Background(), - &db_api.FileQuery{ - Purpose: "", - }, true, 0, 10) - if err != nil { - t.Fatalf("Get should handle empty purpose: %v", err) - } - if len(resItems) != 0 { - t.Fatalf("Expected 0 items for empty purpose, got %d", len(resItems)) - } - }) - - t.Run("Edge cases - Empty fields", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - // Store batch with empty spec and status. - batchID := uuid.New().String() - batch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: batchID, - TenantID: "Tnt1", - Expiry: 0, - Tags: map[string]string{}, - }, - BaseContents: db_api.BaseContents{ - Spec: []byte{}, - Status: []byte{}, - }, - } - err := batchClient.DBStore(context.Background(), batch) - if err != nil { - t.Fatalf("Failed to store batch with empty fields: %v", err) - } - - // Retrieve and verify. - resItems, _, _, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: []string{batchID}, - }, - }, true, 0, 10) - if err != nil { - t.Fatalf("Failed to get batch: %v", err) - } - if len(resItems) != 1 { - t.Fatalf("Expected 1 item, got %d", len(resItems)) - } - if resItems[0].Expiry != 0 { - t.Fatalf("Expected expiry 0, got %d", resItems[0].Expiry) - } - if len(resItems[0].Tags) != 0 { - t.Fatalf("Expected empty tags, got %v", resItems[0].Tags) - } - - // Cleanup. - _, _ = batchClient.DBDelete(context.Background(), []string{batchID}) - }) - - t.Run("Edge cases - Update with empty fields", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - // Store batch. - batchID := uuid.New().String() - batch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: batchID, - TenantID: "EmptyUpdateTnt", - Tags: map[string]string{tagKey1: tagVal1}, - }, - BaseContents: db_api.BaseContents{ - Spec: []byte("spec"), - Status: []byte("status"), - }, - } - err := batchClient.DBStore(context.Background(), batch) - if err != nil { - t.Fatalf("Failed to store batch: %v", err) - } - - // Update with empty status and tags - should do nothing. - updateBatch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: batchID, - Tags: map[string]string{}, - }, - BaseContents: db_api.BaseContents{ - Status: []byte{}, - }, - } - err = batchClient.DBUpdate(context.Background(), updateBatch, nil) - if err != nil { - t.Fatalf("Failed to update batch: %v", err) - } - - // Verify original values unchanged. - resItems, _, _, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: []string{batchID}, - }, - }, true, 0, 10) - if err != nil { - t.Fatalf("Failed to get batch: %v", err) - } - if len(resItems) != 1 { - t.Fatalf("Expected 1 item, got %d", len(resItems)) - } - if !bytes.Equal(resItems[0].Status, []byte("status")) { - t.Fatalf("Status should be unchanged") - } - - // Cleanup. - _, _ = batchClient.DBDelete(context.Background(), []string{batchID}) - }) - - t.Run("Pagination - Batch by tenant", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - // Store 10 items for the same tenant. - nItems := 10 - pageSize := 3 - tenant := "PaginationTenant" - stored := make(map[string]*db_api.BatchItem) - var allIDs []string - for i := 0; i < nItems; i++ { - id := uuid.New().String() - batch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: id, - TenantID: tenant, - Expiry: time.Now().Add(time.Hour).Unix(), - }, - BaseContents: db_api.BaseContents{ - Status: []byte(fmt.Sprintf("status-%d", i)), - }, - } - err := batchClient.DBStore(context.Background(), batch) - if err != nil { - t.Fatalf("Failed to store item: %v", err) - } - stored[id] = batch - allIDs = append(allIDs, id) - } - - // Paginate with small page size. - seen := make(map[string]bool) - cursor := 0 - pages := 0 - expectMore := true - for expectMore { - resItems, cur, em, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - TenantID: tenant, - }, - }, true, cursor, pageSize) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - for _, item := range resItems { - if seen[item.ID] { - t.Fatalf("Duplicate item returned: %s", item.ID) - } - seen[item.ID] = true - if _, ok := stored[item.ID]; !ok { - t.Fatalf("Unexpected item returned: %s", item.ID) - } - } - if len(resItems) > pageSize { - t.Fatalf("Page returned more items than limit: %d > %d", len(resItems), pageSize) - } - expectMore = em - cursor = cur - pages++ - } - if len(seen) != nItems { - t.Fatalf("Expected %d total items, got %d", nItems, len(seen)) - } - expectedPages := (nItems + pageSize - 1) / pageSize - if pages != expectedPages { - t.Fatalf("Expected %d pages, got %d", expectedPages, pages) - } - - // Verify stable ordering: paginate again and compare order. - var firstPassIDs []string - cursor = 0 - expectMore = true - for expectMore { - resItems, cur, em, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - TenantID: tenant, - }, - }, true, cursor, pageSize) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - for _, item := range resItems { - firstPassIDs = append(firstPassIDs, item.ID) - } - expectMore = em - cursor = cur - } - var secondPassIDs []string - cursor = 0 - expectMore = true - for expectMore { - resItems, cur, em, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - TenantID: tenant, - }, - }, true, cursor, pageSize) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - for _, item := range resItems { - secondPassIDs = append(secondPassIDs, item.ID) - } - expectMore = em - cursor = cur - } - if len(firstPassIDs) != len(secondPassIDs) { - t.Fatalf("Pass lengths differ: %d vs %d", len(firstPassIDs), len(secondPassIDs)) - } - for i := range firstPassIDs { - if firstPassIDs[i] != secondPassIDs[i] { - t.Fatalf("Order differs at position %d: %s vs %s", i, firstPassIDs[i], secondPassIDs[i]) - } - } - - // Cleanup. - _, _ = batchClient.DBDelete(context.Background(), allIDs) - }) - - t.Run("Pagination - Batch by tags", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - nItems := 7 - pageSize := 2 - stored := make(map[string]*db_api.BatchItem) - var allIDs []string - for i := 0; i < nItems; i++ { - id := uuid.New().String() - batch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: id, - TenantID: "PagTagTenant", - Expiry: time.Now().Add(time.Hour).Unix(), - Tags: map[string]string{"env": "test-pagination"}, - }, - BaseContents: db_api.BaseContents{ - Status: []byte(fmt.Sprintf("status-%d", i)), - }, - } - err := batchClient.DBStore(context.Background(), batch) - if err != nil { - t.Fatalf("Failed to store item: %v", err) - } - stored[id] = batch - allIDs = append(allIDs, id) - } - - seen := make(map[string]bool) - cursor := 0 - expectMore := true - for expectMore { - resItems, cur, em, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - TagSelectors: db_api.Tags{"env": "test-pagination"}, - TagsLogicalCond: db_api.LogicalCondAnd, - TenantID: "PagTagTenant", - }, - }, true, cursor, pageSize) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - for _, item := range resItems { - if seen[item.ID] { - t.Fatalf("Duplicate item returned: %s", item.ID) - } - seen[item.ID] = true - } - if len(resItems) > pageSize { - t.Fatalf("Page returned more items than limit: %d > %d", len(resItems), pageSize) - } - expectMore = em - cursor = cur - } - if len(seen) != nItems { - t.Fatalf("Expected %d total items, got %d", nItems, len(seen)) - } - - _, _ = batchClient.DBDelete(context.Background(), allIDs) - }) - - t.Run("Pagination - File by purpose", func(t *testing.T) { - t.Parallel() - baseClient, _, fileClient, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - nItems := 9 - pageSize := 4 - stored := make(map[string]*db_api.FileItem) - var allIDs []string - for i := 0; i < nItems; i++ { - id := uuid.New().String() - file := &db_api.FileItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: id, - TenantID: "PagPurposeTenant", - Expiry: time.Now().Add(time.Hour).Unix(), - }, - Purpose: "pagination-test", - BaseContents: db_api.BaseContents{ - Status: []byte(fmt.Sprintf("status-%d", i)), - }, - } - err := fileClient.DBStore(context.Background(), file) - if err != nil { - t.Fatalf("Failed to store item: %v", err) - } - stored[id] = file - allIDs = append(allIDs, id) - } - - seen := make(map[string]bool) - cursor := 0 - expectMore := true - for expectMore { - resItems, cur, em, err := fileClient.DBGet(context.Background(), - &db_api.FileQuery{ - BaseQuery: db_api.BaseQuery{ - TenantID: "PagPurposeTenant", - }, - Purpose: "pagination-test", - }, true, cursor, pageSize) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - for _, item := range resItems { - if seen[item.ID] { - t.Fatalf("Duplicate item returned: %s", item.ID) - } - seen[item.ID] = true - } - if len(resItems) > pageSize { - t.Fatalf("Page returned more items than limit: %d > %d", len(resItems), pageSize) - } - expectMore = em - cursor = cur - } - if len(seen) != nItems { - t.Fatalf("Expected %d total items, got %d", nItems, len(seen)) - } - - _, _ = fileClient.DBDelete(context.Background(), allIDs) - }) - - t.Run("Pagination - Batch by expiry", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - nItems := 8 - pageSize := 3 - var allIDs []string - for i := 0; i < nItems; i++ { - id := uuid.New().String() - batch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: id, - TenantID: "PagExpiryTenant", - Expiry: time.Now().Add(time.Second).Unix(), - }, - BaseContents: db_api.BaseContents{ - Status: []byte(fmt.Sprintf("status-%d", i)), - }, - } - err := batchClient.DBStore(context.Background(), batch) - if err != nil { - t.Fatalf("Failed to store item: %v", err) - } - allIDs = append(allIDs, id) - } - time.Sleep(3 * time.Second) - - seen := make(map[string]bool) - cursor := 0 - expectMore := true - for expectMore { - resItems, cur, em, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - Expired: true, - TenantID: "PagExpiryTenant", - }, - }, true, cursor, pageSize) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - for _, item := range resItems { - if seen[item.ID] { - t.Fatalf("Duplicate item returned: %s", item.ID) - } - seen[item.ID] = true - } - if len(resItems) > pageSize { - t.Fatalf("Page returned more items than limit: %d > %d", len(resItems), pageSize) - } - expectMore = em - cursor = cur - } - if len(seen) != nItems { - t.Fatalf("Expected %d total items, got %d", nItems, len(seen)) - } - - _, _ = batchClient.DBDelete(context.Background(), allIDs) - }) - - t.Run("Expiry query excludes zero-expiry items", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - tenant := "ZeroExpiryTenant" - - // Store items with expiry=0 (no expiry, like batch jobs). - var zeroExpiryIDs []string - for i := 0; i < 5; i++ { - id := uuid.New().String() - batch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: id, - TenantID: tenant, - Expiry: 0, - }, - BaseContents: db_api.BaseContents{ - Status: []byte(fmt.Sprintf("status-%d", i)), - }, - } - err := batchClient.DBStore(context.Background(), batch) - if err != nil { - t.Fatalf("Failed to store item: %v", err) - } - zeroExpiryIDs = append(zeroExpiryIDs, id) - } - - // Store items with a past expiry (truly expired). - var expiredIDs []string - for i := 0; i < 3; i++ { - id := uuid.New().String() - batch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: id, - TenantID: tenant, - Expiry: time.Now().Add(time.Second).Unix(), - }, - BaseContents: db_api.BaseContents{ - Status: []byte(fmt.Sprintf("expired-status-%d", i)), - }, - } - err := batchClient.DBStore(context.Background(), batch) - if err != nil { - t.Fatalf("Failed to store item: %v", err) - } - expiredIDs = append(expiredIDs, id) - } - time.Sleep(3 * time.Second) - - // Query for expired items — should only return the 3 truly expired items, - // not the 5 zero-expiry items. - resItems, _, _, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - Expired: true, - TenantID: tenant, - }, - }, true, 0, 100) - if err != nil { - t.Fatalf("Failed to get items: %v", err) - } - if len(resItems) != 3 { - t.Fatalf("Expected 3 expired items, got %d", len(resItems)) - } - - // Verify returned items are only the truly expired ones. - returnedIDs := make(map[string]bool) - for _, item := range resItems { - returnedIDs[item.ID] = true - } - for _, id := range expiredIDs { - if !returnedIDs[id] { - t.Fatalf("Expected expired item %s to be returned", id) - } - } - for _, id := range zeroExpiryIDs { - if returnedIDs[id] { - t.Fatalf("Zero-expiry item %s should NOT be returned by expiry query", id) - } - } - - _, _ = batchClient.DBDelete(context.Background(), zeroExpiryIDs) - _, _ = batchClient.DBDelete(context.Background(), expiredIDs) - }) - - t.Run("Get by IDs with tenant filter", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - // Store batches with different tenants. - batch1ID := uuid.New().String() - batch1 := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: batch1ID, - TenantID: "TenantA", - }, - BaseContents: db_api.BaseContents{ - Status: []byte("status1"), - }, - } - batch2ID := uuid.New().String() - batch2 := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ - ID: batch2ID, - TenantID: "TenantB", - }, - BaseContents: db_api.BaseContents{ - Status: []byte("status2"), - }, - } - _ = batchClient.DBStore(context.Background(), batch1) - _ = batchClient.DBStore(context.Background(), batch2) - - // Get by IDs with tenant filter. - resItems, _, _, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{ - BaseQuery: db_api.BaseQuery{ - IDs: []string{batch1ID, batch2ID}, - TenantID: "TenantA", - }, - }, true, 0, 10) - if err != nil { - t.Fatalf("Failed to get batches: %v", err) - } - if len(resItems) != 1 { - t.Fatalf("Expected 1 item with TenantA, got %d", len(resItems)) - } - if resItems[0].ID != batch1ID { - t.Fatalf("Expected batch1, got %s", resItems[0].ID) - } - - // Cleanup. - _, _ = batchClient.DBDelete(context.Background(), []string{batch1ID, batch2ID}) - }) - - t.Run("CAS update", func(t *testing.T) { - t.Parallel() - baseClient, batchClient, _, _ := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - tests := []struct { - name string - initialStatus []byte - expectedStatus []byte - newStatus []byte - skipStore bool - wantErr bool - }{ - { - name: "matching expected status succeeds", - initialStatus: []byte("validating"), - expectedStatus: []byte("validating"), - newStatus: []byte("in_progress"), - }, - { - name: "mismatched expected status returns ErrConflict", - initialStatus: []byte("in_progress"), - expectedStatus: []byte("validating"), - newStatus: []byte("failed"), - wantErr: true, - }, - { - name: "non-existent key returns ErrConflict", - expectedStatus: []byte("validating"), - newStatus: []byte("in_progress"), - skipStore: true, - wantErr: true, - }, - { - name: "nil expected status bypasses CAS", - initialStatus: []byte("validating"), - expectedStatus: nil, - newStatus: []byte("in_progress"), - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - id := uuid.New().String() - batch := &db_api.BatchItem{ - BaseIndexes: db_api.BaseIndexes{ID: id, TenantID: "Tnt1"}, - BaseContents: db_api.BaseContents{Status: tt.initialStatus}, - } - if !tt.skipStore { - if err := batchClient.DBStore(context.Background(), batch); err != nil { - t.Fatalf("Failed to store: %v", err) - } - } - - batch.Status = tt.newStatus - err := batchClient.DBUpdate(context.Background(), batch, tt.expectedStatus) - - if tt.wantErr { - if !errors.Is(err, db_api.ErrConflict) { - t.Fatalf("expected ErrConflict, got %v", err) - } - if !tt.skipStore { - // Verify status unchanged. - items, _, _, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{BaseQuery: db_api.BaseQuery{IDs: []string{id}}}, true, 0, 1) - if err != nil { - t.Fatalf("Failed to get: %v", err) - } - if !bytes.Equal(items[0].Status, tt.initialStatus) { - t.Fatalf("status should be unchanged: got %s, want %s", items[0].Status, tt.initialStatus) - } - } - } else { - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - // Verify status updated. - items, _, _, err := batchClient.DBGet(context.Background(), - &db_api.BatchQuery{BaseQuery: db_api.BaseQuery{IDs: []string{id}}}, true, 0, 1) - if err != nil { - t.Fatalf("Failed to get: %v", err) - } - if !bytes.Equal(items[0].Status, tt.newStatus) { - t.Fatalf("status mismatch: got %s, want %s", items[0].Status, tt.newStatus) - } - } - - _, _ = batchClient.DBDelete(context.Background(), []string{id}) - }) - } - }) - - t.Run("InFlight operations", func(t *testing.T) { - t.Parallel() - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) - t.Cleanup(func() { - _ = baseClient.Close() - }) - - t.Run("set and get round-trip", func(t *testing.T) { - err := exchClient.InFlightSet(context.Background(), "job-1", "pod-a") - if err != nil { - t.Fatalf("InFlightSet failed: %v", err) - } - - all, err := exchClient.InFlightGetAll(context.Background()) - if err != nil { - t.Fatalf("InFlightGetAll failed: %v", err) - } - entry, ok := all["job-1"] - if !ok { - t.Fatal("expected job-1 in in-flight entries") - } - if entry.ProcessorID != "pod-a" { - t.Fatalf("expected ProcessorID pod-a, got %s", entry.ProcessorID) - } - if entry.LastSeen <= 0 { - t.Fatalf("expected positive LastSeen, got %d", entry.LastSeen) - } - - // Cleanup. - _ = exchClient.InFlightDelete(context.Background(), "job-1") - }) - - t.Run("set overwrites existing entry", func(t *testing.T) { - _ = exchClient.InFlightSet(context.Background(), "job-2", "pod-a") - _ = exchClient.InFlightSet(context.Background(), "job-2", "pod-b") - - all, err := exchClient.InFlightGetAll(context.Background()) - if err != nil { - t.Fatalf("InFlightGetAll failed: %v", err) - } - if all["job-2"].ProcessorID != "pod-b" { - t.Fatalf("expected overwritten ProcessorID pod-b, got %s", all["job-2"].ProcessorID) - } - - _ = exchClient.InFlightDelete(context.Background(), "job-2") - }) - - t.Run("delete removes entry", func(t *testing.T) { - _ = exchClient.InFlightSet(context.Background(), "job-3", "pod-a") - err := exchClient.InFlightDelete(context.Background(), "job-3") - if err != nil { - t.Fatalf("InFlightDelete failed: %v", err) - } - - all, err := exchClient.InFlightGetAll(context.Background()) - if err != nil { - t.Fatalf("InFlightGetAll failed: %v", err) - } - if _, ok := all["job-3"]; ok { - t.Fatal("expected job-3 to be deleted") - } - }) - - t.Run("delete non-existent key is idempotent", func(t *testing.T) { - err := exchClient.InFlightDelete(context.Background(), "non-existent") - if err != nil { - t.Fatalf("InFlightDelete of non-existent key should not error: %v", err) - } - }) - - t.Run("get all on empty hash returns empty map", func(t *testing.T) { - all, err := exchClient.InFlightGetAll(context.Background()) - if err != nil { - t.Fatalf("InFlightGetAll failed: %v", err) - } - if len(all) != 0 { - t.Fatalf("expected empty map, got %d entries", len(all)) - } - }) - - t.Run("validation rejects empty jobID", func(t *testing.T) { - if err := exchClient.InFlightSet(context.Background(), "", "pod-a"); err == nil { - t.Fatal("expected error for empty jobID") - } - if err := exchClient.InFlightDelete(context.Background(), ""); err == nil { - t.Fatal("expected error for empty jobID") - } - }) - - t.Run("validation rejects empty processorID", func(t *testing.T) { - if err := exchClient.InFlightSet(context.Background(), "job-x", ""); err == nil { - t.Fatal("expected error for empty processorID") - } - }) - }) - t.Run("PQGetIDs", func(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -2578,7 +967,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -2621,7 +1010,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -2656,7 +1045,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -2708,7 +1097,7 @@ func TestRedisDSClient(t *testing.T) { if minirds != nil { t.Skip("Miniredis model") } - baseClient, _, _, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) + baseClient, exchClient := setupRedisDSClients(t, redisUrl, redisCaCert) t.Cleanup(func() { _ = baseClient.Close() }) @@ -2769,76 +1158,3 @@ func isSameEvent(t *testing.T, a, b *db_api.BatchEvent) bool { } return true } - -func sameMembersBatch(t *testing.T, sl []*db_api.BatchItem, mp map[string]*db_api.BatchItem) bool { - t.Helper() - for _, item := range sl { - isEqualBatchItem(t, item, mp[item.ID]) - } - return true -} - -func isEqualBatchItem(t *testing.T, a, b *db_api.BatchItem) bool { - t.Helper() - if a == nil || b == nil { - t.Fatalf("Invalid items to compare") - return false - } - if a.ID != b.ID { - t.Fatalf("Mismatch id %s != %s", a.ID, b.ID) - } - if a.TenantID != b.TenantID { - t.Fatalf("Mismatch TenantID %s != %s", a.TenantID, b.TenantID) - } - if a.Expiry != b.Expiry { - t.Fatalf("Mismatch expiry %d != %d", a.Expiry, b.Expiry) - } - if !maps.Equal(a.Tags, b.Tags) { - t.Fatalf("Mismatch tags %v != %v", a.Tags, b.Tags) - } - if !bytes.Equal(a.Spec, b.Spec) { - t.Fatalf("Mismatch spec %s != %s", a.Spec, b.Spec) - } - if !bytes.Equal(a.Status, b.Status) { - t.Fatalf("Mismatch status %s != %s", a.Spec, b.Spec) - } - return true -} - -func sameMembersFile(t *testing.T, sl []*db_api.FileItem, mp map[string]*db_api.FileItem) bool { - t.Helper() - for _, item := range sl { - isEqualFileItem(t, item, mp[item.ID]) - } - return true -} - -func isEqualFileItem(t *testing.T, a, b *db_api.FileItem) bool { - t.Helper() - if a == nil || b == nil { - t.Fatalf("Invalid items to compare") - return false - } - if a.ID != b.ID { - t.Fatalf("Mismatch id %s != %s", a.ID, b.ID) - } - if a.TenantID != b.TenantID { - t.Fatalf("Mismatch TenantID %s != %s", a.TenantID, b.TenantID) - } - if a.Expiry != b.Expiry { - t.Fatalf("Mismatch expiry %d != %d", a.Expiry, b.Expiry) - } - if !maps.Equal(a.Tags, b.Tags) { - t.Fatalf("Mismatch tags %v != %v", a.Tags, b.Tags) - } - if a.Purpose != b.Purpose { - t.Fatalf("Mismatch purpose %s != %s", a.Purpose, b.Purpose) - } - if !bytes.Equal(a.Spec, b.Spec) { - t.Fatalf("Mismatch spec %s != %s", a.Spec, b.Spec) - } - if !bytes.Equal(a.Status, b.Status) { - t.Fatalf("Mismatch status %s != %s", a.Spec, b.Spec) - } - return true -} diff --git a/internal/database/redis/redis_inflight.go b/internal/database/redis/redis_inflight.go deleted file mode 100644 index 5d0e34147..000000000 --- a/internal/database/redis/redis_inflight.go +++ /dev/null @@ -1,97 +0,0 @@ -/* -Copyright 2026 The llm-d Authors - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -package redis - -import ( - "context" - "encoding/json" - "fmt" - "time" - - db_api "github.com/llm-d/llm-d-batch-gateway/internal/database/api" -) - -func (c *ExchangeDBClientRedis) InFlightSet(ctx context.Context, jobID, processorID string) error { - if ctx == nil { - ctx = context.Background() - } - if jobID == "" { - return fmt.Errorf("jobID is empty") - } - if processorID == "" { - return fmt.Errorf("processorID is empty") - } - - entry := db_api.InFlightEntry{ - ProcessorID: processorID, - LastSeen: time.Now().Unix(), - } - data, err := json.Marshal(entry) - if err != nil { - return fmt.Errorf("failed to marshal in-flight entry: %w", err) - } - - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - return c.redisClient.HSet(cctx, inFlightKeyName, jobID, data).Err() -} - -func (c *ExchangeDBClientRedis) InFlightDelete(ctx context.Context, jobID string) error { - if ctx == nil { - ctx = context.Background() - } - if jobID == "" { - return fmt.Errorf("jobID is empty") - } - - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - return c.redisClient.HDel(cctx, inFlightKeyName, jobID).Err() -} - -func (c *ExchangeDBClientRedis) InFlightGetAll(ctx context.Context) (map[string]*db_api.InFlightEntry, error) { - if ctx == nil { - ctx = context.Background() - } - - result := make(map[string]*db_api.InFlightEntry) - var cursor uint64 - cctx, ccancel := context.WithTimeout(ctx, c.timeout) - defer ccancel() - - for { - entries, nextCursor, err := c.redisClient.HScan(cctx, inFlightKeyName, cursor, "*", 100).Result() - if err != nil { - return nil, fmt.Errorf("failed to scan in-flight hash: %w", err) - } - // HScan returns [field, value, field, value, ...] - for i := 0; i < len(entries); i += 2 { - jobID := entries[i] - var entry db_api.InFlightEntry - if err := json.Unmarshal([]byte(entries[i+1]), &entry); err != nil { - return nil, fmt.Errorf("failed to unmarshal in-flight entry for job %s: %w", jobID, err) - } - result[jobID] = &entry - } - cursor = nextCursor - if cursor == 0 { - break - } - } - - return result, nil -} diff --git a/internal/gc/config/config.go b/internal/gc/config/config.go index 98f80be24..10ab3d1ca 100644 --- a/internal/gc/config/config.go +++ b/internal/gc/config/config.go @@ -40,8 +40,10 @@ const ( // ReconcilerConfig holds the orphan reconciler configuration. type ReconcilerConfig struct { - Enabled bool `yaml:"enabled"` - Interval time.Duration `yaml:"interval"` + Enabled bool `yaml:"enabled"` + Interval time.Duration `yaml:"interval"` + ProcessorLabelSelector string `yaml:"processor_label_selector"` + ProcessorStatefulSet string `yaml:"processor_statefulset"` } // CollectorConfig holds collector-specific settings (interval and concurrency). @@ -111,12 +113,12 @@ func Load(path string) (*Config, error) { } switch cfg.DBClientCfg.Type { - case sharedcfg.DBTypeRedis, sharedcfg.DBTypeValkey, sharedcfg.DBTypePostgreSQL: + case sharedcfg.DBTypePostgreSQL: // valid case "": - return nil, fmt.Errorf("db_client.type is required (must be \"redis\", \"valkey\", or \"postgresql\")") + return nil, fmt.Errorf("db_client.type is required (must be \"postgresql\")") default: - return nil, fmt.Errorf("db_client.type must be \"redis\", \"valkey\", or \"postgresql\", got %q", cfg.DBClientCfg.Type) + return nil, fmt.Errorf("db_client.type must be \"postgresql\", got %q", cfg.DBClientCfg.Type) } switch cfg.FileClientCfg.Type { diff --git a/internal/gc/config/config_test.go b/internal/gc/config/config_test.go index c48f829e8..c09a48833 100644 --- a/internal/gc/config/config_test.go +++ b/internal/gc/config/config_test.go @@ -99,60 +99,6 @@ file_client: } } -func TestLoad_RedisDatabase(t *testing.T) { - path := writeTempConfig(t, ` -db_client: - type: "redis" - redis: - db: 2 - enable_tls: true -file_client: - type: "fs" - fs: - base_path: "/tmp/files" -`) - cfg, err := Load(path) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if cfg.DBClientCfg.Type != "redis" { - t.Errorf("expected db_client.type redis, got %s", cfg.DBClientCfg.Type) - } - if cfg.DBClientCfg.RedisCfg.DB != 2 { - t.Errorf("expected redis db 2, got %d", cfg.DBClientCfg.RedisCfg.DB) - } - if !cfg.DBClientCfg.RedisCfg.EnableTLS { - t.Error("expected redis enable_tls to be true") - } -} - -func TestLoad_ValkeyDatabase(t *testing.T) { - path := writeTempConfig(t, ` -db_client: - type: "valkey" - redis: - db: 2 - enable_tls: true -file_client: - type: "fs" - fs: - base_path: "/tmp/files" -`) - cfg, err := Load(path) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if cfg.DBClientCfg.Type != "valkey" { - t.Errorf("expected db_client.type valkey, got %s", cfg.DBClientCfg.Type) - } - if cfg.DBClientCfg.RedisCfg.DB != 2 { - t.Errorf("expected redis db 2, got %d", cfg.DBClientCfg.RedisCfg.DB) - } - if !cfg.DBClientCfg.RedisCfg.EnableTLS { - t.Error("expected redis enable_tls to be true") - } -} - func TestLoad_PostgreSQLDatabase(t *testing.T) { path := writeTempConfig(t, ` db_client: @@ -521,59 +467,3 @@ file_client: }) } } - -func TestLoad_RedisConfigTuning(t *testing.T) { - path := writeTempConfig(t, ` -db_client: - type: "redis" - redis: - db: 3 - enable_tls: true - insecure: true - timeout: "5s" - max_retries: 5 - min_retry_backoff: "100ms" - max_retry_backoff: "2s" - pool_timeout: "10s" - conn_max_idle_time: "5m" - conn_max_lifetime: "30m" -file_client: - type: "fs" - fs: - base_path: "/tmp/files" -`) - cfg, err := Load(path) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if cfg.DBClientCfg.RedisCfg.DB != 3 { - t.Errorf("expected db 3, got %d", cfg.DBClientCfg.RedisCfg.DB) - } - if !cfg.DBClientCfg.RedisCfg.EnableTLS { - t.Error("expected enable_tls to be true") - } - if !cfg.DBClientCfg.RedisCfg.Insecure { - t.Error("expected insecure to be true") - } - if cfg.DBClientCfg.RedisCfg.Timeout != 5*time.Second { - t.Errorf("expected timeout 5s, got %v", cfg.DBClientCfg.RedisCfg.Timeout) - } - if cfg.DBClientCfg.RedisCfg.MaxRetries != 5 { - t.Errorf("expected max_retries 5, got %d", cfg.DBClientCfg.RedisCfg.MaxRetries) - } - if cfg.DBClientCfg.RedisCfg.MinRetryBackoff != 100*time.Millisecond { - t.Errorf("expected min_retry_backoff 100ms, got %v", cfg.DBClientCfg.RedisCfg.MinRetryBackoff) - } - if cfg.DBClientCfg.RedisCfg.MaxRetryBackoff != 2*time.Second { - t.Errorf("expected max_retry_backoff 2s, got %v", cfg.DBClientCfg.RedisCfg.MaxRetryBackoff) - } - if cfg.DBClientCfg.RedisCfg.PoolTimeout != 10*time.Second { - t.Errorf("expected pool_timeout 10s, got %v", cfg.DBClientCfg.RedisCfg.PoolTimeout) - } - if cfg.DBClientCfg.RedisCfg.ConnMaxIdleTime != 5*time.Minute { - t.Errorf("expected conn_max_idle_time 5m, got %v", cfg.DBClientCfg.RedisCfg.ConnMaxIdleTime) - } - if cfg.DBClientCfg.RedisCfg.ConnMaxLifetime != 30*time.Minute { - t.Errorf("expected conn_max_lifetime 30m, got %v", cfg.DBClientCfg.RedisCfg.ConnMaxLifetime) - } -} diff --git a/internal/gc/podwatcher/watcher.go b/internal/gc/podwatcher/watcher.go new file mode 100644 index 000000000..1aec03db7 --- /dev/null +++ b/internal/gc/podwatcher/watcher.go @@ -0,0 +1,175 @@ +/* +Copyright 2026 The llm-d Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package podwatcher watches the processor StatefulSet via the Kubernetes +// API and notifies the reconciler when the set of ready pods stabilizes. +package podwatcher + +import ( + "context" + "fmt" + "os" + "strings" + + "github.com/go-logr/logr" + "github.com/llm-d/llm-d-batch-gateway/internal/util/logging" + appsv1 "k8s.io/api/apps/v1" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/informers" + "k8s.io/client-go/kubernetes" + "k8s.io/client-go/rest" + "k8s.io/client-go/tools/cache" +) + +// PodEventHandler is called when the StatefulSet is stable (readyReplicas +// == spec.replicas) with the set of ready processor pod names. +type PodEventHandler func(liveProcessors map[string]bool) + +// Watcher watches the processor StatefulSet. When readyReplicas equals +// spec.replicas (stable), it lists the ready pods and calls the handler +// with their names. +type Watcher struct { + clientset kubernetes.Interface + namespace string + stsName string + podLabelSelector string + handler PodEventHandler +} + +// New creates a new Watcher using in-cluster Kubernetes config. +// stsName is the processor StatefulSet name. +// podLabelSelector identifies processor pods for listing. +// The namespace is auto-detected from the service account mount. +func New(stsName, podLabelSelector string, handler PodEventHandler) (*Watcher, error) { + restCfg, err := rest.InClusterConfig() + if err != nil { + return nil, fmt.Errorf("failed to get in-cluster config: %w", err) + } + + cs, err := kubernetes.NewForConfig(restCfg) + if err != nil { + return nil, fmt.Errorf("failed to create kubernetes client: %w", err) + } + + ns, err := detectNamespace() + if err != nil { + return nil, fmt.Errorf("failed to detect namespace: %w", err) + } + + return &Watcher{ + clientset: cs, + namespace: ns, + stsName: stsName, + podLabelSelector: podLabelSelector, + handler: handler, + }, nil +} + +// Run watches the processor StatefulSet and blocks until the context is +// cancelled. On each StatefulSet update where readyReplicas == replicas, +// it lists the ready pods and calls the handler. +func (w *Watcher) Run(ctx context.Context) error { + logger := logr.FromContextOrDiscard(ctx) + + factory := informers.NewSharedInformerFactoryWithOptions( + w.clientset, + 0, + informers.WithNamespace(w.namespace), + informers.WithTweakListOptions(func(opts *metav1.ListOptions) { + opts.FieldSelector = "metadata.name=" + w.stsName + }), + ) + + stsInformer := factory.Apps().V1().StatefulSets().Informer() + if _, err := stsInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{ + AddFunc: func(obj interface{}) { + w.onStatefulSetChange(ctx, logger, obj) + }, + UpdateFunc: func(_, newObj interface{}) { + w.onStatefulSetChange(ctx, logger, newObj) + }, + }); err != nil { + return fmt.Errorf("failed to add event handler: %w", err) + } + + factory.Start(ctx.Done()) + if !cache.WaitForCacheSync(ctx.Done(), stsInformer.HasSynced) { + return fmt.Errorf("failed to sync StatefulSet informer cache") + } + + logger.Info("Pod watcher started", + "namespace", w.namespace, + "statefulSet", w.stsName, + "podLabelSelector", w.podLabelSelector) + + <-ctx.Done() + return ctx.Err() +} + +func (w *Watcher) onStatefulSetChange(ctx context.Context, logger logr.Logger, obj interface{}) { + sts, ok := obj.(*appsv1.StatefulSet) + if !ok || sts.Name != w.stsName { + return + } + + desired := int32(1) + if sts.Spec.Replicas != nil { + desired = *sts.Spec.Replicas + } + + if sts.Status.ReadyReplicas != desired { + logger.V(logging.INFO).Info("StatefulSet not stable", + "ready", sts.Status.ReadyReplicas, "desired", desired) + return + } + + podList, err := w.clientset.CoreV1().Pods(w.namespace).List(ctx, metav1.ListOptions{ + LabelSelector: w.podLabelSelector, + }) + if err != nil { + logger.Error(err, "Failed to list processor pods") + return + } + + live := make(map[string]bool, len(podList.Items)) + for i := range podList.Items { + if isPodReady(&podList.Items[i]) { + live[podList.Items[i].Name] = true + } + } + + logger.Info("StatefulSet stable, updating live processors", + "ready", len(live), "desired", desired) + w.handler(live) +} + +func isPodReady(pod *corev1.Pod) bool { + for _, cond := range pod.Status.Conditions { + if cond.Type == corev1.PodReady && cond.Status == corev1.ConditionTrue { + return true + } + } + return false +} + +func detectNamespace() (string, error) { + data, err := os.ReadFile("/var/run/secrets/kubernetes.io/serviceaccount/namespace") + if err != nil { + return "", err + } + return strings.TrimSpace(string(data)), nil +} diff --git a/internal/gc/podwatcher/watcher_test.go b/internal/gc/podwatcher/watcher_test.go new file mode 100644 index 000000000..a9207f616 --- /dev/null +++ b/internal/gc/podwatcher/watcher_test.go @@ -0,0 +1,255 @@ +/* +Copyright 2026 The llm-d Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package podwatcher + +import ( + "context" + "testing" + "time" + + appsv1 "k8s.io/api/apps/v1" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/kubernetes/fake" + "k8s.io/utils/ptr" +) + +const ( + testNamespace = "default" + testStsName = "processor" + testSelector = "app.kubernetes.io/component=processor" +) + +func readyPod(name string) *corev1.Pod { + return &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: name, + Namespace: testNamespace, + Labels: map[string]string{"app.kubernetes.io/component": "processor"}, + }, + Status: corev1.PodStatus{ + Conditions: []corev1.PodCondition{ + {Type: corev1.PodReady, Status: corev1.ConditionTrue}, + }, + }, + } +} + +func testStatefulSet(replicas int32, readyReplicas int32) *appsv1.StatefulSet { + return &appsv1.StatefulSet{ + ObjectMeta: metav1.ObjectMeta{ + Name: testStsName, + Namespace: testNamespace, + }, + Spec: appsv1.StatefulSetSpec{ + Replicas: ptr.To(replicas), + }, + Status: appsv1.StatefulSetStatus{ + ReadyReplicas: readyReplicas, + }, + } +} + +func waitForLive(t *testing.T, ch <-chan map[string]bool) map[string]bool { + t.Helper() + select { + case live := <-ch: + return live + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for handler call") + return nil + } +} + +func TestWatcher(t *testing.T) { + t.Run("stable StatefulSet calls handler with ready pods", func(t *testing.T) { + cs := fake.NewSimpleClientset( + testStatefulSet(2, 2), + readyPod("processor-0"), + readyPod("processor-1"), + ) + + liveCh := make(chan map[string]bool, 10) + w := &Watcher{ + clientset: cs, + namespace: testNamespace, + stsName: testStsName, + podLabelSelector: testSelector, + handler: func(live map[string]bool) { liveCh <- live }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go func() { _ = w.Run(ctx) }() + + live := waitForLive(t, liveCh) + if len(live) != 2 { + t.Fatalf("expected 2 live pods, got %d: %v", len(live), live) + } + if !live["processor-0"] || !live["processor-1"] { + t.Fatalf("expected processor-0 and processor-1, got %v", live) + } + }) + + t.Run("unstable StatefulSet does not call handler", func(t *testing.T) { + cs := fake.NewSimpleClientset( + testStatefulSet(3, 2), // 3 desired, only 2 ready + readyPod("processor-0"), + readyPod("processor-1"), + ) + + liveCh := make(chan map[string]bool, 10) + w := &Watcher{ + clientset: cs, + namespace: testNamespace, + stsName: testStsName, + podLabelSelector: testSelector, + handler: func(live map[string]bool) { liveCh <- live }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go func() { _ = w.Run(ctx) }() + + select { + case live := <-liveCh: + t.Fatalf("handler should not be called when unstable, got %v", live) + case <-time.After(500 * time.Millisecond): + // Expected — handler not called. + } + }) + + t.Run("scale up triggers handler with expanded set", func(t *testing.T) { + sts := testStatefulSet(1, 1) + cs := fake.NewSimpleClientset( + sts, + readyPod("processor-0"), + ) + + liveCh := make(chan map[string]bool, 10) + w := &Watcher{ + clientset: cs, + namespace: testNamespace, + stsName: testStsName, + podLabelSelector: testSelector, + handler: func(live map[string]bool) { liveCh <- live }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go func() { _ = w.Run(ctx) }() + + // Wait for initial stable state (1 replica). + waitForLive(t, liveCh) + + // Simulate scale up: add a pod and update STS to 2 replicas. + _, err := cs.CoreV1().Pods(testNamespace).Create(ctx, readyPod("processor-1"), metav1.CreateOptions{}) + if err != nil { + t.Fatalf("failed to create pod: %v", err) + } + sts.Spec.Replicas = ptr.To(int32(2)) + sts.Status.ReadyReplicas = 2 + if _, err := cs.AppsV1().StatefulSets(testNamespace).Update(ctx, sts, metav1.UpdateOptions{}); err != nil { + t.Fatalf("failed to update sts: %v", err) + } + + live := waitForLive(t, liveCh) + if len(live) != 2 { + t.Fatalf("expected 2 live pods after scale up, got %d: %v", len(live), live) + } + if !live["processor-0"] || !live["processor-1"] { + t.Fatalf("expected processor-0 and processor-1, got %v", live) + } + }) + + t.Run("scale down triggers handler with reduced set", func(t *testing.T) { + sts := testStatefulSet(2, 2) + cs := fake.NewSimpleClientset( + sts, + readyPod("processor-0"), + readyPod("processor-1"), + ) + + liveCh := make(chan map[string]bool, 10) + w := &Watcher{ + clientset: cs, + namespace: testNamespace, + stsName: testStsName, + podLabelSelector: testSelector, + handler: func(live map[string]bool) { liveCh <- live }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go func() { _ = w.Run(ctx) }() + + // Wait for initial stable state. + waitForLive(t, liveCh) + + // Simulate scale down: update STS to 1 replica and delete processor-1. + sts.Spec.Replicas = ptr.To(int32(1)) + sts.Status.ReadyReplicas = 1 + _, err := cs.AppsV1().StatefulSets(testNamespace).Update(ctx, sts, metav1.UpdateOptions{}) + if err != nil { + t.Fatalf("failed to update sts: %v", err) + } + if err := cs.CoreV1().Pods(testNamespace).Delete(ctx, "processor-1", metav1.DeleteOptions{}); err != nil { + t.Fatalf("failed to delete pod: %v", err) + } + + live := waitForLive(t, liveCh) + if len(live) != 1 { + t.Fatalf("expected 1 live pod after scale down, got %d: %v", len(live), live) + } + if !live["processor-0"] { + t.Fatalf("expected processor-0, got %v", live) + } + }) +} + +func TestIsPodReady(t *testing.T) { + t.Run("ready pod", func(t *testing.T) { + pod := readyPod("test") + if !isPodReady(pod) { + t.Error("expected pod to be ready") + } + }) + + t.Run("not ready pod", func(t *testing.T) { + pod := &corev1.Pod{ + Status: corev1.PodStatus{ + Conditions: []corev1.PodCondition{ + {Type: corev1.PodReady, Status: corev1.ConditionFalse}, + }, + }, + } + if isPodReady(pod) { + t.Error("expected pod to not be ready") + } + }) + + t.Run("no conditions", func(t *testing.T) { + pod := &corev1.Pod{} + if isPodReady(pod) { + t.Error("expected pod with no conditions to not be ready") + } + }) +} diff --git a/internal/gc/reconciler/reconciler.go b/internal/gc/reconciler/reconciler.go index aedb07c28..542e5a150 100644 --- a/internal/gc/reconciler/reconciler.go +++ b/internal/gc/reconciler/reconciler.go @@ -15,7 +15,7 @@ limitations under the License. */ // Package reconciler detects and recovers orphaned batch jobs that are stuck -// in non-terminal states because their processor crashed or lost connectivity. +// in non-terminal states because their processor crashed or was deleted. package reconciler import ( @@ -23,7 +23,7 @@ import ( "encoding/json" "errors" "fmt" - "strconv" + "sync" "time" "github.com/go-logr/logr" @@ -31,40 +31,43 @@ import ( db "github.com/llm-d/llm-d-batch-gateway/internal/database/api" "github.com/llm-d/llm-d-batch-gateway/internal/shared/batch_utils" "github.com/llm-d/llm-d-batch-gateway/internal/shared/openai" - batch_types "github.com/llm-d/llm-d-batch-gateway/internal/shared/types" ) const pageSize = 100 // Result contains the outcome of a single reconciliation cycle. type Result struct { - Cancelled int - Expired int - ReEnqueued int - Failed int - StaleCleanup int - Conflicts int - Errors int - Duration time.Duration + Expired int + ReEnqueued int + Conflicts int + Errors int + Duration time.Duration } -// Reconciler periodically scans for orphaned batch jobs and recovers them. +// Reconciler detects orphaned batch jobs and recovers them. A job is +// considered orphaned when it has a processor_id set but that processor +// is no longer in the live set (maintained by the pod watcher). +// +// The reconciler runs on two triggers: +// - Event-driven: the pod watcher calls Trigger() on pod deletion +// - Periodic: a backstop timer fires every interval as a safety net type Reconciler struct { batchDB db.BatchDBClient queue db.BatchPriorityQueueClient - inflight db.InFlightClient interval time.Duration dryRun bool onCycleComplete func(*Result) + + mu sync.RWMutex + liveProcessors map[string]bool + + triggerCh chan struct{} } // NewReconciler creates a new orphan reconciler. -// interval controls both the scan frequency and the staleness threshold for in-flight entries. -// onCycleComplete, if non-nil, is called after each cycle with the result. func NewReconciler( batchDB db.BatchDBClient, queue db.BatchPriorityQueueClient, - inflight db.InFlightClient, interval time.Duration, dryRun bool, onCycleComplete func(*Result), @@ -75,30 +78,48 @@ func NewReconciler( if queue == nil { return nil, fmt.Errorf("queue client is required") } - if inflight == nil { - return nil, fmt.Errorf("in-flight client is required") - } if interval <= 0 { return nil, fmt.Errorf("interval must be positive, got %v", interval) } return &Reconciler{ batchDB: batchDB, queue: queue, - inflight: inflight, interval: interval, dryRun: dryRun, onCycleComplete: onCycleComplete, + liveProcessors: make(map[string]bool), + triggerCh: make(chan struct{}, 1), }, nil } -// RunLoop runs the reconciler in a continuous loop at the configured interval. +// SetLiveProcessors updates the set of currently alive processor pod names. +// Called by the pod watcher on add/delete events. +func (r *Reconciler) SetLiveProcessors(processors map[string]bool) { + r.mu.Lock() + defer r.mu.Unlock() + r.liveProcessors = processors +} + +// Trigger requests an immediate reconciliation cycle. Non-blocking. +func (r *Reconciler) Trigger() { + select { + case r.triggerCh <- struct{}{}: + default: + } +} + +func (r *Reconciler) isProcessorAlive(processorID string) bool { + r.mu.RLock() + defer r.mu.RUnlock() + return r.liveProcessors[processorID] +} + +// RunLoop runs the reconciler on both event triggers and a periodic timer. // It blocks until the context is cancelled. func (r *Reconciler) RunLoop(ctx context.Context) error { logger := logr.FromContextOrDiscard(ctx) logger.Info("Reconciler: starting loop", "interval", r.interval) - r.run(ctx) - ticker := time.NewTicker(r.interval) defer ticker.Stop() @@ -109,6 +130,8 @@ func (r *Reconciler) RunLoop(ctx context.Context) error { return ctx.Err() case <-ticker.C: r.run(ctx) + case <-r.triggerCh: + r.run(ctx) } } } @@ -122,11 +145,8 @@ func (r *Reconciler) run(ctx context.Context) { defer func() { result.Duration = time.Since(start) logger.Info("Reconciler: cycle completed", - "cancelled", result.Cancelled, "expired", result.Expired, "reEnqueued", result.ReEnqueued, - "failed", result.Failed, - "staleCleanup", result.StaleCleanup, "conflicts", result.Conflicts, "errors", result.Errors, "duration", result.Duration, @@ -141,45 +161,11 @@ func (r *Reconciler) run(ctx context.Context) { return } - inflightEntries, err := r.inflight.InFlightGetAll(ctx) - if err != nil { - logger.Error(err, "Reconciler: failed to get in-flight entries") - result.Errors++ - return - } - - nonTerminalIDs := make(map[string]bool, len(jobs)) - - if len(jobs) > 0 { - queuedIDs, err := r.queue.PQGetIDs(ctx) - if err != nil { - logger.Error(err, "Reconciler: failed to get queued job IDs") - result.Errors++ - return - } - - now := time.Now() - stalenessThreshold := now.Add(-r.interval) - - for _, job := range jobs { - nonTerminalIDs[job.ID] = true - - if queuedIDs[job.ID] { - continue - } - - if entry, ok := inflightEntries[job.ID]; ok { - lastSeen := time.Unix(entry.LastSeen, 0) - if lastSeen.After(stalenessThreshold) { - continue - } - } - + for _, job := range jobs { + if !r.isProcessorAlive(job.ProcessorID) { r.triageOrphan(ctx, job, result) } } - - r.cleanupStaleInflight(ctx, inflightEntries, nonTerminalIDs, result) } func (r *Reconciler) notifyCycle(result *Result) { @@ -188,228 +174,127 @@ func (r *Reconciler) notifyCycle(result *Result) { } } -// fetchNonTerminalJobs retrieves all non-terminal batch jobs via paginated queries. +// fetchNonTerminalJobs paginates through non-terminal batch items that are +// owned by a processor (processor_id IS NOT NULL). Queued jobs (processor_id +// IS NULL) are excluded since they are not orphans. func (r *Reconciler) fetchNonTerminalJobs(ctx context.Context) ([]*db.BatchItem, error) { - query := &db.BatchQuery{NonTerminal: true} - var allJobs []*db.BatchItem + var all []*db.BatchItem cursor := 0 - for { - select { - case <-ctx.Done(): - return nil, ctx.Err() - default: - } - - jobs, nextCursor, expectMore, err := r.batchDB.DBGet(ctx, query, false, cursor, pageSize) + items, nextCursor, more, err := r.batchDB.DBGet(ctx, + &db.BatchQuery{NonTerminal: true, HasProcessorID: true}, + false, cursor, pageSize) if err != nil { - return nil, fmt.Errorf("failed to query non-terminal jobs: %w", err) + return nil, err } - allJobs = append(allJobs, jobs...) - - if !expectMore { + all = append(all, items...) + if !more { break } cursor = nextCursor } - - return allJobs, nil + return all, nil } -// triageOrphan determines the correct recovery action for an orphaned job -// based on its current status and SLO. func (r *Reconciler) triageOrphan(ctx context.Context, job *db.BatchItem, result *Result) { - logger := logr.FromContextOrDiscard(ctx).WithValues("jobId", job.ID) + logger := logr.FromContextOrDiscard(ctx).WithValues("jobId", job.ID, "processorId", job.ProcessorID) var statusInfo openai.BatchStatusInfo if err := json.Unmarshal(job.Status, &statusInfo); err != nil { - logger.Error(err, "Reconciler: failed to unmarshal job status") + logger.Error(err, "Reconciler: failed to unmarshal orphan status") result.Errors++ return } - sloExpired := isSLOExpired(job) - - var ok bool - switch statusInfo.Status { case openai.BatchStatusCancelling: - ok = r.transitionOrphan(ctx, job, &statusInfo, openai.BatchStatusCancelled, result, logger) - - case openai.BatchStatusValidating: - if sloExpired { - ok = r.transitionOrphan(ctx, job, &statusInfo, openai.BatchStatusExpired, result, logger) - } else { - ok = r.reEnqueueOrphan(ctx, job, result, logger) - } - - case openai.BatchStatusInProgress, openai.BatchStatusFinalizing: - if sloExpired { - ok = r.transitionOrphan(ctx, job, &statusInfo, openai.BatchStatusExpired, result, logger) - } else { - ok = r.transitionOrphan(ctx, job, &statusInfo, openai.BatchStatusFailed, result, logger) - } - + r.transitionOrphan(ctx, job, &statusInfo, openai.BatchStatusCancelled, result, logger) default: - logger.Info("Reconciler: orphan in unexpected status, skipping", "status", statusInfo.Status) - result.Errors++ - return - } - - if ok && !r.dryRun { - if err := r.inflight.InFlightDelete(ctx, job.ID); err != nil { - logger.Error(err, "Reconciler: failed to delete in-flight entry for orphan") - result.Errors++ + if isSLOExpired(job) { + r.transitionOrphan(ctx, job, &statusInfo, openai.BatchStatusFailed, result, logger) + } else { + r.reEnqueueOrphan(ctx, job, result, logger) } } } -// transitionOrphan performs a CAS status transition on the orphaned job. -// Returns true if the transition succeeded (or dry-run logged it). -func (r *Reconciler) transitionOrphan( - ctx context.Context, - job *db.BatchItem, - currentStatus *openai.BatchStatusInfo, - newStatus openai.BatchStatus, - result *Result, - logger logr.Logger, -) bool { - updatedStatus, err := batch_utils.BuildUpdatedStatusInfo(currentStatus, newStatus, nil, nil) +// transitionOrphan transitions an orphaned job to the given terminal status. +func (r *Reconciler) transitionOrphan(ctx context.Context, job *db.BatchItem, statusInfo *openai.BatchStatusInfo, target openai.BatchStatus, result *Result, logger logr.Logger) { + updatedStatus, err := batch_utils.BuildUpdatedStatusInfo(statusInfo, target, nil, nil) if err != nil { - logger.Error(err, "Reconciler: failed to build updated status", "newStatus", newStatus) + logger.Error(err, "Reconciler: failed to build target status", "target", target) result.Errors++ - return false + return } updatedBytes, err := json.Marshal(updatedStatus) if err != nil { - logger.Error(err, "Reconciler: failed to marshal updated status") + logger.Error(err, "Reconciler: failed to marshal target status", "target", target) result.Errors++ - return false + return } - if !r.dryRun { - updateItem := &db.BatchItem{ - BaseIndexes: db.BaseIndexes{ID: job.ID}, - BaseContents: db.BaseContents{Status: updatedBytes}, - } - if err := r.batchDB.DBUpdate(ctx, updateItem, job.Status); err != nil { - if errors.Is(err, db.ErrConflict) { - logger.Info("Reconciler: CAS conflict during orphan transition (another actor won the race)", "newStatus", newStatus) - result.Conflicts++ - } else { - logger.Error(err, "Reconciler: failed to transition orphan", "newStatus", newStatus) - result.Errors++ - } - return false - } - logger.Info("Reconciler: orphan transitioned", "from", currentStatus.Status, "to", newStatus) - } else { - logger.Info("Reconciler: dry-run: would transition orphan", "from", currentStatus.Status, "to", newStatus) + if r.dryRun { + logger.Info("Reconciler: dry-run: would transition orphan", "target", target) + result.Expired++ + return } - switch newStatus { - case openai.BatchStatusCancelled: - result.Cancelled++ - case openai.BatchStatusExpired: - result.Expired++ - case openai.BatchStatusFailed: - result.Failed++ + updateItem := &db.BatchItem{ + BaseIndexes: db.BaseIndexes{ID: job.ID}, + BaseContents: db.BaseContents{Status: updatedBytes}, + Epoch: job.Epoch, + BumpEpoch: true, } - return true + if err := r.batchDB.DBUpdate(ctx, updateItem, job.Status); err != nil { + if errors.Is(err, db.ErrConflict) { + logger.Info("Reconciler: CAS conflict during orphan transition", "target", target) + result.Conflicts++ + } else { + logger.Error(err, "Reconciler: failed to transition orphan", "target", target) + result.Errors++ + } + return + } + + logger.Info("Reconciler: orphan transitioned", "target", target) + result.Expired++ } -// reEnqueueOrphan re-enqueues an orphaned validating job with its original SLO. -// Returns true if the re-enqueue succeeded (or dry-run logged it). -func (r *Reconciler) reEnqueueOrphan( - ctx context.Context, - job *db.BatchItem, - result *Result, - logger logr.Logger, -) bool { - slo, err := extractSLO(job) - if err != nil { - logger.Error(err, "Reconciler: cannot re-enqueue orphan with corrupt SLO") +// reEnqueueOrphan re-enqueues an orphaned job whose SLO is still valid. +func (r *Reconciler) reEnqueueOrphan(ctx context.Context, job *db.BatchItem, result *Result, logger logr.Logger) { + if job.Priority <= 0 { + logger.Error(fmt.Errorf("missing priority"), "Reconciler: cannot re-enqueue orphan without priority") result.Errors++ - return false - } - if slo == nil { - logger.Error(fmt.Errorf("missing SLO tag"), "Reconciler: cannot re-enqueue orphan without SLO") - result.Errors++ - return false + return } + slo := time.UnixMicro(job.Priority) + if r.dryRun { logger.Info("Reconciler: dry-run: would re-enqueue orphan", "slo", slo) result.ReEnqueued++ - return true + return } task := &db.BatchJobPriority{ ID: job.ID, - SLO: *slo, + SLO: slo, } if err := r.queue.PQEnqueue(ctx, task); err != nil { logger.Error(err, "Reconciler: failed to re-enqueue orphan") result.Errors++ - return false + return } logger.Info("Reconciler: orphan re-enqueued", "slo", slo) result.ReEnqueued++ - return true -} - -// cleanupStaleInflight removes in-flight entries for jobs that are no longer -// in the non-terminal set (already completed, failed, or deleted from DB). -func (r *Reconciler) cleanupStaleInflight( - ctx context.Context, - inflightEntries map[string]*db.InFlightEntry, - nonTerminalIDs map[string]bool, - result *Result, -) { - logger := logr.FromContextOrDiscard(ctx) - - for jobID := range inflightEntries { - if nonTerminalIDs[jobID] { - continue - } - if r.dryRun { - logger.Info("Reconciler: dry-run: would clean up stale in-flight entry", "jobId", jobID) - result.StaleCleanup++ - continue - } - if err := r.inflight.InFlightDelete(ctx, jobID); err != nil { - logger.Error(err, "Reconciler: failed to clean up stale in-flight entry", "jobId", jobID) - result.Errors++ - continue - } - logger.Info("Reconciler: cleaned up stale in-flight entry", "jobId", jobID) - result.StaleCleanup++ - } } // isSLOExpired checks whether the job's SLO deadline has passed. -// Returns false if the SLO tag is missing or corrupt (caller should check extractSLO separately). func isSLOExpired(job *db.BatchItem) bool { - slo, _ := extractSLO(job) - if slo == nil { + if job.Priority <= 0 { return false } - return time.Now().After(*slo) -} - -// extractSLO parses the SLO tag from the job's tags. -// Returns (nil, nil) if the tag is missing, or (nil, error) if the tag value is corrupt. -func extractSLO(job *db.BatchItem) (*time.Time, error) { - sloStr, ok := job.Tags[batch_types.TagSLO] - if !ok { - return nil, nil - } - sloMicro, err := strconv.ParseInt(sloStr, 10, 64) - if err != nil { - return nil, fmt.Errorf("corrupt SLO tag %q: %w", sloStr, err) - } - slo := time.UnixMicro(sloMicro).UTC() - return &slo, nil + return time.Now().After(time.UnixMicro(job.Priority)) } diff --git a/internal/gc/reconciler/reconciler_test.go b/internal/gc/reconciler/reconciler_test.go index c89b09458..f1f9bf558 100644 --- a/internal/gc/reconciler/reconciler_test.go +++ b/internal/gc/reconciler/reconciler_test.go @@ -19,6 +19,7 @@ package reconciler import ( "context" "encoding/json" + "errors" "fmt" "testing" "time" @@ -26,33 +27,29 @@ import ( db "github.com/llm-d/llm-d-batch-gateway/internal/database/api" "github.com/llm-d/llm-d-batch-gateway/internal/database/mock" "github.com/llm-d/llm-d-batch-gateway/internal/shared/openai" - batch_types "github.com/llm-d/llm-d-batch-gateway/internal/shared/types" ) const testInterval = 60 * time.Minute -func sloTag(slo time.Time) db.Tags { - return db.Tags{batch_types.TagSLO: fmt.Sprintf("%d", slo.UnixMicro())} +func futureSLO() int64 { + return time.Now().Add(24 * time.Hour).UnixMicro() } -func futureSLO() time.Time { - return time.Now().Add(24 * time.Hour) +func expiredSLO() int64 { + return time.Now().Add(-1 * time.Hour).UnixMicro() } -func expiredSLO() time.Time { - return time.Now().Add(-1 * time.Hour) -} - -func newTestBatchItem(id string, status openai.BatchStatus, tags db.Tags) *db.BatchItem { +func newTestBatchItem(id, processorID string, status openai.BatchStatus, priority int64) *db.BatchItem { statusBytes, _ := json.Marshal(openai.BatchStatusInfo{Status: status}) return &db.BatchItem{ BaseIndexes: db.BaseIndexes{ - ID: id, - Tags: tags, + ID: id, }, BaseContents: db.BaseContents{ Status: statusBytes, }, + ProcessorID: processorID, + Priority: priority, } } @@ -60,11 +57,10 @@ func newTestReconciler( t *testing.T, batchDB db.BatchDBClient, queue db.BatchPriorityQueueClient, - inflight db.InFlightClient, ) (*Reconciler, chan *Result) { t.Helper() resultCh := make(chan *Result, 1) - r, err := NewReconciler(batchDB, queue, inflight, testInterval, false, func(res *Result) { + r, err := NewReconciler(batchDB, queue, testInterval, false, func(res *Result) { resultCh <- res }) if err != nil { @@ -73,6 +69,22 @@ func newTestReconciler( return r, resultCh } +func newTestDryRunReconciler( + t *testing.T, + batchDB db.BatchDBClient, + queue db.BatchPriorityQueueClient, +) (*Reconciler, chan *Result) { + t.Helper() + resultCh := make(chan *Result, 1) + r, err := NewReconciler(batchDB, queue, testInterval, true, func(res *Result) { + resultCh <- res + }) + if err != nil { + t.Fatalf("failed to create dry-run reconciler: %v", err) + } + return r, resultCh +} + func storeItems(t *testing.T, batchDB db.BatchDBClient, items ...*db.BatchItem) { t.Helper() ctx := context.Background() @@ -83,256 +95,361 @@ func storeItems(t *testing.T, batchDB db.BatchDBClient, items ...*db.BatchItem) } } -func TestTriageOrphan(t *testing.T) { - ctx := context.Background() - - t.Run("cancelling transitions to cancelled", func(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - - item := newTestBatchItem("job-1", openai.BatchStatusCancelling, sloTag(futureSLO())) - storeItems(t, batchDB, item) - - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) +func assertJobStatus(t *testing.T, batchDB db.BatchDBClient, jobID string, expected openai.BatchStatus) { + t.Helper() + items, _, _, err := batchDB.DBGet(context.Background(), + &db.BatchQuery{BaseQuery: db.BaseQuery{IDs: []string{jobID}}}, false, 0, 10) + if err != nil { + t.Fatalf("failed to get job %s: %v", jobID, err) + } + if len(items) != 1 { + t.Fatalf("expected 1 item for %s, got %d", jobID, len(items)) + } + var info openai.BatchStatusInfo + if err := json.Unmarshal(items[0].Status, &info); err != nil { + t.Fatalf("failed to unmarshal status: %v", err) + } + if info.Status != expected { + t.Errorf("expected status %s for job %s, got %s", expected, jobID, info.Status) + } +} - result := <-resultCh - if result.Cancelled != 1 { - t.Errorf("expected 1 cancelled, got %d", result.Cancelled) +// newMockBatchDB creates a mock batch DB that filters by HasProcessorID when set. +func newMockBatchDB() *mock.MockDBClient[db.BatchItem, db.BatchQuery] { + m := mock.NewMockDBClient[db.BatchItem, db.BatchQuery]( + func(item *db.BatchItem) string { return item.ID }, + func(query *db.BatchQuery) *db.BaseQuery { return &query.BaseQuery }, + ) + m.QueryFilter = func(item *db.BatchItem, query *db.BatchQuery) bool { + if query.HasProcessorID && item.ProcessorID == "" { + return false } - assertJobStatus(t, batchDB, "job-1", openai.BatchStatusCancelled) - }) + return true + } + return m +} - t.Run("validating with expired SLO transitions to expired", func(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() +// casConflictBatchDB is a minimal mock that always returns ErrConflict on DBUpdate. +type casConflictBatchDB struct{} - item := newTestBatchItem("job-1", openai.BatchStatusValidating, sloTag(expiredSLO())) - storeItems(t, batchDB, item) +func (c *casConflictBatchDB) DBStore(_ context.Context, _ *db.BatchItem) error { return nil } +func (c *casConflictBatchDB) DBGet(_ context.Context, _ *db.BatchQuery, _ bool, _, _ int) ([]*db.BatchItem, int, bool, error) { + item := newTestBatchItem("job-cas", "dead-processor", openai.BatchStatusInProgress, expiredSLO()) + return []*db.BatchItem{item}, 1, false, nil +} +func (c *casConflictBatchDB) DBUpdate(_ context.Context, _ *db.BatchItem, _ []byte) error { + return fmt.Errorf("DBUpdate: %w", db.ErrConflict) +} +func (c *casConflictBatchDB) DBDelete(_ context.Context, _ []string) ([]string, error) { + return nil, nil +} +func (c *casConflictBatchDB) Close() error { return nil } +func (c *casConflictBatchDB) GetContext(_ context.Context, _ time.Duration) (context.Context, context.CancelFunc) { + return context.Background(), func() {} +} - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) +func TestTriageOrphan(t *testing.T) { + ctx := context.Background() - result := <-resultCh - if result.Expired != 1 { - t.Errorf("expected 1 expired, got %d", result.Expired) - } - assertJobStatus(t, batchDB, "job-1", openai.BatchStatusExpired) - }) + tests := []struct { + name string + status openai.BatchStatus + priority int64 + wantExpired int + wantReEnqueued int + wantFinalStatus openai.BatchStatus + checkQueue bool + wantInQueue bool + }{ + { + name: "orphan with expired SLO transitions to failed", + status: openai.BatchStatusValidating, + priority: expiredSLO(), + wantExpired: 1, + wantReEnqueued: 0, + wantFinalStatus: openai.BatchStatusFailed, + }, + { + name: "orphan with future SLO is re-enqueued", + status: openai.BatchStatusValidating, + priority: futureSLO(), + wantExpired: 0, + wantReEnqueued: 1, + wantFinalStatus: openai.BatchStatusValidating, // status unchanged on re-enqueue + checkQueue: true, + wantInQueue: true, + }, + { + name: "in_progress orphan with expired SLO transitions to failed", + status: openai.BatchStatusInProgress, + priority: expiredSLO(), + wantExpired: 1, + wantReEnqueued: 0, + wantFinalStatus: openai.BatchStatusFailed, + }, + { + name: "in_progress orphan with future SLO is re-enqueued", + status: openai.BatchStatusInProgress, + priority: futureSLO(), + wantExpired: 0, + wantReEnqueued: 1, + wantFinalStatus: openai.BatchStatusInProgress, + checkQueue: true, + wantInQueue: true, + }, + { + name: "cancelling orphan transitions to cancelled regardless of SLO", + status: openai.BatchStatusCancelling, + priority: futureSLO(), + wantExpired: 1, + wantReEnqueued: 0, + wantFinalStatus: openai.BatchStatusCancelled, + }, + { + name: "cancelling orphan with expired SLO transitions to cancelled", + status: openai.BatchStatusCancelling, + priority: expiredSLO(), + wantExpired: 1, + wantReEnqueued: 0, + wantFinalStatus: openai.BatchStatusCancelled, + }, + } - t.Run("validating with future SLO is re-enqueued", func(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + batchDB := newMockBatchDB() + queue := mock.NewMockBatchPriorityQueueClient() - item := newTestBatchItem("job-1", openai.BatchStatusValidating, sloTag(futureSLO())) - storeItems(t, batchDB, item) + item := newTestBatchItem("job-1", "dead-processor", tc.status, tc.priority) + storeItems(t, batchDB, item) - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) + r, resultCh := newTestReconciler(t, batchDB, queue) + // Mark "dead-processor" as not alive by not including it in live set. + r.SetLiveProcessors(map[string]bool{}) + r.run(ctx) - result := <-resultCh - if result.ReEnqueued != 1 { - t.Errorf("expected 1 re-enqueued, got %d", result.ReEnqueued) - } + result := <-resultCh + if result.Expired != tc.wantExpired { + t.Errorf("expected %d expired, got %d", tc.wantExpired, result.Expired) + } + if result.ReEnqueued != tc.wantReEnqueued { + t.Errorf("expected %d re-enqueued, got %d", tc.wantReEnqueued, result.ReEnqueued) + } - queuedIDs, _ := queue.PQGetIDs(ctx) - if !queuedIDs["job-1"] { - t.Error("expected job-1 to be in queue after re-enqueue") - } - }) + assertJobStatus(t, batchDB, "job-1", tc.wantFinalStatus) - t.Run("in_progress with expired SLO transitions to expired", func(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() + if tc.checkQueue { + queuedIDs, _ := queue.PQGetIDs(ctx) + if tc.wantInQueue && !queuedIDs["job-1"] { + t.Error("expected job-1 to be in queue after re-enqueue") + } + if !tc.wantInQueue && queuedIDs["job-1"] { + t.Error("expected job-1 NOT to be in queue") + } + } + }) + } +} - item := newTestBatchItem("job-1", openai.BatchStatusInProgress, sloTag(expiredSLO())) - storeItems(t, batchDB, item) +func TestSkipNonOrphans(t *testing.T) { + ctx := context.Background() - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) + tests := []struct { + name string + processorID string + liveSet map[string]bool + }{ + { + name: "job with no processor_id is skipped", + processorID: "", + liveSet: map[string]bool{}, + }, + { + name: "job with alive processor is skipped", + processorID: "alive-processor", + liveSet: map[string]bool{"alive-processor": true}, + }, + } - result := <-resultCh - if result.Expired != 1 { - t.Errorf("expected 1 expired, got %d", result.Expired) - } - assertJobStatus(t, batchDB, "job-1", openai.BatchStatusExpired) - }) + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + batchDB := newMockBatchDB() + queue := mock.NewMockBatchPriorityQueueClient() - t.Run("in_progress with future SLO transitions to failed", func(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() + item := newTestBatchItem("job-1", tc.processorID, openai.BatchStatusValidating, futureSLO()) + storeItems(t, batchDB, item) - item := newTestBatchItem("job-1", openai.BatchStatusInProgress, sloTag(futureSLO())) - storeItems(t, batchDB, item) + r, resultCh := newTestReconciler(t, batchDB, queue) + r.SetLiveProcessors(tc.liveSet) + r.run(ctx) - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) + result := <-resultCh + if result.Expired != 0 || result.ReEnqueued != 0 || result.Conflicts != 0 || result.Errors != 0 { + t.Errorf("expected no actions for non-orphan job, got %+v", result) + } + }) + } +} - result := <-resultCh - if result.Failed != 1 { - t.Errorf("expected 1 failed, got %d", result.Failed) - } - assertJobStatus(t, batchDB, "job-1", openai.BatchStatusFailed) - }) +func TestRunCycleMixedJobs(t *testing.T) { + ctx := context.Background() - t.Run("finalizing with expired SLO transitions to expired", func(t *testing.T) { + t.Run("only dead-processor jobs are triaged", func(t *testing.T) { batchDB := newMockBatchDB() queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - item := newTestBatchItem("job-1", openai.BatchStatusFinalizing, sloTag(expiredSLO())) - storeItems(t, batchDB, item) - - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) + storeItems(t, batchDB, + // Queued (no processor_id) — should be skipped. + newTestBatchItem("queued-job", "", openai.BatchStatusValidating, futureSLO()), + // Alive processor — should be skipped. + newTestBatchItem("alive-job", "processor-0", openai.BatchStatusInProgress, futureSLO()), + // Dead processor, valid SLO — should be re-enqueued. + newTestBatchItem("dead-valid", "processor-1", openai.BatchStatusInProgress, futureSLO()), + // Dead processor, expired SLO — should be failed. + newTestBatchItem("dead-expired", "processor-2", openai.BatchStatusInProgress, expiredSLO()), + ) + + r, resultCh := newTestReconciler(t, batchDB, queue) + r.SetLiveProcessors(map[string]bool{"processor-0": true}) r.run(ctx) result := <-resultCh + if result.ReEnqueued != 1 { + t.Errorf("expected 1 re-enqueued, got %d", result.ReEnqueued) + } if result.Expired != 1 { t.Errorf("expected 1 expired, got %d", result.Expired) } - }) - - t.Run("finalizing with future SLO transitions to failed", func(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - - item := newTestBatchItem("job-1", openai.BatchStatusFinalizing, sloTag(futureSLO())) - storeItems(t, batchDB, item) - - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) - - result := <-resultCh - if result.Failed != 1 { - t.Errorf("expected 1 failed, got %d", result.Failed) + if result.Errors != 0 { + t.Errorf("expected 0 errors, got %d", result.Errors) } + + // Verify the queued and alive jobs were not touched. + assertJobStatus(t, batchDB, "queued-job", openai.BatchStatusValidating) + assertJobStatus(t, batchDB, "alive-job", openai.BatchStatusInProgress) }) } -func TestSkipNonOrphans(t *testing.T) { +func TestCASConflict(t *testing.T) { ctx := context.Background() - t.Run("job in queue is not treated as orphan", func(t *testing.T) { - batchDB := newMockBatchDB() + t.Run("CAS conflict is counted as conflict not error", func(t *testing.T) { + batchDB := &casConflictBatchDB{} queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - - slo := futureSLO() - item := newTestBatchItem("job-1", openai.BatchStatusValidating, sloTag(slo)) - storeItems(t, batchDB, item) - _ = queue.PQEnqueue(ctx, &db.BatchJobPriority{ID: "job-1", SLO: slo}) - - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) - - result := <-resultCh - if result.ReEnqueued != 0 || result.Failed != 0 || result.Expired != 0 || result.Cancelled != 0 { - t.Errorf("expected no actions for queued job, got %+v", result) + resultCh := make(chan *Result, 1) + r, err := NewReconciler(batchDB, queue, testInterval, false, func(res *Result) { + resultCh <- res + }) + if err != nil { + t.Fatalf("failed to create reconciler: %v", err) } - }) - - t.Run("job with fresh in-flight entry is not treated as orphan", func(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - - item := newTestBatchItem("job-1", openai.BatchStatusInProgress, sloTag(futureSLO())) - storeItems(t, batchDB, item) - - _ = inflight.InFlightSet(ctx, "job-1", "processor-1") - - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) + // Ensure "dead-processor" is not in live set so the job is orphaned. + r.SetLiveProcessors(map[string]bool{}) r.run(ctx) result := <-resultCh - if result.Failed != 0 { - t.Errorf("expected no failures for fresh in-flight job, got %d", result.Failed) + if result.Conflicts != 1 { + t.Errorf("expected 1 conflict from CAS, got %d", result.Conflicts) + } + if result.Errors != 0 { + t.Errorf("expected 0 errors (CAS is a conflict, not an error), got %d", result.Errors) + } + if result.Expired != 0 { + t.Errorf("expected 0 expired (CAS failed), got %d", result.Expired) } }) } -func TestStaleInflightCleanup(t *testing.T) { +func TestEpochFencing(t *testing.T) { ctx := context.Background() - t.Run("removes in-flight entry for job not in non-terminal set", func(t *testing.T) { + t.Run("zombie write fails after GC bumps epoch", func(t *testing.T) { + // Simulate: processor owns job at epoch=3, GC transitions to failed + // (bumping epoch to 4), then zombie tries to write with epoch=3. batchDB := newMockBatchDB() queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - // Stale in-flight entry for a job that no longer appears in the - // non-terminal query (e.g. it already reached a terminal state - // or was deleted, but its in-flight entry was not cleaned up). - _ = inflight.InFlightSet(ctx, "job-stale", "processor-1") + item := newTestBatchItem("job-epoch", "dead-processor", openai.BatchStatusInProgress, expiredSLO()) + item.Epoch = 3 + storeItems(t, batchDB, item) - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) + r, resultCh := newTestReconciler(t, batchDB, queue) + r.SetLiveProcessors(map[string]bool{}) r.run(ctx) result := <-resultCh - if result.StaleCleanup != 1 { - t.Errorf("expected 1 stale cleanup, got %d", result.StaleCleanup) + if result.Expired != 1 { + t.Fatalf("expected 1 expired, got %d", result.Expired) } - entries, _ := inflight.InFlightGetAll(ctx) - if _, ok := entries["job-stale"]; ok { - t.Error("expected stale in-flight entry to be removed") - } + // Verify GC set BumpEpoch — the mock doesn't actually increment, + // but in production the epoch would be bumped. Verify the job + // was transitioned to failed. + assertJobStatus(t, batchDB, "job-epoch", openai.BatchStatusFailed) }) - t.Run("preserves in-flight entry for non-terminal job", func(t *testing.T) { + t.Run("processor write with matching epoch succeeds", func(t *testing.T) { batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - item := newTestBatchItem("job-active", openai.BatchStatusInProgress, sloTag(futureSLO())) + statusBytes, _ := json.Marshal(openai.BatchStatusInfo{Status: openai.BatchStatusInProgress}) + item := &db.BatchItem{ + BaseIndexes: db.BaseIndexes{ID: "job-match"}, + BaseContents: db.BaseContents{Status: statusBytes}, + ProcessorID: "processor-0", + Epoch: 5, + } storeItems(t, batchDB, item) - // Fresh in-flight entry — should NOT be cleaned up (it's non-terminal - // and recently seen, so it's treated as actively processing). - _ = inflight.InFlightSet(ctx, "job-active", "processor-1") - - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) - - result := <-resultCh - if result.StaleCleanup != 0 { - t.Errorf("expected 0 stale cleanup, got %d", result.StaleCleanup) + // Update with matching epoch should succeed. + newStatus, _ := json.Marshal(openai.BatchStatusInfo{Status: openai.BatchStatusFinalizing}) + err := batchDB.DBUpdate(ctx, &db.BatchItem{ + BaseIndexes: db.BaseIndexes{ID: "job-match"}, + BaseContents: db.BaseContents{Status: newStatus}, + Epoch: 5, + }, nil) + if err != nil { + t.Fatalf("expected epoch-matched update to succeed, got %v", err) } - entries, _ := inflight.InFlightGetAll(ctx) - if _, ok := entries["job-active"]; !ok { - t.Error("expected in-flight entry to be preserved for non-terminal job") - } + assertJobStatus(t, batchDB, "job-match", openai.BatchStatusFinalizing) }) -} - -func TestCASConflict(t *testing.T) { - ctx := context.Background() - - t.Run("CAS conflict is counted as conflict not error", func(t *testing.T) { - batchDB := &casConflictBatchDB{} - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) + t.Run("processor write with stale epoch fails", func(t *testing.T) { + batchDB := newMockBatchDB() - result := <-resultCh - if result.Conflicts != 1 { - t.Errorf("expected 1 conflict from CAS, got %d", result.Conflicts) + statusBytes, _ := json.Marshal(openai.BatchStatusInfo{Status: openai.BatchStatusInProgress}) + item := &db.BatchItem{ + BaseIndexes: db.BaseIndexes{ID: "job-stale"}, + BaseContents: db.BaseContents{Status: statusBytes}, + ProcessorID: "processor-0", + Epoch: 5, } - if result.Errors != 0 { - t.Errorf("expected 0 errors (CAS is a conflict, not an error), got %d", result.Errors) + storeItems(t, batchDB, item) + + // Simulate GC bumping epoch to 6 by directly updating the stored item. + items, _, _, _ := batchDB.DBGet(ctx, + &db.BatchQuery{BaseQuery: db.BaseQuery{IDs: []string{"job-stale"}}}, + true, 0, 1) + items[0].Epoch = 6 + _ = batchDB.DBUpdate(ctx, items[0], nil) + + // Zombie write with stale epoch=5 should fail. + newStatus, _ := json.Marshal(openai.BatchStatusInfo{Status: openai.BatchStatusCompleted}) + err := batchDB.DBUpdate(ctx, &db.BatchItem{ + BaseIndexes: db.BaseIndexes{ID: "job-stale"}, + BaseContents: db.BaseContents{Status: newStatus}, + Epoch: 5, + }, nil) + if err == nil { + t.Fatal("expected epoch-mismatched update to fail") } - if result.Cancelled != 0 { - t.Errorf("expected 0 cancelled (CAS failed), got %d", result.Cancelled) + if !errors.Is(err, db.ErrConflict) { + t.Fatalf("expected ErrConflict, got %v", err) } + + // Job should still be in_progress (zombie write rejected). + assertJobStatus(t, batchDB, "job-stale", openai.BatchStatusInProgress) }) } @@ -340,10 +457,9 @@ func TestRunLoop(t *testing.T) { t.Run("stops on context cancel", func(t *testing.T) { batchDB := newMockBatchDB() queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() ran := make(chan struct{}, 1) - r, err := NewReconciler(batchDB, queue, inflight, testInterval, false, func(*Result) { + r, err := NewReconciler(batchDB, queue, testInterval, false, func(*Result) { select { case ran <- struct{}{}: default: @@ -358,6 +474,10 @@ func TestRunLoop(t *testing.T) { done := make(chan error, 1) go func() { done <- r.RunLoop(ctx) }() + // The first cycle is deferred until Trigger() is called (normally + // by the pod watcher after cache sync). + r.Trigger() + <-ran cancel() @@ -367,181 +487,120 @@ func TestRunLoop(t *testing.T) { }) } -func TestTriageEdgeCases(t *testing.T) { - ctx := context.Background() - - t.Run("validating orphan without SLO tag errors", func(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - - item := newTestBatchItem("job-1", openai.BatchStatusValidating, db.Tags{}) - storeItems(t, batchDB, item) - - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) - - result := <-resultCh - if result.Errors != 1 { - t.Errorf("expected 1 error for missing SLO, got %d", result.Errors) - } - if result.ReEnqueued != 0 { - t.Errorf("expected 0 re-enqueued, got %d", result.ReEnqueued) - } - }) - - t.Run("validating orphan with corrupt SLO tag errors", func(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() +func TestNewReconcilerValidation(t *testing.T) { + batchDB := newMockBatchDB() + queue := mock.NewMockBatchPriorityQueueClient() - item := newTestBatchItem("job-1", openai.BatchStatusValidating, db.Tags{batch_types.TagSLO: "not-a-number"}) - storeItems(t, batchDB, item) + tests := []struct { + name string + batchDB db.BatchDBClient + queue db.BatchPriorityQueueClient + interval time.Duration + }{ + { + name: "nil batchDB", + batchDB: nil, + queue: queue, + interval: testInterval, + }, + { + name: "nil queue", + batchDB: batchDB, + queue: nil, + interval: testInterval, + }, + { + name: "zero interval", + batchDB: batchDB, + queue: queue, + interval: 0, + }, + { + name: "negative interval", + batchDB: batchDB, + queue: queue, + interval: -time.Minute, + }, + } - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, err := NewReconciler(tc.batchDB, tc.queue, tc.interval, false, nil) + if err == nil { + t.Fatalf("expected error for %s", tc.name) + } + }) + } +} - result := <-resultCh - if result.Errors != 1 { - t.Errorf("expected 1 error for corrupt SLO, got %d", result.Errors) - } - if result.ReEnqueued != 0 { - t.Errorf("expected 0 re-enqueued, got %d", result.ReEnqueued) - } - }) +func TestRunCycle_DBFailure(t *testing.T) { + ctx := context.Background() - t.Run("malformed status JSON errors", func(t *testing.T) { - batchDB := newMockBatchDB() + t.Run("DB error is counted and cycle returns early", func(t *testing.T) { queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - - item := &db.BatchItem{ - BaseIndexes: db.BaseIndexes{ID: "job-1", Tags: sloTag(futureSLO())}, - BaseContents: db.BaseContents{Status: []byte(`{{invalid json`)}, - } - storeItems(t, batchDB, item) + failDB := &failGetBatchDB{err: fmt.Errorf("connection refused")} - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) + r, resultCh := newTestReconciler(t, failDB, queue) + r.SetLiveProcessors(map[string]bool{}) r.run(ctx) result := <-resultCh if result.Errors != 1 { - t.Errorf("expected 1 error for malformed status, got %d", result.Errors) + t.Errorf("expected 1 error, got %d", result.Errors) } - }) - - t.Run("stale in-flight entry triggers triage", func(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - - item := newTestBatchItem("job-1", openai.BatchStatusInProgress, sloTag(futureSLO())) - storeItems(t, batchDB, item) - - _ = inflight.InFlightSet(ctx, "job-1", "processor-1") - // Backdate the LastSeen to make it stale (older than the reconciler interval). - staleTime := time.Now().Add(-2 * testInterval).Unix() - inflight.SetLastSeen("job-1", staleTime) - - r, resultCh := newTestReconciler(t, batchDB, queue, inflight) - r.run(ctx) - - result := <-resultCh - if result.Failed != 1 { - t.Errorf("expected 1 failed for stale in-flight job, got %d", result.Failed) + if result.Expired != 0 || result.ReEnqueued != 0 { + t.Errorf("expected no actions on DB failure, got expired=%d reEnqueued=%d", result.Expired, result.ReEnqueued) } - assertJobStatus(t, batchDB, "job-1", openai.BatchStatusFailed) }) } -func TestNewReconcilerValidation(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - - t.Run("nil batchDB", func(t *testing.T) { - _, err := NewReconciler(nil, queue, inflight, testInterval, false, nil) - if err == nil { - t.Fatal("expected error for nil batchDB") - } - }) - - t.Run("nil queue", func(t *testing.T) { - _, err := NewReconciler(batchDB, nil, inflight, testInterval, false, nil) - if err == nil { - t.Fatal("expected error for nil queue") - } - }) - - t.Run("nil inflight", func(t *testing.T) { - _, err := NewReconciler(batchDB, queue, nil, testInterval, false, nil) - if err == nil { - t.Fatal("expected error for nil inflight") - } - }) - - t.Run("zero interval", func(t *testing.T) { - _, err := NewReconciler(batchDB, queue, inflight, 0, false, nil) - if err == nil { - t.Fatal("expected error for zero interval") - } - }) - - t.Run("negative interval", func(t *testing.T) { - _, err := NewReconciler(batchDB, queue, inflight, -time.Minute, false, nil) - if err == nil { - t.Fatal("expected error for negative interval") - } - }) +// failGetBatchDB always returns an error on DBGet. +type failGetBatchDB struct { + err error } -func newTestDryRunReconciler( - t *testing.T, - batchDB db.BatchDBClient, - queue db.BatchPriorityQueueClient, - inflight db.InFlightClient, -) (*Reconciler, chan *Result) { - t.Helper() - resultCh := make(chan *Result, 1) - r, err := NewReconciler(batchDB, queue, inflight, testInterval, true, func(res *Result) { - resultCh <- res - }) - if err != nil { - t.Fatalf("failed to create dry-run reconciler: %v", err) - } - return r, resultCh +func (f *failGetBatchDB) DBStore(_ context.Context, _ *db.BatchItem) error { return nil } +func (f *failGetBatchDB) DBGet(_ context.Context, _ *db.BatchQuery, _ bool, _, _ int) ([]*db.BatchItem, int, bool, error) { + return nil, 0, false, f.err +} +func (f *failGetBatchDB) DBUpdate(_ context.Context, _ *db.BatchItem, _ []byte) error { return nil } +func (f *failGetBatchDB) DBDelete(_ context.Context, _ []string) ([]string, error) { return nil, nil } +func (f *failGetBatchDB) Close() error { return nil } +func (f *failGetBatchDB) GetContext(_ context.Context, _ time.Duration) (context.Context, context.CancelFunc) { + return context.Background(), func() {} } func TestDryRun(t *testing.T) { ctx := context.Background() - t.Run("transition is counted but DB is not mutated", func(t *testing.T) { + t.Run("expire is counted but DB is not mutated", func(t *testing.T) { batchDB := newMockBatchDB() queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - item := newTestBatchItem("job-1", openai.BatchStatusCancelling, sloTag(futureSLO())) + item := newTestBatchItem("job-1", "dead-processor", openai.BatchStatusInProgress, expiredSLO()) storeItems(t, batchDB, item) - r, resultCh := newTestDryRunReconciler(t, batchDB, queue, inflight) + r, resultCh := newTestDryRunReconciler(t, batchDB, queue) + r.SetLiveProcessors(map[string]bool{}) r.run(ctx) result := <-resultCh - if result.Cancelled != 1 { - t.Errorf("expected 1 cancelled, got %d", result.Cancelled) + if result.Expired != 1 { + t.Errorf("expected 1 expired, got %d", result.Expired) } - assertJobStatus(t, batchDB, "job-1", openai.BatchStatusCancelling) + // In dry-run mode the DB status should remain unchanged. + assertJobStatus(t, batchDB, "job-1", openai.BatchStatusInProgress) }) t.Run("re-enqueue is counted but queue is not mutated", func(t *testing.T) { batchDB := newMockBatchDB() queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() - item := newTestBatchItem("job-1", openai.BatchStatusValidating, sloTag(futureSLO())) + item := newTestBatchItem("job-1", "dead-processor", openai.BatchStatusValidating, futureSLO()) storeItems(t, batchDB, item) - r, resultCh := newTestDryRunReconciler(t, batchDB, queue, inflight) + r, resultCh := newTestDryRunReconciler(t, batchDB, queue) + r.SetLiveProcessors(map[string]bool{}) r.run(ctx) result := <-resultCh @@ -554,101 +613,103 @@ func TestDryRun(t *testing.T) { t.Error("expected job-1 NOT to be in queue in dry-run mode") } }) +} - t.Run("stale cleanup is counted but in-flight entry is preserved", func(t *testing.T) { - batchDB := newMockBatchDB() - queue := mock.NewMockBatchPriorityQueueClient() - inflight := mock.NewMockInFlightClient() +func TestSetLiveProcessors(t *testing.T) { + tests := []struct { + name string + liveSet map[string]bool + processor string + wantAlive bool + }{ + { + name: "processor in live set is alive", + liveSet: map[string]bool{"proc-1": true, "proc-2": true}, + processor: "proc-1", + wantAlive: true, + }, + { + name: "processor not in live set is not alive", + liveSet: map[string]bool{"proc-1": true}, + processor: "proc-2", + wantAlive: false, + }, + { + name: "empty live set means no processor is alive", + liveSet: map[string]bool{}, + processor: "proc-1", + wantAlive: false, + }, + { + name: "updating live set replaces previous set", + liveSet: map[string]bool{"proc-new": true}, + processor: "proc-old", + wantAlive: false, + }, + } - _ = inflight.InFlightSet(ctx, "job-stale", "processor-1") + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + batchDB := newMockBatchDB() + queue := mock.NewMockBatchPriorityQueueClient() - r, resultCh := newTestDryRunReconciler(t, batchDB, queue, inflight) - r.run(ctx) + r, _ := newTestReconciler(t, batchDB, queue) - result := <-resultCh - if result.StaleCleanup != 1 { - t.Errorf("expected 1 stale cleanup, got %d", result.StaleCleanup) - } + // Set an initial live set then overwrite it. + r.SetLiveProcessors(map[string]bool{"proc-old": true}) + r.SetLiveProcessors(tc.liveSet) - entries, _ := inflight.InFlightGetAll(ctx) - if _, ok := entries["job-stale"]; !ok { - t.Error("expected stale in-flight entry to be preserved in dry-run mode") - } - }) + got := r.isProcessorAlive(tc.processor) + if got != tc.wantAlive { + t.Errorf("isProcessorAlive(%q) = %v, want %v", tc.processor, got, tc.wantAlive) + } + }) + } } -func TestTerminalStatusesSync(t *testing.T) { - // Verify that every status returned by TerminalStatuses() is actually final, - // and every final status is included in TerminalStatuses(). - terminalSet := make(map[openai.BatchStatus]bool) - for _, s := range openai.TerminalStatuses() { - if !s.IsTerminal() { - t.Errorf("TerminalStatuses() contains %q which is not IsTerminal()", s) - } - terminalSet[s] = true - } +func TestTrigger(t *testing.T) { + t.Run("trigger causes immediate reconciliation", func(t *testing.T) { + batchDB := newMockBatchDB() + queue := mock.NewMockBatchPriorityQueueClient() - allStatuses := []openai.BatchStatus{ - openai.BatchStatusValidating, - openai.BatchStatusFailed, - openai.BatchStatusInProgress, - openai.BatchStatusFinalizing, - openai.BatchStatusCompleted, - openai.BatchStatusExpired, - openai.BatchStatusCancelling, - openai.BatchStatusCancelled, - } - for _, s := range allStatuses { - if s.IsTerminal() && !terminalSet[s] { - t.Errorf("status %q is IsTerminal() but missing from TerminalStatuses()", s) + // Use a very long interval so the ticker never fires. + resultCh := make(chan *Result, 10) + r, err := NewReconciler(batchDB, queue, 24*time.Hour, false, func(res *Result) { + resultCh <- res + }) + if err != nil { + t.Fatalf("failed to create reconciler: %v", err) } - } -} -// --- Helpers --- + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() -func assertJobStatus(t *testing.T, batchDB db.BatchDBClient, jobID string, expected openai.BatchStatus) { - t.Helper() - items, _, _, err := batchDB.DBGet(context.Background(), - &db.BatchQuery{BaseQuery: db.BaseQuery{IDs: []string{jobID}}}, false, 0, 10) - if err != nil { - t.Fatalf("failed to get job %s: %v", jobID, err) - } - if len(items) != 1 { - t.Fatalf("expected 1 item for %s, got %d", jobID, len(items)) - } - var info openai.BatchStatusInfo - if err := json.Unmarshal(items[0].Status, &info); err != nil { - t.Fatalf("failed to unmarshal status: %v", err) - } - if info.Status != expected { - t.Errorf("expected status %s for job %s, got %s", expected, jobID, info.Status) - } -} + done := make(chan error, 1) + go func() { done <- r.RunLoop(ctx) }() -// newMockBatchDB creates a mock batch DB that always returns all items for NonTerminal queries. -func newMockBatchDB() *mock.MockDBClient[db.BatchItem, db.BatchQuery] { - return mock.NewMockDBClient[db.BatchItem, db.BatchQuery]( - func(item *db.BatchItem) string { return item.ID }, - func(query *db.BatchQuery) *db.BaseQuery { return &query.BaseQuery }, - ) -} + r.Trigger() -// casConflictBatchDB is a minimal mock that always returns ErrConflict on DBUpdate. -type casConflictBatchDB struct{} + select { + case <-resultCh: + // Trigger caused the cycle — success. + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for triggered reconciliation cycle") + } -func (c *casConflictBatchDB) DBStore(_ context.Context, _ *db.BatchItem) error { return nil } -func (c *casConflictBatchDB) DBGet(_ context.Context, _ *db.BatchQuery, _ bool, _, _ int) ([]*db.BatchItem, int, bool, error) { - item := newTestBatchItem("job-cas", openai.BatchStatusCancelling, sloTag(futureSLO())) - return []*db.BatchItem{item}, 1, false, nil -} -func (c *casConflictBatchDB) DBUpdate(_ context.Context, _ *db.BatchItem, _ []byte) error { - return fmt.Errorf("DBUpdate: %w", db.ErrConflict) -} -func (c *casConflictBatchDB) DBDelete(_ context.Context, _ []string) ([]string, error) { - return nil, nil -} -func (c *casConflictBatchDB) Close() error { return nil } -func (c *casConflictBatchDB) GetContext(_ context.Context, _ time.Duration) (context.Context, context.CancelFunc) { - return context.Background(), func() {} + cancel() + <-done + }) + + t.Run("trigger is non-blocking", func(t *testing.T) { + batchDB := newMockBatchDB() + queue := mock.NewMockBatchPriorityQueueClient() + + r, _ := newTestReconciler(t, batchDB, queue) + + // Multiple triggers should not block even without a consumer. + r.Trigger() + r.Trigger() + r.Trigger() + // If we reach here without deadlocking, the test passes. + }) } diff --git a/internal/processor/config/config.go b/internal/processor/config/config.go index 67773da9e..b4e99b57c 100644 --- a/internal/processor/config/config.go +++ b/internal/processor/config/config.go @@ -204,14 +204,6 @@ type ProcessorConfig struct { // FileClient holds configuration for the shared file storage client (fs or s3). FileClientCfg sharedcfg.FileClientConfig `yaml:"file_client"` - // HeartbeatInterval controls how often the processor refreshes its in-flight - // entry for a running job. The orphan reconciler uses staleness (no heartbeat - // for > reconciler interval) to detect abandoned jobs. - // Must be shorter than the reconciler's interval so that active jobs are - // never mistaken for orphans. - // Zero means use the default (5 minutes). - HeartbeatInterval time.Duration `yaml:"heartbeat_interval"` - // DispatchMode selects the inference dispatch backend. // "sync" (default): direct HTTP via InferenceClient. // "async": submit via llm-d-async producer, collect from result queue. @@ -347,7 +339,7 @@ func NewConfig() *ProcessorConfig { ShutdownTimeout: 30 * time.Second, WorkDir: "/var/lib/batch-gateway/processor", DBClientCfg: sharedcfg.DBClientConfig{ - Type: sharedcfg.DBTypeRedis, + Type: sharedcfg.DBTypePostgreSQL, }, FileClientCfg: sharedcfg.FileClientConfig{ Type: sharedcfg.FileTypeMock, @@ -360,8 +352,7 @@ func NewConfig() *ProcessorConfig { DefaultOutputExpirationSeconds: 90 * 24 * 60 * 60, // 90 days ProgressTTLSeconds: 24 * 60 * 60, // 24 hours - HeartbeatInterval: 5 * time.Minute, - DispatchMode: DispatchModeSync, + DispatchMode: DispatchModeSync, AsyncDispatchConfig: AsyncDispatchConfig{ ResultPollTimeout: 5 * time.Second, }, @@ -389,9 +380,6 @@ func (c *ProcessorConfig) Validate() error { if c.ShutdownTimeout <= 0 { return fmt.Errorf("shutdown_timeout must be > 0") } - if c.HeartbeatInterval <= 0 { - return fmt.Errorf("heartbeat_interval must be > 0") - } if c.Addr == "" { return fmt.Errorf("addr cannot be empty") } diff --git a/internal/processor/config/config_test.go b/internal/processor/config/config_test.go index 008c5bd90..e91224d44 100644 --- a/internal/processor/config/config_test.go +++ b/internal/processor/config/config_test.go @@ -74,8 +74,8 @@ func TestNewConfig_Defaults(t *testing.T) { if c.WorkDir == "" { t.Fatalf("WorkDir should not be empty") } - if c.DBClientCfg.Type != "redis" { - t.Fatalf("DBClientCfg.Type = %q, want %q", c.DBClientCfg.Type, "redis") + if c.DBClientCfg.Type != "postgresql" { + t.Fatalf("DBClientCfg.Type = %q, want %q", c.DBClientCfg.Type, "postgresql") } if c.Concurrency.Recovery != 5 { t.Fatalf("Concurrency.Recovery = %d, want %d", c.Concurrency.Recovery, 5) diff --git a/internal/processor/worker/job_runner.go b/internal/processor/worker/job_runner.go index fc27a3a1c..4c7d24d19 100644 --- a/internal/processor/worker/job_runner.go +++ b/internal/processor/worker/job_runner.go @@ -48,15 +48,7 @@ import ( // Declared as var (not const) so tests can shorten it. var panicRecoveryTimeout = time.Minute -// defaultHeartbeatInterval is the fallback heartbeat interval used when -// the config value is zero. Matches ProcessorConfig.HeartbeatInterval default. -const defaultHeartbeatInterval = 5 * time.Minute - func (p *Processor) runJob(ctx context.Context, params *jobExecutionParams) { - // Clean up in-flight entry on exit (first defer = last to run via LIFO), - // ensuring the entry is removed regardless of how runJob terminates. - defer p.deleteInFlight(context.Background(), params.jobItem.ID) - // Restore parent trace context propagated from the apiserver via Redis tags if len(params.jobInfo.TraceContext) > 0 { propagator := otel.GetTextMapPropagator() @@ -162,16 +154,6 @@ func (p *Processor) runJob(ctx context.Context, params *jobExecutionParams) { params.eventWatcher = eventWatcher go p.watchCancel(ctx, params) - // Start heartbeat: periodically refreshes the in-flight entry so the - // orphan reconciler knows this job is still being actively processed. - // On each tick it also checks the DB status — if the reconciler acted - // (terminal status or reverted to validating), it aborts (neutral - // context.Canceled) to stop all in-flight requests. The processor's terminal - // CAS write will then fail with ErrConflict, and the processor yields. - heartbeatCtx, heartbeatCancel := context.WithCancel(ctx) - defer heartbeatCancel() - go p.heartbeat(heartbeatCtx, params.jobItem.ID, func() { abortCause(context.Canceled) }) - // ingestion: pre-process job (rejects unregistered-model requests early) ingestCtx, ingestSpan := uotel.StartSpan(abortCtx, "ingest-and-plan") err = p.preProcessJob(ingestCtx, params.jobInfo) diff --git a/internal/processor/worker/recovery.go b/internal/processor/worker/recovery.go index 422041919..e5608bb79 100644 --- a/internal/processor/worker/recovery.go +++ b/internal/processor/worker/recovery.go @@ -22,7 +22,6 @@ import ( "fmt" "os" "path/filepath" - "strconv" "time" "github.com/go-logr/logr" @@ -60,55 +59,47 @@ type recoveryResult struct { cancelPhase string // non-empty → RecordCancellation is called } -// recoverStaleJobs scans the workdir for leftover job directories from a previous -// container execution and performs phase-aware recovery for each discovered job. +// recoverOwnedJobs queries the DB for non-terminal jobs owned by this processor +// (via processor_id) and recovers them. This handles both container-level crashes +// (where emptyDir survives) and pod-level restarts within a StatefulSet (where +// the processor identity is preserved across restarts). // -// This handles container-level crashes (OOM kill, process panic) where K8s restarts -// the container within the same pod and emptyDir survives. Pod-level failures -// (node eviction, pod deletion) destroy emptyDir and are out of scope. -// -// Runs once at startup before the polling loop. Individual job recovery failures -// do not prevent the processor from starting. -func (p *Processor) recoverStaleJobs(ctx context.Context) { +// Runs once at startup before the polling loop. +func (p *Processor) recoverOwnedJobs(ctx context.Context) { logger := logr.FromContextOrDiscard(ctx) - dirs, err := p.discoverStaleJobDirs() + items, _, _, err := p.batchDB.DBGet(ctx, &db.BatchQuery{ + ProcessorID: p.processorID, + NonTerminal: true, + }, true, 0, 1000) if err != nil { - logger.Error(err, "Startup recovery: failed to scan workdir") + logger.Error(err, "Startup recovery: failed to query owned jobs") return } - if len(dirs) == 0 { - logger.V(logging.DEBUG).Info("Startup recovery: no stale job directories found") + if len(items) == 0 { + logger.V(logging.DEBUG).Info("Startup recovery: no owned jobs found") return } - logger.V(logging.INFO).Info("Startup recovery: found stale job directories", "count", len(dirs)) + logger.V(logging.INFO).Info("Startup recovery: found owned jobs", "count", len(items)) var grp errgroup.Group grp.SetLimit(p.cfg.Concurrency.Recovery) - for _, dir := range dirs { - jobID := filepath.Base(dir) + for _, item := range items { grp.Go(func() error { - jlogger := logger.WithValues("jobId", jobID) + jlogger := logger.WithValues("jobId", item.ID) jctx := logr.NewContext(ctx, jlogger) - if err := p.recoverJob(jctx, jobID); err != nil { - jlogger.Error(err, "Startup recovery: failed to recover job") + if recoverErr := p.recoverJob(jctx, item.ID); recoverErr != nil { + jlogger.Error(recoverErr, "Startup recovery: failed to recover owned job") } - return nil // individual failures shouldn't block other recoveries + return nil }) } _ = grp.Wait() } -// discoverStaleJobDirs returns paths to job directories left over from a previous execution. -// The workdir layout is //jobs//. -func (p *Processor) discoverStaleJobDirs() ([]string, error) { - pattern := filepath.Join(p.cfg.WorkDir, "*", jobsDirName, "*") - return filepath.Glob(pattern) -} - // recoverJob is the single routing point for startup recovery. Each recover* // function returns (*recoveryResult, nil) on success, or (*recoveryResult, error) // on failure — where the result may be non-nil (carrying partial file IDs for @@ -408,16 +399,10 @@ func (p *Processor) extractRequestCounts(dbItem *db.BatchItem) *openai.BatchRequ } // extractRecoverySLO recovers the exact SLO deadline for queue re-enqueue. -// Prefer the stored microsecond tag so later CancelBatch can reconstruct the same queue score. func (p *Processor) extractRecoverySLO(dbItem *db.BatchItem, jobInfo *batch_types.JobInfo) (*time.Time, error) { - if dbItem != nil { - if sloStr, ok := dbItem.Tags[batch_types.TagSLO]; ok { - sloMicro, err := strconv.ParseInt(sloStr, 10, 64) - if err == nil { - slo := time.UnixMicro(sloMicro).UTC() - return &slo, nil - } - } + if dbItem != nil && dbItem.Priority > 0 { + slo := time.UnixMicro(dbItem.Priority).UTC() + return &slo, nil } if jobInfo.BatchJob.ExpiresAt != nil { diff --git a/internal/processor/worker/recovery_test.go b/internal/processor/worker/recovery_test.go index 043d54423..55c928bbc 100644 --- a/internal/processor/worker/recovery_test.go +++ b/internal/processor/worker/recovery_test.go @@ -16,7 +16,6 @@ import ( mockfiles "github.com/llm-d/llm-d-batch-gateway/internal/files_store/mock" "github.com/llm-d/llm-d-batch-gateway/internal/processor/config" "github.com/llm-d/llm-d-batch-gateway/internal/shared/openai" - batch_types "github.com/llm-d/llm-d-batch-gateway/internal/shared/types" "github.com/llm-d/llm-d-batch-gateway/internal/util/clientset" "github.com/llm-d/llm-d-batch-gateway/pkg/clients/inference" ) @@ -39,7 +38,6 @@ func newRecoveryTestProcessor(t *testing.T, workDir string) (*Processor, db.Batc Queue: spyQueue, Status: statusClient, Event: mockdb.NewMockBatchEventChannelClient(), - InFlight: mockdb.NewMockInFlightClient(), Inference: inference.NewSingleClientResolver(&fakeInferenceClient{}), }, "test-processor", testLogger(t)) if err != nil { @@ -82,14 +80,12 @@ func seedDBJobWithStatusAndSLO(t *testing.T, dbClient db.BatchDBClient, jobID, t BaseIndexes: db.BaseIndexes{ ID: jobID, TenantID: tenantID, - Tags: db.Tags{ - batch_types.TagSLO: fmt.Sprintf("%d", slo.UTC().UnixMicro()), - }, }, BaseContents: db.BaseContents{ Status: statusBytes, Spec: specBytes, }, + Priority: slo.UTC().UnixMicro(), } if err := dbClient.DBStore(context.Background(), item); err != nil { t.Fatalf("seed DB job: %v", err) @@ -153,39 +149,147 @@ func getDBJobStatus(t *testing.T, dbClient db.BatchDBClient, jobID string) opena // --- Tests --- -func TestRecoverStaleJobs_NoStaleJobs(t *testing.T) { - workDir := t.TempDir() - p, _, _ := newRecoveryTestProcessor(t, workDir) +// newRecoveryTestProcessorWithQueryFilter is like newRecoveryTestProcessor but +// installs a QueryFilter on the mock BatchDB that respects ProcessorID and +// HasProcessorID query fields — matching the real Postgres behavior. +func newRecoveryTestProcessorWithQueryFilter(t *testing.T, workDir string) (*Processor, *mockdb.MockDBClient[db.BatchItem, db.BatchQuery], *spyPQ) { + t.Helper() - p.recoverStaleJobs(testLoggerCtx(t)) -} + batchDB := mockdb.NewMockDBClient[db.BatchItem, db.BatchQuery]( + func(b *db.BatchItem) string { return b.ID }, + func(q *db.BatchQuery) *db.BaseQuery { return &q.BaseQuery }, + ) + batchDB.QueryFilter = func(item *db.BatchItem, query *db.BatchQuery) bool { + if query.ProcessorID != "" && item.ProcessorID != query.ProcessorID { + return false + } + if query.HasProcessorID && item.ProcessorID == "" { + return false + } + return true + } -func TestRecoverStaleJobs_DiscoversDirs(t *testing.T) { - workDir := t.TempDir() - p, _, _ := newRecoveryTestProcessor(t, workDir) + pq := mockdb.NewMockBatchPriorityQueueClient() + spyQueue := &spyPQ{inner: pq} + statusClient := mockdb.NewMockBatchStatusClient() + + cfg := config.NewConfig() + cfg.WorkDir = workDir - dirs, err := p.discoverStaleJobDirs() + p, err := NewProcessor(cfg, &clientset.Clientset{ + BatchDB: batchDB, + FileDB: newMockFileDBClient(), + File: mockfiles.NewMockBatchFilesClient(t.TempDir()), + Queue: spyQueue, + Status: statusClient, + Event: mockdb.NewMockBatchEventChannelClient(), + Inference: inference.NewSingleClientResolver(&fakeInferenceClient{}), + }, "test-processor", testLogger(t)) if err != nil { - t.Fatalf("discoverStaleJobDirs: %v", err) - } - if len(dirs) != 0 { - t.Fatalf("expected 0 dirs, got %d", len(dirs)) + t.Fatalf("NewProcessor: %v", err) } + p.poller = NewPoller(spyQueue, batchDB) + p.updater = NewStatusUpdater(batchDB, statusClient, 86400) - tenantDir := filepath.Join(workDir, "t-abc123", jobsDirName, "job-1") - if err := os.MkdirAll(tenantDir, 0o755); err != nil { - t.Fatalf("MkdirAll: %v", err) - } + return p, batchDB, spyQueue +} - dirs, err = p.discoverStaleJobDirs() - if err != nil { - t.Fatalf("discoverStaleJobDirs: %v", err) +func TestRecoverOwnedJobs(t *testing.T) { + t.Run("recovers jobs owned by this processor", func(t *testing.T) { + workDir := t.TempDir() + p, batchDB, spyQueue := newRecoveryTestProcessorWithQueryFilter(t, workDir) + + // Seed two jobs owned by this processor (in_progress) and one owned by another. + seedDBJobWithStatus(t, batchDB, "owned-1", "tenant-1", openai.BatchStatusInProgress, nil) + seedDBJobWithStatus(t, batchDB, "owned-2", "tenant-1", openai.BatchStatusFinalizing, nil) + seedDBJobWithStatus(t, batchDB, "other-1", "tenant-1", openai.BatchStatusInProgress, nil) + + // Set processor_id on owned items. + setProcessorID(t, batchDB, "owned-1", p.processorID) + setProcessorID(t, batchDB, "owned-2", p.processorID) + setProcessorID(t, batchDB, "other-1", "other-processor") + + // Create job dirs so recovery can find artifacts. + createJobDir(t, p, "owned-1", "tenant-1") + createJobDir(t, p, "owned-2", "tenant-1") + + ctx := testLoggerCtx(t) + p.recoverOwnedJobs(ctx) + + // owned-1 was in_progress with no local artifacts → re-enqueued. + // owned-2 was finalizing with no output files → transitions to failed. + // other-1 should not be touched. + if spyQueue.EnqueueCalls() < 1 { + t.Errorf("expected at least 1 re-enqueue for owned jobs, got %d", spyQueue.EnqueueCalls()) + } + + // Verify other-1 is untouched (still in_progress). + otherStatus := getDBJobStatus(t, batchDB, "other-1") + if otherStatus != openai.BatchStatusInProgress { + t.Errorf("expected other-1 to remain in_progress, got %s", otherStatus) + } + }) + + t.Run("no owned jobs is a no-op", func(t *testing.T) { + workDir := t.TempDir() + p, batchDB, spyQueue := newRecoveryTestProcessorWithQueryFilter(t, workDir) + + // Seed a job owned by a different processor. + seedDBJobWithStatus(t, batchDB, "other-1", "tenant-1", openai.BatchStatusInProgress, nil) + setProcessorID(t, batchDB, "other-1", "other-processor") + + ctx := testLoggerCtx(t) + p.recoverOwnedJobs(ctx) + + if spyQueue.EnqueueCalls() != 0 { + t.Errorf("expected 0 enqueue calls, got %d", spyQueue.EnqueueCalls()) + } + }) + + t.Run("DB error is handled gracefully", func(t *testing.T) { + workDir := t.TempDir() + p, _, _ := newRecoveryTestProcessorWithQueryFilter(t, workDir) + + // Replace batchDB with one that always fails on DBGet. + p.batchDB = &failOnGetDB{err: fmt.Errorf("connection refused")} + + ctx := testLoggerCtx(t) + // Should not panic — just logs the error and returns. + p.recoverOwnedJobs(ctx) + }) +} + +// setProcessorID updates a stored BatchItem's ProcessorID in the mock DB. +func setProcessorID(t *testing.T, dbClient *mockdb.MockDBClient[db.BatchItem, db.BatchQuery], jobID, processorID string) { + t.Helper() + items, _, _, err := dbClient.DBGet(context.Background(), + &db.BatchQuery{BaseQuery: db.BaseQuery{IDs: []string{jobID}}}, + true, 0, 1) + if err != nil || len(items) == 0 { + t.Fatalf("setProcessorID: failed to get job %s: %v", jobID, err) } - if len(dirs) != 1 { - t.Fatalf("expected 1 dir, got %d", len(dirs)) + items[0].ProcessorID = processorID + if err := dbClient.DBUpdate(context.Background(), items[0], nil); err != nil { + t.Fatalf("setProcessorID: failed to update job %s: %v", jobID, err) } } +// failOnGetDB is a minimal mock that returns an error on DBGet. +type failOnGetDB struct { + err error +} + +func (f *failOnGetDB) DBStore(_ context.Context, _ *db.BatchItem) error { return nil } +func (f *failOnGetDB) DBGet(_ context.Context, _ *db.BatchQuery, _ bool, _, _ int) ([]*db.BatchItem, int, bool, error) { + return nil, 0, false, f.err +} +func (f *failOnGetDB) DBUpdate(_ context.Context, _ *db.BatchItem, _ []byte) error { return nil } +func (f *failOnGetDB) DBDelete(_ context.Context, _ []string) ([]string, error) { return nil, nil } +func (f *failOnGetDB) Close() error { return nil } +func (f *failOnGetDB) GetContext(_ context.Context, _ time.Duration) (context.Context, context.CancelFunc) { + return context.Background(), func() {} +} + func TestRecoverJob_Finalizing(t *testing.T) { workDir := t.TempDir() p, dbClient, _ := newRecoveryTestProcessor(t, workDir) @@ -587,7 +691,6 @@ func newRecoveryTestProcessorWithFailDB(t *testing.T, workDir string, failOn int Queue: pq, Status: statusClient, Event: mockdb.NewMockBatchEventChannelClient(), - InFlight: mockdb.NewMockInFlightClient(), Inference: inference.NewSingleClientResolver(&fakeInferenceClient{}), }, "test-processor", testLogger(t)) if err != nil { @@ -627,7 +730,6 @@ func TestRecoverJob_Cancelling_AllUpdatesFail_ReturnsError(t *testing.T) { Queue: pq, Status: statusClient, Event: mockdb.NewMockBatchEventChannelClient(), - InFlight: mockdb.NewMockInFlightClient(), Inference: inference.NewSingleClientResolver(&fakeInferenceClient{}), }, "test-processor", testLogger(t)) if err != nil { @@ -672,7 +774,6 @@ func TestRecoverJob_Validating_EnqueueFails_FallsBackToFailed(t *testing.T) { Queue: pq, Status: statusClient, Event: mockdb.NewMockBatchEventChannelClient(), - InFlight: mockdb.NewMockInFlightClient(), Inference: inference.NewSingleClientResolver(&fakeInferenceClient{}), }, "test-processor", testLogger(t)) if err != nil { @@ -721,7 +822,6 @@ func TestRecoverJob_InProgressReEnqueue_EnqueueFails_FallsBackToFailed(t *testin Queue: pq, Status: statusClient, Event: mockdb.NewMockBatchEventChannelClient(), - InFlight: mockdb.NewMockInFlightClient(), Inference: inference.NewSingleClientResolver(&fakeInferenceClient{}), }, "test-processor", testLogger(t)) if err != nil { @@ -835,7 +935,7 @@ func (s *slowBatchDBClient) peakConcurrency() int { return s.maxActive } -func TestRecoverStaleJobs_RunsConcurrently(t *testing.T) { +func TestRecoverOwnedJobs_RunsConcurrently(t *testing.T) { workDir := t.TempDir() innerDB := newMockBatchDBClient() @@ -857,7 +957,6 @@ func TestRecoverStaleJobs_RunsConcurrently(t *testing.T) { Queue: pq, Status: statusClient, Event: mockdb.NewMockBatchEventChannelClient(), - InFlight: mockdb.NewMockInFlightClient(), Inference: inference.NewSingleClientResolver(&fakeInferenceClient{}), }, "test-processor", testLogger(t)) if err != nil { @@ -866,18 +965,23 @@ func TestRecoverStaleJobs_RunsConcurrently(t *testing.T) { p.poller = NewPoller(pq, slowDB) p.updater = NewStatusUpdater(slowDB, statusClient, 86400) - // Create 5 stale job directories with terminal status (completed) so + // Create 5 owned jobs with terminal status (completed) so // recovery just cleans them up after the DB lookup. numJobs := 5 tenantID := "tenant-conc" for i := 0; i < numJobs; i++ { jobID := fmt.Sprintf("job-conc-%d", i) seedDBJobWithStatus(t, innerDB, jobID, tenantID, openai.BatchStatusCompleted, nil) + // Set processor_id so recoverOwnedJobs finds them. + items, _, _, _ := innerDB.DBGet(context.Background(), + &db.BatchQuery{BaseQuery: db.BaseQuery{IDs: []string{jobID}}}, true, 0, 1) + items[0].ProcessorID = p.processorID + _ = innerDB.DBUpdate(context.Background(), items[0], nil) createJobDir(t, p, jobID, tenantID) } ctx := testLoggerCtx(t) - p.recoverStaleJobs(ctx) + p.recoverOwnedJobs(ctx) // With concurrency=5 and 5 jobs at 50ms delay each, parallel should // complete in ~50ms. Sequential would take ~250ms. @@ -1034,7 +1138,6 @@ func TestRecoverJob_ExpiredWriteFails_FallbackToFailed(t *testing.T) { Queue: pq, Status: statusClient, Event: mockdb.NewMockBatchEventChannelClient(), - InFlight: mockdb.NewMockInFlightClient(), Inference: inference.NewSingleClientResolver(&fakeInferenceClient{}), }, "test-processor", testLogger(t)) if err != nil { diff --git a/internal/processor/worker/status_updater.go b/internal/processor/worker/status_updater.go index 42dda402c..581c433da 100644 --- a/internal/processor/worker/status_updater.go +++ b/internal/processor/worker/status_updater.go @@ -122,6 +122,7 @@ func (s *StatusUpdater) UpdatePersistentStatus( BaseContents: db.BaseContents{ Status: statusBytes, }, + Epoch: dbJob.Epoch, }, nil); err != nil { return err } diff --git a/internal/processor/worker/test_helpers_test.go b/internal/processor/worker/test_helpers_test.go index 0fe9cb574..6f767148b 100644 --- a/internal/processor/worker/test_helpers_test.go +++ b/internal/processor/worker/test_helpers_test.go @@ -313,9 +313,6 @@ func (f *failOnStatusDB) Close() error { return f.inner.Close() } func mustNewProcessor(t *testing.T, cfg *config.ProcessorConfig, clients *clientset.Clientset) *Processor { t.Helper() - if clients.InFlight == nil { - clients.InFlight = mockdb.NewMockInFlightClient() - } p, err := NewProcessor(cfg, clients, "test-processor", testLogger(t)) if err != nil { t.Fatalf("NewProcessor: %v", err) @@ -344,7 +341,6 @@ func validProcessorClients(t testing.TB) *clientset.Clientset { Queue: mockdb.NewMockBatchPriorityQueueClient(), Status: mockdb.NewMockBatchStatusClient(), Event: mockdb.NewMockBatchEventChannelClient(), - InFlight: mockdb.NewMockInFlightClient(), Inference: inference.NewSingleClientResolver(&fakeInferenceClient{}), } } @@ -373,7 +369,6 @@ func newTestProcessorEnv(t *testing.T, cfg *config.ProcessorConfig, inferClient Queue: pqClient, Status: statusClient, Event: mockdb.NewMockBatchEventChannelClient(), - InFlight: mockdb.NewMockInFlightClient(), Inference: inference.NewSingleClientResolver(inferClient), }, "test-processor", testLogger(t)) if err != nil { @@ -722,29 +717,3 @@ func uniqueTestFolder(t *testing.T, base string) string { testName := strings.ReplaceAll(t.Name(), "/", "_") return filepath.Join(base, testName, fmt.Sprintf("%d", time.Now().UnixNano())) } - -type countingInFlightClient struct { - inner *mockdb.MockInFlightClient - setCount atomic.Int32 -} - -func newCountingInFlightClient() *countingInFlightClient { - return &countingInFlightClient{inner: mockdb.NewMockInFlightClient()} -} - -func (c *countingInFlightClient) InFlightSet(ctx context.Context, jobID, processorID string) error { - c.setCount.Add(1) - return c.inner.InFlightSet(ctx, jobID, processorID) -} - -func (c *countingInFlightClient) InFlightDelete(ctx context.Context, jobID string) error { - return c.inner.InFlightDelete(ctx, jobID) -} - -func (c *countingInFlightClient) InFlightGetAll(ctx context.Context) (map[string]*db.InFlightEntry, error) { - return c.inner.InFlightGetAll(ctx) -} - -func (c *countingInFlightClient) Close() error { - return c.inner.Close() -} diff --git a/internal/processor/worker/worker.go b/internal/processor/worker/worker.go index fd68bb794..d63397962 100644 --- a/internal/processor/worker/worker.go +++ b/internal/processor/worker/worker.go @@ -19,7 +19,6 @@ package worker import ( "context" - "encoding/json" "fmt" "sync" "time" @@ -70,9 +69,8 @@ type Processor struct { poller *Poller updater *StatusUpdater - batchDB db.BatchDBClient // job status lookups (heartbeat DB check) + batchDB db.BatchDBClient // job status lookups event db.BatchEventChannelClient // cancel-event subscription - inflight db.InFlightClient // in-flight job tracking for orphan recovery inference *inference.GatewayResolver // model → gateway routing (sync) asyncInference *inference.AsyncGatewayResolver // model → async client routing broadcasters *broadcasterRegistry // per-model result broadcasters (async only) @@ -100,7 +98,6 @@ func NewProcessor( updater: updater, batchDB: clients.BatchDB, event: clients.Event, - inflight: clients.InFlight, inference: clients.Inference, asyncInference: clients.AsyncInference, files: newFileManager(clients.File, clients.FileDB), @@ -115,7 +112,7 @@ func (p *Processor) Run(ctx context.Context, onReady func()) error { return err } - p.recoverStaleJobs(ctx) + p.recoverOwnedJobs(ctx) if onReady != nil { onReady() @@ -272,13 +269,6 @@ func (p *Processor) runPollingLoop(pollingCtx, jobBaseCtx context.Context) error continue } - // Record in-flight entry immediately after dequeue so the orphan - // reconciler can track this job. Non-fatal on error: the reconciler - // can still detect orphans via DB + queue cross-reference. - if err := p.inflight.InFlightSet(pollingCtx, task.ID, p.processorID); err != nil { - logr.FromContextOrDiscard(pollingCtx).Error(err, "Failed to set in-flight entry", "jobId", task.ID) - } - // Pre-launch: use pollingCtx so guard cancel / SIGTERM interrupts // DB fetch and validation promptly. jobBaseCtx is only used once // we commit to launching the job goroutine. @@ -296,10 +286,8 @@ func (p *Processor) runPollingLoop(pollingCtx, jobBaseCtx context.Context) error bgCtx, bgSpan := uotel.DetachedContext(pollCtx, "re-enqueue-fetch-failure") if reEnqueueErr := p.poller.enqueueOne(bgCtx, task); reEnqueueErr != nil { pollLogger.Error(reEnqueueErr, "Failed to re-enqueue the job to the queue") - p.deleteInFlight(bgCtx, task.ID) metrics.RecordJobProcessed(metrics.ResultFailed, metrics.ReasonSystemError) } else { - p.deleteInFlight(bgCtx, task.ID) metrics.RecordJobProcessed(metrics.ResultReEnqueued, metrics.ReasonDBTransient) pollLogger.V(logging.INFO).Info("Re-enqueued the job to the queue") } @@ -311,7 +299,6 @@ func (p *Processor) runPollingLoop(pollingCtx, jobBaseCtx context.Context) error if jobItem == nil { pollLogger.Error(fmt.Errorf("job item is not found in the DB"), "Ignoring job (data inconsistency)") p.releaseForNextPoll() - p.deleteInFlight(pollCtx, task.ID) metrics.RecordJobProcessed(metrics.ResultSkipped, metrics.ReasonDBInconsistency) continue } @@ -326,7 +313,6 @@ func (p *Processor) runPollingLoop(pollingCtx, jobBaseCtx context.Context) error if failErr := p.handleFailed(pollCtx, p.updater, jobItem, nil, nil); failErr != nil { pollLogger.Error(failErr, "Failed to mark malformed job as failed") } - p.deleteInFlight(pollCtx, task.ID) continue } @@ -353,7 +339,6 @@ func (p *Processor) runPollingLoop(pollingCtx, jobBaseCtx context.Context) error } p.releaseForNextPoll() - p.deleteInFlight(pollCtx, task.ID) recordE2ELatency(jobInfo, metrics.E2EStatusExpired) metrics.RecordJobProcessed(metrics.ResultExpired, metrics.ReasonExpiredDequeue) continue @@ -371,7 +356,6 @@ func (p *Processor) runPollingLoop(pollingCtx, jobBaseCtx context.Context) error continue } p.releaseForNextPoll() - p.deleteInFlight(pollCtx, task.ID) recordE2ELatency(jobInfo, metrics.E2EStatusCancelled) metrics.RecordCancellation(metrics.CancelPhaseQueued) metrics.RecordJobProcessed(metrics.ResultSuccess, metrics.ReasonNone) @@ -381,7 +365,6 @@ func (p *Processor) runPollingLoop(pollingCtx, jobBaseCtx context.Context) error pollLogger.V(logging.INFO).Info("job is not in processible state. skipping this job.", "status", jobInfo.BatchJob.Status) p.releaseForNextPoll() - p.deleteInFlight(pollCtx, task.ID) metrics.RecordJobProcessed(metrics.ResultSkipped, metrics.ReasonNotRunnableState) continue } @@ -400,10 +383,8 @@ func (p *Processor) runPollingLoop(pollingCtx, jobBaseCtx context.Context) error if failErr := p.handleFailed(bgCtx, p.updater, jobItem, nil, jobInfo); failErr != nil { pollLogger.Error(failErr, "Failed to mark dequeued job as failed after re-enqueue failure") } - p.deleteInFlight(bgCtx, task.ID) metrics.RecordJobProcessed(metrics.ResultFailed, metrics.ReasonGuardShutdown) } else { - p.deleteInFlight(bgCtx, task.ID) metrics.RecordJobProcessed(metrics.ResultReEnqueued, metrics.ReasonGuardShutdown) pollLogger.V(logging.INFO).Info("Re-enqueued the job to the queue during graceful shutdown") } @@ -450,70 +431,6 @@ func (p *Processor) releaseForNextPoll() { p.release() } -func (p *Processor) deleteInFlight(ctx context.Context, jobID string) { - if err := p.inflight.InFlightDelete(ctx, jobID); err != nil { - logr.FromContextOrDiscard(ctx).Error(err, "Failed to delete in-flight entry", "jobId", jobID) - } -} - -func (p *Processor) heartbeat(ctx context.Context, jobID string, abortFn context.CancelFunc) { - logger := logr.FromContextOrDiscard(ctx).WithValues("jobId", jobID) - logger.V(logging.INFO).Info("Heartbeat: started") - - interval := p.cfg.HeartbeatInterval - if interval <= 0 { - interval = defaultHeartbeatInterval - } - ticker := time.NewTicker(interval) - defer ticker.Stop() - - for { - select { - case <-ctx.Done(): - logger.V(logging.INFO).Info("Heartbeat: stopped") - return - case <-ticker.C: - if err := p.inflight.InFlightSet(ctx, jobID, p.processorID); err != nil { - logger.Error(err, "Heartbeat: failed to refresh in-flight entry") - } else { - logger.V(logging.INFO).Info("Heartbeat: refreshed") - } - - if p.checkReconcilerActed(ctx, jobID, logger) { - logger.Info("Heartbeat: reconciler acted on job, aborting") - abortFn() - return - } - } - } -} - -func (p *Processor) checkReconcilerActed(ctx context.Context, jobID string, logger logr.Logger) bool { - query := &db.BatchQuery{BaseQuery: db.BaseQuery{IDs: []string{jobID}}} - items, _, _, err := p.batchDB.DBGet(ctx, query, false, 0, 1) - if err != nil { - logger.Error(err, "Heartbeat: DB status check failed") - return false - } - if len(items) == 0 { - logger.Info("Heartbeat: job not found in DB, reconciler may have acted") - return true - } - - var statusInfo openai.BatchStatusInfo - if err := json.Unmarshal(items[0].Status, &statusInfo); err != nil { - logger.Error(err, "Heartbeat: failed to unmarshal job status") - return false - } - - if statusInfo.Status.IsTerminal() || statusInfo.Status == openai.BatchStatusValidating { - logger.Info("Heartbeat: unexpected DB status, reconciler acted", "dbStatus", statusInfo.Status) - return true - } - - return false -} - // pre-flight check func (p *Processor) prepare(ctx context.Context) error { logger := logr.FromContextOrDiscard(ctx) @@ -545,9 +462,6 @@ func (p *Processor) validate() error { if p.event == nil { return fmt.Errorf("event channel client is missing") } - if p.inflight == nil { - return fmt.Errorf("in-flight client is missing") - } if p.inference == nil && p.asyncInference == nil { return fmt.Errorf("inference client is missing") } diff --git a/internal/processor/worker/worker_test.go b/internal/processor/worker/worker_test.go index 17a7c6c5d..00696819b 100644 --- a/internal/processor/worker/worker_test.go +++ b/internal/processor/worker/worker_test.go @@ -2,13 +2,10 @@ package worker import ( "context" - "encoding/json" "testing" "time" - db "github.com/llm-d/llm-d-batch-gateway/internal/database/api" "github.com/llm-d/llm-d-batch-gateway/internal/processor/config" - "github.com/llm-d/llm-d-batch-gateway/internal/shared/openai" "github.com/llm-d/llm-d-batch-gateway/internal/util/clientset" "github.com/llm-d/llm-d-batch-gateway/internal/util/semaphore" "github.com/llm-d/llm-d-batch-gateway/pkg/clients/inference" @@ -152,74 +149,6 @@ func TestSemaphoreGuard_JobBaseCtxSurvives(t *testing.T) { } } -func TestHeartbeat_StopsOnContextCancel(t *testing.T) { - mock := newCountingInFlightClient() - cfg := config.NewConfig() - cfg.HeartbeatInterval = 10 * time.Millisecond - p := mustNewProcessor(t, cfg, validProcessorClients(t)) - p.inflight = mock - - statusBytes, _ := json.Marshal(openai.BatchStatusInfo{Status: openai.BatchStatusInProgress}) - _ = p.batchDB.DBStore(context.Background(), &db.BatchItem{ - BaseIndexes: db.BaseIndexes{ID: "job-1"}, - BaseContents: db.BaseContents{Status: statusBytes}, - }) - - ctx, cancel := context.WithCancel(testLoggerCtx(t)) - done := make(chan struct{}) - go func() { - p.heartbeat(ctx, "job-1", func() {}) - close(done) - }() - - // Let a few heartbeats fire. - time.Sleep(50 * time.Millisecond) - cancel() - - select { - case <-done: - case <-time.After(time.Second): - t.Fatal("heartbeat goroutine did not stop after context cancel") - } - - countAtStop := mock.setCount.Load() - if countAtStop == 0 { - t.Fatal("expected at least one InFlightSet call") - } - - // Verify no more calls after cancel. - time.Sleep(30 * time.Millisecond) - if mock.setCount.Load() != countAtStop { - t.Fatal("InFlightSet called after context was cancelled") - } -} - -func TestHeartbeat_AbortsWhenReconcilerActs(t *testing.T) { - cfg := config.NewConfig() - cfg.HeartbeatInterval = 10 * time.Millisecond - p := mustNewProcessor(t, cfg, validProcessorClients(t)) - - statusBytes, _ := json.Marshal(openai.BatchStatusInfo{Status: openai.BatchStatusFailed}) - _ = p.batchDB.DBStore(context.Background(), &db.BatchItem{ - BaseIndexes: db.BaseIndexes{ID: "job-reconciled"}, - BaseContents: db.BaseContents{Status: statusBytes}, - }) - - aborted := make(chan struct{}) - abortFn := func() { close(aborted) } - - ctx, cancel := context.WithCancel(testLoggerCtx(t)) - defer cancel() - - go p.heartbeat(ctx, "job-reconciled", abortFn) - - select { - case <-aborted: - case <-time.After(time.Second): - t.Fatal("heartbeat did not call abortFn when DB status was terminal") - } -} - func TestProcessorTokenHelpers(t *testing.T) { cfg := config.NewConfig() cfg.NumWorkers = 1 diff --git a/internal/shared/config/config.go b/internal/shared/config/config.go index d262076e0..9b739320b 100644 --- a/internal/shared/config/config.go +++ b/internal/shared/config/config.go @@ -27,8 +27,6 @@ import ( // Database backend types. const ( - DBTypeRedis = "redis" - DBTypeValkey = "valkey" DBTypePostgreSQL = "postgresql" DBTypeMock = "mock" ) @@ -42,7 +40,7 @@ const ( // DBClientConfig holds database client configuration shared by all components. type DBClientConfig struct { - // Type specifies the database backend: DBTypeRedis, DBTypeValkey, or DBTypePostgreSQL. + // Type specifies the database backend: DBTypePostgreSQL. Type string `yaml:"type"` // PostgreSQLCfg holds PostgreSQL connection settings (used when Type is "postgresql"). PostgreSQLCfg postgresql.PostgreSQLConfig `yaml:"postgresql"` diff --git a/internal/shared/converter/batch_converter.go b/internal/shared/converter/batch_converter.go index 1cf4c6f7a..e9b82544d 100644 --- a/internal/shared/converter/batch_converter.go +++ b/internal/shared/converter/batch_converter.go @@ -19,6 +19,7 @@ package converter import ( "encoding/json" "fmt" + "time" "github.com/llm-d/llm-d-batch-gateway/internal/database/api" "github.com/llm-d/llm-d-batch-gateway/internal/shared/openai" @@ -44,6 +45,16 @@ func BatchToDBItem(batch *openai.Batch, tenantID string, tags api.Tags) (*api.Ba expiry = *batch.ExpiresAt } + var priority int64 + if batch.CompletionWindow != "" && batch.CreatedAt > 0 { + d, err := time.ParseDuration(batch.CompletionWindow) + if err != nil { + return nil, fmt.Errorf("failed to parse completion window %q: %w", batch.CompletionWindow, err) + } + slo := time.Unix(batch.CreatedAt, 0).Add(d) + priority = slo.UnixMicro() + } + item := &api.BatchItem{} item.ID = batch.ID item.TenantID = tenantID @@ -51,6 +62,7 @@ func BatchToDBItem(batch *openai.Batch, tenantID string, tags api.Tags) (*api.Ba item.Tags = tags item.Spec = specData item.Status = statusData + item.Priority = priority return item, nil } diff --git a/internal/shared/types/job.go b/internal/shared/types/job.go index 953521953..24d5bcfd1 100644 --- a/internal/shared/types/job.go +++ b/internal/shared/types/job.go @@ -24,7 +24,6 @@ import ( const ( TagPrefixPassThroughHeader = "pth:" TagPrefixOTel = "otel:" - TagSLO = "slo_unix_micro" TagOutputExpiresAfterAnchor = "output_expires_after_anchor" TagOutputExpiresAfterSeconds = "output_expires_after_seconds" ) diff --git a/internal/util/clientset/clientset.go b/internal/util/clientset/clientset.go index c4329874f..9771bf1cb 100644 --- a/internal/util/clientset/clientset.go +++ b/internal/util/clientset/clientset.go @@ -24,6 +24,7 @@ import ( "errors" "fmt" "maps" + "os" "github.com/go-logr/logr" dbapi "github.com/llm-d/llm-d-batch-gateway/internal/database/api" @@ -48,7 +49,6 @@ type Clientset struct { Queue dbapi.BatchPriorityQueueClient Event dbapi.BatchEventChannelClient Status dbapi.BatchStatusClient - InFlight dbapi.InFlightClient Inference *inference.GatewayResolver AsyncInference *inference.AsyncGatewayResolver } @@ -93,31 +93,6 @@ func NewS3FileClient(ctx context.Context, cfg *s3client.Config) (fsapi.BatchFile return c, nil } -// NewRedisDBClients creates Redis-backed batch and file database clients. -// It reads the Redis URL from the mounted secrets when not set in the config. -func NewRedisDBClients(ctx context.Context, cfg *uredis.RedisClientConfig) (dbapi.BatchDBClient, dbapi.FileDBClient, error) { - if cfg == nil { - return nil, nil, fmt.Errorf("redis config cannot be nil") - } - if cfg.Url == "" { - redisURL, err := ucom.ReadSecretFile(ucom.SecretKeyRedisURL) - if err != nil { - return nil, nil, err - } - cfg.Url = redisURL - } - batchDB, err := dbRedis.NewBatchDBClientRedis(ctx, nil, cfg, 0) - if err != nil { - return nil, nil, fmt.Errorf("failed to create redis batch-db client: %w", err) - } - fileDB, err := dbRedis.NewFileDBClientRedis(ctx, nil, cfg, 0) - if err != nil { - return nil, nil, fmt.Errorf("failed to create redis file-db client: %w", err) - } - logr.FromContextOrDiscard(ctx).Info("Redis-based database client created") - return batchDB, fileDB, nil -} - // NewPostgreSQLDBClients creates PostgreSQL-backed batch and file database clients. // It reads the URL from the mounted secrets when not set in the config. func NewPostgreSQLDBClients(ctx context.Context, cfg *postgresql.PostgreSQLConfig) (dbapi.BatchDBClient, dbapi.FileDBClient, error) { @@ -194,7 +169,8 @@ func WithAsyncInference(cfg inference.AsyncClientConfig) Option { } // NewClientset creates the clients specified by the given options. -func NewClientset(ctx context.Context, component ucom.Component, opts ...Option) (*Clientset, error) { +// On error, any clients already created are closed before returning. +func NewClientset(ctx context.Context, component ucom.Component, opts ...Option) (_ *Clientset, retErr error) { logger := logr.FromContextOrDiscard(ctx) cfg := &clientsetConfig{} @@ -203,6 +179,13 @@ func NewClientset(ctx context.Context, component ucom.Component, opts ...Option) } cs := &Clientset{} + defer func() { + if retErr != nil { + if closeErr := cs.Close(); closeErr != nil { + logger.Error(closeErr, "failed to close partially constructed clientset") + } + } + }() // build redis exchange client if cfg.exchangeRedisCfg != nil { @@ -224,7 +207,6 @@ func NewClientset(ctx context.Context, component ucom.Component, opts ...Option) cs.Queue = redisClient cs.Event = redisClient cs.Status = redisClient - cs.InFlight = redisClient } // build file store client @@ -254,14 +236,6 @@ func NewClientset(ctx context.Context, component ucom.Component, opts ...Option) // build database client if cfg.dbCfg != nil { switch cfg.dbCfg.Type { - case sharedcfg.DBTypeRedis, sharedcfg.DBTypeValkey: - redisCfg := &cfg.dbCfg.RedisCfg - batchDB, fileDB, err := NewRedisDBClients(ctx, redisCfg) - if err != nil { - return nil, err - } - cs.BatchDB = batchDB - cs.FileDB = fileDB case sharedcfg.DBTypePostgreSQL: batchDB, fileDB, err := NewPostgreSQLDBClients(ctx, &cfg.dbCfg.PostgreSQLCfg) if err != nil { @@ -269,8 +243,23 @@ func NewClientset(ctx context.Context, component ucom.Component, opts ...Option) } cs.BatchDB = batchDB cs.FileDB = fileDB + + var processorID string + if component == ucom.ComponentProcessor { + processorID, err = os.Hostname() + if err != nil { + return nil, fmt.Errorf("failed to get hostname for processor ID: %w", err) + } + } + queueClient, err := postgresql.NewPostgresBatchQueueClient(ctx, &cfg.dbCfg.PostgreSQLCfg, processorID) + if err != nil { + return nil, fmt.Errorf("failed to create postgres queue client: %w", err) + } + // Postgres queue intentionally replaces the Redis queue set above. + // Redis is retained only for Event and Status channels. + cs.Queue = queueClient default: - return nil, fmt.Errorf("unsupported database.type: %s (supported values: redis, valkey, postgresql)", cfg.dbCfg.Type) + return nil, fmt.Errorf("unsupported database.type: %s (supported values: postgresql)", cfg.dbCfg.Type) } } @@ -340,11 +329,6 @@ func (cs *Clientset) Close() error { errs = append(errs, err) } } - if cs.InFlight != nil { - if err := cs.InFlight.Close(); err != nil { - errs = append(errs, err) - } - } if cs.Inference != nil { if err := cs.Inference.Close(); err != nil { errs = append(errs, err) diff --git a/scripts/dev-deploy-dispatcher.sh b/scripts/dev-deploy-dispatcher.sh index c119c4ea5..ca3d93b5f 100755 --- a/scripts/dev-deploy-dispatcher.sh +++ b/scripts/dev-deploy-dispatcher.sh @@ -199,8 +199,8 @@ helm upgrade "${HELM_RELEASE}" "${REPO_ROOT}/charts/batch-gateway" \ rm -f "${REUSED_VALUES}" step "Restarting processor to pick up new config..." -kubectl rollout restart deployment/"${HELM_RELEASE}-processor" --namespace "${NAMESPACE}" -kubectl rollout status deployment/"${HELM_RELEASE}-processor" --namespace "${NAMESPACE}" --timeout=60s +kubectl rollout restart statefulset/"${HELM_RELEASE}-processor" --namespace "${NAMESPACE}" +kubectl rollout status statefulset/"${HELM_RELEASE}-processor" --namespace "${NAMESPACE}" --timeout=60s log "Processor reconfigured for async dispatch." diff --git a/scripts/dev-deploy.sh b/scripts/dev-deploy.sh index 9bff6b8b6..306ac6112 100755 --- a/scripts/dev-deploy.sh +++ b/scripts/dev-deploy.sh @@ -1150,7 +1150,10 @@ install_batch_gateway() { kubectl rollout restart deployment \ -l "app.kubernetes.io/instance=${HELM_RELEASE}" \ -n "${NAMESPACE}" - # rollout status blocks until new ReplicaSet pods are Ready. + kubectl rollout restart statefulset \ + -l "app.kubernetes.io/instance=${HELM_RELEASE}" \ + -n "${NAMESPACE}" + # rollout status blocks until new ReplicaSet/StatefulSet pods are Ready. # wait_for_deployment (condition=Available) is insufficient here because # the old ReplicaSet satisfies Available immediately after restart. wait_for_rollout "${HELM_RELEASE}-apiserver" "${NAMESPACE}" 120s @@ -1175,31 +1178,49 @@ verify_deployment() { } # wait_for_deployment -# Suitable for initial install where no old ReplicaSet exists. +# Suitable for initial install where no old ReplicaSet/StatefulSet exists. +# Automatically detects whether the resource is a Deployment or StatefulSet. wait_for_deployment() { local name="$1" local ns="$2" local timeout="${3:-120s}" - step "Waiting for deployment '${name}' to be ready..." - if ! kubectl wait deployment/"${name}" \ - -n "${ns}" --for=condition=Available --timeout="${timeout}"; then - die "Deployment '${name}' did not become ready within ${timeout}" + local kind="deployment" + if kubectl get statefulset/"${name}" -n "${ns}" >/dev/null 2>&1; then + kind="statefulset" fi - log "Deployment '${name}' is ready." + + step "Waiting for ${kind}/${name} to be ready..." + if [ "${kind}" = "statefulset" ]; then + # StatefulSets don't have an Available condition; use rollout status. + if ! kubectl rollout status "${kind}/${name}" \ + -n "${ns}" --timeout="${timeout}"; then + die "${kind}/${name} did not become ready within ${timeout}" + fi + else + if ! kubectl wait "${kind}/${name}" \ + -n "${ns}" --for=condition=Available --timeout="${timeout}"; then + die "${kind}/${name} did not become ready within ${timeout}" + fi + fi + log "${kind}/${name} is ready." } # wait_for_rollout -# Blocks until the latest rollout (new ReplicaSet) is fully complete. -# Use after rollout restart; condition=Available can pass prematurely -# when the old ReplicaSet still satisfies the Available condition. +# Blocks until the latest rollout is fully complete. +# Automatically detects whether the resource is a Deployment or StatefulSet. wait_for_rollout() { local name="$1" local ns="$2" local timeout="${3:-120s}" - step "Waiting for rollout of '${name}' to complete..." - if ! kubectl rollout status deployment/"${name}" \ + local kind="deployment" + if kubectl get statefulset/"${name}" -n "${ns}" >/dev/null 2>&1; then + kind="statefulset" + fi + + step "Waiting for rollout of ${kind}/${name} to complete..." + if ! kubectl rollout status "${kind}/${name}" \ -n "${ns}" --timeout="${timeout}"; then die "Rollout of '${name}' did not complete within ${timeout}" fi diff --git a/test/e2e/flow_control_test.go b/test/e2e/flow_control_test.go index 68efdd421..61e2d7a40 100644 --- a/test/e2e/flow_control_test.go +++ b/test/e2e/flow_control_test.go @@ -675,14 +675,14 @@ func assertRequestErrors(t *testing.T, model string) { func getProcessorLogsSince(t *testing.T, sinceTime string) string { t.Helper() - deployment := fmt.Sprintf("%s-processor", testHelmRelease) + sts := fmt.Sprintf("%s-processor", testHelmRelease) out, err := exec.Command("kubectl", "logs", - fmt.Sprintf("deployment/%s", deployment), + fmt.Sprintf("statefulset/%s", sts), "-n", testNamespace, fmt.Sprintf("--since-time=%s", sinceTime), ).CombinedOutput() if err != nil { - t.Fatalf("kubectl logs for %s failed: %v\n%s", deployment, err, out) + t.Fatalf("kubectl logs for %s failed: %v\n%s", sts, err, out) } return string(out) } diff --git a/test/e2e/helm_upgrade_test.go b/test/e2e/helm_upgrade_test.go index 86aff76ad..6f1a0070e 100644 --- a/test/e2e/helm_upgrade_test.go +++ b/test/e2e/helm_upgrade_test.go @@ -83,7 +83,7 @@ func testHelmUpgrade(t *testing.T) { rollCtx, rollCancel := context.WithTimeout(context.Background(), helmCmdTimeout) defer rollCancel() out, err = exec.CommandContext(rollCtx, "kubectl", "rollout", "status", - fmt.Sprintf("deployment/%s-processor", testHelmRelease), "-n", testNamespace, "--timeout=180s", + fmt.Sprintf("statefulset/%s-processor", testHelmRelease), "-n", testNamespace, "--timeout=180s", ).CombinedOutput() if err != nil { t.Errorf("cleanup: rollout wait failed: %v%s\n%s", err, execContextFailureHint(err), out) @@ -194,21 +194,28 @@ func kubectlGetConfigMap(t *testing.T, name string) string { return result } -func waitForRollout(t *testing.T, deployment string) { +func waitForRollout(t *testing.T, name string) { t.Helper() - t.Logf("waiting for rollout of %s...", deployment) + // Auto-detect whether the resource is a Deployment or StatefulSet. + kind := "deployment" + if out, err := exec.Command("kubectl", "get", "statefulset/"+name, "-n", testNamespace).CombinedOutput(); err == nil && len(out) > 0 { + kind = "statefulset" + } + + resource := fmt.Sprintf("%s/%s", kind, name) + t.Logf("waiting for rollout of %s...", resource) ctx, cancel := context.WithTimeout(context.Background(), helmCmdTimeout) defer cancel() out, err := exec.CommandContext(ctx, "kubectl", "rollout", "status", - fmt.Sprintf("deployment/%s", deployment), + resource, "-n", testNamespace, "--timeout=180s", ).CombinedOutput() if err != nil { - t.Fatalf("rollout of %s failed: %v%s\n%s", deployment, err, execContextFailureHint(err), out) + t.Fatalf("rollout of %s failed: %v%s\n%s", resource, err, execContextFailureHint(err), out) } - t.Logf("rollout of %s complete", deployment) + t.Logf("rollout of %s complete", resource) } // parseModelGatewayMaxRetries returns model_gateways[model].max_retries from processor config.yaml. diff --git a/test/e2e/helpers_test.go b/test/e2e/helpers_test.go index 891cfca79..bdb64f405 100644 --- a/test/e2e/helpers_test.go +++ b/test/e2e/helpers_test.go @@ -965,9 +965,11 @@ func startProcessorObsPortForward(t *testing.T) (*processorObsPortForward, error select { case line, ok := <-lineCh: if !ok { - if err := <-scanDone; err != nil { - return nil, fmt.Errorf("read processor port-forward output: %w", err) + // Drain scanDone here so the defer cleanup doesn't block on it. + if scanErr := <-scanDone; scanErr != nil { + return nil, fmt.Errorf("read processor port-forward output: %w", scanErr) } + scanDone = nil // prevent defer from reading it again return nil, fmt.Errorf("processor port-forward exited before reporting a local port:\n%s", strings.Join(output, "\n")) } output = append(output, line) @@ -1006,11 +1008,19 @@ func (pf *processorObsPortForward) Close() { if pf.reader != nil { _ = pf.reader.Close() } + // Use a timeout to prevent hanging if the subprocess doesn't exit cleanly. + closeTimeout := time.After(5 * time.Second) if pf.waitDone != nil { - <-pf.waitDone + select { + case <-pf.waitDone: + case <-closeTimeout: + } } if pf.scanDone != nil { - <-pf.scanDone + select { + case <-pf.scanDone: + case <-closeTimeout: + } } } @@ -1053,6 +1063,21 @@ func waitForProcessorReady(t *testing.T, timeout time.Duration) { } }() + // Wait for the processor pod to be Running and Ready before attempting + // a port-forward. This avoids the race where kubectl port-forward fails + // because the pod is still Pending after a delete/restart. + podLabel := fmt.Sprintf("app.kubernetes.io/instance=%s,app.kubernetes.io/component=processor", testHelmRelease) + waitCtx, waitCancel := context.WithDeadline(context.Background(), deadline) + defer waitCancel() + if out, err := exec.CommandContext(waitCtx, "kubectl", "wait", "pod", + "-l", podLabel, + "-n", testNamespace, + "--for=condition=Ready", + fmt.Sprintf("--timeout=%ds", int(time.Until(deadline).Seconds())), + ).CombinedOutput(); err != nil { + t.Fatalf("processor pod not ready after %v: %v\n%s", timeout, err, out) + } + for { if pf == nil { var err error diff --git a/test/e2e/orphan_recovery_test.go b/test/e2e/orphan_recovery_test.go index f6db93d40..2d4b2bb63 100644 --- a/test/e2e/orphan_recovery_test.go +++ b/test/e2e/orphan_recovery_test.go @@ -29,8 +29,8 @@ import ( // when no processor is running to handle a stranded job. // // Unlike testProcessorGracefulShutdown which tests SIGTERM with a replacement -// pod available, these tests scale the processor deployment to 0 replicas -// before killing the pod, ensuring no processor can pick up the job. +// pod available, these tests scale the processor StatefulSet to 0 replicas, +// ensuring no processor can self-recover the job. // The reconciler (running in the GC pod) then detects the stale in-flight // entry and transitions the orphaned job to a terminal state. // @@ -43,19 +43,16 @@ func testOrphanRecovery(t *testing.T) { } // doTestHardCrashOrphanRecovery submits a batch with long-running requests, -// force-kills the processor pod and scales the deployment to 0, then verifies -// the GC reconciler transitions the orphaned job to failed. -// -// Since the errShutdown handler does NOT re-enqueue, the job stays in_progress -// in the DB regardless of whether SIGTERM or SIGKILL kills the process. Scaling -// the deployment to 0 ensures no replacement pod can interfere. +// scales the processor StatefulSet to 0 (which kills the pod), then verifies +// the GC reconciler re-enqueues the orphaned job (SLO is still valid with a +// 24h window). After scaling back to 1, the new processor picks up the +// re-enqueued job and completes it. // // Timeline: // 1. Submit batch → wait for in_progress -// 2. Force-kill pod + scale to 0 (no replacement, no re-enqueue) -// 3. Reconciler detects orphan (staleness threshold = reconciler interval) -// 4. in_progress + stale/missing in-flight → reconciler transitions to failed -// 5. Scale processor back to 1 (cleanup for subsequent tests) +// 2. Scale to 0 (kills pod, no replacement available) +// 3. Reconciler detects orphan → re-enqueues (SLO valid) +// 4. Scale back to 1 → processor picks up job and completes it func doTestHardCrashOrphanRecovery(t *testing.T) { t.Helper() @@ -63,10 +60,10 @@ func doTestHardCrashOrphanRecovery(t *testing.T) { t.Skip("kubectl not available, skipping orphan recovery test") } - deployment := fmt.Sprintf("%s-processor", testHelmRelease) + sts := fmt.Sprintf("%s-processor", testHelmRelease) var lines []string - for i := 1; i <= 50; i++ { + for i := 1; i <= 10; i++ { lines = append(lines, fmt.Sprintf( `{"custom_id":"orphan-%d","method":"POST","url":"/v1/chat/completions","body":{"model":"%s","max_tokens":200,"messages":[{"role":"user","content":"slow %d"}]}}`, i, testSimModel, i)) } @@ -76,18 +73,23 @@ func doTestHardCrashOrphanRecovery(t *testing.T) { _, _ = waitForBatchStatus(t, batchID, 2*time.Minute, openai.BatchStatusInProgress) time.Sleep(2 * time.Second) - killAndScaleDown(t, deployment) - t.Cleanup(func() { scaleUp(t, deployment) }) + killAndScaleDown(t, sts) + + // The reconciler detects the orphan and re-enqueues it (SLO still valid). + // Wait for the job to go back to validating (queued, no processor). + waitForBatchStatus(t, batchID, 3*time.Minute, openai.BatchStatusValidating) + t.Log("orphan re-enqueued to validating") - // Wait for the reconciler to detect the orphan and transition it to failed. - // With reconciler interval=60s (dev-deploy): - // - Staleness threshold: 60s after last heartbeat (or immediate if in-flight - // entry was deleted by the processor's defer before process death) - // - Next cycle: up to 60s after staleness - // - Total: ~2m + buffer - finalBatch := waitForOrphanTerminal(t, batchID, 5*time.Minute, openai.BatchStatusFailed) + // Scale back to 1 so the processor picks up the re-enqueued job. + scaleUp(t, sts) - t.Logf("orphan recovery: batch %s reached %s", batchID, finalBatch.Status) + finalBatch := waitForOrphanTerminal(t, batchID, 5*time.Minute, openai.BatchStatusCompleted) + + t.Logf("orphan recovery: batch %s reached %s (completed=%d, failed=%d, total=%d)", + batchID, finalBatch.Status, + finalBatch.RequestCounts.Completed, + finalBatch.RequestCounts.Failed, + finalBatch.RequestCounts.Total) } // doTestCancellingOrphanRecovery submits a batch, waits for it to reach @@ -110,7 +112,7 @@ func doTestCancellingOrphanRecovery(t *testing.T) { t.Skip("kubectl not available, skipping cancelling orphan recovery test") } - deployment := fmt.Sprintf("%s-processor", testHelmRelease) + sts := fmt.Sprintf("%s-processor", testHelmRelease) var lines []string for i := 1; i <= 5; i++ { @@ -138,8 +140,8 @@ func doTestCancellingOrphanRecovery(t *testing.T) { } t.Logf("batch %s is now cancelling", batchID) - killAndScaleDown(t, deployment) - t.Cleanup(func() { scaleUp(t, deployment) }) + killAndScaleDown(t, sts) + t.Cleanup(func() { scaleUp(t, sts) }) finalBatch := waitForOrphanTerminal(t, batchID, 5*time.Minute, openai.BatchStatusCancelled) @@ -151,36 +153,24 @@ func doTestCancellingOrphanRecovery(t *testing.T) { } } -// killAndScaleDown force-kills all processor pods and scales the deployment -// to 0 replicas. Since the processor does NOT re-enqueue on SIGTERM (the -// errShutdown handler is a no-op), a simple force-kill is sufficient — the -// scale-to-0 just prevents replacement pods from interfering with the -// reconciler's orphan detection. -func killAndScaleDown(t *testing.T, deployment string) { +// killAndScaleDown force-kills all processor pods and scales the StatefulSet +// to 0 replicas. The scale-to-0 prevents replacement pods from self-recovering +// the orphaned job, ensuring the GC reconciler handles it. +func killAndScaleDown(t *testing.T, sts string) { t.Helper() podSelector := fmt.Sprintf("app.kubernetes.io/instance=%s,app.kubernetes.io/component=processor", testHelmRelease) - killOut, killErr := exec.Command("kubectl", "delete", "pod", - "-l", podSelector, - "-n", testNamespace, - "--grace-period=0", "--force", - ).CombinedOutput() - if killErr != nil { - t.Logf("force-kill pods (may be already gone): %v\n%s", killErr, killOut) - } else { - t.Logf("force-killed processor pods: %s", strings.TrimSpace(string(killOut))) - } - + // Scale to 0 first so the StatefulSet controller doesn't recreate the pod. scaleOut, scaleErr := exec.Command("kubectl", "scale", - fmt.Sprintf("deployment/%s", deployment), + fmt.Sprintf("statefulset/%s", sts), "--replicas=0", "-n", testNamespace, ).CombinedOutput() if scaleErr != nil { t.Fatalf("kubectl scale --replicas=0 failed: %v\n%s", scaleErr, scaleOut) } - t.Logf("scaled %s to 0: %s", deployment, strings.TrimSpace(string(scaleOut))) + t.Logf("scaled %s to 0: %s", sts, strings.TrimSpace(string(scaleOut))) waitCtx, waitCancel := context.WithTimeout(context.Background(), 2*time.Minute) defer waitCancel() @@ -196,13 +186,13 @@ func killAndScaleDown(t *testing.T, deployment string) { } } -// scaleUp scales the given deployment back to 1 replica and waits for it +// scaleUp scales the given StatefulSet back to 1 replica and waits for it // to become ready. Used in t.Cleanup to restore the processor for subsequent tests. -func scaleUp(t *testing.T, deployment string) { +func scaleUp(t *testing.T, sts string) { t.Helper() out, err := exec.Command("kubectl", "scale", - fmt.Sprintf("deployment/%s", deployment), + fmt.Sprintf("statefulset/%s", sts), "--replicas=1", "-n", testNamespace, ).CombinedOutput() @@ -210,7 +200,7 @@ func scaleUp(t *testing.T, deployment string) { t.Logf("kubectl scale --replicas=1 failed (cleanup): %v\n%s", err, out) return } - t.Logf("scaled %s back to 1: %s", deployment, strings.TrimSpace(string(out))) + t.Logf("scaled %s back to 1: %s", sts, strings.TrimSpace(string(out))) waitCtx, waitCancel := context.WithTimeout(context.Background(), 2*time.Minute) defer waitCancel() diff --git a/test/e2e/processor_graceful_shutdown_test.go b/test/e2e/processor_graceful_shutdown_test.go index 22c0fe72f..78291acf1 100644 --- a/test/e2e/processor_graceful_shutdown_test.go +++ b/test/e2e/processor_graceful_shutdown_test.go @@ -83,11 +83,11 @@ func doTestPodDeleteMidJob(t *testing.T) { waitForProcessorReady(t, 2*time.Minute) t.Log("new processor pod is ready") - // The orphan reconciler detects the stranded in_progress job (not in - // queue, stale or missing in-flight entry) and transitions it to failed. - // Use waitForOrphanTerminal because the reconciler's transition preserves - // whatever request counts existed at crash time and does not upload files. - finalBatch := waitForOrphanTerminal(t, batchID, 5*time.Minute, openai.BatchStatusFailed) + // The processor is a StatefulSet — the restarted pod has the same identity + // (batch-gateway-processor-0) and recovers its own orphaned jobs via + // recoverOwnedJobs at startup. Since the SLO is still valid (24h window), + // the job is re-enqueued, re-processed, and completed. + finalBatch := waitForOrphanTerminal(t, batchID, 5*time.Minute, openai.BatchStatusCompleted) t.Logf("pod delete: batch %s reached %s (completed=%d, failed=%d, total=%d)", batchID, finalBatch.Status, @@ -97,9 +97,8 @@ func doTestPodDeleteMidJob(t *testing.T) { } // doTestRollingRestartReEnqueue submits a batch, triggers a rolling restart of -// the processor deployment, and verifies the GC reconciler transitions the -// orphaned job to failed. Same SIGTERM -> orphan -> reconciler path as -// PodDeleteMidJob, different trigger. +// the processor StatefulSet, and verifies the restarted processor self-recovers +// the orphaned job via recoverOwnedJobs and completes it. // // Rolling restart delivers SIGTERM only after the new pod is Ready (~12s). // To guarantee requests are still in-flight when SIGTERM arrives, we set the @@ -132,12 +131,12 @@ func doTestRollingRestartReEnqueue(t *testing.T) { _, _ = waitForBatchStatus(t, batchID, 2*time.Minute, openai.BatchStatusInProgress) time.Sleep(2 * time.Second) - deployment := fmt.Sprintf("%s-processor", testHelmRelease) - t.Logf("triggering rolling restart of %s...", deployment) + sts := fmt.Sprintf("%s-processor", testHelmRelease) + t.Logf("triggering rolling restart of statefulset/%s...", sts) ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) defer cancel() out, err := exec.CommandContext(ctx, "kubectl", "rollout", "restart", - fmt.Sprintf("deployment/%s", deployment), + fmt.Sprintf("statefulset/%s", sts), "-n", testNamespace, ).CombinedOutput() if err != nil { @@ -145,13 +144,17 @@ func doTestRollingRestartReEnqueue(t *testing.T) { } t.Logf("rollout restart triggered: %s", strings.TrimSpace(string(out))) - waitForRollout(t, deployment) + waitForRollout(t, sts) waitForProcessorReady(t, 2*time.Minute) t.Log("processor rollout complete and ready") - // Same reconciler path as PodDeleteMidJob: the orphaned job is detected - // and transitioned to failed. - finalBatch := waitForOrphanTerminal(t, batchID, 5*time.Minute, openai.BatchStatusFailed) + // Restore normal latency so the re-processed job completes quickly. + setSimAdminConfig(t, testSimService, `{"inter-token-latency":"100ms"}`) + + // The processor is a StatefulSet — the restarted pod has the same identity + // and recovers its own orphaned jobs via recoverOwnedJobs. Since the SLO + // is still valid (24h window), the job is re-enqueued and completed. + finalBatch := waitForOrphanTerminal(t, batchID, 5*time.Minute, openai.BatchStatusCompleted) t.Logf("rolling restart: batch %s reached %s (completed=%d, failed=%d, total=%d)", batchID, finalBatch.Status, diff --git a/test/integration/setup_test.go b/test/integration/setup_test.go index 8df0ee08b..c69073f6d 100644 --- a/test/integration/setup_test.go +++ b/test/integration/setup_test.go @@ -20,12 +20,14 @@ package integration import ( "bytes" + "context" "encoding/json" "io" "mime/multipart" "net/http" "net/http/httptest" "testing" + "time" "github.com/llm-d/llm-d-batch-gateway/internal/apiserver/batch" "github.com/llm-d/llm-d-batch-gateway/internal/apiserver/common" @@ -46,12 +48,11 @@ type testServer struct { URL string Client *http.Client - batchDB *dbmock.MockDBClient[dbapi.BatchItem, dbapi.BatchQuery] - fileDB *dbmock.MockDBClient[dbapi.FileItem, dbapi.FileQuery] - queue *dbmock.MockBatchPriorityQueueClient - event *dbmock.MockBatchEventChannelClient - status *dbmock.MockBatchStatusClient - inFlight *dbmock.MockInFlightClient + batchDB *dbmock.MockDBClient[dbapi.BatchItem, dbapi.BatchQuery] + fileDB *dbmock.MockDBClient[dbapi.FileItem, dbapi.FileQuery] + queue *dbmock.MockBatchPriorityQueueClient + event *dbmock.MockBatchEventChannelClient + status *dbmock.MockBatchStatusClient } func newTestServer(t *testing.T) *testServer { @@ -68,7 +69,29 @@ func newTestServer(t *testing.T) *testServer { queue := dbmock.NewMockBatchPriorityQueueClient() event := dbmock.NewMockBatchEventChannelClient() statusClient := dbmock.NewMockBatchStatusClient() - inFlight := dbmock.NewMockInFlightClient() + + // Mirror Postgres PQDelete: atomically transition cancelled jobs in the DB. + queue.OnDelete = func(ctx context.Context, id string) error { + items, _, _, err := batchDB.DBGet(ctx, + &dbapi.BatchQuery{BaseQuery: dbapi.BaseQuery{IDs: []string{id}}}, + true, 0, 1) + if err != nil || len(items) == 0 { + return err + } + item := items[0] + var statusInfo map[string]any + if err := json.Unmarshal(item.Status, &statusInfo); err != nil { + return err + } + statusInfo["status"] = "cancelled" + statusInfo["cancelled_at"] = time.Now().UTC().Unix() + newStatus, err := json.Marshal(statusInfo) + if err != nil { + return err + } + item.Status = newStatus + return batchDB.DBUpdate(ctx, item, nil) + } filesClient, err := fsclient.New(t.TempDir()) if err != nil { @@ -76,13 +99,12 @@ func newTestServer(t *testing.T) *testServer { } clients := &clientset.Clientset{ - File: filesClient, - BatchDB: batchDB, - FileDB: fileDB, - Queue: queue, - Event: event, - Status: statusClient, - InFlight: inFlight, + File: filesClient, + BatchDB: batchDB, + FileDB: fileDB, + Queue: queue, + Event: event, + Status: statusClient, } config := &common.ServerConfig{ @@ -119,14 +141,13 @@ func newTestServer(t *testing.T) *testServer { }) return &testServer{ - URL: srv.URL, - Client: srv.Client(), - batchDB: batchDB, - fileDB: fileDB, - queue: queue, - event: event, - status: statusClient, - inFlight: inFlight, + URL: srv.URL, + Client: srv.Client(), + batchDB: batchDB, + fileDB: fileDB, + queue: queue, + event: event, + status: statusClient, } }