Skip to content

Commit 36913be

Browse files
lioraronlioraron
authored
feat: in-flight tracking, CAS on DBUpdate, processor instrumentation (llm-d#437)
* feat: add in-flight tracking, CAS on DBUpdate, and processor instrumentation Introduce infrastructure for orphan job recovery (issue llm-d#434, PR 1 of 2). In-flight tracker: Redis hash records (jobID → processorID, timestamp) on dequeue and removes on terminal transition. Heartbeat goroutine refreshes every 5 minutes. CAS on DBUpdate: adds expectedStatus parameter to DBUpdate across all implementations (Redis Lua script, PostgreSQL WHERE clause). Returns ErrConflict sentinel when current status doesn't match expected. Callers that don't need CAS pass nil. PQGetIDs: new method on BatchPriorityQueueClient returns all job IDs in the priority queue via server-side Lua ZSCAN. Used by the reconciler (PR 2). Processor instrumentation: InFlightSet after dequeue, InFlightDelete on all terminal and re-enqueue paths, heartbeat in runJob. Zero behavioral risk — writes breadcrumbs but nothing reads them yet. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * fix: heartbeat context lifetime and CAS argv flattening Address review feedback: - Use dedicated heartbeatCtx with defer cancel in runJob() so the heartbeat goroutine stops when the job finishes, not only on SIGTERM - Flatten expectedStatus + fields into a single args slice before passing to the CAS Lua script, fixing incorrect ARGV shape - Add tests: heartbeat stops on context cancel, CAS with non-nil expectedStatus, PQGetIDs behavior Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * fix: review findings from code analysis - Add missing deleteInFlight on re-enqueue failure path in polling loop - CAS Lua script: distinguish NOT_FOUND from CONFLICT for deleted keys - CAS Go code: verify "OK" explicitly, error on unexpected script results - Add InFlight to Clientset.Close() for defensive resource cleanup - Add InFlight Redis integration tests (set, get, delete, validation) - Add CAS test case for non-existent key - Add copyright headers to new Lua scripts Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * refactor(db): clean up CAS error messages, comments, and naming - Use consistent "DBUpdate:" prefix for all CAS error messages - Remove caller-referencing comments from API declarations - Add PostgreSQL CAS unit tests (match + conflict) - Rename redis_cas_update.lua → redis_update_cas.lua for consistency - Change InFlightEntry JSON tag from "ts" to "ls" for LastSeen Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * refactor: move inFlightKeyName constant, remove stray blank line Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> --------- Signed-off-by: Lior Aronovich <lioraronpr@gmail.com>
1 parent 04e1714 commit 36913be

31 files changed

Lines changed: 1105 additions & 68 deletions

‎cmd/batch-processor/main.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ func run() error {
117117

118118
// init processor
119119
logger.V(logging.INFO).Info("Initializing worker processor", "maxWorkers", cfg.NumWorkers)
120-
proc, err := worker.NewProcessor(cfg, procClients, logger)
120+
proc, err := worker.NewProcessor(cfg, procClients, hostname, logger)
121121
if err != nil {
122122
logger.Error(err, "Failed to create processor")
123123
return err

‎docs/design/orphan-recovery.md‎

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
# Orphan Job Recovery Design
2+
3+
- **Revision**: 1
4+
- **Last Updated**: 2026-05-15
5+
- **Issue**: [#434](https://github.com/llm-d-incubation/batch-gateway/issues/434)
6+
7+
---
8+
9+
## Problem
10+
11+
When a batch-processor pod crashes (OOM kill, node eviction, pod deletion), jobs that were dequeued from the Redis priority queue can become **orphaned** — stuck in a non-terminal DB status with no processor working on them.
12+
13+
The existing startup recovery (`recovery.go`) only handles **container restarts within the same pod** (it scans the local `emptyDir` workdir). If the pod is **replaced** (e.g., node eviction destroys `emptyDir`), those jobs are lost forever.
14+
15+
### Orphan scenarios
16+
17+
1. **Pod replacement**: Processor dequeues a job, then the pod is replaced. Workdir is gone; startup recovery on new pods can't find it.
18+
2. **Silent abandonment**: A worker fails to re-enqueue or transition a job to a terminal state (e.g., DB unreachable during error handling), logs the error and moves on. The pod is healthy, so startup recovery won't run.
19+
20+
### The orphan window
21+
22+
After `PQDequeue` atomically removes a job from Redis, if the pod crashes before the worker creates a workdir and transitions the status:
23+
- **Queue**: Job removed (won't be re-dequeued)
24+
- **Workdir**: Doesn't exist yet (no startup recovery can find it)
25+
- **DB**: Stuck with non-terminal status forever
26+
27+
---
28+
29+
## Design Overview
30+
31+
Three mechanisms work together to detect orphans, recover them, and guarantee data consistency:
32+
33+
1. **In-flight tracker with heartbeat** — processor records job ownership in Redis, refreshes every 5 minutes
34+
2. **Orphan reconciler** (PR 2) — batch-gc periodically cross-references DB, queue, and in-flight hash to find and recover orphans
35+
3. **Compare-and-swap (CAS) on DB status transitions** — prevents a reconciled job from being overwritten by a stale processor
36+
37+
### Data flow
38+
39+
```
40+
Normal path:
41+
PQDequeue → InFlightSet(jobID, podID) → process (heartbeat every 5m)
42+
→ CAS terminal write → InFlightDelete(jobID)
43+
44+
Crash recovery:
45+
PQDequeue → InFlightSet(jobID, podID) → heartbeat → [POD CRASHES]
46+
↓
47+
[Reconciler cycle, ~60 min later]
48+
Query DB: non-terminal jobs
49+
Query queue: pending job IDs → jobID not in queue
50+
Query in-flight hash → heartbeat stale (>60 min old)
51+
Triage: re-enqueue or CAS to terminal
52+
InFlightDelete(jobID) → cleanup
53+
54+
False positive recovery (processor still alive but reconciler acted):
55+
Processor heartbeat reads DB status → sees unexpected status → cancels job context
56+
OR: Processor tries CAS terminal write → ErrConflict → processor aborts
57+
```
58+
59+
---
60+
61+
## Mechanism 1: In-Flight Tracker with Heartbeat
62+
63+
### Redis data structure
64+
65+
**Key**: `llmd_batch:inflight` (single Redis hash)
66+
**Fields**: job IDs
67+
**Values**: JSON `{"pid":"<hostname>","ts":<unix_seconds>}`
68+
69+
Why a hash: `HSCAN` lets the reconciler iterate all in-flight jobs without knowing IDs upfront. Single key, no key-space pollution, `HSET`/`HDEL` are O(1).
70+
71+
### Interface
72+
73+
File: `internal/database/api/database.go`
74+
75+
```go
76+
type InFlightEntry struct {
77+
ProcessorID string `json:"pid"`
78+
LastSeen int64 `json:"ts"`
79+
}
80+
81+
type InFlightClient interface {
82+
store.BatchClientAdmin
83+
InFlightSet(ctx context.Context, jobID, processorID string) error
84+
InFlightDelete(ctx context.Context, jobID string) error
85+
InFlightGetAll(ctx context.Context) (map[string]*InFlightEntry, error)
86+
}
87+
```
88+
89+
Separate from `BatchPriorityQueueClient` so the reconciler gets `InFlightGetAll` without queue mutation methods.
90+
91+
### Redis implementation
92+
93+
File: `internal/database/redis/redis_inflight.go`
94+
95+
- `InFlightSet`: `HSET llmd_batch:inflight <jobID> <json>` — O(1). Serves as both initial registration and heartbeat refresh.
96+
- `InFlightDelete`: `HDEL llmd_batch:inflight <jobID>` — O(1).
97+
- `InFlightGetAll`: `HSCAN llmd_batch:inflight 0 * 100` — paginated iteration.
98+
99+
### Processor instrumentation
100+
101+
**Registration** — after dequeue (`worker.go`):
102+
Call `InFlightSet(jobID, processorID)` immediately after `PQDequeue` succeeds. Non-fatal on error.
103+
104+
> **Why not atomic with PQDequeue**: `PQDequeue` uses `BZMPOP` (blocking pop), which Redis forbids inside Lua scripts. The gap between the two calls is microseconds, while the reconciler runs every 60 minutes. Even if the reconciler scans in between, it would see a `validating` job and re-enqueue it — CAS prevents data corruption when a second processor picks up the duplicate.
105+
106+
**Heartbeat** — in `runJob()` (`job_runner.go`):
107+
A goroutine with a 5-minute ticker calls `InFlightSet` to refresh the timestamp. The goroutine stops when the job context is cancelled.
108+
109+
**Cleanup** — deferred in `runJob()`:
110+
`InFlightDelete` as the first defer (LIFO = executes last, after terminal transitions and panic recovery).
111+
112+
**Re-enqueue paths** (`worker.go`):
113+
All pre-launch paths that re-enqueue a job (DB fetch failure, malformed job, expired, cancelling, etc.) also call `InFlightDelete`.
114+
115+
### Processor identity
116+
117+
`os.Hostname()` = Kubernetes pod name. Already used in `cmd/batch-processor/main.go` for logging.
118+
119+
---
120+
121+
## Mechanism 2: Orphan Reconciler
122+
123+
> Implemented in PR 2.
124+
125+
Runs in **batch-gc** as a periodic scan. Default interval: 60 minutes (also used as the staleness threshold — ≥ 12× the heartbeat interval).
126+
127+
### Algorithm
128+
129+
1. Fetch non-terminal jobs from DB (paginated).
130+
2. Fetch queue membership via `PQGetIDs()` → `map[string]bool`.
131+
3. Fetch in-flight entries via `InFlightGetAll()`.
132+
4. For each non-terminal job:
133+
- In queue → skip (waiting to be picked up).
134+
- In-flight and not stale → skip (actively being processed).
135+
- Otherwise → orphan (triage).
136+
5. Triage:
137+
- `cancelling` → CAS to `cancelled`.
138+
- SLO expired → CAS to `expired`.
139+
- `validating` → re-enqueue.
140+
- `in_progress` / `finalizing` → CAS to `failed` (local files are lost).
141+
6. Cleanup stale in-flight entries.
142+
7. Cleanup terminal in-flight entries (processor crashed between CAS write and `InFlightDelete`).
143+
144+
### PQGetIDs
145+
146+
New method on `BatchPriorityQueueClient`. Implemented as a Lua script that iterates the sorted set server-side via `ZSCAN`, extracts IDs with `cjson.decode`, returns the ID list. Go side converts to `map[string]bool`.
147+
148+
---
149+
150+
## Mechanism 3: Compare-and-Swap (CAS) on DB Status Transitions
151+
152+
### Problem
153+
154+
Without CAS, the reconciler and processor can race: the processor overwrites the reconciler's `failed` with `completed`, leading to status flicker or orphaned file records.
155+
156+
### Design
157+
158+
`DBUpdate` accepts an optional `expectedStatus []byte`. If non-nil, the update only proceeds if the current status matches; otherwise returns `ErrConflict`.
159+
160+
- **PostgreSQL**: `UPDATE ... SET status = $new WHERE id = $id AND status = $expected` — 0 rows affected = conflict.
161+
- **Redis**: Lua script checks status field before writing. Single atomic call.
162+
163+
Callers that don't need CAS pass `nil` (e.g., apiserver cancel handler).
164+
165+
### Conflict handling
166+
167+
Both processor and reconciler use CAS, so regardless of which writes first, the other detects the conflict and yields.
168+
169+
---
170+
171+
## Risks and Mitigations
172+
173+
| Risk | Mitigation |
174+
|------|------------|
175+
| False positive (reconciler recovers a job still being processed) | Heartbeat makes this unlikely (requires ~60 min of sustained Redis unreachability). CAS prevents data corruption. |
176+
| InFlightSet fails (Redis down at dequeue time) | Non-fatal. Reconciler detects orphan via DB + queue cross-reference. |
177+
| CAS conflict during normal operation | Only if reconciler and processor race. Processor sees `ErrConflict`, logs, aborts — no data loss. |
178+
| Large DB scan | Paginated (existing GC pattern). Active batch count bounded by throughput. |
179+
| Race: reconciler vs. startup recovery | Startup recovery runs before polling. `PQEnqueue` uses `ZADD NX` (idempotent). CAS prevents conflicting writes. |
180+
| In-flight hash unbounded | Entries ~100 bytes. Reconciler clears stale entries. |
181+
182+
---
183+
184+
## Implementation Plan
185+
186+
### PR 1: In-flight tracking + CAS + processor instrumentation
187+
188+
Adds tracking infrastructure and processor instrumentation. Writes breadcrumbs but nothing reads them yet — zero behavioral risk.
189+
190+
### PR 2: Orphan reconciler + GC wiring
191+
192+
Builds reconciler logic and wires it into batch-gc. Activates orphan detection and recovery.
193+
194+
### Follow-up: Hardening and observability
195+
196+
- Prometheus metrics for reconciler
197+
- Reconciler dry-run mode validation

‎internal/apiserver/batch/batch_handler.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -569,7 +569,7 @@ func (c *BatchAPIHandler) CancelBatch(w http.ResponseWriter, r *http.Request) {
569569
return
570570
}
571571

572-
if err := c.clients.BatchDB.DBUpdate(ctx, dbItem); err != nil {
572+
if err := c.clients.BatchDB.DBUpdate(ctx, dbItem, nil); err != nil {
573573
logger.Error(err, "failed to update batch in database")
574574
common.WriteInternalServerError(w, r)
575575
return

‎internal/database/api/database.go‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,20 @@ package api
2020

2121
import (
2222
"context"
23+
"errors"
2324
"fmt"
2425
"time"
2526

2627
"github.com/llm-d-incubation/batch-gateway/internal/shared/store"
2728
)
2829

30+
// ErrConflict is returned by DBUpdate when expectedStatus is non-nil and
31+
// the current status in the database does not match. Callers can use
32+
// errors.Is(err, ErrConflict) to detect CAS failures.
33+
var ErrConflict = errors.New("status conflict")
34+
35+
// -- Metadata storage --
36+
2937
// DBClient is a generic interface for managing database items in persistent storage.
3038
//
3139
// Each domain type (e.g., BatchItem, FileItem) gets its own typed DBClient implementation.
@@ -66,7 +74,12 @@ type DBClient[T any, Q any] interface {
6674
// The function will update in the item's record in the database - all the dynamic fields of the item which are not empty
6775
// in the given item object.
6876
// Any dynamic field that is empty in the given item object - will not be updated in the item's record in the database.
69-
DBUpdate(ctx context.Context, item *T) (err error)
77+
//
78+
// When expectedStatus is non-nil, the update is conditional (CAS): it only
79+
// succeeds if the current status in the database matches expectedStatus exactly.
80+
// Returns ErrConflict if the status has changed since it was read.
81+
// Pass nil to skip the CAS check (unconditional update).
82+
DBUpdate(ctx context.Context, item *T, expectedStatus []byte) (err error)
7083

7184
// DBDelete removes items by their IDs.
7285
DBDelete(ctx context.Context, IDs []string) (deletedIDs []string, err error)
@@ -131,6 +144,9 @@ type BatchPriorityQueueClient interface {
131144
// It returns the number of deleted objects.
132145
// An error is returned only if the deletion operation failed.
133146
PQDelete(ctx context.Context, jobPriority *BatchJobPriority) (nDeleted int, err error)
147+
148+
// PQGetIDs returns the set of all job IDs currently in the priority queue.
149+
PQGetIDs(ctx context.Context) (map[string]bool, error)
134150
}
135151

136152
// -- Batch jobs events and channels --
@@ -199,3 +215,28 @@ type BatchStatusClient interface {
199215
// StatusDelete deletes the status data for a job.
200216
StatusDelete(ctx context.Context, ID string) (nDeleted int, err error)
201217
}
218+
219+
// -- In-flight job tracking --
220+
221+
// InFlightEntry records which processor owns a dequeued job and when it last
222+
// reported liveness.
223+
type InFlightEntry struct {
224+
ProcessorID string `json:"pid"`
225+
LastSeen int64 `json:"ls"`
226+
}
227+
228+
// InFlightClient tracks jobs that have been dequeued and are being processed.
229+
type InFlightClient interface {
230+
store.BatchClientAdmin
231+
232+
// InFlightSet records or refreshes the in-flight entry for a job.
233+
// Called after dequeue and periodically as a heartbeat.
234+
InFlightSet(ctx context.Context, jobID, processorID string) error
235+
236+
// InFlightDelete removes the in-flight entry for a job.
237+
// Called when the job reaches a terminal state.
238+
InFlightDelete(ctx context.Context, jobID string) error
239+
240+
// InFlightGetAll returns all in-flight entries keyed by job ID.
241+
InFlightGetAll(ctx context.Context) (map[string]*InFlightEntry, error)
242+
}

‎internal/database/mock/mock_db_client.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ func (m *MockDBClient[T, Q]) DBGet(
106106
return
107107
}
108108

109-
func (m *MockDBClient[T, Q]) DBUpdate(ctx context.Context, item *T) (err error) {
109+
func (m *MockDBClient[T, Q]) DBUpdate(ctx context.Context, item *T, expectedStatus []byte) (err error) {
110110
if item == nil {
111111
return fmt.Errorf("item is nil")
112112
}

‎internal/database/mock/mock_inflight_client.go‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
/*
2+
Copyright 2026 The llm-d Authors
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package mock
18+
19+
import (
20+
"context"
21+
"fmt"
22+
"sync"
23+
"time"
24+
25+
"github.com/llm-d-incubation/batch-gateway/internal/database/api"
26+
)
27+
28+
var _ api.InFlightClient = (*MockInFlightClient)(nil)
29+
30+
type MockInFlightClient struct {
31+
mu sync.Mutex
32+
entries map[string]*api.InFlightEntry
33+
}
34+
35+
func NewMockInFlightClient() *MockInFlightClient {
36+
return &MockInFlightClient{
37+
entries: make(map[string]*api.InFlightEntry),
38+
}
39+
}
40+
41+
func (m *MockInFlightClient) InFlightSet(_ context.Context, jobID, processorID string) error {
42+
if jobID == "" {
43+
return fmt.Errorf("jobID is empty")
44+
}
45+
if processorID == "" {
46+
return fmt.Errorf("processorID is empty")
47+
}
48+
49+
m.mu.Lock()
50+
defer m.mu.Unlock()
51+
m.entries[jobID] = &api.InFlightEntry{
52+
ProcessorID: processorID,
53+
LastSeen: time.Now().Unix(),
54+
}
55+
return nil
56+
}
57+
58+
func (m *MockInFlightClient) InFlightDelete(_ context.Context, jobID string) error {
59+
if jobID == "" {
60+
return fmt.Errorf("jobID is empty")
61+
}
62+
63+
m.mu.Lock()
64+
defer m.mu.Unlock()
65+
delete(m.entries, jobID)
66+
return nil
67+
}
68+
69+
func (m *MockInFlightClient) InFlightGetAll(_ context.Context) (map[string]*api.InFlightEntry, error) {
70+
m.mu.Lock()
71+
defer m.mu.Unlock()
72+
73+
result := make(map[string]*api.InFlightEntry, len(m.entries))
74+
for k, v := range m.entries {
75+
copied := *v
76+
result[k] = &copied
77+
}
78+
return result, nil
79+
}
80+
81+
func (m *MockInFlightClient) Close() error {
82+
m.mu.Lock()
83+
defer m.mu.Unlock()
84+
m.entries = nil
85+
return nil
86+
}

0 commit comments

Comments
 (0)