From 9321285f57e6acd7ac38547cc62590afc84d21ec Mon Sep 17 00:00:00 2001 From: seemeroland Date: Mon, 5 Oct 2026 16:09:31 -0700 Subject: [PATCH 1/4] Scope worker cancellation and closure acknowledgements to executions Co-Authored-By: Warp --- internal/types/messages.go | 15 +- internal/worker/backend.go | 11 +- internal/worker/command_integration_test.go | 2 +- internal/worker/direct.go | 10 + internal/worker/dispatch_test.go | 20 +- internal/worker/docker.go | 34 +- .../worker/execution_cancellation_test.go | 317 ++++++++++++++++++ .../worker/execution_closure_backends_test.go | 193 +++++++++++ internal/worker/kubernetes.go | 48 ++- internal/worker/worker.go | 246 ++++++++++---- internal/worker/worker_test.go | 45 +-- 11 files changed, 833 insertions(+), 108 deletions(-) create mode 100644 internal/worker/execution_cancellation_test.go create mode 100644 internal/worker/execution_closure_backends_test.go diff --git a/internal/types/messages.go b/internal/types/messages.go index 7332600..56f49e8 100644 --- a/internal/types/messages.go +++ b/internal/types/messages.go @@ -15,6 +15,7 @@ const ( MessageTypeTaskFailed MessageType = "task_failed" MessageTypeTaskRejected MessageType = "task_rejected" MessageTypeTaskCancellation MessageType = "task_cancellation" + MessageTypeExecutionClosed MessageType = "execution_closed" MessageTypeHeartbeat MessageType = "heartbeat" ) @@ -131,8 +132,9 @@ type TaskAssignmentMessage struct { // TaskClaimedMessage is sent from worker to server after successfully claiming a task type TaskClaimedMessage struct { - TaskID string `json:"task_id"` - WorkerID string `json:"worker_id"` + TaskID string `json:"task_id"` + WorkerID string `json:"worker_id"` + ExecutionID string `json:"execution_id,omitempty"` } // TaskCompletedMessage tells the server to end the active run execution after a successful agent process exit. @@ -166,7 +168,14 @@ type TaskRejectedMessage struct { // TaskCancellationMessage is sent from server to worker to cancel an active task. type TaskCancellationMessage struct { - TaskID string `json:"task_id"` + TaskID string `json:"task_id"` + ExecutionID string `json:"execution_id,omitempty"` +} + +// ExecutionClosedMessage confirms backend closure without changing the logical run state. +type ExecutionClosedMessage struct { + TaskID string `json:"task_id"` + ExecutionID string `json:"execution_id"` } // TaskState is the serialized terminal task state accepted by warp-server. diff --git a/internal/worker/backend.go b/internal/worker/backend.go index 6e77844..496c997 100644 --- a/internal/worker/backend.go +++ b/internal/worker/backend.go @@ -176,9 +176,9 @@ type Backend interface { // failures are surfaced as ExecuteOutcomeError with an error that is (or // wraps) *TaskFailure. ExecuteTask(ctx context.Context, params *TaskParams) ExecuteResult - // CancelTask makes a best-effort attempt to cancel a task. The worker - // invokes it for every task cancellation, alongside cancelling the - // ExecuteTask context. + // CancelTask makes a best-effort attempt to cancel a task after the + // ExecuteTask context is cancelled and ExecuteTask returns. A nil error + // is not proof that asynchronous backend resources have stopped. CancelTask(ctx context.Context, params *CancelParams) error // PreservesTasksOnShutdown reports whether active task execution units can // safely outlive the worker process during shutdown. @@ -195,6 +195,11 @@ type CancelParams struct { ExecutionID string } +// executionClosureVerifier confirms closure after ExecuteTask and CancelTask have returned. +type executionClosureVerifier interface { + ConfirmTaskClosed(ctx context.Context, params *CancelParams) (bool, error) +} + // TaskFailure is the structured error backends return from ExecuteTask when // task execution fails. Backends record only the facts they observe (metrics // labels and the failing process's exit status); the worker adds lifecycle diff --git a/internal/worker/command_integration_test.go b/internal/worker/command_integration_test.go index 7829987..7178590 100644 --- a/internal/worker/command_integration_test.go +++ b/internal/worker/command_integration_test.go @@ -94,7 +94,7 @@ func TestIntegrationCommandBackendDispatchSuppressesTerminalMessage(t *testing.T waitFor(t, 5*time.Second, func() bool { w.tasksMutex.Lock() defer w.tasksMutex.Unlock() - task, ok := w.activeTasks["task-1"] + task, ok := w.activeTasks[taskExecution{taskID: "task-1", executionID: "exec-1"}] return ok && task.spawned }) diff --git a/internal/worker/direct.go b/internal/worker/direct.go index 3c176e1..fa28f12 100644 --- a/internal/worker/direct.go +++ b/internal/worker/direct.go @@ -173,6 +173,16 @@ func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe } else { // Create per-task workspace directory. workspaceDir = filepath.Join(b.config.WorkspaceRoot, taskID) + if params.ExecutionID != "" { + if err := validateTaskIDForPath(params.ExecutionID); err != nil { + return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonWorkspaceSetup, fmt.Errorf("invalid execution ID for workspace path: %w", err))) + } + var err error + workspaceDir, err = os.MkdirTemp(b.config.WorkspaceRoot, taskID+"-"+params.ExecutionID+"-") + if err != nil { + return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonWorkspaceSetup, fmt.Errorf("failed to create execution workspace: %w", err))) + } + } if err := os.MkdirAll(workspaceDir, 0700); err != nil { return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonWorkspaceSetup, fmt.Errorf("failed to create workspace directory: %w", err))) } diff --git a/internal/worker/dispatch_test.go b/internal/worker/dispatch_test.go index 988337c..59798a0 100644 --- a/internal/worker/dispatch_test.go +++ b/internal/worker/dispatch_test.go @@ -37,7 +37,7 @@ func newDispatchWorker(backend Backend) *Worker { ctx: context.Background(), config: Config{}, outbound: newOutboundQueue(4), - activeTasks: map[string]activeTask{"task-1": {cancel: func() {}, executionID: "exec-1"}}, + activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "exec-1"}: {cancel: func() {}, executionID: "exec-1"}}, backend: backend, } } @@ -61,7 +61,7 @@ func TestExecuteTaskDispatchedSuppressesTerminalMessage(t *testing.T) { t.Fatalf("expected no terminal message after dispatch, got %q", msg.Type) } w.tasksMutex.Lock() - task, ok := w.activeTasks["task-1"] + task, ok := w.activeTasks[taskExecution{taskID: "task-1", executionID: "exec-1"}] w.tasksMutex.Unlock() if !ok { t.Fatal("spawned task should remain in activeTasks") @@ -104,7 +104,7 @@ func TestHandleTaskCancellationRoutesToBackendCancelTask(t *testing.T) { w := &Worker{ ctx: context.Background(), outbound: newOutboundQueue(1), - activeTasks: map[string]activeTask{"task-1": spawnedActiveTask("exec-1")}, + activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "exec-1"}: spawnedActiveTask("exec-1")}, backend: backend, } @@ -120,10 +120,10 @@ func TestHandleTaskCancellationRoutesToBackendCancelTask(t *testing.T) { } w.tasksMutex.Lock() - _, ok := w.activeTasks["task-1"] + _, ok := w.activeTasks[taskExecution{taskID: "task-1", executionID: "exec-1"}] w.tasksMutex.Unlock() - if ok { - t.Error("spawned task should be removed after cancellation is routed") + if !ok { + t.Error("spawned task should remain tracked until closure is confirmed") } } @@ -134,7 +134,7 @@ func TestHandleTaskCancellationRunningTaskCancelsContextAndBackend(t *testing.T) w := &Worker{ ctx: context.Background(), outbound: newOutboundQueue(1), - activeTasks: map[string]activeTask{"task-1": { + activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "exec-1"}: { ctx: taskCtx, cancel: taskCancel, executionID: "exec-1", @@ -156,9 +156,9 @@ func TestHandleTaskCancellationRunningTaskCancelsContextAndBackend(t *testing.T) t.Fatal("CancelTask was not invoked for a running task") } - // The entry stays until executeTask's deferred cleanup removes it. + // An unverified cancellation must stay available for retry. w.tasksMutex.Lock() - task, ok := w.activeTasks["task-1"] + task, ok := w.activeTasks[taskExecution{taskID: "task-1", executionID: "exec-1"}] w.tasksMutex.Unlock() if !ok { t.Fatal("running task should remain in activeTasks until executeTask returns") @@ -172,7 +172,7 @@ func TestHandleTaskCancellationSpawnedNoopCancelEmitsNoMessage(t *testing.T) { w := &Worker{ ctx: context.Background(), outbound: newOutboundQueue(1), - activeTasks: map[string]activeTask{"task-1": spawnedActiveTask("exec-1")}, + activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "exec-1"}: spawnedActiveTask("exec-1")}, backend: &dispatchBackend{}, } diff --git a/internal/worker/docker.go b/internal/worker/docker.go index 5f530b0..ce4c4d5 100644 --- a/internal/worker/docker.go +++ b/internal/worker/docker.go @@ -67,6 +67,7 @@ type DockerBackend struct { platform string // Docker daemon platform (e.g., "linux/amd64" or "linux/arm64") platformSpec ocispec.Platform sidecarVolumes sync.Map // volume name -> *sidecarVolumeState + taskContainers sync.Map // taskExecution -> container ID; retained after cleanup failures } // NewDockerBackend creates a new Docker backend, connecting to the Docker daemon. @@ -172,6 +173,9 @@ func (b *DockerBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe } doneContainerStart := params.SetupEvents.startPhase(ctx, SetupEventContainerStart) + key := taskExecution{params.TaskID, params.ExecutionID} + // A failed create response cannot prove that the daemon created no container. + b.taskContainers.Store(key, "") resp, err := dockerClient.ContainerCreate(ctx, client.ContainerCreateOptions{ Config: containerConfig, HostConfig: hostConfig, @@ -182,14 +186,17 @@ func (b *DockerBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe } containerID := resp.ID + b.taskContainers.Store(key, containerID) log.Debugf(ctx, "Created Docker container: %s", containerID) defer func() { if containerID != "" && !b.config.NoCleanup { cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), BackendShutdownTimeout) defer cleanupCancel() - if _, removeErr := dockerClient.ContainerRemove(cleanupCtx, containerID, client.ContainerRemoveOptions{Force: true}); removeErr != nil { + if _, removeErr := dockerClient.ContainerRemove(cleanupCtx, containerID, client.ContainerRemoveOptions{Force: true}); removeErr != nil && !cerrdefs.IsNotFound(removeErr) { log.Debugf(ctx, "Container %s already removed or removal failed: %v", containerID, removeErr) + } else { + b.taskContainers.Delete(key) } } }() @@ -261,9 +268,28 @@ func dockerResourcesForShape(shape *types.InstanceShape) container.Resources { return res } -// CancelTask is a no-op: cancelling the ExecuteTask context fully stops a -// Docker-backend task. -func (b *DockerBackend) CancelTask(context.Context, *CancelParams) error { return nil } +// CancelTask removes the captured execution container, including when normal cleanup failed. +func (b *DockerBackend) CancelTask(ctx context.Context, params *CancelParams) error { + key := taskExecution{params.TaskID, params.ExecutionID} + containerID, ok := b.taskContainers.Load(key) + if !ok { + return nil + } + if containerID == "" { + return fmt.Errorf("container creation outcome is unknown for execution %s", params.ExecutionID) + } + if _, err := b.dockerClient.ContainerRemove(ctx, containerID.(string), client.ContainerRemoveOptions{Force: true}); err != nil && !cerrdefs.IsNotFound(err) { + return err + } + b.taskContainers.Delete(key) + return nil +} + +// ConfirmTaskClosed reports whether the execution has no remaining container. +func (b *DockerBackend) ConfirmTaskClosed(_ context.Context, params *CancelParams) (bool, error) { + _, exists := b.taskContainers.Load(taskExecution{params.TaskID, params.ExecutionID}) + return !exists, nil +} // Shutdown closes the Docker client. func (b *DockerBackend) Shutdown(ctx context.Context) { diff --git a/internal/worker/execution_cancellation_test.go b/internal/worker/execution_cancellation_test.go new file mode 100644 index 0000000..b62eec4 --- /dev/null +++ b/internal/worker/execution_cancellation_test.go @@ -0,0 +1,317 @@ +package worker + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +type executionTestBackend struct { + dispatchBackend + started chan string + release map[string]chan struct{} + cancelled chan CancelParams + cancelGate chan struct{} + failCancel atomic.Bool +} + +func TestClaimIncludesAssignedExecution(t *testing.T) { + for _, id := range []string{"", "old", "new"} { + t.Run(id, func(t *testing.T) { + w := &Worker{config: Config{WorkerID: "worker"}, outbound: newOutboundQueue(1)} + if err := w.sendTaskClaimed("run", id); err != nil { + t.Fatal(err) + } + msg := readWebSocketMessage(t, w.outbound.messages) + var claimed types.TaskClaimedMessage + if err := json.Unmarshal(msg.Data, &claimed); err != nil { + t.Fatal(err) + } + if msg.Type != types.MessageTypeTaskClaimed || claimed.TaskID != "run" || claimed.WorkerID != "worker" || claimed.ExecutionID != id { + t.Fatalf("incorrect claim: type %q, data %+v", msg.Type, claimed) + } + if id == "" && strings.Contains(string(msg.Data), "execution_id") { + t.Fatal("legacy empty execution ID should be omitted") + } + }) + } +} + +func (b *executionTestBackend) ExecuteTask(ctx context.Context, params *TaskParams) ExecuteResult { + b.started <- params.ExecutionID + select { + case <-b.release[params.ExecutionID]: + return executeCompleted() + case <-ctx.Done(): + return executeError(ctx.Err()) + } +} + +func (b *executionTestBackend) CancelTask(ctx context.Context, params *CancelParams) error { + b.cancelled <- *params + if b.cancelGate != nil { + select { + case <-b.cancelGate: + case <-ctx.Done(): + return ctx.Err() + } + } + if b.failCancel.Load() { + return errors.New("cleanup unavailable") + } + return nil +} + +func (b *executionTestBackend) ConfirmTaskClosed(context.Context, *CancelParams) (bool, error) { + return true, nil +} + +func newExecutionTestWorker(t *testing.T) (*Worker, *executionTestBackend) { + t.Helper() + backend := &executionTestBackend{ + started: make(chan string, 4), + release: map[string]chan struct{}{"old": make(chan struct{}), "new": make(chan struct{})}, + cancelled: make(chan CancelParams, 4), + } + w := &Worker{ + ctx: context.Background(), + backend: backend, + activeTasks: make(map[taskExecution]activeTask), + outbound: newOutboundQueue(32), + } + t.Cleanup(func() { + for _, release := range backend.release { + select { + case <-release: + default: + close(release) + } + } + w.taskWG.Wait() + }) + return w, backend +} + +func executionAssignment(id string) *types.TaskAssignmentMessage { + return &types.TaskAssignmentMessage{TaskID: "run", ExecutionID: id, Task: &types.Task{ID: "run"}} +} + +func awaitExecutionStarted(t *testing.T, backend *executionTestBackend, id string) { + t.Helper() + select { + case got := <-backend.started: + if got != id { + t.Fatalf("started %q, want %q", got, id) + } + case <-time.After(2 * time.Second): + t.Fatal("execution did not start") + } +} + +func awaitExecutionClosed(t *testing.T, w *Worker, id string) { + t.Helper() + timeout := time.NewTimer(2 * time.Second) + defer timeout.Stop() + for { + var msg types.WebSocketMessage + select { + case data := <-w.outbound.messages: + if err := json.Unmarshal(data, &msg); err != nil { + t.Fatal(err) + } + case <-timeout.C: + t.Fatal("execution closure was not acknowledged") + } + if msg.Type != types.MessageTypeExecutionClosed { + continue + } + var closed types.ExecutionClosedMessage + if err := json.Unmarshal(msg.Data, &closed); err != nil { + t.Fatal(err) + } + if closed.TaskID != "run" || closed.ExecutionID != id { + t.Fatalf("wrong closure identity: %+v", closed) + } + return + } +} + +func TestExecutionScopedCancellation(t *testing.T) { + t.Run("delayed predecessor cancellation leaves successor running and replays closure", func(t *testing.T) { + w, backend := newExecutionTestWorker(t) + w.handleTaskAssignment(executionAssignment("old")) + awaitExecutionStarted(t, backend, "old") + w.handleTaskAssignment(executionAssignment("new")) + awaitExecutionStarted(t, backend, "new") + request := &types.TaskCancellationMessage{TaskID: "run", ExecutionID: "old"} + w.handleTaskCancellation(request) + awaitExecutionClosed(t, w, "old") + w.tasksMutex.Lock() + newTask, exists := w.activeTasks[taskExecution{"run", "new"}] + w.tasksMutex.Unlock() + if !exists || newTask.ctx.Err() != nil { + t.Fatal("predecessor cancellation affected successor") + } + if params := <-backend.cancelled; params.ExecutionID != "old" { + t.Fatalf("cancelled wrong execution: %+v", params) + } + w.handleTaskCancellation(request) + awaitExecutionClosed(t, w, "old") + select { + case params := <-backend.cancelled: + t.Fatalf("replay called backend again: %+v", params) + default: + } + }) + + t.Run("old completion cannot remove successor and remains closable", func(t *testing.T) { + w, backend := newExecutionTestWorker(t) + for _, id := range []string{"old", "new"} { + w.handleTaskAssignment(executionAssignment(id)) + awaitExecutionStarted(t, backend, id) + } + close(backend.release["old"]) + waitFor(t, 2*time.Second, func() bool { + w.tasksMutex.Lock() + defer w.tasksMutex.Unlock() + _, completed := w.completedTasks[taskExecution{"run", "old"}] + return completed + }) + if count := w.activeTaskCount(); count != 1 { + t.Fatalf("active executions = %d, want 1", count) + } + w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "run", ExecutionID: "old"}) + awaitExecutionClosed(t, w, "old") + if count := w.activeTaskCount(); count != 1 { + t.Fatalf("cleanup removed successor: active executions = %d", count) + } + }) + + t.Run("unknown execution and duplicate assignment are no-ops", func(t *testing.T) { + w, backend := newExecutionTestWorker(t) + w.handleTaskAssignment(executionAssignment("new")) + awaitExecutionStarted(t, backend, "new") + drainMessages(t, w.outbound.messages) + w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "run", ExecutionID: "unknown"}) + w.handleTaskAssignment(executionAssignment("new")) + if len(w.outbound.messages) != 0 || len(backend.cancelled) != 0 || len(backend.started) != 0 { + t.Fatal("unknown cancellation or duplicate assignment produced effects") + } + }) + + t.Run("legacy task-only cancellation cancels all matching executions", func(t *testing.T) { + w, backend := newExecutionTestWorker(t) + for _, id := range []string{"old", "new"} { + w.handleTaskAssignment(executionAssignment(id)) + awaitExecutionStarted(t, backend, id) + } + w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "run"}) + waitFor(t, 2*time.Second, func() bool { return w.activeTaskCount() == 0 }) + got := map[string]bool{} + for i := 0; i < 2; i++ { + params := <-backend.cancelled + got[params.ExecutionID] = true + } + if !got["old"] || !got["new"] { + t.Fatalf("legacy cancellation targets = %v", got) + } + }) + + t.Run("queued cleanup is not proof and failed cleanup remains retryable", func(t *testing.T) { + w, backend := newExecutionTestWorker(t) + backend.cancelGate = make(chan struct{}) + backend.failCancel.Store(true) + w.handleTaskAssignment(executionAssignment("old")) + awaitExecutionStarted(t, backend, "old") + request := &types.TaskCancellationMessage{TaskID: "run", ExecutionID: "old"} + w.handleTaskCancellation(request) + select { + case <-backend.cancelled: + case <-time.After(2 * time.Second): + t.Fatal("backend cleanup did not start") + } + for _, msg := range drainMessages(t, w.outbound.messages) { + if msg.Type == types.MessageTypeExecutionClosed { + t.Fatal("acknowledged queued cleanup") + } + } + w.handleTaskCancellation(request) + if len(backend.cancelled) != 0 { + t.Fatal("concurrent retry duplicated cleanup") + } + close(backend.cancelGate) + waitFor(t, 2*time.Second, func() bool { + w.tasksMutex.Lock() + defer w.tasksMutex.Unlock() + task, tracked := w.activeTasks[taskExecution{"run", "old"}] + return tracked && !task.cancelling + }) + if len(w.outbound.messages) != 0 { + t.Fatal("failed cleanup emitted closure") + } + backend.failCancel.Store(false) + w.handleTaskCancellation(request) + awaitExecutionClosed(t, w, "old") + }) + + t.Run("waits for dispatch before backend cancellation", func(t *testing.T) { + w, backend := newExecutionTestWorker(t) + key := taskExecution{"run", "old"} + done := make(chan struct{}) + ctx, cancel := context.WithCancel(context.Background()) + w.activeTasks[key] = activeTask{ctx: ctx, cancel: cancel, done: done} + w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "run", ExecutionID: "old"}) + if len(backend.cancelled) != 0 || len(w.outbound.messages) != 0 { + t.Fatal("cancel acknowledged before dispatch returned") + } + close(done) + awaitExecutionClosed(t, w, "old") + }) +} + +func TestExecutionClosureReplayCacheBounded(t *testing.T) { + w := &Worker{} + w.rememberCompletedLocked(taskExecution{taskID: "legacy-run"}, true) + if len(w.completedTasks) != 0 { + t.Fatal("legacy completion must not suppress a later run-only assignment") + } + for i := 0; i < 5000; i++ { + w.rememberCompletedLocked(taskExecution{taskID: "run", executionID: time.Unix(int64(i), 0).String()}, true) + } + if len(w.completedTasks) != 4096 || len(w.closedExecutions) != 4096 { + t.Fatalf("cache unbounded: %d records, %d keys", len(w.completedTasks), len(w.closedExecutions)) + } +} + +func TestExecutionCancellationCapabilityHandshake(t *testing.T) { + capability := make(chan string, 1) + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + capability <- r.Header.Get("X-Warp-Worker-Capabilities") + upgrader := websocket.Upgrader{} + conn, err := upgrader.Upgrade(rw, r, nil) + if err != nil { + t.Error(err) + return + } + _ = conn.Close() + })) + defer server.Close() + w := &Worker{ctx: context.Background(), config: Config{WebSocketURL: "ws" + strings.TrimPrefix(server.URL, "http")}} + conn, err := w.connect() + if err != nil { + t.Fatal(err) + } + defer func() { _ = conn.Close() }() + if got := <-capability; got != "execution-scoped-cancellation-v1" { + t.Fatalf("capability = %q", got) + } +} diff --git a/internal/worker/execution_closure_backends_test.go b/internal/worker/execution_closure_backends_test.go new file mode 100644 index 0000000..8cd8a7b --- /dev/null +++ b/internal/worker/execution_closure_backends_test.go @@ -0,0 +1,193 @@ +package worker + +import ( + "context" + "errors" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/moby/moby/client" + batchv1 "k8s.io/api/batch/v1" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + "k8s.io/client-go/kubernetes/fake" + k8stesting "k8s.io/client-go/testing" +) + +func TestDirectExecutionWorkspaceIsolation(t *testing.T) { + for _, predecessor := range []string{"", "old"} { + t.Run("predecessor="+predecessor, func(t *testing.T) { + root := t.TempDir() + ozPath := filepath.Join(root, "oz") + if err := os.WriteFile(ozPath, []byte("#!/bin/sh\nprintf '%s' \"$PWD\" > \"$CAPTURE\"\n"), 0o700); err != nil { + t.Fatal(err) + } + backend, err := NewDirectBackend(context.Background(), DirectBackendConfig{ + WorkspaceRoot: filepath.Join(root, "workspaces"), OzPath: ozPath, NoCleanup: true, + }) + if err != nil { + t.Fatal(err) + } + var workspaces []string + for _, id := range []string{predecessor, "new"} { + capture := filepath.Join(root, "capture-"+id) + result := backend.ExecuteTask(context.Background(), &TaskParams{ + TaskID: "run", ExecutionID: id, EnvVars: []string{"CAPTURE=" + capture}, + }) + if result.Error != nil { + t.Fatal(result.Error) + } + workspace, err := os.ReadFile(capture) + if err != nil { + t.Fatal(err) + } + workspaces = append(workspaces, string(workspace)) + } + backend.cleanup(context.Background(), "run", workspaces[0], "") + if _, err := os.Stat(workspaces[1]); err != nil { + t.Fatalf("predecessor cleanup removed successor workspace: %v", err) + } + }) + } +} + +type closureDockerTransport func(*http.Request) (*http.Response, error) + +func (f closureDockerTransport) RoundTrip(r *http.Request) (*http.Response, error) { + return f(r) +} + +func TestDockerExecutionClosure(t *testing.T) { + t.Run("failed removal retains exact resource for retry", func(t *testing.T) { + fail := true + var removed []string + transport := closureDockerTransport(func(r *http.Request) (*http.Response, error) { + if r.URL.Path == "/_ping" { + return mockEngineResponse(r, http.StatusOK, "text/plain", "OK"), nil + } + removed = append(removed, r.URL.Path) + if r.Method != http.MethodDelete || !strings.HasSuffix(r.URL.Path, "/containers/old-container") { + t.Fatalf("unexpected cleanup request: %s %s", r.Method, r.URL.Path) + } + if fail { + return mockEngineErrorResponse(t, r, http.StatusInternalServerError, "unavailable"), nil + } + return mockEngineResponse(r, http.StatusNoContent, "application/json", ""), nil + }) + dockerClient, err := client.New(client.WithHTTPClient(&http.Client{Transport: transport})) + if err != nil { + t.Fatal(err) + } + defer func() { _ = dockerClient.Close() }() + backend := &DockerBackend{dockerClient: dockerClient} + old := taskExecution{"run", "old"} + successor := taskExecution{"run", "new"} + backend.taskContainers.Store(old, "old-container") + backend.taskContainers.Store(successor, "new-container") + params := &CancelParams{TaskID: "run", ExecutionID: "old"} + if err := backend.CancelTask(context.Background(), params); err == nil { + t.Fatal("expected retryable removal failure") + } + if closed, _ := backend.ConfirmTaskClosed(context.Background(), params); closed { + t.Fatal("failed removal reported closure") + } + fail = false + if err := backend.CancelTask(context.Background(), params); err != nil { + t.Fatal(err) + } + if closed, err := backend.ConfirmTaskClosed(context.Background(), params); !closed || err != nil { + t.Fatalf("closed=%v, error=%v", closed, err) + } + if _, found := backend.taskContainers.Load(successor); !found { + t.Fatal("successor tracking removed") + } + if len(removed) != 2 { + t.Fatalf("removal attempts = %d", len(removed)) + } + }) + + t.Run("unknown create outcome cannot acknowledge absence", func(t *testing.T) { + backend := &DockerBackend{} + backend.taskContainers.Store(taskExecution{"run", "old"}, "") + if err := backend.CancelTask(context.Background(), &CancelParams{TaskID: "run", ExecutionID: "old"}); err == nil { + t.Fatal("ambiguous creation treated as closure") + } + }) +} + +func TestKubernetesExecutionClosure(t *testing.T) { + t.Run("deleting job is insufficient while pods remain", func(t *testing.T) { + backend := &KubernetesBackend{config: KubernetesBackendConfig{Namespace: "agents"}} + params := &CancelParams{TaskID: "run", ExecutionID: "old"} + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{ + Name: "old-pod", Namespace: "agents", Labels: backend.baseLabels("run", "old"), + }} + clientset := fake.NewSimpleClientset(pod) + backend.clientset = clientset + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + clientset.PrependReactor("list", "pods", func(k8stesting.Action) (bool, runtime.Object, error) { + cancel() + return false, nil, nil + }) + if closed, err := backend.ConfirmTaskClosed(ctx, params); closed || !errors.Is(err, context.Canceled) { + t.Fatalf("remaining pod yielded closed=%v, err=%v", closed, err) + } + if err := clientset.CoreV1().Pods("agents").Delete(context.Background(), pod.Name, metav1.DeleteOptions{}); err != nil { + t.Fatal(err) + } + if closed, err := backend.ConfirmTaskClosed(context.Background(), params); !closed || err != nil { + t.Fatalf("removed pod yielded closed=%v, err=%v", closed, err) + } + }) + + t.Run("API errors are not absence", func(t *testing.T) { + clientset := fake.NewSimpleClientset() + clientset.PrependReactor("get", "jobs", func(k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, apierrors.NewForbidden(schema.GroupResource{Resource: "jobs"}, "job", errors.New("denied")) + }) + backend := &KubernetesBackend{config: KubernetesBackendConfig{Namespace: "agents"}, clientset: clientset} + if closed, err := backend.ConfirmTaskClosed(context.Background(), &CancelParams{TaskID: "run", ExecutionID: "old"}); closed || err == nil { + t.Fatalf("API failure yielded closed=%v, err=%v", closed, err) + } + }) + + t.Run("name suffix collision never deletes successor", func(t *testing.T) { + backend := &KubernetesBackend{config: KubernetesBackendConfig{Namespace: "agents"}} + oldID, newID := "old-12345678", "new-12345678" + jobName := kubernetesTaskJobName("run", newID) + if jobName != kubernetesTaskJobName("run", oldID) { + t.Fatal("test requires colliding names") + } + clientset := fake.NewSimpleClientset(&batchv1.Job{ObjectMeta: metav1.ObjectMeta{ + Name: jobName, Namespace: "agents", Labels: backend.baseLabels("run", newID), + }}) + backend.clientset = clientset + if err := backend.CancelTask(context.Background(), &CancelParams{TaskID: "run", ExecutionID: oldID}); err == nil { + t.Fatal("expected execution mismatch") + } + if _, err := clientset.BatchV1().Jobs("agents").Get(context.Background(), jobName, metav1.GetOptions{}); err != nil { + t.Fatalf("successor removed: %v", err) + } + }) + + t.Run("existing job cannot be acknowledged closed", func(t *testing.T) { + backend := &KubernetesBackend{config: KubernetesBackendConfig{Namespace: "agents"}} + job := &batchv1.Job{ObjectMeta: metav1.ObjectMeta{ + Name: kubernetesTaskJobName("run", "old"), Namespace: "agents", Labels: backend.baseLabels("run", "old"), + }} + backend.clientset = fake.NewSimpleClientset(job) + ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second)) + defer cancel() + if closed, err := backend.ConfirmTaskClosed(ctx, &CancelParams{TaskID: "run", ExecutionID: "old"}); closed || err == nil { + t.Fatalf("live Job yielded closed=%v, err=%v", closed, err) + } + }) +} diff --git a/internal/worker/kubernetes.go b/internal/worker/kubernetes.go index 3b06ed5..f1736b9 100644 --- a/internal/worker/kubernetes.go +++ b/internal/worker/kubernetes.go @@ -525,14 +525,58 @@ func (b *KubernetesBackend) CancelTask(ctx context.Context, params *CancelParams if params == nil { return fmt.Errorf("cancel params are required") } - jobName := kubernetesTaskJobName(params.TaskID, executionIDOrTaskID(params.TaskID, params.ExecutionID)) + executionID := executionIDOrTaskID(params.TaskID, params.ExecutionID) + jobName := kubernetesTaskJobName(params.TaskID, executionID) + job, err := b.getTaskJob(ctx, jobName) + if apierrors.IsNotFound(err) { + return nil + } + if err != nil { + return err + } + if job.Labels[kubernetesExecutionHashLabel] != kubernetesLabelHash(executionID) || + job.Labels[kubernetesTaskHashLabel] != kubernetesLabelHash(params.TaskID) { + return fmt.Errorf("job %s belongs to another execution", jobName) + } log.Infof(ctx, "Deleting Kubernetes Job %s after task cancellation", jobName) - if err := b.deleteJob(ctx, jobName); err != nil { + if err := b.deleteTaskJob(ctx, jobName, &job.UID); err != nil { return fmt.Errorf("failed to delete Kubernetes Job %s: %w", jobName, err) } return nil } +// ConfirmTaskClosed waits for both the Job and its Pods to disappear. +func (b *KubernetesBackend) ConfirmTaskClosed(ctx context.Context, params *CancelParams) (bool, error) { + executionID := executionIDOrTaskID(params.TaskID, params.ExecutionID) + jobName := kubernetesTaskJobName(params.TaskID, executionID) + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + job, err := b.getTaskJob(ctx, jobName) + if err != nil && !apierrors.IsNotFound(err) { + return false, err + } + if err == nil && (job.Labels[kubernetesExecutionHashLabel] != kubernetesLabelHash(executionID) || + job.Labels[kubernetesTaskHashLabel] != kubernetesLabelHash(params.TaskID)) { + return false, fmt.Errorf("job %s belongs to another execution", jobName) + } + if apierrors.IsNotFound(err) { + pods, err := b.listTaskPods(ctx, executionID) + if err != nil { + return false, err + } + if len(pods) == 0 { + return true, nil + } + } + select { + case <-ctx.Done(): + return false, ctx.Err() + case <-ticker.C: + } + } +} + // Shutdown intentionally does not delete task Jobs. // // Kubernetes Jobs are the durable execution unit for this backend. During diff --git a/internal/worker/worker.go b/internal/worker/worker.go index b6833b7..f73f3fa 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -31,7 +31,8 @@ const ( WriteWait = 10 * time.Second BackendShutdownTimeout = 10 * time.Second - warpServerRootURLEnv = "WARP_SERVER_ROOT_URL" + warpServerRootURLEnv = "WARP_SERVER_ROOT_URL" + executionCancellationCapability = "execution-scoped-cancellation-v1" ) type Config struct { @@ -62,18 +63,24 @@ type Config struct { Kubernetes *KubernetesBackendConfig Command *CommandBackendConfig } +type taskExecution struct { + taskID string + executionID string +} type Worker struct { - config Config - ctx context.Context - outbound *outboundQueue - activeTasks map[string]activeTask - tasksMutex sync.Mutex - taskWG sync.WaitGroup - shuttingDown bool - oneShot oneShotState - backend Backend - taskSemaphore *semaphore.Weighted // nil when unlimited + config Config + ctx context.Context + outbound *outboundQueue + activeTasks map[taskExecution]activeTask + completedTasks map[taskExecution]bool + closedExecutions []taskExecution + tasksMutex sync.Mutex + taskWG sync.WaitGroup + shuttingDown bool + oneShot oneShotState + backend Backend + taskSemaphore *semaphore.Weighted // nil when unlimited // heartbeatInterval is how often the worker pings the server. It defaults // to HeartbeatInterval and is overridable in tests. heartbeatInterval time.Duration @@ -144,7 +151,9 @@ type activeTask struct { // spawned marks a task whose backend returned ExecuteOutcomeSpawned: it no // longer executes locally, but the entry is kept so a later cancellation // can be routed to the backend's CancelTask. - spawned bool + spawned bool + done chan struct{} + cancelling bool } func New(ctx context.Context, config Config) (*Worker, error) { @@ -199,7 +208,7 @@ func New(ctx context.Context, config Config) (*Worker, error) { config: config, ctx: ctx, outbound: newOutboundQueue(256), - activeTasks: make(map[string]activeTask), + activeTasks: make(map[taskExecution]activeTask), oneShot: newOneShotState(), backend: backend, taskSemaphore: taskSemaphore, @@ -292,6 +301,7 @@ func (w *Worker) connect() (*websocket.Conn, error) { headers := make(map[string][]string) headers["Authorization"] = []string{fmt.Sprintf("Bearer %s", w.config.APIKey)} + headers["X-Warp-Worker-Capabilities"] = []string{executionCancellationCapability} log.Infof(w.ctx, "Connecting to %s", u.String()) conn, resp, err := websocket.DefaultDialer.Dial(u.String(), headers) @@ -489,55 +499,137 @@ func (w *Worker) handleMessage(message []byte) { func (w *Worker) handleTaskCancellation(cancellation *types.TaskCancellationMessage) { w.tasksMutex.Lock() - task, ok := w.activeTasks[cancellation.TaskID] - if ok { - if task.cancellationSource == "" { - task.cancellationSource = taskCancellationSourceUser - w.activeTasks[cancellation.TaskID] = task + key := taskExecution{cancellation.TaskID, cancellation.ExecutionID} + if cancellation.ExecutionID == "" { + // A task-only request is explicitly legacy run-wide cancellation. + for candidate := range w.activeTasks { + if candidate.taskID == cancellation.TaskID { + w.startTaskCancellationLocked(candidate) + } } - if task.spawned { - // executeTask has already returned for a spawned task, so no - // deferred cleanup will remove the entry; drop it now that the - // cancellation is being routed to the backend. - delete(w.activeTasks, cancellation.TaskID) + w.tasksMutex.Unlock() + return + } + if closed, known := w.completedTasks[key]; known { + if closed { + w.tasksMutex.Unlock() + if err := w.sendExecutionClosed(key); err != nil { + log.Warnf(w.ctx, "Failed to replay execution closure: %v", err) + } + return + } + if _, active := w.activeTasks[key]; !active { + done := make(chan struct{}) + close(done) + w.activeTasks[key] = activeTask{ctx: context.Background(), cancel: func() {}, executionID: key.executionID, done: done} } } + w.startTaskCancellationLocked(key) w.tasksMutex.Unlock() +} - if !ok { - log.Warnf(w.ctx, "Received cancellation for inactive task: taskID=%s", cancellation.TaskID) +func (w *Worker) startTaskCancellationLocked(key taskExecution) { + task, ok := w.activeTasks[key] + if !ok || task.cancelling { return } - - log.Infof(w.ctx, "Cancelling task from server request: taskID=%s", cancellation.TaskID) + task.cancellationSource = taskCancellationSourceUser + task.cancelling = true + w.activeTasks[key] = task metrics.AddTaskEvent(task.ctx, "task.cancellation_requested", attribute.String("source", "server"), - attribute.String("task.id", cancellation.TaskID), + attribute.String("task.id", key.taskID), ) - // Every backend gets the same cancellation contract: its cancelation hook - // is invoked explicitly, and then its execution context is canceled.. - w.cancelTaskOnBackend(&CancelParams{TaskID: cancellation.TaskID, ExecutionID: task.executionID}) task.cancel() + go w.cancelTaskOnBackend(key, task.done) } -// cancelTaskOnBackend makes a best-effort attempt to cancel a task via the -// backend's CancelTask. -func (w *Worker) cancelTaskOnBackend(params *CancelParams) { - log.Infof(w.ctx, "Requesting backend cancellation for task %s", params.TaskID) - go func() { - ctx, cancel := context.WithTimeout(context.WithoutCancel(w.ctx), BackendShutdownTimeout) - defer cancel() - if err := w.backend.CancelTask(ctx, params); err != nil { - log.Warnf(w.ctx, "Backend cancellation failed for task %s: %v", params.TaskID, err) - metrics.AddTaskEvent(ctx, "cancel.failed", - attribute.String("reason", string(metrics.TaskFailureReasonCancelCommand)), - attribute.String("task.id", params.TaskID), - ) +func (w *Worker) cancelTaskOnBackend(key taskExecution, done <-chan struct{}) { + ctx, cancel := context.WithTimeout(context.WithoutCancel(w.ctx), BackendShutdownTimeout) + defer cancel() + closed, err := w.closeExecution(ctx, key, done) + w.tasksMutex.Lock() + task, tracked := w.activeTasks[key] + if !tracked { + w.tasksMutex.Unlock() + return + } + task.cancelling = false + w.activeTasks[key] = task + if err != nil { + w.tasksMutex.Unlock() + log.Warnf(w.ctx, "Backend cancellation failed for task %s execution %s: %v", key.taskID, key.executionID, err) + return + } + if closed { + w.rememberCompletedLocked(key, true) + delete(w.activeTasks, key) + } + w.tasksMutex.Unlock() + if closed && key.executionID != "" { + if err := w.sendExecutionClosed(key); err != nil { + log.Warnf(w.ctx, "Failed to send execution closure: %v", err) } - }() + } +} + +func (w *Worker) closeExecution(ctx context.Context, key taskExecution, done <-chan struct{}) (bool, error) { + // A cancel racing provisioning must not confirm absence before creation finishes. + if done != nil { + select { + case <-done: + case <-ctx.Done(): + return false, ctx.Err() + } + } + params := &CancelParams{TaskID: key.taskID, ExecutionID: key.executionID} + if err := w.backend.CancelTask(ctx, params); err != nil { + return false, err + } + if verifier, ok := w.backend.(executionClosureVerifier); ok { + closed, err := verifier.ConfirmTaskClosed(ctx, params) + if err != nil { + return false, err + } + if closed { + return true, nil + } + } + return false, fmt.Errorf("backend has not confirmed execution closure") +} + +func (w *Worker) rememberCompletedLocked(key taskExecution, closed bool) { + if key.executionID == "" { + return + } + if w.completedTasks == nil { + w.completedTasks = make(map[taskExecution]bool) + } + w.completedTasks[key] = closed + if closed { + w.closedExecutions = append(w.closedExecutions, key) + const closedExecutionCacheSize = 4096 + if len(w.closedExecutions) > closedExecutionCacheSize { + delete(w.completedTasks, w.closedExecutions[0]) + w.closedExecutions = w.closedExecutions[1:] + } + } +} + +func (w *Worker) sendExecutionClosed(key taskExecution) error { + data, err := json.Marshal(types.ExecutionClosedMessage{TaskID: key.taskID, ExecutionID: key.executionID}) + if err != nil { + return err + } + message, err := json.Marshal(types.WebSocketMessage{Type: types.MessageTypeExecutionClosed, Data: data}) + if err != nil { + return err + } + return w.sendMessage(message) } func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage) { + key := taskExecution{assignment.TaskID, assignment.ExecutionID} receivedAt := time.Now() log.Infof(w.ctx, "Received task assignment: taskID=%s, title=%s", assignment.TaskID, assignment.Task.Title) taskCtx, span := metrics.StartTaskSpan(w.ctx, assignment.TaskID, assignment.Task.Title) @@ -548,6 +640,13 @@ func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage) { ) w.tasksMutex.Lock() + _, active := w.activeTasks[key] + _, completed := w.completedTasks[key] + if active || completed { + w.tasksMutex.Unlock() + span.End() + return + } rejectionReason := "" rejectionMetric := "" switch { @@ -600,16 +699,29 @@ func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage) { w.rejectTaskAssignment(taskCtx, span, assignment.TaskID, "worker is shutting down", metrics.RejectReasonShuttingDown) return } - w.activeTasks[assignment.TaskID] = activeTask{ + _, active = w.activeTasks[key] + _, completed = w.completedTasks[key] + if active || completed { + w.tasksMutex.Unlock() + taskCancel() + span.End() + if w.taskSemaphore != nil { + w.taskSemaphore.Release(1) + } + w.taskWG.Done() + return + } + w.activeTasks[key] = activeTask{ ctx: taskCtx, cancel: taskCancel, executionID: assignment.ExecutionID, + done: make(chan struct{}), } w.tasksMutex.Unlock() // It's important to update the task state to claimed as the task lifecycle // treats this as a dependency to advance to further states. - if err := w.sendTaskClaimed(assignment.TaskID); err != nil { + if err := w.sendTaskClaimed(assignment.TaskID, assignment.ExecutionID); err != nil { log.Errorf(w.ctx, "Failed to send task claimed message: %v", err) } metrics.RecordTaskClaim() @@ -787,6 +899,7 @@ func (w *Worker) defaultImageForTask(assignmentImage string, task *types.Task) s } func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc, span trace.Span, assignment *types.TaskAssignmentMessage, receivedAt time.Time) { + key := taskExecution{assignment.TaskID, assignment.ExecutionID} start := time.Now() result := metrics.TaskResultSucceeded // One-shot-capable backends always wait for a terminal outcome. Register @@ -798,10 +911,14 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc, taskCancel() span.End() w.tasksMutex.Lock() - // Spawned tasks stay in activeTasks so a later cancellation can be - // routed to the backend's CancelTask; everything else is done. - if task, tracked := w.activeTasks[assignment.TaskID]; !tracked || !task.spawned { - delete(w.activeTasks, assignment.TaskID) + if task, tracked := w.activeTasks[key]; tracked { + if task.done != nil { + close(task.done) + } + if !task.spawned && task.cancellationSource != taskCancellationSourceUser { + w.rememberCompletedLocked(key, false) + delete(w.activeTasks, key) + } } w.tasksMutex.Unlock() @@ -826,7 +943,7 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc, executeResult := w.backend.ExecuteTask(ctx, params) if executeResult.Error != nil { err := executeResult.Error - if ctx.Err() == context.Canceled && w.cancellationSource(taskID) == taskCancellationSourceUser { + if ctx.Err() == context.Canceled && w.cancellationSource(key) == taskCancellationSourceUser { result = metrics.TaskResultCancelled metrics.AddTaskEvent(ctx, "task.cancelled", attribute.String("source", string(taskCancellationSourceUser)), @@ -845,7 +962,7 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc, // Reclassify failures caused by a graceful worker shutdown (task // cancelled, or agent killed by the shutdown's SIGTERM) as // graceful_shutdown. - if w.cancellationSource(taskID) == taskCancellationSourceShutdown && + if w.cancellationSource(key) == taskCancellationSourceShutdown && (metricsReason == metrics.TaskFailureReasonTaskCancelled || exitCode == sigtermExitCode) { metricsReason = metrics.TaskFailureReasonGracefulShutdown } @@ -871,9 +988,9 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc, // implementation later. result = metrics.TaskResultDispatched w.tasksMutex.Lock() - if task, tracked := w.activeTasks[taskID]; tracked && task.cancellationSource == "" { + if task, tracked := w.activeTasks[key]; tracked { task.spawned = true - w.activeTasks[taskID] = task + w.activeTasks[key] = task } w.tasksMutex.Unlock() metrics.AddTaskEvent(ctx, "task.dispatched") @@ -897,20 +1014,21 @@ func (w *Worker) signalOneShotComplete() { w.oneShot.complete() } -func (w *Worker) cancellationSource(taskID string) taskCancellationSource { +func (w *Worker) cancellationSource(key taskExecution) taskCancellationSource { w.tasksMutex.Lock() defer w.tasksMutex.Unlock() - task, ok := w.activeTasks[taskID] + task, ok := w.activeTasks[key] if !ok { return "" } return task.cancellationSource } -func (w *Worker) sendTaskClaimed(taskID string) error { +func (w *Worker) sendTaskClaimed(taskID, executionID string) error { claimed := types.TaskClaimedMessage{ - TaskID: taskID, - WorkerID: w.config.WorkerID, + TaskID: taskID, + WorkerID: w.config.WorkerID, + ExecutionID: executionID, } data, err := json.Marshal(claimed) @@ -1055,15 +1173,15 @@ func (w *Worker) shutdownTasks() { log.Infof(w.ctx, "Preserving %d active tasks during worker shutdown", activeTaskCount) } else if activeTaskCount > 0 { log.Infof(w.ctx, "Cancelling %d active tasks", activeTaskCount) - for taskID, task := range w.activeTasks { + for key, task := range w.activeTasks { if task.cancellationSource == "" { task.cancellationSource = taskCancellationSourceShutdown - w.activeTasks[taskID] = task + w.activeTasks[key] = task } - log.Debugf(w.ctx, "Cancelling task: %s", taskID) + log.Debugf(w.ctx, "Cancelling task: %s", key.taskID) metrics.AddTaskEvent(task.ctx, "task.cancellation_requested", attribute.String("source", "signal"), - attribute.String("task.id", taskID), + attribute.String("task.id", key.taskID), ) task.cancel() } diff --git a/internal/worker/worker_test.go b/internal/worker/worker_test.go index 177e5ac..4f6a564 100644 --- a/internal/worker/worker_test.go +++ b/internal/worker/worker_test.go @@ -143,7 +143,7 @@ func TestOneShotAcceptsOnlyOneTaskAndWaitsForCompletion(t *testing.T) { }, ctx: ctx, outbound: newOutboundQueue(8), - activeTasks: make(map[string]activeTask), + activeTasks: make(map[taskExecution]activeTask), oneShot: newOneShotState(), backend: backend, taskSemaphore: semaphore.NewWeighted(1), @@ -342,7 +342,7 @@ func TestExecuteTaskReportsGracefulShutdownOnWorkerShutdown(t *testing.T) { ctx: context.Background(), config: Config{}, outbound: newOutboundQueue(1), - activeTasks: map[string]activeTask{"task-1": { + activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "execution-1"}: { cancel: func() {}, cancellationSource: taskCancellationSourceShutdown, }}, @@ -375,7 +375,7 @@ func TestExecuteTaskReportsTaskCancelledOnUserCancellation(t *testing.T) { config: Config{OneShot: true}, outbound: newOutboundQueue(1), oneShot: newOneShotState(), - activeTasks: map[string]activeTask{"task-1": { + activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "execution-1"}: { cancel: func() {}, cancellationSource: taskCancellationSourceUser, }}, @@ -408,8 +408,8 @@ func TestExecuteTaskReportsTaskCancelledOnUserCancellation(t *testing.T) { if completed.Message != "Task cancelled by user request." { t.Errorf("message = %q, want %q", completed.Message, "Task cancelled by user request.") } - if _, ok := w.activeTasks["task-1"]; ok { - t.Fatal("task should be removed from active tasks") + if _, ok := w.activeTasks[taskExecution{taskID: "task-1", executionID: "execution-1"}]; !ok { + t.Fatal("cancelled task must remain tracked until backend closure is confirmed") } assertOneShotDone(t, w) } @@ -419,7 +419,7 @@ func TestExecuteTaskDoesNotReportTaskCancelledOnBackendCancellationError(t *test ctx: context.Background(), config: Config{}, outbound: newOutboundQueue(1), - activeTasks: map[string]activeTask{"task-1": {cancel: func() {}}}, + activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "execution-1"}: {cancel: func() {}}}, backend: &recordingBackend{err: fmt.Errorf("backend request failed: %w", context.Canceled)}, } @@ -442,8 +442,8 @@ func TestHandleMessageCancelsActiveTask(t *testing.T) { w := &Worker{ ctx: context.Background(), outbound: newOutboundQueue(1), - activeTasks: map[string]activeTask{ - "task-1": { + activeTasks: map[taskExecution]activeTask{ + {taskID: "task-1", executionID: ""}: { ctx: taskCtx, cancel: taskCancel, }, @@ -468,7 +468,10 @@ func TestHandleMessageCancelsActiveTask(t *testing.T) { if taskCtx.Err() != context.Canceled { t.Fatalf("task context error = %v, want %v", taskCtx.Err(), context.Canceled) } - if task := w.activeTasks["task-1"]; task.cancellationSource != taskCancellationSourceUser { + w.tasksMutex.Lock() + task := w.activeTasks[taskExecution{taskID: "task-1", executionID: ""}] + w.tasksMutex.Unlock() + if task.cancellationSource != taskCancellationSourceUser { t.Fatalf("task cancellation source = %q, want %q", task.cancellationSource, taskCancellationSourceUser) } } @@ -478,7 +481,7 @@ func TestExecuteTaskReportsTaskCompletedOnSuccess(t *testing.T) { ctx: context.Background(), config: Config{OneShot: true}, outbound: newOutboundQueue(1), - activeTasks: map[string]activeTask{"task-1": {cancel: func() {}}}, + activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "execution-1"}: {cancel: func() {}}}, oneShot: newOneShotState(), backend: &recordingBackend{}, } @@ -507,7 +510,7 @@ func TestExecuteTaskReportsTaskCompletedOnSuccess(t *testing.T) { if completed.Message != "Task completed successfully" { t.Errorf("message = %q, want %q", completed.Message, "Task completed successfully") } - if _, ok := w.activeTasks["task-1"]; ok { + if _, ok := w.activeTasks[taskExecution{taskID: "task-1", executionID: "execution-1"}]; ok { t.Fatal("task should be removed from active tasks") } assertOneShotDone(t, w) @@ -518,7 +521,7 @@ func TestExecuteTaskReportsTaskFailedOnBackendError(t *testing.T) { ctx: context.Background(), config: Config{OneShot: true}, outbound: newOutboundQueue(1), - activeTasks: map[string]activeTask{"task-1": {cancel: func() {}}}, + activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: ""}: {cancel: func() {}}}, oneShot: newOneShotState(), backend: &recordingBackend{err: errors.New("boom")}, } @@ -543,7 +546,7 @@ func TestExecuteTaskReportsTaskFailedOnBackendError(t *testing.T) { if failed.Message != "Failed to execute task: boom" { t.Errorf("message = %q, want %q", failed.Message, "Failed to execute task: boom") } - if _, ok := w.activeTasks["task-1"]; ok { + if _, ok := w.activeTasks[taskExecution{taskID: "task-1", executionID: ""}]; ok { t.Fatal("task should be removed from active tasks") } assertOneShotDone(t, w) @@ -554,7 +557,7 @@ func TestExecuteTaskReportsUserFriendlyMessageOnDeadlineExceeded(t *testing.T) { ctx: context.Background(), config: Config{}, outbound: newOutboundQueue(1), - activeTasks: map[string]activeTask{"task-1": {cancel: func() {}}}, + activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: ""}: {cancel: func() {}}}, backend: &recordingBackend{err: context.DeadlineExceeded}, } @@ -674,7 +677,7 @@ func TestRunServerCancellationClosesProtocolAndBackend(t *testing.T) { }, ctx: serverCtx, outbound: newOutboundQueue(8), - activeTasks: make(map[string]activeTask), + activeTasks: make(map[taskExecution]activeTask), oneShot: newOneShotState(), backend: backend, heartbeatInterval: HeartbeatInterval, @@ -1020,7 +1023,7 @@ func TestRunHeartbeatAndWritesAreConcurrencySafe(t *testing.T) { config: Config{}, ctx: context.Background(), outbound: newOutboundQueue(256), - activeTasks: make(map[string]activeTask), + activeTasks: make(map[taskExecution]activeTask), oneShot: newOneShotState(), backend: &recordingBackend{}, heartbeatInterval: time.Millisecond, @@ -1521,7 +1524,7 @@ func TestShutdownBackendUsesFreshContextForCleanup(t *testing.T) { backend := &shutdownRecordingBackend{} w := &Worker{ ctx: context.Background(), - activeTasks: make(map[string]activeTask), + activeTasks: make(map[taskExecution]activeTask), backend: backend, } w.shutdownBackend() @@ -1539,8 +1542,8 @@ func TestWorkerShutdownPreservesActiveTasksForPreservingBackend(t *testing.T) { cancelledTask := false w := &Worker{ ctx: context.Background(), - activeTasks: map[string]activeTask{ - "task-1": {cancel: func() { + activeTasks: map[taskExecution]activeTask{ + {taskID: "task-1", executionID: ""}: {cancel: func() { cancelledTask = true }}, }, @@ -1565,7 +1568,7 @@ func TestHandleTaskAssignmentDoesNotStartTaskAfterShutdownDuringClaim(t *testing ctx: workerCtx, config: Config{}, outbound: newOutboundQueue(1), - activeTasks: make(map[string]activeTask), + activeTasks: make(map[taskExecution]activeTask), backend: &preservingShutdownRecordingBackend{}, } @@ -1584,7 +1587,7 @@ func TestHandleTaskAssignmentRejectsAfterShutdownStarts(t *testing.T) { ctx: context.Background(), config: Config{}, outbound: newOutboundQueue(1), - activeTasks: make(map[string]activeTask), + activeTasks: make(map[taskExecution]activeTask), shuttingDown: true, backend: &recordingBackend{}, } From e38237cb1ef50e18ee8b59068d327af3e8525471 Mon Sep 17 00:00:00 2001 From: seemeroland Date: Tue, 6 Oct 2026 15:09:30 -0700 Subject: [PATCH 2/4] Track shutdown requests for specific executions --- internal/types/messages.go | 7 - internal/worker/backend.go | 11 +- internal/worker/dispatch_test.go | 21 +-- internal/worker/docker.go | 6 - ...> execution_cancellation_backends_test.go} | 66 ++++--- .../worker/execution_cancellation_test.go | 169 ++++++++++++------ internal/worker/kubernetes.go | 32 ---- internal/worker/worker.go | 165 ++++++++--------- 8 files changed, 230 insertions(+), 247 deletions(-) rename internal/worker/{execution_closure_backends_test.go => execution_cancellation_backends_test.go} (71%) diff --git a/internal/types/messages.go b/internal/types/messages.go index 56f49e8..fc614df 100644 --- a/internal/types/messages.go +++ b/internal/types/messages.go @@ -15,7 +15,6 @@ const ( MessageTypeTaskFailed MessageType = "task_failed" MessageTypeTaskRejected MessageType = "task_rejected" MessageTypeTaskCancellation MessageType = "task_cancellation" - MessageTypeExecutionClosed MessageType = "execution_closed" MessageTypeHeartbeat MessageType = "heartbeat" ) @@ -172,12 +171,6 @@ type TaskCancellationMessage struct { ExecutionID string `json:"execution_id,omitempty"` } -// ExecutionClosedMessage confirms backend closure without changing the logical run state. -type ExecutionClosedMessage struct { - TaskID string `json:"task_id"` - ExecutionID string `json:"execution_id"` -} - // TaskState is the serialized terminal task state accepted by warp-server. type TaskState string diff --git a/internal/worker/backend.go b/internal/worker/backend.go index 496c997..ccd6936 100644 --- a/internal/worker/backend.go +++ b/internal/worker/backend.go @@ -176,9 +176,9 @@ type Backend interface { // failures are surfaced as ExecuteOutcomeError with an error that is (or // wraps) *TaskFailure. ExecuteTask(ctx context.Context, params *TaskParams) ExecuteResult - // CancelTask makes a best-effort attempt to cancel a task after the - // ExecuteTask context is cancelled and ExecuteTask returns. A nil error - // is not proof that asynchronous backend resources have stopped. + // CancelTask requests cancellation after ExecuteTask returns. A nil error + // means the request was accepted, not that asynchronous resources stopped. + // Failed requests may be retried and must be execution-scoped and idempotent. CancelTask(ctx context.Context, params *CancelParams) error // PreservesTasksOnShutdown reports whether active task execution units can // safely outlive the worker process during shutdown. @@ -195,11 +195,6 @@ type CancelParams struct { ExecutionID string } -// executionClosureVerifier confirms closure after ExecuteTask and CancelTask have returned. -type executionClosureVerifier interface { - ConfirmTaskClosed(ctx context.Context, params *CancelParams) (bool, error) -} - // TaskFailure is the structured error backends return from ExecuteTask when // task execution fails. Backends record only the facts they observe (metrics // labels and the failing process's exit status); the worker adds lifecycle diff --git a/internal/worker/dispatch_test.go b/internal/worker/dispatch_test.go index 59798a0..16a6e41 100644 --- a/internal/worker/dispatch_test.go +++ b/internal/worker/dispatch_test.go @@ -119,11 +119,9 @@ func TestHandleTaskCancellationRoutesToBackendCancelTask(t *testing.T) { t.Fatal("CancelTask was not invoked for a spawned task") } - w.tasksMutex.Lock() - _, ok := w.activeTasks[taskExecution{taskID: "task-1", executionID: "exec-1"}] - w.tasksMutex.Unlock() - if !ok { - t.Error("spawned task should remain tracked until closure is confirmed") + w.cancellationWG.Wait() + if w.activeTaskCount() != 0 { + t.Error("accepted cancellation should remove spawned task tracking") } } @@ -156,15 +154,9 @@ func TestHandleTaskCancellationRunningTaskCancelsContextAndBackend(t *testing.T) t.Fatal("CancelTask was not invoked for a running task") } - // An unverified cancellation must stay available for retry. - w.tasksMutex.Lock() - task, ok := w.activeTasks[taskExecution{taskID: "task-1", executionID: "exec-1"}] - w.tasksMutex.Unlock() - if !ok { - t.Fatal("running task should remain in activeTasks until executeTask returns") - } - if task.cancellationSource != taskCancellationSourceUser { - t.Fatalf("cancellation source = %q, want %q", task.cancellationSource, taskCancellationSourceUser) + w.cancellationWG.Wait() + if w.activeTaskCount() != 0 { + t.Fatal("accepted cancellation should remove task tracking") } } @@ -179,6 +171,7 @@ func TestHandleTaskCancellationSpawnedNoopCancelEmitsNoMessage(t *testing.T) { // Must not panic and must not emit any task status message when the // backend's CancelTask is a no-op. w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "task-1"}) + w.cancellationWG.Wait() if len(w.outbound.messages) != 0 { t.Fatalf("expected no message for no-op-cancel spawned task, got %d", len(w.outbound.messages)) diff --git a/internal/worker/docker.go b/internal/worker/docker.go index ce4c4d5..742250f 100644 --- a/internal/worker/docker.go +++ b/internal/worker/docker.go @@ -285,12 +285,6 @@ func (b *DockerBackend) CancelTask(ctx context.Context, params *CancelParams) er return nil } -// ConfirmTaskClosed reports whether the execution has no remaining container. -func (b *DockerBackend) ConfirmTaskClosed(_ context.Context, params *CancelParams) (bool, error) { - _, exists := b.taskContainers.Load(taskExecution{params.TaskID, params.ExecutionID}) - return !exists, nil -} - // Shutdown closes the Docker client. func (b *DockerBackend) Shutdown(ctx context.Context) { if b.dockerClient != nil { diff --git a/internal/worker/execution_closure_backends_test.go b/internal/worker/execution_cancellation_backends_test.go similarity index 71% rename from internal/worker/execution_closure_backends_test.go rename to internal/worker/execution_cancellation_backends_test.go index 8cd8a7b..667d3a6 100644 --- a/internal/worker/execution_closure_backends_test.go +++ b/internal/worker/execution_cancellation_backends_test.go @@ -8,7 +8,6 @@ import ( "path/filepath" "strings" "testing" - "time" "github.com/moby/moby/client" batchv1 "k8s.io/api/batch/v1" @@ -58,17 +57,17 @@ func TestDirectExecutionWorkspaceIsolation(t *testing.T) { } } -type closureDockerTransport func(*http.Request) (*http.Response, error) +type cancellationDockerTransport func(*http.Request) (*http.Response, error) -func (f closureDockerTransport) RoundTrip(r *http.Request) (*http.Response, error) { +func (f cancellationDockerTransport) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } -func TestDockerExecutionClosure(t *testing.T) { +func TestDockerExecutionCancellation(t *testing.T) { t.Run("failed removal retains exact resource for retry", func(t *testing.T) { fail := true var removed []string - transport := closureDockerTransport(func(r *http.Request) (*http.Response, error) { + transport := cancellationDockerTransport(func(r *http.Request) (*http.Response, error) { if r.URL.Path == "/_ping" { return mockEngineResponse(r, http.StatusOK, "text/plain", "OK"), nil } @@ -95,15 +94,15 @@ func TestDockerExecutionClosure(t *testing.T) { if err := backend.CancelTask(context.Background(), params); err == nil { t.Fatal("expected retryable removal failure") } - if closed, _ := backend.ConfirmTaskClosed(context.Background(), params); closed { - t.Fatal("failed removal reported closure") + if id, found := backend.taskContainers.Load(old); !found || id != "old-container" { + t.Fatal("failed removal lost resource tracking") } fail = false if err := backend.CancelTask(context.Background(), params); err != nil { t.Fatal(err) } - if closed, err := backend.ConfirmTaskClosed(context.Background(), params); !closed || err != nil { - t.Fatalf("closed=%v, error=%v", closed, err) + if err := backend.CancelTask(context.Background(), params); err != nil { + t.Fatal(err) } if _, found := backend.taskContainers.Load(successor); !found { t.Fatal("successor tracking removed") @@ -113,38 +112,36 @@ func TestDockerExecutionClosure(t *testing.T) { } }) - t.Run("unknown create outcome cannot acknowledge absence", func(t *testing.T) { + t.Run("unknown create outcome remains an observable failure", func(t *testing.T) { backend := &DockerBackend{} backend.taskContainers.Store(taskExecution{"run", "old"}, "") if err := backend.CancelTask(context.Background(), &CancelParams{TaskID: "run", ExecutionID: "old"}); err == nil { - t.Fatal("ambiguous creation treated as closure") + t.Fatal("ambiguous creation treated as accepted cancellation") } }) } -func TestKubernetesExecutionClosure(t *testing.T) { - t.Run("deleting job is insufficient while pods remain", func(t *testing.T) { +func TestKubernetesExecutionCancellation(t *testing.T) { + t.Run("accepted deletion succeeds even while job and pods remain", func(t *testing.T) { backend := &KubernetesBackend{config: KubernetesBackendConfig{Namespace: "agents"}} params := &CancelParams{TaskID: "run", ExecutionID: "old"} pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{ Name: "old-pod", Namespace: "agents", Labels: backend.baseLabels("run", "old"), }} - clientset := fake.NewSimpleClientset(pod) + job := &batchv1.Job{ObjectMeta: metav1.ObjectMeta{ + Name: kubernetesTaskJobName("run", "old"), Namespace: "agents", Labels: backend.baseLabels("run", "old"), + }} + clientset := fake.NewSimpleClientset(job, pod) backend.clientset = clientset - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - clientset.PrependReactor("list", "pods", func(k8stesting.Action) (bool, runtime.Object, error) { - cancel() - return false, nil, nil + clientset.PrependReactor("delete", "jobs", func(k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, nil }) - if closed, err := backend.ConfirmTaskClosed(ctx, params); closed || !errors.Is(err, context.Canceled) { - t.Fatalf("remaining pod yielded closed=%v, err=%v", closed, err) - } - if err := clientset.CoreV1().Pods("agents").Delete(context.Background(), pod.Name, metav1.DeleteOptions{}); err != nil { + if err := backend.CancelTask(context.Background(), params); err != nil { t.Fatal(err) } - if closed, err := backend.ConfirmTaskClosed(context.Background(), params); !closed || err != nil { - t.Fatalf("removed pod yielded closed=%v, err=%v", closed, err) + actions := clientset.Actions() + if len(actions) != 2 || actions[0].GetVerb() != "get" || actions[1].GetVerb() != "delete" { + t.Fatalf("cancellation should only get and delete the job: %v", actions) } }) @@ -154,8 +151,8 @@ func TestKubernetesExecutionClosure(t *testing.T) { return true, nil, apierrors.NewForbidden(schema.GroupResource{Resource: "jobs"}, "job", errors.New("denied")) }) backend := &KubernetesBackend{config: KubernetesBackendConfig{Namespace: "agents"}, clientset: clientset} - if closed, err := backend.ConfirmTaskClosed(context.Background(), &CancelParams{TaskID: "run", ExecutionID: "old"}); closed || err == nil { - t.Fatalf("API failure yielded closed=%v, err=%v", closed, err) + if err := backend.CancelTask(context.Background(), &CancelParams{TaskID: "run", ExecutionID: "old"}); err == nil { + t.Fatal("API failure treated as accepted cancellation") } }) @@ -178,16 +175,13 @@ func TestKubernetesExecutionClosure(t *testing.T) { } }) - t.Run("existing job cannot be acknowledged closed", func(t *testing.T) { + t.Run("missing job cancellation is idempotent", func(t *testing.T) { backend := &KubernetesBackend{config: KubernetesBackendConfig{Namespace: "agents"}} - job := &batchv1.Job{ObjectMeta: metav1.ObjectMeta{ - Name: kubernetesTaskJobName("run", "old"), Namespace: "agents", Labels: backend.baseLabels("run", "old"), - }} - backend.clientset = fake.NewSimpleClientset(job) - ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second)) - defer cancel() - if closed, err := backend.ConfirmTaskClosed(ctx, &CancelParams{TaskID: "run", ExecutionID: "old"}); closed || err == nil { - t.Fatalf("live Job yielded closed=%v, err=%v", closed, err) + backend.clientset = fake.NewSimpleClientset() + for i := 0; i < 2; i++ { + if err := backend.CancelTask(context.Background(), &CancelParams{TaskID: "run", ExecutionID: "old"}); err != nil { + t.Fatal(err) + } } }) } diff --git a/internal/worker/execution_cancellation_test.go b/internal/worker/execution_cancellation_test.go index b62eec4..e9d4ae7 100644 --- a/internal/worker/execution_cancellation_test.go +++ b/internal/worker/execution_cancellation_test.go @@ -71,10 +71,6 @@ func (b *executionTestBackend) CancelTask(ctx context.Context, params *CancelPar return nil } -func (b *executionTestBackend) ConfirmTaskClosed(context.Context, *CancelParams) (bool, error) { - return true, nil -} - func newExecutionTestWorker(t *testing.T) (*Worker, *executionTestBackend) { t.Helper() backend := &executionTestBackend{ @@ -97,7 +93,10 @@ func newExecutionTestWorker(t *testing.T) (*Worker, *executionTestBackend) { } } w.taskWG.Wait() + w.shutdownTasks() + w.cancellationWG.Wait() }) + return w, backend } @@ -117,36 +116,19 @@ func awaitExecutionStarted(t *testing.T, backend *executionTestBackend, id strin } } -func awaitExecutionClosed(t *testing.T, w *Worker, id string) { +func awaitCancellationAccepted(t *testing.T, w *Worker, id string) { t.Helper() - timeout := time.NewTimer(2 * time.Second) - defer timeout.Stop() - for { - var msg types.WebSocketMessage - select { - case data := <-w.outbound.messages: - if err := json.Unmarshal(data, &msg); err != nil { - t.Fatal(err) - } - case <-timeout.C: - t.Fatal("execution closure was not acknowledged") - } - if msg.Type != types.MessageTypeExecutionClosed { - continue - } - var closed types.ExecutionClosedMessage - if err := json.Unmarshal(msg.Data, &closed); err != nil { - t.Fatal(err) - } - if closed.TaskID != "run" || closed.ExecutionID != id { - t.Fatalf("wrong closure identity: %+v", closed) - } - return - } + waitFor(t, 2*time.Second, func() bool { + w.tasksMutex.Lock() + defer w.tasksMutex.Unlock() + key := taskExecution{"run", id} + _, active := w.activeTasks[key] + return !active && w.completedTasks[key] + }) } func TestExecutionScopedCancellation(t *testing.T) { - t.Run("delayed predecessor cancellation leaves successor running and replays closure", func(t *testing.T) { + t.Run("delayed predecessor cancellation leaves successor running and is idempotent", func(t *testing.T) { w, backend := newExecutionTestWorker(t) w.handleTaskAssignment(executionAssignment("old")) awaitExecutionStarted(t, backend, "old") @@ -154,7 +136,7 @@ func TestExecutionScopedCancellation(t *testing.T) { awaitExecutionStarted(t, backend, "new") request := &types.TaskCancellationMessage{TaskID: "run", ExecutionID: "old"} w.handleTaskCancellation(request) - awaitExecutionClosed(t, w, "old") + awaitCancellationAccepted(t, w, "old") w.tasksMutex.Lock() newTask, exists := w.activeTasks[taskExecution{"run", "new"}] w.tasksMutex.Unlock() @@ -165,7 +147,16 @@ func TestExecutionScopedCancellation(t *testing.T) { t.Fatalf("cancelled wrong execution: %+v", params) } w.handleTaskCancellation(request) - awaitExecutionClosed(t, w, "old") + w.handleTaskAssignment(executionAssignment("old")) + w.cancellationWG.Wait() + for _, msg := range drainMessages(t, w.outbound.messages) { + if msg.Type != types.MessageTypeTaskClaimed && msg.Type != types.MessageTypeTaskCompleted { + t.Fatalf("unexpected cancellation message: %s", msg.Type) + } + } + if len(backend.started) != 0 { + t.Fatal("accepted cancellation did not suppress duplicate assignment") + } select { case params := <-backend.cancelled: t.Fatalf("replay called backend again: %+v", params) @@ -173,7 +164,7 @@ func TestExecutionScopedCancellation(t *testing.T) { } }) - t.Run("old completion cannot remove successor and remains closable", func(t *testing.T) { + t.Run("old completion cannot remove successor and remains cancellable", func(t *testing.T) { w, backend := newExecutionTestWorker(t) for _, id := range []string{"old", "new"} { w.handleTaskAssignment(executionAssignment(id)) @@ -190,7 +181,7 @@ func TestExecutionScopedCancellation(t *testing.T) { t.Fatalf("active executions = %d, want 1", count) } w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "run", ExecutionID: "old"}) - awaitExecutionClosed(t, w, "old") + awaitCancellationAccepted(t, w, "old") if count := w.activeTaskCount(); count != 1 { t.Fatalf("cleanup removed successor: active executions = %d", count) } @@ -226,7 +217,7 @@ func TestExecutionScopedCancellation(t *testing.T) { } }) - t.Run("queued cleanup is not proof and failed cleanup remains retryable", func(t *testing.T) { + t.Run("failed cancellation retries locally without duplicating in-progress or accepted requests", func(t *testing.T) { w, backend := newExecutionTestWorker(t) backend.cancelGate = make(chan struct{}) backend.failCancel.Store(true) @@ -239,28 +230,37 @@ func TestExecutionScopedCancellation(t *testing.T) { case <-time.After(2 * time.Second): t.Fatal("backend cleanup did not start") } - for _, msg := range drainMessages(t, w.outbound.messages) { - if msg.Type == types.MessageTypeExecutionClosed { - t.Fatal("acknowledged queued cleanup") - } - } w.handleTaskCancellation(request) if len(backend.cancelled) != 0 { t.Fatal("concurrent retry duplicated cleanup") } close(backend.cancelGate) - waitFor(t, 2*time.Second, func() bool { - w.tasksMutex.Lock() - defer w.tasksMutex.Unlock() - task, tracked := w.activeTasks[taskExecution{"run", "old"}] - return tracked && !task.cancelling - }) - if len(w.outbound.messages) != 0 { - t.Fatal("failed cleanup emitted closure") + select { + case <-backend.cancelled: + case <-time.After(2 * time.Second): + t.Fatal("failed cancellation did not retry locally") } + w.handleTaskAssignment(executionAssignment("new")) + awaitExecutionStarted(t, backend, "new") backend.failCancel.Store(false) + awaitCancellationAccepted(t, w, "old") + w.cancellationWG.Wait() + w.tasksMutex.Lock() + successor := w.activeTasks[taskExecution{"run", "new"}] + w.tasksMutex.Unlock() + if successor.ctx.Err() != nil { + t.Fatal("retry cancelled successor") + } + for len(backend.cancelled) > 0 { + if params := <-backend.cancelled; params.ExecutionID != "old" { + t.Fatalf("retry targeted successor: %+v", params) + } + } w.handleTaskCancellation(request) - awaitExecutionClosed(t, w, "old") + w.cancellationWG.Wait() + if len(backend.cancelled) != 0 { + t.Fatal("accepted request was retried") + } }) t.Run("waits for dispatch before backend cancellation", func(t *testing.T) { @@ -271,24 +271,83 @@ func TestExecutionScopedCancellation(t *testing.T) { w.activeTasks[key] = activeTask{ctx: ctx, cancel: cancel, done: done} w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "run", ExecutionID: "old"}) if len(backend.cancelled) != 0 || len(w.outbound.messages) != 0 { - t.Fatal("cancel acknowledged before dispatch returned") + t.Fatal("backend cancellation started before dispatch returned") } close(done) - awaitExecutionClosed(t, w, "old") + awaitCancellationAccepted(t, w, "old") }) } -func TestExecutionClosureReplayCacheBounded(t *testing.T) { +func TestCancellationStopsOnShutdown(t *testing.T) { + t.Run("provisioning wait respects worker context", func(t *testing.T) { + w, backend := newExecutionTestWorker(t) + ctx, stop := context.WithCancel(context.Background()) + defer stop() + w.ctx = ctx + w.activeTasks[taskExecution{"run", "old"}] = activeTask{ + ctx: context.Background(), cancel: func() {}, done: make(chan struct{}), + } + w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "run", ExecutionID: "old"}) + stop() + w.cancellationWG.Wait() + if len(backend.cancelled) != 0 { + t.Fatal("cancellation bypassed unfinished provisioning") + } + }) + t.Run("in-flight backend request respects worker context", func(t *testing.T) { + w, backend := newExecutionTestWorker(t) + ctx, stop := context.WithCancel(context.Background()) + defer stop() + w.ctx = ctx + backend.cancelGate = make(chan struct{}) + w.handleTaskAssignment(executionAssignment("old")) + awaitExecutionStarted(t, backend, "old") + w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "run", ExecutionID: "old"}) + select { + case <-backend.cancelled: + case <-time.After(2 * time.Second): + t.Fatal("cancellation did not start") + } + stop() + w.cancellationWG.Wait() + if w.completedTasks[taskExecution{"run", "old"}] { + t.Fatal("interrupted cancellation was marked accepted") + } + }) + t.Run("failed request", func(t *testing.T) { + w, backend := newExecutionTestWorker(t) + backend.failCancel.Store(true) + w.handleTaskAssignment(executionAssignment("old")) + awaitExecutionStarted(t, backend, "old") + w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "run", ExecutionID: "old"}) + select { + case <-backend.cancelled: + case <-time.After(2 * time.Second): + t.Fatal("cancellation did not start") + } + w.shutdownTasks() + w.cancellationWG.Wait() + w.tasksMutex.Lock() + accepted := w.completedTasks[taskExecution{"run", "old"}] + w.tasksMutex.Unlock() + if accepted { + t.Fatal("failed cancellation was marked accepted") + } + }) +} +func TestCompletedExecutionCacheBounded(t *testing.T) { w := &Worker{} w.rememberCompletedLocked(taskExecution{taskID: "legacy-run"}, true) if len(w.completedTasks) != 0 { t.Fatal("legacy completion must not suppress a later run-only assignment") } for i := 0; i < 5000; i++ { - w.rememberCompletedLocked(taskExecution{taskID: "run", executionID: time.Unix(int64(i), 0).String()}, true) + key := taskExecution{taskID: "run", executionID: time.Unix(int64(i), 0).String()} + w.rememberCompletedLocked(key, false) + w.rememberCompletedLocked(key, true) } - if len(w.completedTasks) != 4096 || len(w.closedExecutions) != 4096 { - t.Fatalf("cache unbounded: %d records, %d keys", len(w.completedTasks), len(w.closedExecutions)) + if len(w.completedTasks) != 4096 || len(w.completedExecutions) != 4096 { + t.Fatalf("cache unbounded: %d records, %d keys", len(w.completedTasks), len(w.completedExecutions)) } } diff --git a/internal/worker/kubernetes.go b/internal/worker/kubernetes.go index f1736b9..2a03a1e 100644 --- a/internal/worker/kubernetes.go +++ b/internal/worker/kubernetes.go @@ -545,38 +545,6 @@ func (b *KubernetesBackend) CancelTask(ctx context.Context, params *CancelParams return nil } -// ConfirmTaskClosed waits for both the Job and its Pods to disappear. -func (b *KubernetesBackend) ConfirmTaskClosed(ctx context.Context, params *CancelParams) (bool, error) { - executionID := executionIDOrTaskID(params.TaskID, params.ExecutionID) - jobName := kubernetesTaskJobName(params.TaskID, executionID) - ticker := time.NewTicker(100 * time.Millisecond) - defer ticker.Stop() - for { - job, err := b.getTaskJob(ctx, jobName) - if err != nil && !apierrors.IsNotFound(err) { - return false, err - } - if err == nil && (job.Labels[kubernetesExecutionHashLabel] != kubernetesLabelHash(executionID) || - job.Labels[kubernetesTaskHashLabel] != kubernetesLabelHash(params.TaskID)) { - return false, fmt.Errorf("job %s belongs to another execution", jobName) - } - if apierrors.IsNotFound(err) { - pods, err := b.listTaskPods(ctx, executionID) - if err != nil { - return false, err - } - if len(pods) == 0 { - return true, nil - } - } - select { - case <-ctx.Done(): - return false, ctx.Err() - case <-ticker.C: - } - } -} - // Shutdown intentionally does not delete task Jobs. // // Kubernetes Jobs are the durable execution unit for this backend. During diff --git a/internal/worker/worker.go b/internal/worker/worker.go index f73f3fa..55f86b2 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -26,10 +26,12 @@ const ( MaxReconnectDelay = 60 * time.Second ReconnectBackoffRate = 2.0 - HeartbeatInterval = 30 * time.Second - PongWait = 60 * time.Second - WriteWait = 10 * time.Second - BackendShutdownTimeout = 10 * time.Second + HeartbeatInterval = 30 * time.Second + PongWait = 60 * time.Second + WriteWait = 10 * time.Second + BackendShutdownTimeout = 10 * time.Second + initialCancellationRetryDelay = 100 * time.Millisecond + maxCancellationRetryDelay = 10 * time.Second warpServerRootURLEnv = "WARP_SERVER_ROOT_URL" executionCancellationCapability = "execution-scoped-cancellation-v1" @@ -69,18 +71,21 @@ type taskExecution struct { } type Worker struct { - config Config - ctx context.Context - outbound *outboundQueue - activeTasks map[taskExecution]activeTask - completedTasks map[taskExecution]bool - closedExecutions []taskExecution - tasksMutex sync.Mutex - taskWG sync.WaitGroup - shuttingDown bool - oneShot oneShotState - backend Backend - taskSemaphore *semaphore.Weighted // nil when unlimited + config Config + ctx context.Context + outbound *outboundQueue + activeTasks map[taskExecution]activeTask + // The value distinguishes accepted cancellation from ordinary completion, + // which may still need backend cleanup after a failed execution or watch. + completedTasks map[taskExecution]bool + completedExecutions []taskExecution + tasksMutex sync.Mutex + taskWG sync.WaitGroup + cancellationWG sync.WaitGroup + shuttingDown bool + oneShot oneShotState + backend Backend + taskSemaphore *semaphore.Weighted // nil when unlimited // heartbeatInterval is how often the worker pings the server. It defaults // to HeartbeatInterval and is overridable in tests. heartbeatInterval time.Duration @@ -151,9 +156,9 @@ type activeTask struct { // spawned marks a task whose backend returned ExecuteOutcomeSpawned: it no // longer executes locally, but the entry is kept so a later cancellation // can be routed to the backend's CancelTask. - spawned bool - done chan struct{} - cancelling bool + spawned bool + done chan struct{} + stopCancellation context.CancelFunc } func New(ctx context.Context, config Config) (*Worker, error) { @@ -499,6 +504,10 @@ func (w *Worker) handleMessage(message []byte) { func (w *Worker) handleTaskCancellation(cancellation *types.TaskCancellationMessage) { w.tasksMutex.Lock() + if w.shuttingDown || w.ctx.Err() != nil { + w.tasksMutex.Unlock() + return + } key := taskExecution{cancellation.TaskID, cancellation.ExecutionID} if cancellation.ExecutionID == "" { // A task-only request is explicitly legacy run-wide cancellation. @@ -510,12 +519,9 @@ func (w *Worker) handleTaskCancellation(cancellation *types.TaskCancellationMess w.tasksMutex.Unlock() return } - if closed, known := w.completedTasks[key]; known { - if closed { + if cancellationAccepted, known := w.completedTasks[key]; known { + if cancellationAccepted { w.tasksMutex.Unlock() - if err := w.sendExecutionClosed(key); err != nil { - log.Warnf(w.ctx, "Failed to replay execution closure: %v", err) - } return } if _, active := w.activeTasks[key]; !active { @@ -530,102 +536,76 @@ func (w *Worker) handleTaskCancellation(cancellation *types.TaskCancellationMess func (w *Worker) startTaskCancellationLocked(key taskExecution) { task, ok := w.activeTasks[key] - if !ok || task.cancelling { + if !ok || task.stopCancellation != nil { return } task.cancellationSource = taskCancellationSourceUser - task.cancelling = true + ctx, cancel := context.WithCancel(w.ctx) + task.stopCancellation = cancel w.activeTasks[key] = task metrics.AddTaskEvent(task.ctx, "task.cancellation_requested", attribute.String("source", "server"), attribute.String("task.id", key.taskID), ) task.cancel() - go w.cancelTaskOnBackend(key, task.done) -} - -func (w *Worker) cancelTaskOnBackend(key taskExecution, done <-chan struct{}) { - ctx, cancel := context.WithTimeout(context.WithoutCancel(w.ctx), BackendShutdownTimeout) - defer cancel() - closed, err := w.closeExecution(ctx, key, done) - w.tasksMutex.Lock() - task, tracked := w.activeTasks[key] - if !tracked { - w.tasksMutex.Unlock() - return - } - task.cancelling = false - w.activeTasks[key] = task - if err != nil { - w.tasksMutex.Unlock() - log.Warnf(w.ctx, "Backend cancellation failed for task %s execution %s: %v", key.taskID, key.executionID, err) - return - } - if closed { - w.rememberCompletedLocked(key, true) - delete(w.activeTasks, key) - } - w.tasksMutex.Unlock() - if closed && key.executionID != "" { - if err := w.sendExecutionClosed(key); err != nil { - log.Warnf(w.ctx, "Failed to send execution closure: %v", err) - } - } + w.cancellationWG.Add(1) + go func() { + defer w.cancellationWG.Done() + defer cancel() + w.cancelTaskOnBackend(ctx, key, task.done) + }() } -func (w *Worker) closeExecution(ctx context.Context, key taskExecution, done <-chan struct{}) (bool, error) { - // A cancel racing provisioning must not confirm absence before creation finishes. +func (w *Worker) cancelTaskOnBackend(ctx context.Context, key taskExecution, done <-chan struct{}) { + // Cancelling before provisioning returns can miss resources still being created. if done != nil { select { case <-done: case <-ctx.Done(): - return false, ctx.Err() + return } } params := &CancelParams{TaskID: key.taskID, ExecutionID: key.executionID} - if err := w.backend.CancelTask(ctx, params); err != nil { - return false, err - } - if verifier, ok := w.backend.(executionClosureVerifier); ok { - closed, err := verifier.ConfirmTaskClosed(ctx, params) - if err != nil { - return false, err + delay := initialCancellationRetryDelay + for ctx.Err() == nil { + attemptCtx, cancel := context.WithTimeout(ctx, BackendShutdownTimeout) + err := w.backend.CancelTask(attemptCtx, params) + cancel() + if err == nil { + w.tasksMutex.Lock() + w.rememberCompletedLocked(key, true) + delete(w.activeTasks, key) + w.tasksMutex.Unlock() + return } - if closed { - return true, nil + log.Warnf(w.ctx, "Backend cancellation failed for task %s execution %s; retrying in %v: %v", key.taskID, key.executionID, delay, err) + timer := time.NewTimer(delay) + select { + case <-ctx.Done(): + timer.Stop() + return + case <-timer.C: } + delay = min(delay*2, maxCancellationRetryDelay) } - return false, fmt.Errorf("backend has not confirmed execution closure") } -func (w *Worker) rememberCompletedLocked(key taskExecution, closed bool) { +func (w *Worker) rememberCompletedLocked(key taskExecution, cancellationAccepted bool) { if key.executionID == "" { return } if w.completedTasks == nil { w.completedTasks = make(map[taskExecution]bool) } - w.completedTasks[key] = closed - if closed { - w.closedExecutions = append(w.closedExecutions, key) - const closedExecutionCacheSize = 4096 - if len(w.closedExecutions) > closedExecutionCacheSize { - delete(w.completedTasks, w.closedExecutions[0]) - w.closedExecutions = w.closedExecutions[1:] + if _, known := w.completedTasks[key]; !known { + w.completedExecutions = append(w.completedExecutions, key) + const completedExecutionCacheSize = 4096 + if len(w.completedExecutions) > completedExecutionCacheSize { + delete(w.completedTasks, w.completedExecutions[0]) + w.completedExecutions = w.completedExecutions[1:] } } -} - -func (w *Worker) sendExecutionClosed(key taskExecution) error { - data, err := json.Marshal(types.ExecutionClosedMessage{TaskID: key.taskID, ExecutionID: key.executionID}) - if err != nil { - return err - } - message, err := json.Marshal(types.WebSocketMessage{Type: types.MessageTypeExecutionClosed, Data: data}) - if err != nil { - return err - } - return w.sendMessage(message) + w.completedTasks[key] = cancellationAccepted } func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage) { @@ -1168,6 +1148,11 @@ func (w *Worker) shutdownTasks() { return } w.shuttingDown = true + for _, task := range w.activeTasks { + if task.stopCancellation != nil { + task.stopCancellation() + } + } activeTaskCount := len(w.activeTasks) if activeTaskCount > 0 && preserveActiveTasks { log.Infof(w.ctx, "Preserving %d active tasks during worker shutdown", activeTaskCount) @@ -1203,6 +1188,8 @@ func (w *Worker) shutdownTasks() { } func (w *Worker) shutdownBackend() { + w.shutdownTasks() + w.cancellationWG.Wait() backendShutdownCtx, backendShutdownCancel := context.WithTimeout(context.Background(), BackendShutdownTimeout) defer backendShutdownCancel() From 0369f4cf95aca37da206a182c1b78b282b6ea098 Mon Sep 17 00:00:00 2001 From: seemeroland Date: Wed, 7 Oct 2026 16:08:08 -0700 Subject: [PATCH 3/4] Small docker improvement --- internal/worker/docker.go | 5 +- .../execution_cancellation_backends_test.go | 72 +++++++++++++++++++ 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/internal/worker/docker.go b/internal/worker/docker.go index 742250f..4e33373 100644 --- a/internal/worker/docker.go +++ b/internal/worker/docker.go @@ -174,13 +174,16 @@ func (b *DockerBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe doneContainerStart := params.SetupEvents.startPhase(ctx, SetupEventContainerStart) key := taskExecution{params.TaskID, params.ExecutionID} - // A failed create response cannot prove that the daemon created no container. + // Transport and server failures cannot prove that the daemon created no container. b.taskContainers.Store(key, "") resp, err := dockerClient.ContainerCreate(ctx, client.ContainerCreateOptions{ Config: containerConfig, HostConfig: hostConfig, }) if err != nil { + if cerrdefs.IsInvalidArgument(err) || cerrdefs.IsNotFound(err) || cerrdefs.IsConflict(err) { + b.taskContainers.Delete(key) + } doneContainerStart(true) return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonContainerCreate, fmt.Errorf("failed to create container: %w", err))) } diff --git a/internal/worker/execution_cancellation_backends_test.go b/internal/worker/execution_cancellation_backends_test.go index 667d3a6..eefc386 100644 --- a/internal/worker/execution_cancellation_backends_test.go +++ b/internal/worker/execution_cancellation_backends_test.go @@ -63,6 +63,78 @@ func (f cancellationDockerTransport) RoundTrip(r *http.Request) (*http.Response, return f(r) } +func TestDockerCreateFailureTracking(t *testing.T) { + for _, tc := range []struct { + name string + status int + body string + transportError error + expectedTracked bool + }{ + {name: "invalid configuration", status: http.StatusBadRequest}, + {name: "missing image", status: http.StatusNotFound}, + {name: "name conflict", status: http.StatusConflict}, + {name: "server failure", status: http.StatusInternalServerError, expectedTracked: true}, + {name: "lost response", transportError: errors.New("response lost"), expectedTracked: true}, + {name: "cancelled request", transportError: context.Canceled, expectedTracked: true}, + {name: "malformed success response", status: http.StatusCreated, body: "{", expectedTracked: true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv("DOCKER_CONFIG", t.TempDir()) + createCalled := false + transport := cancellationDockerTransport(func(r *http.Request) (*http.Response, error) { + switch { + case r.URL.Path == "/_ping": + return mockEngineResponse(r, http.StatusOK, "text/plain", "OK"), nil + case r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/images/task-image/json"): + return mockImageInspectResponse(t, r, "linux", "amd64", ""), nil + case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/containers/create"): + createCalled = true + if tc.transportError != nil { + return nil, tc.transportError + } + if tc.body != "" { + return mockEngineResponse(r, tc.status, "application/json", tc.body), nil + } + return mockEngineErrorResponse(t, r, tc.status, tc.name), nil + default: + t.Fatalf("unexpected request: %s %s", r.Method, r.URL.Path) + return nil, nil + } + }) + dockerClient, err := client.New(client.WithHTTPClient(&http.Client{Transport: transport})) + if err != nil { + t.Fatal(err) + } + defer func() { _ = dockerClient.Close() }() + backend := &DockerBackend{ + dockerClient: dockerClient, + config: DockerBackendConfig{ImagePullPolicy: PullPolicyNever}, + platform: "linux/amd64", + } + key := taskExecution{"run", "old"} + successor := taskExecution{"run", "new"} + backend.taskContainers.Store(successor, "successor-container") + result := backend.ExecuteTask(context.Background(), &TaskParams{ + TaskID: key.taskID, ExecutionID: key.executionID, DockerImage: "task-image", + }) + if !createCalled || result.Error == nil { + t.Fatalf("expected create failure: create called = %t, result = %+v", createCalled, result) + } + if value, tracked := backend.taskContainers.Load(key); tracked != tc.expectedTracked || (tracked && value != "") { + t.Fatalf("create tracking = (%v, %t), expected tracked = %t", value, tracked, tc.expectedTracked) + } + err = backend.CancelTask(context.Background(), &CancelParams{TaskID: key.taskID, ExecutionID: key.executionID}) + if (err != nil) != tc.expectedTracked { + t.Fatalf("cancellation error = %v, expected uncertain outcome = %t", err, tc.expectedTracked) + } + if value, tracked := backend.taskContainers.Load(successor); !tracked || value != "successor-container" { + t.Fatal("create failure affected successor tracking") + } + }) + } +} + func TestDockerExecutionCancellation(t *testing.T) { t.Run("failed removal retains exact resource for retry", func(t *testing.T) { fail := true From f1d7bfea72b11b15d1dc7ee63bef04748448aef0 Mon Sep 17 00:00:00 2001 From: seemeroland Date: Thu, 15 Oct 2026 15:04:17 -0700 Subject: [PATCH 4/4] Address comments --- internal/worker/direct.go | 6 +- internal/worker/docker.go | 5 ++ .../execution_cancellation_backends_test.go | 72 +++++++++++++++++++ internal/worker/worker.go | 2 + 4 files changed, 83 insertions(+), 2 deletions(-) diff --git a/internal/worker/direct.go b/internal/worker/direct.go index 21c3221..da5bb71 100644 --- a/internal/worker/direct.go +++ b/internal/worker/direct.go @@ -169,20 +169,22 @@ func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonWorkspaceSetup, fmt.Errorf("invalid task ID for workspace path: %w", err))) } - // Determine working directory: shared target dir or per-task workspace. + // Determine working directory: shared target dir or isolated workspace. var workspaceDir string usingTargetDir := b.config.TargetDir != "" if usingTargetDir { workspaceDir = b.config.TargetDir } else { - // Create per-task workspace directory. + // Legacy assignments use a per-task workspace. workspaceDir = filepath.Join(b.config.WorkspaceRoot, taskID) if params.ExecutionID != "" { if err := validateTaskIDForPath(params.ExecutionID); err != nil { return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonWorkspaceSetup, fmt.Errorf("invalid execution ID for workspace path: %w", err))) } var err error + // Give each execution a fresh workspace so overlapping executions + // cannot share files or remove each other's workspace during cleanup. workspaceDir, err = os.MkdirTemp(b.config.WorkspaceRoot, taskID+"-"+params.ExecutionID+"-") if err != nil { return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonWorkspaceSetup, fmt.Errorf("failed to create execution workspace: %w", err))) diff --git a/internal/worker/docker.go b/internal/worker/docker.go index b576385..018333d 100644 --- a/internal/worker/docker.go +++ b/internal/worker/docker.go @@ -231,6 +231,11 @@ func (b *DockerBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe } case status := <-waitResult.Result: log.Debugf(ctx, "Container exited with status code: %d", status.StatusCode) + if b.config.NoCleanup && status.Error == nil { + // Preserve the exited container for debugging, but not its tracking. + // Later cancellation intentionally leaves this stopped container alone. + b.taskContainers.Delete(key) + } if params.Logs == nil && (zerolog.GlobalLevel() <= zerolog.DebugLevel || status.StatusCode != 0) { logOutput, logErr := b.getContainerLogs(ctx, dockerClient, containerID) diff --git a/internal/worker/execution_cancellation_backends_test.go b/internal/worker/execution_cancellation_backends_test.go index e7b7a45..9f323a7 100644 --- a/internal/worker/execution_cancellation_backends_test.go +++ b/internal/worker/execution_cancellation_backends_test.go @@ -135,6 +135,78 @@ func TestDockerCreateFailureTracking(t *testing.T) { } } +func TestDockerNoCleanupTracking(t *testing.T) { + for _, tc := range []struct { + name string + waitBody string + waitFails bool + expectedError bool + expectedTracked bool + }{ + {name: "successful exit", waitBody: `{"StatusCode":0}`}, + {name: "failed exit", waitBody: `{"StatusCode":1}`, expectedError: true}, + {name: "wait failure", waitFails: true, expectedError: true, expectedTracked: true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv("DOCKER_CONFIG", t.TempDir()) + transport := cancellationDockerTransport(func(r *http.Request) (*http.Response, error) { + switch { + case r.URL.Path == "/_ping": + return mockEngineResponse(r, http.StatusOK, "text/plain", "OK"), nil + case r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/images/task-image/json"): + return mockImageInspectResponse(t, r, "linux", "amd64", ""), nil + case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/containers/create"): + return mockEngineResponse(r, http.StatusCreated, "application/json", `{"Id":"task-container"}`), nil + case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/containers/task-container/start"): + return mockEngineResponse(r, http.StatusNoContent, "application/json", ""), nil + case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/containers/task-container/wait"): + if tc.waitFails { + return mockEngineErrorResponse(t, r, http.StatusInternalServerError, "unavailable"), nil + } + return mockEngineResponse(r, http.StatusOK, "application/json", tc.waitBody), nil + case r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/containers/task-container/logs"): + return mockEngineResponse(r, http.StatusOK, "application/octet-stream", ""), nil + case r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/containers/task-container/json"): + return mockEngineResponse(r, http.StatusOK, "application/json", `{"State":{"OOMKilled":false}}`), nil + default: + t.Errorf("unexpected request (NoCleanup must preserve the container): %s %s", r.Method, r.URL.Path) + return mockEngineErrorResponse(t, r, http.StatusInternalServerError, "unexpected request"), nil + } + }) + dockerClient, err := client.New(client.WithHTTPClient(&http.Client{Transport: transport})) + if err != nil { + t.Fatal(err) + } + defer func() { _ = dockerClient.Close() }() + backend := &DockerBackend{ + dockerClient: dockerClient, + config: DockerBackendConfig{NoCleanup: true, ImagePullPolicy: PullPolicyNever}, + platform: "linux/amd64", + } + key := taskExecution{"run", "old"} + successor := taskExecution{"run", "new"} + backend.taskContainers.Store(successor, "successor-container") + result := backend.ExecuteTask(context.Background(), &TaskParams{ + TaskID: key.taskID, ExecutionID: key.executionID, DockerImage: "task-image", + }) + if (result.Error != nil) != tc.expectedError { + t.Fatalf("execution error = %v, expected error = %t", result.Error, tc.expectedError) + } + if id, tracked := backend.taskContainers.Load(key); tracked != tc.expectedTracked || (tracked && id != "task-container") { + t.Fatalf("container tracking = (%v, %t), expected tracked = %t", id, tracked, tc.expectedTracked) + } + if id, tracked := backend.taskContainers.Load(successor); !tracked || id != "successor-container" { + t.Fatal("execution affected successor tracking") + } + if !tc.expectedTracked { + if err := backend.CancelTask(context.Background(), &CancelParams{TaskID: key.taskID, ExecutionID: key.executionID}); err != nil { + t.Fatalf("late cancellation of preserved exited container: %v", err) + } + } + }) + } +} + func TestDockerExecutionCancellation(t *testing.T) { t.Run("failed removal retains exact resource for retry", func(t *testing.T) { fail := true diff --git a/internal/worker/worker.go b/internal/worker/worker.go index 5b9da10..a6de76b 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -512,6 +512,8 @@ func (w *Worker) handleTaskCancellation(cancellation *types.TaskCancellationMess key := taskExecution{cancellation.TaskID, cancellation.ExecutionID} if cancellation.ExecutionID == "" { // A task-only request is explicitly legacy run-wide cancellation. + // TODO: Remove this legacy branch once execution-scoped cancellation is + // fully rolled out on the server and pre-rollout messages have drained. for candidate := range w.activeTasks { if candidate.taskID == cancellation.TaskID { w.startTaskCancellationLocked(candidate)