diff --git a/internal/types/messages.go b/internal/types/messages.go index 529aa3e..6aad91b 100644 --- a/internal/types/messages.go +++ b/internal/types/messages.go @@ -143,8 +143,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. @@ -178,7 +179,8 @@ 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"` } // TaskState is the serialized terminal task state accepted by warp-server. diff --git a/internal/worker/backend.go b/internal/worker/backend.go index 1034e07..91e585a 100644 --- a/internal/worker/backend.go +++ b/internal/worker/backend.go @@ -178,9 +178,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 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. 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 4f54811..da5bb71 100644 --- a/internal/worker/direct.go +++ b/internal/worker/direct.go @@ -169,15 +169,27 @@ 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))) + } + } 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..16a6e41 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, } @@ -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["task-1"] - w.tasksMutex.Unlock() - if ok { - t.Error("spawned task should be removed after cancellation is routed") + w.cancellationWG.Wait() + if w.activeTaskCount() != 0 { + t.Error("accepted cancellation should remove spawned task tracking") } } @@ -134,7 +132,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,15 +154,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. - w.tasksMutex.Lock() - task, ok := w.activeTasks["task-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") } } @@ -172,13 +164,14 @@ 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{}, } // 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 5cc1604..018333d 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. @@ -181,24 +182,33 @@ func (b *DockerBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe } doneContainerStart := params.SetupEvents.startPhase(ctx, SetupEventContainerStart) + key := taskExecution{params.TaskID, params.ExecutionID} + // 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))) } 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) } } }() @@ -221,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) @@ -272,9 +287,22 @@ 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 +} // Shutdown closes the Docker client. func (b *DockerBackend) Shutdown(ctx context.Context) { diff --git a/internal/worker/execution_cancellation_backends_test.go b/internal/worker/execution_cancellation_backends_test.go new file mode 100644 index 0000000..9f323a7 --- /dev/null +++ b/internal/worker/execution_cancellation_backends_test.go @@ -0,0 +1,331 @@ +package worker + +import ( + "context" + "errors" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + + "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], "", nil) + if _, err := os.Stat(workspaces[1]); err != nil { + t.Fatalf("predecessor cleanup removed successor workspace: %v", err) + } + }) + } +} + +type cancellationDockerTransport func(*http.Request) (*http.Response, error) + +func (f cancellationDockerTransport) RoundTrip(r *http.Request) (*http.Response, error) { + 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 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 + var removed []string + transport := cancellationDockerTransport(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 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 err := backend.CancelTask(context.Background(), params); err != nil { + t.Fatal(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 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 accepted cancellation") + } + }) +} + +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"), + }} + 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 + clientset.PrependReactor("delete", "jobs", func(k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, nil + }) + if err := backend.CancelTask(context.Background(), params); err != nil { + t.Fatal(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) + } + }) + + 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 err := backend.CancelTask(context.Background(), &CancelParams{TaskID: "run", ExecutionID: "old"}); err == nil { + t.Fatal("API failure treated as accepted cancellation") + } + }) + + 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("missing job cancellation is idempotent", func(t *testing.T) { + backend := &KubernetesBackend{config: KubernetesBackendConfig{Namespace: "agents"}} + 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 new file mode 100644 index 0000000..e9d4ae7 --- /dev/null +++ b/internal/worker/execution_cancellation_test.go @@ -0,0 +1,376 @@ +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 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() + w.shutdownTasks() + w.cancellationWG.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 awaitCancellationAccepted(t *testing.T, w *Worker, id string) { + t.Helper() + 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 is idempotent", 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) + awaitCancellationAccepted(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) + 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) + default: + } + }) + + 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)) + 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"}) + awaitCancellationAccepted(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("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) + 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") + } + w.handleTaskCancellation(request) + if len(backend.cancelled) != 0 { + t.Fatal("concurrent retry duplicated cleanup") + } + close(backend.cancelGate) + 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) + 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) { + 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("backend cancellation started before dispatch returned") + } + close(done) + awaitCancellationAccepted(t, w, "old") + }) +} + +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++ { + 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.completedExecutions) != 4096 { + t.Fatalf("cache unbounded: %d records, %d keys", len(w.completedTasks), len(w.completedExecutions)) + } +} + +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/kubernetes.go b/internal/worker/kubernetes.go index a04c4b4..ef81f72 100644 --- a/internal/worker/kubernetes.go +++ b/internal/worker/kubernetes.go @@ -572,9 +572,21 @@ 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 diff --git a/internal/worker/task_logs_test.go b/internal/worker/task_logs_test.go index 3b05132..c38dbca 100644 --- a/internal/worker/task_logs_test.go +++ b/internal/worker/task_logs_test.go @@ -76,14 +76,18 @@ func TestSetupEnvironmentCredentialsRedactedInAgentOutput(t *testing.T) { if err := os.WriteFile(ozPath, []byte(script), 0700); err != nil { t.Fatal(err) } - backend := &DirectBackend{ozPath: ozPath, config: DirectBackendConfig{ + backend, err := NewDirectBackend(t.Context(), DirectBackendConfig{ + OzPath: ozPath, WorkspaceRoot: filepath.Join(dir, "workspaces"), SetupCommand: `printf 'CUSTOM_CREDENTIAL=generated-credential\n' >"$OZ_ENVIRONMENT_FILE"`, - }} + }) + if err != nil { + t.Fatal(err) + } ctx, cancel := context.WithCancel(t.Context()) defer cancel() w := &Worker{config: Config{BackendType: "direct", ServerRootURL: server.URL, CollectTaskLogs: true}, - ctx: ctx, backend: backend, activeTasks: make(map[string]activeTask), outbound: newOutboundQueue(8)} + ctx: ctx, backend: backend, activeTasks: make(map[taskExecution]activeTask), outbound: newOutboundQueue(8)} _, span := noop.NewTracerProvider().Tracer("test").Start(ctx, "task") w.executeTask(ctx, cancel, span, logAssignment(server.URL), time.Now()) text := logPayloads(payloads) @@ -242,15 +246,17 @@ func TestDirectTaskLogLifecycle(t *testing.T) { if err := os.WriteFile(ozPath, []byte("#!/bin/sh\nprintf 'agent credential:%s\\n' \"$CUSTOM_CREDENTIAL\"\nprintf 'error output\\n' >&2\n"), 0700); err != nil { t.Fatal(err) } - backend := &DirectBackend{ - ozPath: ozPath, - config: DirectBackendConfig{WorkspaceRoot: filepath.Join(dir, "workspaces"), SetupCommand: tc.setup, Env: map[string]string{"CUSTOM_CREDENTIAL": "private-backend-value"}}, + backend, err := NewDirectBackend(t.Context(), DirectBackendConfig{ + OzPath: ozPath, WorkspaceRoot: filepath.Join(dir, "workspaces"), SetupCommand: tc.setup, Env: map[string]string{"CUSTOM_CREDENTIAL": "private-backend-value"}, + }) + if err != nil { + t.Fatal(err) } ctx, cancel := context.WithCancel(t.Context()) defer cancel() w := &Worker{ config: Config{BackendType: "direct", ServerRootURL: server.URL, CollectTaskLogs: !tc.disabled}, - ctx: ctx, backend: backend, activeTasks: make(map[string]activeTask), outbound: newOutboundQueue(8), + ctx: ctx, backend: backend, activeTasks: make(map[taskExecution]activeTask), outbound: newOutboundQueue(8), } _, span := noop.NewTracerProvider().Tracer("test").Start(ctx, "task") assignment := logAssignment(server.URL) @@ -314,16 +320,16 @@ func TestTaskLogsFlushOnCancellationAndPreservingShutdown(t *testing.T) { backend := &cancellableLoggingBackend{started: make(chan struct{}), preserve: preserve} w := &Worker{ config: Config{ServerRootURL: server.URL, BackendType: "direct", CollectTaskLogs: true}, - ctx: ctx, backend: backend, activeTasks: make(map[string]activeTask), outbound: newOutboundQueue(8), oneShot: newOneShotState(), + ctx: ctx, backend: backend, activeTasks: make(map[taskExecution]activeTask), outbound: newOutboundQueue(8), oneShot: newOneShotState(), } w.handleTaskAssignment(logAssignment(server.URL)) <-backend.started if preserve { w.shutdownTasks() - if w.activeTasks["run"].ctx.Err() != nil { + if w.activeTasks[taskExecution{"run", "execution"}].ctx.Err() != nil { t.Error("reporter shutdown cancelled a preserved task") } - w.activeTasks["run"].cancel() + w.activeTasks[taskExecution{"run", "execution"}].cancel() w.taskWG.Wait() } else { w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "run"}) diff --git a/internal/worker/worker.go b/internal/worker/worker.go index 071aad8..a6de76b 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -27,12 +27,15 @@ const ( MaxReconnectDelay = 60 * time.Second ReconnectBackoffRate = 2.0 - HeartbeatInterval = 30 * time.Second - PongWait = 60 * time.Second - WriteWait = 10 * time.Second - BackendShutdownTimeout = 10 * time.Second - - warpServerRootURLEnv = "WARP_SERVER_ROOT_URL" + 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" ) type Config struct { @@ -64,18 +67,27 @@ 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 + // 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 @@ -146,8 +158,10 @@ 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 - logs *tasklogs.Reporter + spawned bool + done chan struct{} + stopCancellation context.CancelFunc + logs *tasklogs.Reporter } func New(ctx context.Context, config Config) (*Worker, error) { @@ -202,7 +216,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, @@ -295,6 +309,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) @@ -490,55 +505,114 @@ 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 + 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. + // 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) + } } - 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 cancellationAccepted, known := w.completedTasks[key]; known { + if cancellationAccepted { + w.tasksMutex.Unlock() + 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.stopCancellation != nil { return } - - log.Infof(w.ctx, "Cancelling task from server request: taskID=%s", cancellation.TaskID) + task.cancellationSource = taskCancellationSourceUser + 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", 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() -} - -// 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) + w.cancellationWG.Add(1) go func() { - ctx, cancel := context.WithTimeout(context.WithoutCancel(w.ctx), BackendShutdownTimeout) + defer w.cancellationWG.Done() 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), - ) - } + w.cancelTaskOnBackend(ctx, key, task.done) }() } +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 + } + } + params := &CancelParams{TaskID: key.taskID, ExecutionID: key.executionID} + 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 + } + 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) + } +} + +func (w *Worker) rememberCompletedLocked(key taskExecution, cancellationAccepted bool) { + if key.executionID == "" { + return + } + if w.completedTasks == nil { + w.completedTasks = make(map[taskExecution]bool) + } + 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:] + } + } + w.completedTasks[key] = cancellationAccepted +} + 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) @@ -549,6 +623,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 { @@ -601,16 +682,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() @@ -788,6 +882,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 @@ -799,13 +894,17 @@ 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) - } else { - task.logs = nil - w.activeTasks[assignment.TaskID] = task + 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) + } else { + task.logs = nil + w.activeTasks[key] = task + } } w.tasksMutex.Unlock() @@ -839,9 +938,9 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc, params.SetupEvents.logs = reporter w.tasksMutex.Lock() shuttingDown := w.shuttingDown - if task, ok := w.activeTasks[taskID]; ok { + if task, ok := w.activeTasks[key]; ok { task.logs = reporter - w.activeTasks[taskID] = task + w.activeTasks[key] = task } w.tasksMutex.Unlock() if shuttingDown { @@ -862,7 +961,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)), @@ -882,7 +981,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 } @@ -909,9 +1008,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") @@ -936,20 +1035,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) @@ -1089,6 +1189,11 @@ func (w *Worker) shutdownTasks() { return } w.shuttingDown = true + for _, task := range w.activeTasks { + if task.stopCancellation != nil { + task.stopCancellation() + } + } activeTaskCount := len(w.activeTasks) var reporters []*tasklogs.Reporter for _, task := range w.activeTasks { @@ -1100,15 +1205,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() } @@ -1145,6 +1250,8 @@ func (w *Worker) shutdownTasks() { } func (w *Worker) shutdownBackend() { + w.shutdownTasks() + w.cancellationWG.Wait() backendShutdownCtx, backendShutdownCancel := context.WithTimeout(context.Background(), BackendShutdownTimeout) defer backendShutdownCancel() 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{}, }