| 1 | package taskmonitor |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "fmt" |
| 7 | "strings" |
| 8 | "sync" |
| 9 | "time" |
| 10 | ) |
| 11 | |
| 12 | // JobKiller routes a control request to the runtime that owns a task. SessionID |
| 13 | // is part of the target because local job IDs are only unique within a |
| 14 | // controller. Stop and cancel fail closed when no owner accepts the request. |
| 15 | type JobKiller interface { |
| 16 | Kill(sessionID, jobID string) bool |
| 17 | } |
| 18 | |
| 19 | // runtimeJobID returns the jobs.Manager-local ID for runtime control. JobID is |
| 20 | // present on new snapshots. The prefix fallback keeps snapshots written by the |
| 21 | // short-lived namespaced-ID implementation controllable after an upgrade, while |
| 22 | // legacy snapshots continue to use their unnamespaced TaskID. |
| 23 | func runtimeJobID(snap *TaskSnapshot) string { |
| 24 | if snap == nil { |
| 25 | return "" |
| 26 | } |
| 27 | if snap.JobID != "" { |
| 28 | return snap.JobID |
| 29 | } |
| 30 | if snap.SessionID != "" { |
| 31 | prefix := monitorTaskID(snap.SessionID, "") |
| 32 | if strings.HasPrefix(snap.TaskID, prefix) && len(snap.TaskID) > len(prefix) { |
| 33 | return strings.TrimPrefix(snap.TaskID, prefix) |
| 34 | } |
| 35 | } |
| 36 | return snap.TaskID |
| 37 | } |
| 38 | |
| 39 | // ControlResult is the unified response for all task control operations. |
| 40 | type ControlResult struct { |
| 41 | SchemaVersion int `json:"schema_version"` |
| 42 | Command string `json:"command"` |
| 43 | TaskID string `json:"task_id"` |
| 44 | SessionID string `json:"session_id"` |
| 45 | State TaskState `json:"state"` |
| 46 | RuntimeState RuntimeState `json:"runtime_state,omitempty"` |
| 47 | Version uint64 `json:"version"` |
| 48 | Accepted bool `json:"accepted"` |
| 49 | Idempotent bool `json:"idempotent"` |
| 50 | Error *CtrlError `json:"error,omitempty"` |
| 51 | } |
| 52 | |
| 53 | // CtrlError carries a stable machine-readable code and message. |
| 54 | type CtrlError struct { |
| 55 | Code string `json:"code"` |
| 56 | Message string `json:"message"` |
| 57 | } |
| 58 | |
| 59 | const ( |
| 60 | ErrTaskNotFound = "task_not_found" |
| 61 | ErrTaskScopeMismatch = "task_scope_mismatch" |
| 62 | ErrTaskVersionConflict = "task_version_conflict" |
| 63 | ErrTaskInvalidTransition = "task_invalid_transition" |
| 64 | ErrTaskNotRequeueable = "task_not_requeueable" |
| 65 | ErrTaskAlreadyTerminal = "task_already_terminal" |
| 66 | ErrTaskInProgress = "task_operation_in_progress" |
| 67 | ErrTaskPermissionDenied = "task_permission_denied" |
| 68 | ErrTaskIdempotencyConflict = "task_idempotency_conflict" |
| 69 | ErrTaskAuditFailed = "task_audit_failed" |
| 70 | ErrTaskRuntimeUnavailable = "task_runtime_unavailable" |
| 71 | ) |
| 72 | |
| 73 | // ControlService provides atomic control operations on tasks. |
| 74 | type ControlService struct { |
| 75 | mu sync.Mutex |
| 76 | store WriteStore |
| 77 | } |
| 78 | |
| 79 | // NewControlService returns a ControlService backed by store. |
| 80 | func NewControlService(store WriteStore) *ControlService { |
| 81 | return &ControlService{store: store} |
| 82 | } |
| 83 | |
| 84 | func (cs *ControlService) StopTask(ctx context.Context, projectDir, taskID string, expectedVersion uint64, reason, idemKey string) (ControlResult, error) { |
| 85 | return cs.StopTaskWithKiller(ctx, projectDir, taskID, expectedVersion, reason, idemKey, nil) |
| 86 | } |
| 87 | |
| 88 | // StopTaskWithKiller binds the live runtime target to this control operation. |
| 89 | // Keeping the killer call-scoped prevents concurrent clients from overwriting a |
| 90 | // shared killer and cancelling a same-named job in another session. |
| 91 | func (cs *ControlService) StopTaskWithKiller(ctx context.Context, projectDir, taskID string, expectedVersion uint64, reason, idemKey string, killer JobKiller) (ControlResult, error) { |
| 92 | return cs.controlOp(ctx, projectDir, taskID, expectedVersion, "stop", TaskStateCancelled, reason, idemKey, killer) |
| 93 | } |
| 94 | |
| 95 | func (cs *ControlService) CancelTask(ctx context.Context, projectDir, taskID string, expectedVersion uint64, reason, idemKey string) (ControlResult, error) { |
| 96 | return cs.CancelTaskWithKiller(ctx, projectDir, taskID, expectedVersion, reason, idemKey, nil) |
| 97 | } |
| 98 | |
| 99 | // CancelTaskWithKiller is the call-scoped-killer form of CancelTask. |
| 100 | func (cs *ControlService) CancelTaskWithKiller(ctx context.Context, projectDir, taskID string, expectedVersion uint64, reason, idemKey string, killer JobKiller) (ControlResult, error) { |
| 101 | return cs.controlOp(ctx, projectDir, taskID, expectedVersion, "cancel", TaskStateCancelled, reason, idemKey, killer) |
| 102 | } |
| 103 | |
| 104 | // RequeueTask moves a failed or stale task back to queued. It does not start a |
| 105 | // new runtime; RuntimeState therefore remains exited (or unknown for legacy |
| 106 | // data) until a scheduler starts the task and records a new lifecycle. |
| 107 | func (cs *ControlService) RequeueTask(ctx context.Context, projectDir, taskID string, expectedVersion uint64, idemKey string) (ControlResult, error) { |
| 108 | return cs.controlOp(ctx, projectDir, taskID, expectedVersion, "requeue", TaskStateQueued, "", idemKey, nil) |
| 109 | } |
| 110 | |
| 111 | func (cs *ControlService) OpenTaskSession(ctx context.Context, projectDir, taskID string) (ControlResult, error) { |
| 112 | snap, err := cs.store.GetTask(ctx, projectDir, taskID) |
| 113 | if err != nil { |
| 114 | return ControlResult{}, err |
| 115 | } |
| 116 | if snap == nil { |
| 117 | return ControlResult{ |
| 118 | SchemaVersion: 1, Command: "open_session", TaskID: taskID, |
| 119 | Error: &CtrlError{Code: ErrTaskNotFound, Message: "task not found"}, |
| 120 | }, nil |
| 121 | } |
| 122 | return ControlResult{ |
| 123 | SchemaVersion: 1, Command: "open_session", |
| 124 | TaskID: snap.TaskID, SessionID: snap.SessionID, |
| 125 | State: snap.State, RuntimeState: snap.RuntimeState, |
| 126 | Version: snap.Version, Accepted: true, |
| 127 | }, nil |
| 128 | } |
| 129 | |
| 130 | func (cs *ControlService) controlOp(ctx context.Context, projectDir, taskID string, expectedVersion uint64, cmd string, targetState TaskState, reason, idemKey string, killer JobKiller) (ControlResult, error) { |
| 131 | if err := ctx.Err(); err != nil { |
| 132 | return ControlResult{}, err |
| 133 | } |
| 134 | |
| 135 | cs.mu.Lock() |
| 136 | defer cs.mu.Unlock() |
| 137 | |
| 138 | var claimer IdempotencyClaimer |
| 139 | claimed := false |
| 140 | releaseClaim := func() { |
| 141 | if claimed && claimer != nil { |
| 142 | _ = claimer.ReleaseIdempotency(ctx, projectDir, idemKey) |
| 143 | claimed = false |
| 144 | } |
| 145 | } |
| 146 | |
| 147 | // ── idempotency check (persisted) ── |
| 148 | if idemKey != "" { |
| 149 | var rec *IdempotencyRecord |
| 150 | var err error |
| 151 | if c, ok := cs.store.(IdempotencyClaimer); ok { |
| 152 | claimer = c |
| 153 | rec, err = claimer.ClaimIdempotency(ctx, projectDir, IdempotencyRecord{Key: idemKey, Op: cmd, TaskID: taskID, Version: expectedVersion}) |
| 154 | if err != nil { |
| 155 | return ControlResult{}, fmt.Errorf("claim idempotency: %w", err) |
| 156 | } |
| 157 | claimed = rec == nil |
| 158 | } else { |
| 159 | rec, err = cs.store.CheckIdempotency(ctx, projectDir, idemKey) |
| 160 | if err != nil { |
| 161 | return ControlResult{}, fmt.Errorf("check idempotency: %w", err) |
| 162 | } |
| 163 | } |
| 164 | if rec != nil { |
| 165 | // Must match exactly |
| 166 | if rec.Op != cmd || rec.TaskID != taskID || rec.Version != expectedVersion { |
| 167 | return ControlResult{ |
| 168 | SchemaVersion: 1, Command: cmd, TaskID: taskID, |
| 169 | Error: &CtrlError{Code: ErrTaskIdempotencyConflict, Message: "idempotency key reused with different parameters"}, |
| 170 | }, nil |
| 171 | } |
| 172 | if rec.Pending { |
| 173 | return ControlResult{SchemaVersion: 1, Command: cmd, TaskID: taskID, |
| 174 | Error: &CtrlError{Code: ErrTaskInProgress, Message: "idempotency key is in progress"}}, nil |
| 175 | } |
| 176 | // Replay: fetch current state |
| 177 | snap, err := cs.store.GetTask(ctx, projectDir, taskID) |
| 178 | if err != nil { |
| 179 | return ControlResult{}, err |
| 180 | } |
| 181 | if snap == nil { |
| 182 | return ControlResult{ |
| 183 | SchemaVersion: 1, Command: cmd, TaskID: taskID, |
| 184 | Error: &CtrlError{Code: ErrTaskNotFound, Message: "task not found"}, |
| 185 | }, nil |
| 186 | } |
| 187 | return ControlResult{ |
| 188 | SchemaVersion: 1, Command: cmd, TaskID: taskID, SessionID: snap.SessionID, |
| 189 | State: snap.State, RuntimeState: snap.RuntimeState, |
| 190 | Version: snap.Version, Accepted: true, Idempotent: true, |
| 191 | }, nil |
| 192 | } |
| 193 | } |
| 194 | |
| 195 | // ── fetch + validate ── |
| 196 | snap, err := cs.store.GetTask(ctx, projectDir, taskID) |
| 197 | if err != nil { |
| 198 | releaseClaim() |
| 199 | return ControlResult{}, err |
| 200 | } |
| 201 | if snap == nil { |
| 202 | releaseClaim() |
| 203 | return ControlResult{ |
| 204 | SchemaVersion: 1, Command: cmd, TaskID: taskID, |
| 205 | Error: &CtrlError{Code: ErrTaskNotFound, Message: "task not found"}, |
| 206 | }, nil |
| 207 | } |
| 208 | if expectedVersion != snap.Version { |
| 209 | releaseClaim() |
| 210 | return ControlResult{ |
| 211 | SchemaVersion: 1, Command: cmd, TaskID: taskID, SessionID: snap.SessionID, |
| 212 | State: snap.State, RuntimeState: snap.RuntimeState, Version: snap.Version, |
| 213 | Error: &CtrlError{Code: ErrTaskVersionConflict, Message: "version mismatch"}, |
| 214 | }, nil |
| 215 | } |
| 216 | requeue := cmd == "requeue" |
| 217 | requeueable := requeue && (snap.State == TaskStateFailed || snap.State == TaskStateStale) |
| 218 | if requeue && !requeueable { |
| 219 | releaseClaim() |
| 220 | return ControlResult{ |
| 221 | SchemaVersion: 1, Command: cmd, TaskID: taskID, SessionID: snap.SessionID, |
| 222 | State: snap.State, RuntimeState: snap.RuntimeState, Version: snap.Version, |
| 223 | Error: &CtrlError{Code: ErrTaskNotRequeueable, Message: "task is not failed or stale"}, |
| 224 | }, nil |
| 225 | } |
| 226 | if requeueable && snap.RuntimeState.Effective() == RuntimeStateAlive { |
| 227 | releaseClaim() |
| 228 | return ControlResult{ |
| 229 | SchemaVersion: 1, Command: cmd, TaskID: taskID, SessionID: snap.SessionID, |
| 230 | State: snap.State, RuntimeState: snap.RuntimeState, Version: snap.Version, |
| 231 | Error: &CtrlError{Code: ErrTaskInProgress, Message: "task runtime is still alive"}, |
| 232 | }, nil |
| 233 | } |
| 234 | if snap.State.Terminal() && !requeueable { |
| 235 | releaseClaim() |
| 236 | return ControlResult{ |
| 237 | SchemaVersion: 1, Command: cmd, TaskID: taskID, SessionID: snap.SessionID, |
| 238 | State: snap.State, RuntimeState: snap.RuntimeState, Version: snap.Version, |
| 239 | Error: &CtrlError{Code: ErrTaskAlreadyTerminal, Message: "task is terminal"}, |
| 240 | }, nil |
| 241 | } |
| 242 | if !requeueable && !snap.State.ValidTransition(targetState) { |
| 243 | releaseClaim() |
| 244 | return ControlResult{ |
| 245 | SchemaVersion: 1, Command: cmd, TaskID: taskID, SessionID: snap.SessionID, |
| 246 | State: snap.State, RuntimeState: snap.RuntimeState, Version: snap.Version, |
| 247 | Error: &CtrlError{Code: ErrTaskInvalidTransition, Message: "invalid transition"}, |
| 248 | }, nil |
| 249 | } |
| 250 | runtimeControl := cmd == "stop" || cmd == "cancel" |
| 251 | if runtimeControl && killer == nil { |
| 252 | releaseClaim() |
| 253 | return ControlResult{ |
| 254 | SchemaVersion: 1, Command: cmd, TaskID: taskID, SessionID: snap.SessionID, |
| 255 | State: snap.State, RuntimeState: snap.RuntimeState, Version: snap.Version, |
| 256 | Error: &CtrlError{Code: ErrTaskRuntimeUnavailable, Message: "task runtime owner is unavailable"}, |
| 257 | }, nil |
| 258 | } |
| 259 | if runtimeControl && !killer.Kill(snap.SessionID, runtimeJobID(snap)) { |
| 260 | releaseClaim() |
| 261 | return ControlResult{ |
| 262 | SchemaVersion: 1, Command: cmd, TaskID: taskID, SessionID: snap.SessionID, |
| 263 | State: snap.State, RuntimeState: snap.RuntimeState, Version: snap.Version, |
| 264 | Error: &CtrlError{Code: ErrTaskRuntimeUnavailable, Message: "task runtime owner rejected control request"}, |
| 265 | }, nil |
| 266 | } |
| 267 | |
| 268 | if runtimeControl { |
| 269 | // The runtime owner accepted the request. From this point the idempotency |
| 270 | // claim must not be released: a retry must never repeat an admitted runtime |
| 271 | // side effect merely because persistence or audit reporting failed. |
| 272 | claimed = false |
| 273 | } |
| 274 | |
| 275 | // ── 1. SaveTask (state mutation) ── |
| 276 | // Kill admission races the recorder's terminal completion and the runtime |
| 277 | // heartbeat. Retry those expected version advances. If RecordDone already |
| 278 | // persisted the requested terminal state, use that snapshot as the result. |
| 279 | const maxControlSaveAttempts = 4 |
| 280 | for attempt := 0; attempt < maxControlSaveAttempts; attempt++ { |
| 281 | next := *snap |
| 282 | next.Version++ |
| 283 | next.State = targetState |
| 284 | next.UpdatedAt = timeNow() |
| 285 | if requeueable { |
| 286 | next.RuntimeLeaseUntil = time.Time{} |
| 287 | next.RuntimeOwnerID = "" |
| 288 | } |
| 289 | if runtimeControl && next.RuntimeState.Effective() == RuntimeStateAlive && next.RuntimeLeaseUntil.IsZero() { |
| 290 | // A successful kill request is only an admission signal: the runtime |
| 291 | // may still be exiting. Preserve an existing owner lease, and give |
| 292 | // legacy lease-less snapshots a bounded deadline so observers can |
| 293 | // eventually reconcile alive to exited if RecordDone never arrives. |
| 294 | next.RuntimeLeaseUntil = next.UpdatedAt.Add(runtimeLeaseTTL) |
| 295 | } |
| 296 | if err := cs.store.SaveTask(ctx, projectDir, next); err == nil { |
| 297 | snap = &next |
| 298 | claimed = false |
| 299 | break |
| 300 | } else if !errors.Is(err, ErrStoreVersionConflict) { |
| 301 | releaseClaim() |
| 302 | return ControlResult{}, fmt.Errorf("save task control state: %w", err) |
| 303 | } |
| 304 | |
| 305 | latest, getErr := cs.store.GetTask(ctx, projectDir, taskID) |
| 306 | if getErr != nil { |
| 307 | releaseClaim() |
| 308 | return ControlResult{}, getErr |
| 309 | } |
| 310 | if latest == nil { |
| 311 | releaseClaim() |
| 312 | return ControlResult{}, fmt.Errorf("save task control state: task disappeared") |
| 313 | } |
| 314 | if latest.State == targetState { |
| 315 | snap = latest |
| 316 | claimed = false |
| 317 | break |
| 318 | } |
| 319 | if latest.State.Terminal() || !latest.State.ValidTransition(targetState) { |
| 320 | releaseClaim() |
| 321 | return ControlResult{ |
| 322 | SchemaVersion: 1, Command: cmd, TaskID: taskID, SessionID: latest.SessionID, |
| 323 | State: latest.State, RuntimeState: latest.RuntimeState, Version: latest.Version, |
| 324 | Error: &CtrlError{Code: ErrTaskVersionConflict, Message: "task changed concurrently after runtime accepted control"}, |
| 325 | }, nil |
| 326 | } |
| 327 | snap = latest |
| 328 | if attempt == maxControlSaveAttempts-1 { |
| 329 | releaseClaim() |
| 330 | return ControlResult{ |
| 331 | SchemaVersion: 1, Command: cmd, TaskID: taskID, SessionID: latest.SessionID, |
| 332 | State: latest.State, RuntimeState: latest.RuntimeState, Version: latest.Version, |
| 333 | Error: &CtrlError{Code: ErrTaskVersionConflict, Message: "task kept changing after runtime accepted control"}, |
| 334 | }, nil |
| 335 | } |
| 336 | } |
| 337 | |
| 338 | // ── 2. AppendAuditEvent (atomic sequence + write) ── |
| 339 | auditEv := TaskEvent{ |
| 340 | Sequence: 0, // assigned atomically by store |
| 341 | Timestamp: timeNow(), |
| 342 | EventType: "control_" + cmd, |
| 343 | TaskID: taskID, |
| 344 | SessionID: snap.SessionID, |
| 345 | State: targetState, |
| 346 | RuntimeState: snap.RuntimeState, |
| 347 | } |
| 348 | if err := cs.store.AppendAuditEvent(ctx, projectDir, auditEv); err != nil { |
| 349 | // State is committed but audit is missing. This is a degraded |
| 350 | // but not silent state — the caller receives an error. |
| 351 | return ControlResult{ |
| 352 | SchemaVersion: 1, Command: cmd, TaskID: taskID, |
| 353 | Error: &CtrlError{Code: ErrTaskAuditFailed, Message: "state saved but audit event failed"}, |
| 354 | }, fmt.Errorf("append audit event: %w", err) |
| 355 | } |
| 356 | |
| 357 | // ── 3. RecordIdempotency (claim key after successful mutation) ── |
| 358 | if idemKey != "" { |
| 359 | rec := IdempotencyRecord{Key: idemKey, Op: cmd, TaskID: taskID, Version: expectedVersion} |
| 360 | var err error |
| 361 | if claimer != nil { |
| 362 | err = claimer.FinalizeIdempotency(ctx, projectDir, rec) |
| 363 | claimed = false |
| 364 | } else { |
| 365 | err = cs.store.RecordIdempotency(ctx, projectDir, rec) |
| 366 | } |
| 367 | if err != nil { |
| 368 | return ControlResult{}, fmt.Errorf("record idempotency: %w", err) |
| 369 | } |
| 370 | } |
| 371 | |
| 372 | return ControlResult{ |
| 373 | SchemaVersion: 1, Command: cmd, TaskID: taskID, SessionID: snap.SessionID, |
| 374 | State: snap.State, RuntimeState: snap.RuntimeState, |
| 375 | Version: snap.Version, Accepted: true, |
| 376 | }, nil |
| 377 | } |
| 378 | |
| 379 | var timeNow = func() time.Time { return time.Now() } |
| 380 |