fix worker capacity and startup cancellation races

[luke-lombardi]

Summary by cubic

Fixes two races: controller worker updates no longer overwrite live scheduler capacity, and container startup/stop now honor cancellation and escalation reliably. Improves runtime stop handling with bounded timeouts, verification, and safe worker restart if the runtime is wedged.

• Bug Fixes
  • Preserve live capacity on existing workers using an atomic Redis script that updates controller-owned fields and recomputes free resources from “used” values.
  • Ignore STOPPING container states when checking for outstanding work during worker rollout; capacity is still accounted for by the runtime.
  • Bound runtime control calls with stopContainerWithContext; add retry + verify + force-delete path, and restart the worker if the runtime remains live (without releasing capacity).
  • Re-arm SIGTERM→SIGKILL escalation when a container starts after a pre-runtime stop request.
  • Propagate request context to spawn and checkpoint restore so startup cancels correctly and doesn’t proceed after cancellation.
  • Return context error when image load is canceled.
  • Correct memory free calculation on resize (test expectation updated).
  • Tests added for rollout behavior, stale-capacity protection, bounded stop, force-delete, stuck-runtime recovery, and checkpoint-restore cancellation.

^{Written for commit 9ea8544. Summary will update on new commits.}

[cubic-dev-ai]

4 issues found across 7 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="pkg/worker/lifecycle.go">

<violation number="1" location="pkg/worker/lifecycle.go:487">
P1: At the startup-timeout or shutdown handoff boundary, one container can release worker capacity twice. Passing the cancelable startup context to `spawn` lets it finalize after `failContainerRequest` has already cleared the same request; the handoff should make finalization idempotent or ensure only one path owns cleanup.</violation>

<violation number="2" location="pkg/worker/lifecycle.go:1975">
P1: A container started after a stop request can be force-killed at the old stop request's deadline instead of receiving its full new grace period. Resetting the boolean does not invalidate the already waiting escalation goroutine; a generation or cancelable timer is needed when re-arming escalation.</violation>
</file>

<file name="pkg/repository/worker_redis.go">

<violation number="1" location="pkg/repository/worker_redis.go:335">
P2: If the worker hash expires between the initial read and this script, `AddWorker` reports success and creates dangling indexes instead of recreating or reporting the missing worker. Return an error for the missing-key branch, or check the script result before adding indexes.</violation>

<violation number="2" location="pkg/repository/worker_redis.go:420">
P2: An existing worker can be updated successfully yet disappear from scheduling when the later index pipeline fails, because the Lua state update is executed outside the transaction that updates the indexes. Queue the script on the same `TxPipeline` so the state and index changes succeed or fail together.</violation>
</file>

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[cubic-dev-ai]

P1: A container started after a stop request can be force-killed at the old stop request's deadline instead of receiving its full new grace period. Resetting the boolean does not invalidate the already waiting escalation goroutine; a generation or cancelable timer is needed when re-arming escalation.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At pkg/worker/lifecycle.go, line 1975:

<comment>A container started after a stop request can be force-killed at the old stop request's deadline instead of receiving its full new grace period. Resetting the boolean does not invalidate the already waiting escalation goroutine; a generation or cancelable timer is needed when re-arming escalation.</comment>

<file context>
@@ -1936,7 +1968,15 @@ func (s *Worker) markContainerRunning(ctx context.Context, request *types.Contai
+			// A grace timer may have fired before the runtime existed. Starting a
+			// runtime is a new enforcement point, so it must receive its own
+			// SIGTERM-to-SIGKILL escalation.
+			instance.StopEscalationStarted.Store(false)
+			s.containerInstances.Set(containerId, instance)
+		}
</file context>

[cubic-dev-ai]

P1: At the startup-timeout or shutdown handoff boundary, one container can release worker capacity twice. Passing the cancelable startup context to spawn lets it finalize after failContainerRequest has already cleared the same request; the handoff should make finalization idempotent or ensure only one path owns cleanup.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At pkg/worker/lifecycle.go, line 487:

<comment>At the startup-timeout or shutdown handoff boundary, one container can release worker capacity twice. Passing the cancelable startup context to `spawn` lets it finalize after `failContainerRequest` has already cleared the same request; the handoff should make finalization idempotent or ensure only one path owns cleanup.</comment>

<file context>
@@ -461,7 +484,7 @@ func (s *Worker) RunContainer(ctx context.Context, request *types.ContainerReque
 		filesystemRestoreHandedOff = true
 		s.containerWg.Add(1)
-		go s.spawn(request, spec, outputLogger, opts)
+		go s.spawn(ctx, request, spec, outputLogger, opts)
 		metrics.RecordWorkerStartupPhase("spawn_enqueue", time.Since(phaseStart), request, nil)
 	}
</file context>

[cubic-dev-ai]

P2: If the worker hash expires between the initial read and this script, AddWorker reports success and creates dangling indexes instead of recreating or reporting the missing worker. Return an error for the missing-key branch, or check the script result before adding indexes.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At pkg/repository/worker_redis.go, line 335:

<comment>If the worker hash expires between the initial read and this script, `AddWorker` reports success and creates dangling indexes instead of recreating or reporting the missing worker. Return an error for the missing-key branch, or check the script result before adding indexes.</comment>

<file context>
@@ -327,6 +327,39 @@ local version = redis.call("HINCRBY", KEYS[1], "resource_version", ARGV[4])
+// preserving the resources already in use if the machine's totals changed.
+var updateExistingWorkerScript = redis.NewScript(`
+if redis.call("EXISTS", KEYS[1]) == 0 then
+    return 0
+end
+
</file context>

Suggested change

	return 0
	return redis.error_reply("worker state disappeared")

[cubic-dev-ai]

P2: An existing worker can be updated successfully yet disappear from scheduling when the later index pipeline fails, because the Lua state update is executed outside the transaction that updates the indexes. Queue the script on the same TxPipeline so the state and index changes succeed or fail together.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At pkg/repository/worker_redis.go, line 420:

<comment>An existing worker can be updated successfully yet disappear from scheduling when the later index pipeline fails, because the Lua state update is executed outside the transaction that updates the indexes. Queue the script on the same `TxPipeline` so the state and index changes succeed or fail together.</comment>

<file context>
@@ -368,10 +401,26 @@ func (r *WorkerRedisRepository) AddWorker(worker *types.Worker) error {
+		}
+		args := []interface{}{worker.TotalCpu, worker.TotalMemory, worker.TotalGpuCount}
+		args = append(args, fields...)
+		if err := updateExistingWorkerScript.Run(ctx, r.rdb, []string{stateKey}, args...).Err(); err != nil {
+			return fmt.Errorf("failed to update existing worker <%s>: %w", stateKey, err)
+		}
</file context>

Suggested change

	if err := updateExistingWorkerScript.Run(ctx, r.rdb, []string{stateKey}, args...).Err(); err != nil {
	return fmt.Errorf("failed to update existing worker <%s>: %w", stateKey, err)
	}
	updateExistingWorkerScript.Eval(ctx, pipe, []string{stateKey}, args...)