From 864488c67de2f274abbc651753e03b23dfe8c9be Mon Sep 17 00:00:00 2001 From: JBD Date: Sat, 26 Sep 2026 00:13:55 -0700 Subject: [PATCH] Replace Redis Streams queue and controller with direct execution and resource locking Transition the AX architecture from asynchronous event queues (Redis Streams and the `ax-controller` worker pool) to direct, synchronous reconciliation in `ax-server` guarded by fine-grained distributed locks. Key changes: - Implement per-resource locking (`lock:task`, `lock:workspace`, `lock:model`) using Redis distributed locks (`SET NX PX` with token validation) and an in-memory locker for testing. - Move Substrate reconciliation directly into `ax-server` for Task, Workspace, and Model lifecycle methods (`Create`, `Resume`, `Suspend`, `Delete`). - Decommission `ax-controller`, `deploy/ax-controller.yaml`, and Redis Stream event publishing/consuming interfaces. - Simplify `ax delete` and `ax watch` in the CLI to eliminate two-phase polling loops now that operations complete synchronously. - Update `demo.sh`, build configurations (`Makefile`, `.ko.yaml`), and documentation to reflect the single-binary control plane architecture. --- .dockerignore | 1 - .github/workflows/go.yml | 1 - .gitignore | 1 - .ko.yaml | 4 +- CONTRIBUTING.md | 4 +- DESIGN.md | 34 +--- Makefile | 15 +- README.md | 6 +- cmd/ax-controller/main.go | 111 ----------- cmd/ax-server/main.go | 55 +++++- cmd/ax-task-runner/main.go | 2 +- cmd/ax/main.go | 55 +----- cmd/ax/main_test.go | 72 +++++++ demo.sh | 55 ++---- deploy/ax-controller.yaml | 119 ----------- deploy/ax-server.yaml | 72 +++++++ docs/concepts.md | 2 +- docs/development.md | 2 +- docs/networking.md | 4 +- docs/roadmap.md | 2 +- docs/runner.md | 12 +- examples/simple.yaml | 2 +- internal/controller/reconciler.go | 4 +- internal/controller/reconciler_test.go | 113 +++++++++++ internal/controller/worker.go | 163 --------------- internal/controller/worker_test.go | 233 ---------------------- internal/lock/lock.go | 237 ++++++++++++++++++++++ internal/lock/lock_test.go | 96 +++++++++ internal/server/server.go | 261 +++++++++++++++++++++---- internal/server/server_test.go | 215 ++++++++++++++++++-- internal/store/memory/store.go | 82 -------- internal/store/redis/store.go | 164 +--------------- internal/store/store.go | 38 +--- internal/substrate/client.go | 86 ++++++-- internal/substrate/client_test.go | 142 ++++++++++++++ 35 files changed, 1333 insertions(+), 1132 deletions(-) delete mode 100644 cmd/ax-controller/main.go delete mode 100644 deploy/ax-controller.yaml delete mode 100644 internal/controller/worker.go delete mode 100644 internal/controller/worker_test.go create mode 100644 internal/lock/lock.go create mode 100644 internal/lock/lock_test.go create mode 100644 internal/substrate/client_test.go diff --git a/.dockerignore b/.dockerignore index 96a55e46..c687fdc2 100644 --- a/.dockerignore +++ b/.dockerignore @@ -12,4 +12,3 @@ bin/* ax ax-server ax-task-runner -ax-controller diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index 2bc48a6e..c2b65d83 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -42,7 +42,6 @@ jobs: - name: Build local binaries run: | go build -v ./cmd/ax - go build -v ./cmd/ax-controller go build -v ./cmd/ax-server go build -v ./cmd/ax-task-runner diff --git a/.gitignore b/.gitignore index 22061f52..5bb6ff24 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,5 @@ bin/ /ax /ax-server /ax-task-runner -/ax-controller cmd/ax-task-runner/__pycache__/antigravity_bootstrap.cpython-313.pyc diff --git a/.ko.yaml b/.ko.yaml index 18e06c85..3a3c0bc8 100644 --- a/.ko.yaml +++ b/.ko.yaml @@ -16,8 +16,8 @@ defaultBaseImage: cgr.dev/chainguard/static:latest baseImageOverrides: github.com/google/ax/cmd/ax-task-runner: alpine/git:latest builds: - - id: ax-controller - main: ./cmd/ax-controller + - id: ax-server + main: ./cmd/ax-server env: - CGO_ENABLED=0 flags: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d1ee05e8..aa14f69d 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -33,7 +33,7 @@ All submissions, including submissions by project members, require review. We us ### Building Binaries -Build all local binaries (`bin/ax`, `bin/ax-controller`, `bin/ax-server`): +Build all local binaries (`bin/ax`, `bin/ax-server`): ```bash make build @@ -81,7 +81,7 @@ make push-task-runner TASK_RUNNER_REPO=gcr.io//ax-task-runner ### Deploying to Kubernetes -Deploy Redis, controller, and server components to your cluster in the `ax-system` namespace: +Deploy Redis and server components to your cluster in the `ax-system` namespace: ```bash make deploy AX_IMAGE_REPO=gcr.io//ax-images diff --git a/DESIGN.md b/DESIGN.md index 849a0a83..84f93c87 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -2,7 +2,7 @@ ## Architecture -Storing millions of short-lived tasks as Kubernetes CRDs pushes etcd past its comfort zone (single-digit GB storage limits, write-rate bottlenecks, control plane degradation). AX keeps its state in Redis and uses Redis Streams as the work queue between the API server and a horizontally scaled pool of controllers. +Storing millions of short-lived tasks as Kubernetes CRDs pushes etcd past its comfort zone (single-digit GB storage limits, write-rate bottlenecks, control plane degradation). AX stores its state in Redis and reconciles directly with Agent Substrate under fine-grained distributed locks. ``` ax apply -f task.yaml @@ -11,27 +11,14 @@ Storing millions of short-lived tasks as Kubernetes CRDs pushes etcd past its co ax-server (gRPC API + /healthz) │ - store & publish event - │ - ▼ - Redis - (Task Hashes + Event Streams + PubSub) - │ - XREADGROUP (Streams) - │ - ▼ - ax-controller - (Horizontally Scaled Workers) - │ - gRPC (Control API) - │ - ▼ - Agent Substrate - ┌───────────────────────────────┐ - │ • Atespace Provisioning │ - │ • Actor Creation & Activation │ - │ • Worker Assignment │ - └───────────────────────────────┘ + ┌──────────────────────┴──────────────────────┐ + ▼ ▼ + Redis Agent Substrate + (Resource Store, ┌───────────────────────────────┐ + Locks, PubSub) │ • Atespace Provisioning │ + │ • Actor Creation & Activation │ + │ • Worker Assignment │ + └───────────────────────────────┘ ``` ## Components @@ -39,8 +26,7 @@ Storing millions of short-lived tasks as Kubernetes CRDs pushes etcd past its co | Binary | Role | |---|---| | `ax` | Developer CLI. Applies manifests, inspects and watches resources, tunnels to the cluster. | -| `ax-server` | Stateless gRPC API on port 8080. Validates manifests, persists to Redis, publishes events. | -| `ax-controller` | Reconciliation workers. Consume the Redis stream, provision atespaces and actors on Agent Substrate, and drive tasks toward desired state. Scale by adding replicas. | +| `ax-server` | Direct-execution gRPC API on port 8080. Validates manifests, manages distributed locks, reconciles directly with Agent Substrate, and persists state to Redis. | | `ax-task-runner` | Entrypoint inside every task container. Bootstraps the workspace, serves metadata, and runs the agent command. A thin wrapper over the `runner` package, which custom images can embed directly. | ## API reference diff --git a/Makefile b/Makefile index cdfac40c..0ffba309 100644 --- a/Makefile +++ b/Makefile @@ -19,7 +19,7 @@ AX_IMAGE_REPO ?= gcr.io/ax-substrate/ate-images TASK_RUNNER_REPO ?= $(AX_IMAGE_REPO)/ax-task-runner CONTAINER_CLI ?= $(shell which podman 2>/dev/null || which docker 2>/dev/null) -.PHONY: all build build-binaries build-task-runner install push push-task-runner deploy deploy-controller deploy-server deploy-redis apply-example test clean +.PHONY: all build build-binaries build-task-runner install push push-task-runner deploy deploy-server deploy-redis apply-example test clean all: build @@ -27,14 +27,13 @@ all: build ## Build Targets ## -------------------------------------- -# Build all local binaries (ax CLI, controller, server) +# Build all local binaries (ax CLI, server) build: build-binaries build-binaries: - @echo "==> Building local binaries (ax, ax-controller, ax-server)..." + @echo "==> Building local binaries (ax, ax-server)..." @mkdir -p bin go build -trimpath -ldflags="-s -w" -o bin/ax ./cmd/ax - go build -trimpath -ldflags="-s -w" -o bin/ax-controller ./cmd/ax-controller go build -trimpath -ldflags="-s -w" -o bin/ax-server ./cmd/ax-server # Install the ax CLI into $(go env GOPATH)/bin @@ -64,17 +63,13 @@ push: push-task-runner ## Deployment Targets ## -------------------------------------- -# Deploy all AX components to Kubernetes (Redis, ax-controller, ax-server) -deploy: deploy-redis deploy-controller deploy-server +# Deploy all AX components to Kubernetes (Redis, ax-server) +deploy: deploy-redis deploy-server deploy-redis: @echo "==> Deploying Redis to ax-system namespace..." kubectl apply -f deploy/redis.yaml -deploy-controller: - @echo "==> Building and deploying ax-controller using ko..." - KO_DOCKER_REPO=$(AX_IMAGE_REPO) ko apply -f deploy/ax-controller.yaml - deploy-server: @echo "==> Building and deploying ax-server using ko..." KO_DOCKER_REPO=$(AX_IMAGE_REPO) ko apply -f deploy/ax-server.yaml diff --git a/README.md b/README.md index 22b385a1..ebe78bf5 100644 --- a/README.md +++ b/README.md @@ -4,9 +4,9 @@ > [!WARNING] -> We are still actively refining our core concepts, protocols, -> and specifications. We will likely to introduce major breaking -> changes prior to a stable release. +> AX and several of its features are in heavy development. We are actively +> refining our core concepts, protocols, and specifications, and will likely +> introduce major breaking changes prior to a stable release. **Declare an agentic task with workspaces and model specifications. AX sandboxes it, wires up its workspace, and helps running it at scale.** diff --git a/cmd/ax-controller/main.go b/cmd/ax-controller/main.go deleted file mode 100644 index 9d1020ec..00000000 --- a/cmd/ax-controller/main.go +++ /dev/null @@ -1,111 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package main - -import ( - "context" - "flag" - "log/slog" - "os" - "os/signal" - "syscall" - - "github.com/google/ax/internal/controller" - "github.com/google/ax/internal/store/redis" - "github.com/google/ax/internal/substrate" - goredis "github.com/redis/go-redis/v9" -) - -func main() { - var ( - substrateEndpoint string - substrateAuthority string - substrateTokenFile string - substrateCAFile string - substrateInsecureTLS bool - substratePlaintext bool - defaultTemplate string - defaultTemplateAtespace string - redisAddr string - redisPassword string - redisGroup string - redisConsumer string - ) - - flag.StringVar(&redisAddr, "redis-addr", "localhost:6379", "Redis server address (e.g. localhost:6379)") - flag.StringVar(&redisPassword, "redis-password", "", "Redis password") - flag.StringVar(&redisGroup, "redis-group", "ax-controllers", "Redis stream consumer group") - flag.StringVar(&redisConsumer, "redis-consumer", "", "Redis stream consumer ID (defaults to hostname)") - flag.StringVar(&substrateEndpoint, "substrate-endpoint", "api.ate-system.svc.cluster.local:443", "Agent Substrate Control API endpoint") - flag.StringVar(&substrateAuthority, "substrate-authority", "api.ate-system.svc", "Authority / TLS ServerName for Substrate endpoint") - flag.StringVar(&substrateTokenFile, "substrate-token-file", "", "Path to bearer token file for Substrate auth") - flag.StringVar(&substrateCAFile, "substrate-ca-file", "", "Path to CA PEM file for Substrate TLS") - flag.BoolVar(&substrateInsecureTLS, "substrate-insecure-tls", false, "Skip Substrate TLS verification") - flag.BoolVar(&substratePlaintext, "substrate-plaintext", false, "Use insecure plaintext gRPC connection to Substrate") - flag.StringVar(&defaultTemplate, "template", "default-template", "Default Substrate ActorTemplate name") - flag.StringVar(&defaultTemplateAtespace, "template-atespace", "ax-system", "Default Substrate ActorTemplate atespace") - flag.Parse() - - if envRedis := os.Getenv("REDIS_ADDR"); envRedis != "" { - redisAddr = envRedis - } - if envPass := os.Getenv("REDIS_PASSWORD"); envPass != "" && redisPassword == "" { - redisPassword = envPass - } - - logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo})) - slog.SetDefault(logger) - - slog.Info("starting ax-controller", - "redisAddr", redisAddr, - "group", redisGroup, - "substrateEndpoint", substrateEndpoint, - "authority", substrateAuthority, - "templateAtespace", defaultTemplateAtespace, - "template", defaultTemplate, - ) - - subClient, err := substrate.NewClientWithOptions(substrate.ClientOptions{ - Target: substrateEndpoint, - Authority: substrateAuthority, - TokenFile: substrateTokenFile, - CAFile: substrateCAFile, - InsecureTLS: substrateInsecureTLS, - Plaintext: substratePlaintext, - }) - if err != nil { - slog.Error("failed to initialize substrate client", "error", err) - os.Exit(1) - } - defer subClient.Close() - - reconciler := controller.NewTaskReconciler(subClient, defaultTemplate, defaultTemplateAtespace) - - ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) - defer cancel() - - rClient := goredis.NewClient(&goredis.Options{ - Addr: redisAddr, - Password: redisPassword, - }) - defer rClient.Close() - - rStore := redis.NewStore(rClient, redis.Options{}) - worker := controller.NewWorker(rStore, reconciler, redisGroup, redisConsumer) - if err := worker.Run(ctx); err != nil && err != context.Canceled { - slog.Error("redis worker stopped with error", "error", err) - os.Exit(1) - } -} diff --git a/cmd/ax-server/main.go b/cmd/ax-server/main.go index 8e1a98ea..4f224fda 100644 --- a/cmd/ax-server/main.go +++ b/cmd/ax-server/main.go @@ -24,21 +24,40 @@ import ( "syscall" "time" + "github.com/google/ax/internal/controller" + "github.com/google/ax/internal/lock" "github.com/google/ax/internal/server" "github.com/google/ax/internal/store/redis" + "github.com/google/ax/internal/substrate" goredis "github.com/redis/go-redis/v9" ) func main() { var ( - listenAddr string - redisAddr string - redisPassword string + listenAddr string + redisAddr string + redisPassword string + substrateEndpoint string + substrateAuthority string + substrateTokenFile string + substrateCAFile string + substrateInsecureTLS bool + substratePlaintext bool + defaultTemplate string + defaultTemplateAtespace string ) flag.StringVar(&listenAddr, "addr", ":8080", "HTTP listen address") flag.StringVar(&redisAddr, "redis-addr", "localhost:6379", "Redis server address") flag.StringVar(&redisPassword, "redis-password", "", "Redis password") + flag.StringVar(&substrateEndpoint, "substrate-endpoint", "api.ate-system.svc.cluster.local:443", "Agent Substrate Control API endpoint") + flag.StringVar(&substrateAuthority, "substrate-authority", "api.ate-system.svc", "Authority / TLS ServerName for Substrate endpoint") + flag.StringVar(&substrateTokenFile, "substrate-token-file", "", "Path to bearer token file for Substrate auth") + flag.StringVar(&substrateCAFile, "substrate-ca-file", "", "Path to CA PEM file for Substrate TLS") + flag.BoolVar(&substrateInsecureTLS, "substrate-insecure-tls", false, "Skip Substrate TLS verification") + flag.BoolVar(&substratePlaintext, "substrate-plaintext", false, "Use insecure plaintext gRPC connection to Substrate") + flag.StringVar(&defaultTemplate, "template", "default-template", "Default Substrate ActorTemplate name") + flag.StringVar(&defaultTemplateAtespace, "template-atespace", "ax-system", "Default Substrate ActorTemplate atespace") flag.Parse() if envAddr := os.Getenv("ADDR"); envAddr != "" { @@ -54,7 +73,12 @@ func main() { logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo})) slog.SetDefault(logger) - slog.Info("starting ax-server", "listenAddr", listenAddr, "redisAddr", redisAddr) + slog.Info("starting ax-server", + "listenAddr", listenAddr, + "redisAddr", redisAddr, + "substrateEndpoint", substrateEndpoint, + "template", defaultTemplate, + ) rClient := goredis.NewClient(&goredis.Options{ Addr: redisAddr, @@ -63,7 +87,28 @@ func main() { defer rClient.Close() rStore := redis.NewStore(rClient, redis.Options{}) - srv := server.NewServer(rStore) + rLocker := lock.NewRedisLocker(rClient, lock.RedisLockerOptions{}) + + var reconciler server.Reconciler + subClient, err := substrate.NewClientWithOptions(substrate.ClientOptions{ + Target: substrateEndpoint, + Authority: substrateAuthority, + TokenFile: substrateTokenFile, + CAFile: substrateCAFile, + InsecureTLS: substrateInsecureTLS, + Plaintext: substratePlaintext, + }) + if err != nil { + slog.Warn("could not initialize substrate client; running without substrate reconciliation", "error", err) + } else { + defer subClient.Close() + reconciler = controller.NewTaskReconciler(subClient, defaultTemplate, defaultTemplateAtespace) + } + + srv := server.NewServer(rStore, server.Options{ + Locker: rLocker, + Reconciler: reconciler, + }) httpServer := &http.Server{ Addr: listenAddr, diff --git a/cmd/ax-task-runner/main.go b/cmd/ax-task-runner/main.go index e94564e9..cb278cb0 100644 --- a/cmd/ax-task-runner/main.go +++ b/cmd/ax-task-runner/main.go @@ -16,7 +16,7 @@ // loads the Task and Workspace specs and hands them to the runner package, // which does everything else. // -// The controller delivers the Task as YAML in AX_TASK_YAML and the bound +// AX delivers the Task as YAML in AX_TASK_YAML and the bound // Workspaces as a multi-document YAML stream in AX_WORKSPACES_YAML. For local // runs the specs can be read from files instead with --task-file and one or // more --workspace-file flags. diff --git a/cmd/ax/main.go b/cmd/ax/main.go index 0875177c..a7bfb3af 100644 --- a/cmd/ax/main.go +++ b/cmd/ax/main.go @@ -712,7 +712,7 @@ func runWatch(serverURL, atespace string, args []string) error { actor, workerIP, ) - if phase == "Running" || phase == "Completed" || phase == "Failed" { + if phase == "Completed" || phase == "Failed" { fmt.Printf("Task reached terminal phase %q.\n", phase) break } @@ -737,11 +737,9 @@ func runDelete(serverURL, atespace string, args []string) error { } defer conn.Close() - // Task deletion waits for the controller to tear down the actor, which can take - // a while; the other kinds are removed synchronously. - timeout := 15 * time.Second + timeout := 30 * time.Second if kind == v1alpha1.KindTask { - timeout = deleteTaskTimeout + timeout = 2 * time.Minute } ctx, cancel := context.WithTimeout(context.Background(), timeout) defer cancel() @@ -749,13 +747,7 @@ func runDelete(serverURL, atespace string, args []string) error { return deleteResource(ctx, client, kind, atespace, args[1]) } -const ( - deleteTaskTimeout = 5 * time.Minute - deletePollInterval = 500 * time.Millisecond -) - -// deleteResource requests deletion, then blocks until the resource is really gone -// and prints a kubectl-style confirmation. +// deleteResource requests deletion and prints a confirmation. func deleteResource(ctx context.Context, client v1alpha1.AXClient, kind, atespace, name string) error { lower := strings.ToLower(kind) @@ -774,49 +766,10 @@ func deleteResource(ctx context.Context, client v1alpha1.AXClient, kind, atespac return fmt.Errorf("deleting %s %s/%s: %w", lower, atespace, name, err) } - if err := waitForDeletion(ctx, client, kind, atespace, name); err != nil { - return err - } fmt.Printf("%s.ax.io/%s deleted\n", lower, name) return nil } -// waitForDeletion polls until the resource returns NotFound or ctx expires. -func waitForDeletion(ctx context.Context, client v1alpha1.AXClient, kind, atespace, name string) error { - lookup := func() error { - var err error - switch kind { - case v1alpha1.KindTask: - _, err = client.GetTask(ctx, &v1alpha1.GetTaskRequest{Atespace: atespace, Name: name}) - case v1alpha1.KindWorkspace: - _, err = client.GetWorkspace(ctx, &v1alpha1.GetWorkspaceRequest{Atespace: atespace, Name: name}) - case v1alpha1.KindModel: - _, err = client.GetModel(ctx, &v1alpha1.GetModelRequest{Atespace: atespace, Name: name}) - } - return err - } - - announced := false - for { - err := lookup() - if status.Code(err) == codes.NotFound { - return nil - } - if err != nil { - return fmt.Errorf("checking %s %s/%s after delete: %w", strings.ToLower(kind), atespace, name, err) - } - if !announced { - fmt.Fprintf(os.Stderr, "waiting for %s %s/%s to be deleted...\n", strings.ToLower(kind), atespace, name) - announced = true - } - select { - case <-ctx.Done(): - return fmt.Errorf("timed out waiting for %s %s/%s to be deleted; it is still being torn down (check `ax describe` and the controller logs)", strings.ToLower(kind), atespace, name) - case <-time.After(deletePollInterval): - } - } -} - // normalizeKind maps user-typed kinds ("task", "tasks", "Task") to the canonical // manifest kind, rejecting anything unknown. func normalizeKind(kind string) (string, error) { diff --git a/cmd/ax/main_test.go b/cmd/ax/main_test.go index d1c9d031..2040aa56 100644 --- a/cmd/ax/main_test.go +++ b/cmd/ax/main_test.go @@ -16,13 +16,16 @@ package main import ( "context" + "errors" "net" "os" "path/filepath" "strings" "testing" + "time" "github.com/google/ax/internal/server" + "github.com/google/ax/internal/store" "github.com/google/ax/internal/store/memory" "github.com/google/ax/pkg/apis/v1alpha1" "google.golang.org/grpc/codes" @@ -206,3 +209,72 @@ func TestRunGetResourceAliases(t *testing.T) { } } } + +type fakeMainReconciler struct { + deleted chan struct{} +} + +func (f *fakeMainReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, workspaces ...*v1alpha1.Workspace) (*v1alpha1.Task, error) { + return task, nil +} + +func (f *fakeMainReconciler) ReconcileDelete(ctx context.Context, atespace, taskName string) error { + time.Sleep(100 * time.Millisecond) + close(f.deleted) + return nil +} + +func TestRunDeleteTask(t *testing.T) { + s := memory.NewStore() + meta := &v1alpha1.ObjectMeta{Name: "task-to-delete", Atespace: "default"} + if err := s.SaveTask(context.Background(), &v1alpha1.Task{Metadata: meta}); err != nil { + t.Fatal(err) + } + + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + rec := &fakeMainReconciler{deleted: make(chan struct{})} + srv := server.NewServer(s, server.Options{Reconciler: rec}).GRPCServer() + t.Cleanup(srv.Stop) + go func() { _ = srv.Serve(listener) }() + + output, err := os.CreateTemp(t.TempDir(), "stdout") + if err != nil { + t.Fatal(err) + } + stdout := os.Stdout + os.Stdout = output + t.Cleanup(func() { + os.Stdout = stdout + _ = output.Close() + }) + + runErr := runDelete(listener.Addr().String(), "default", []string{"task", "task-to-delete"}) + if runErr != nil { + t.Fatalf("runDelete failed: %v", runErr) + } + + // Verify actor deletion finished before runDelete returned + select { + case <-rec.deleted: + default: + t.Fatal("expected actor deletion to complete before runDelete returned") + } + + data, err := os.ReadFile(output.Name()) + if err != nil { + t.Fatal(err) + } + wantOutput := "task.ax.io/task-to-delete deleted\n" + if string(data) != wantOutput { + t.Fatalf("expected output %q, got %q", wantOutput, string(data)) + } + + // Verify task is deleted from store + _, err = s.GetTask(context.Background(), "default", "task-to-delete") + if !errors.Is(err, store.ErrNotFound) { + t.Fatalf("expected task to be NotFound after deletion, got %v", err) + } +} diff --git a/demo.sh b/demo.sh index 5ab756a8..5f41ed5f 100755 --- a/demo.sh +++ b/demo.sh @@ -36,7 +36,7 @@ ax() { "${AX_BIN}" "$@"; } ATESPACE="${ATESPACE:-default}" TASK_NAME="demo-task" WORKSPACE_NAME="demo-workspace" -TASK_IMAGE="${AX_TASK_IMAGE:-${AX_IMAGE_REPO:-gcr.io/ax-substrate/ate-images}/ax-task-runner@sha256:464c5a53c68c67e929dbbb5450f1eb41f99b2742efcf42b721c09825a58397f1}" +TASK_IMAGE="${AX_TASK_IMAGE:-${AX_IMAGE_REPO:-gcr.io/dberkov-gke-dev3}/ax-task-runner@sha256:c81cfd6eaf0a96b7bad8d4e650ab25ca8d632e2e9f522d518f0a8c526fd93943}" # --------------------------------------------------------------------------- # Presentation helpers @@ -69,43 +69,6 @@ in_sandbox() { ok() { printf '%s✔ %s%s\n' "${GREEN}" "$*" "${RESET}"; } note() { printf '%s%s%s\n' "${YELLOW}" "$*" "${RESET}"; } err() { printf '%s✘ %s%s\n' "${RED}" "$*" "${RESET}"; } -task_field() { - ax describe task "${TASK_NAME}" -a "${ATESPACE}" 2>/dev/null | awk -v key="$1" '$1 == key {print $2}' -} - -# wait_for PHASE [READY] polls the task until it reaches PHASE (and Ready=READY -# when given), printing a dot per poll and the elapsed time when it gets there. -wait_for() { - local want_phase="$1" want_ready="${2:-}" timeout="${3:-180}" - local start phase ready elapsed - start=$(date +%s) - printf '%swaiting for Phase=%s' "${DIM}" "${want_phase}" - [[ -n "${want_ready}" ]] && printf ' Ready=%s' "${want_ready}" - printf '%s ' "${RESET}" - while :; do - phase=$(task_field "Phase:") - ready=$(task_field "Ready") - if [[ "${phase}" == "${want_phase}" && ( -z "${want_ready}" || "${ready}" == "${want_ready}" ) ]]; then - elapsed=$(( $(date +%s) - start )) - printf ' %s%ds%s\n' "${GREEN}" "${elapsed}" "${RESET}" - return 0 - fi - if [[ "${phase}" == "Failed" && "${want_phase}" != "Failed" ]]; then - printf '\n' - err "Task entered Failed phase! Last seen Ready=${ready:-?}" - ax describe task "${TASK_NAME}" -a "${ATESPACE}" || true - return 1 - fi - if (( $(date +%s) - start > timeout )); then - printf '\n' - note "Gave up after ${timeout}s. Last seen Phase=${phase:-?} Ready=${ready:-?}" - ax describe task "${TASK_NAME}" -a "${ATESPACE}" || true - return 1 - fi - printf '.' - sleep 0.2 - done -} # --------------------------------------------------------------------------- # Demo @@ -137,8 +100,16 @@ ok "Agent Substrate Control API found in namespace ${SUBSTRATE_NAMESPACE}" step "Clean up any previous demo run" CLEANED=0 -ax delete task "${TASK_NAME}" -a "${ATESPACE}" >/dev/null 2>&1 && { ok "removed old task"; CLEANED=1; } || true -ax delete workspace "${WORKSPACE_NAME}" -a "${ATESPACE}" >/dev/null 2>&1 && { ok "removed old workspace"; CLEANED=1; } || true +if ax get task "${TASK_NAME}" -a "${ATESPACE}" >/dev/null 2>&1; then + run ax delete task "${TASK_NAME}" -a "${ATESPACE}" + ok "removed old task" + CLEANED=1 +fi +if ax get workspace "${WORKSPACE_NAME}" -a "${ATESPACE}" >/dev/null 2>&1; then + run ax delete workspace "${WORKSPACE_NAME}" -a "${ATESPACE}" + ok "removed old workspace" + CLEANED=1 +fi (( CLEANED )) || echo "nothing to clean up" step "Declare a Workspace and a Task" @@ -169,10 +140,9 @@ YAML printf '%s' "${DIM}"; sed 's/^/ /' "${DEMO_YAML}"; printf '%s\n\n' "${RESET}" run ax apply -f "${DEMO_YAML}" -step "Resume the task and watch it come up" +step "Resume the task" note "New tasks are created Suspended by default. Resuming creates the worker on Agent Substrate and initializes /workspace." run ax resume task "${TASK_NAME}" -a "${ATESPACE}" -wait_for "Running" "True" ok "${TASK_NAME} is Running and Ready" echo run ax get tasks -a "${ATESPACE}" @@ -189,7 +159,6 @@ in_sandbox 'curl -s "$AX_METADATA_URL/metadata/v1alpha1/ax/task" | head -20' step "Suspend the task" run ax suspend task "${TASK_NAME}" -a "${ATESPACE}" -wait_for "Suspended" ok "${TASK_NAME} is Suspended. The workspace has been checkpointed and the sandbox is gone." echo run ax get tasks -a "${ATESPACE}" diff --git a/deploy/ax-controller.yaml b/deploy/ax-controller.yaml deleted file mode 100644 index 66c2c44b..00000000 --- a/deploy/ax-controller.yaml +++ /dev/null @@ -1,119 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -apiVersion: v1 -kind: Namespace -metadata: - name: ax-system ---- -apiVersion: v1 -kind: ServiceAccount -metadata: - name: ax-controller - namespace: ax-system ---- -apiVersion: rbac.authorization.k8s.io/v1 -kind: ClusterRole -metadata: - name: ax-controller -rules: - - apiGroups: [""] - resources: ["secrets"] - verbs: ["get", "list", "watch"] ---- -apiVersion: rbac.authorization.k8s.io/v1 -kind: ClusterRoleBinding -metadata: - name: ax-controller -subjects: - - kind: ServiceAccount - name: ax-controller - namespace: ax-system -roleRef: - apiGroup: rbac.authorization.k8s.io - kind: ClusterRole - name: ax-controller ---- -apiVersion: apps/v1 -kind: Deployment -metadata: - name: ax-controller - namespace: ax-system - labels: - app.kubernetes.io/name: ax-controller - app.kubernetes.io/part-of: ax -spec: - replicas: 1 - selector: - matchLabels: - app.kubernetes.io/name: ax-controller - template: - metadata: - labels: - app.kubernetes.io/name: ax-controller - spec: - serviceAccountName: ax-controller - containers: - - name: controller - image: ko://github.com/google/ax/cmd/ax-controller - imagePullPolicy: IfNotPresent - args: - - "--redis-addr=ax-redis.ax-system.svc.cluster.local:6379" - - "--substrate-endpoint=api.ate-system.svc.cluster.local:443" - - "--substrate-authority=api.ate-system.svc" - - "--substrate-token-file=/var/run/secrets/ateapi/token" - - "--substrate-ca-file=/run/servicedns-ca/trust-bundle.pem" - - "--template=default-template" - - "--template-atespace=ax-system" - env: - - name: ATENET_ROUTER_ADDR - value: "atenet-router.ate-system.svc.cluster.local:80" - - name: AX_SNAPSHOTS_BUCKET - value: "gs://snapshot-substrate-test-ax-substrate/ate-env/" - resources: - requests: - cpu: "100m" - memory: "128Mi" - limits: - cpu: "500m" - memory: "512Mi" - securityContext: - readOnlyRootFilesystem: true - allowPrivilegeEscalation: false - volumeMounts: - - mountPath: /var/run/secrets/ateapi - name: ate-token - readOnly: true - - mountPath: /run/servicedns-ca - name: servicedns-ca - readOnly: true - volumes: - - name: ate-token - projected: - defaultMode: 420 - sources: - - serviceAccountToken: - audience: api.ate-system.svc - expirationSeconds: 7200 - path: token - - name: servicedns-ca - projected: - defaultMode: 420 - sources: - - clusterTrustBundle: - labelSelector: - matchLabels: - podcert.ate.dev/canarying: live - path: trust-bundle.pem - signerName: servicedns.podcert.ate.dev/identity diff --git a/deploy/ax-server.yaml b/deploy/ax-server.yaml index 658a9dd9..1e577d3b 100644 --- a/deploy/ax-server.yaml +++ b/deploy/ax-server.yaml @@ -12,6 +12,39 @@ # See the License for the specific language governing permissions and # limitations under the License. +apiVersion: v1 +kind: Namespace +metadata: + name: ax-system +--- +apiVersion: v1 +kind: ServiceAccount +metadata: + name: ax-server + namespace: ax-system +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: ax-server +rules: + - apiGroups: [""] + resources: ["secrets"] + verbs: ["get", "list", "watch"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: ax-server +subjects: + - kind: ServiceAccount + name: ax-server + namespace: ax-system +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: ax-server +--- apiVersion: apps/v1 kind: Deployment metadata: @@ -30,6 +63,7 @@ spec: labels: app.kubernetes.io/name: ax-server spec: + serviceAccountName: ax-server containers: - name: ax-server image: ko://github.com/google/ax/cmd/ax-server @@ -37,6 +71,17 @@ spec: args: - "--addr=:8080" - "--redis-addr=ax-redis.ax-system.svc.cluster.local:6379" + - "--substrate-endpoint=api.ate-system.svc.cluster.local:443" + - "--substrate-authority=api.ate-system.svc" + - "--substrate-token-file=/var/run/secrets/ateapi/token" + - "--substrate-ca-file=/run/servicedns-ca/trust-bundle.pem" + - "--template=default-template" + - "--template-atespace=ax-system" + env: + - name: ATENET_ROUTER_ADDR + value: "atenet-router.ate-system.svc.cluster.local:80" + - name: AX_SNAPSHOTS_BUCKET + value: "gs://dberkov-gke-dev3/ate-env/" ports: - containerPort: 8080 name: http @@ -59,6 +104,32 @@ spec: port: 8080 initialDelaySeconds: 5 periodSeconds: 10 + volumeMounts: + - mountPath: /var/run/secrets/ateapi + name: ate-token + readOnly: true + - mountPath: /run/servicedns-ca + name: servicedns-ca + readOnly: true + volumes: + - name: ate-token + projected: + defaultMode: 420 + sources: + - serviceAccountToken: + audience: api.ate-system.svc + expirationSeconds: 7200 + path: token + - name: servicedns-ca + projected: + defaultMode: 420 + sources: + - clusterTrustBundle: + labelSelector: + matchLabels: + podcert.ate.dev/canarying: live + path: trust-bundle.pem + signerName: servicedns.podcert.ate.dev/identity --- apiVersion: v1 kind: Service @@ -75,3 +146,4 @@ spec: name: http selector: app.kubernetes.io/name: ax-server + diff --git a/docs/concepts.md b/docs/concepts.md index ca309ff2..6de3235a 100644 --- a/docs/concepts.md +++ b/docs/concepts.md @@ -17,7 +17,7 @@ The unit is deliberately small. An agent is not one process that runs to complet | `WorkspaceReady` | Every workspace has finished setting up. Stays True afterwards. | | `Ready` | The task is running and `WorkspaceReady` is True. This is the one to wait on. | -Two transitions are worth knowing. Suspending a task sets `Ready` to False with reason `TaskSuspended`; resuming sets it back. Deleting a task moves it to `Terminating` while the controller tears down the sandbox, then removes the record entirely. `ax delete` blocks until that has happened. +Two transitions are worth knowing. Suspending a task sets `Ready` to False with reason `TaskSuspended`; resuming sets it back. Deleting a task tears down the sandbox on Agent Substrate and removes the record. `ax delete` blocks until that has completed. ## Workspace diff --git a/docs/development.md b/docs/development.md index c765b854..0425b9aa 100644 --- a/docs/development.md +++ b/docs/development.md @@ -10,7 +10,7 @@ ## Build ```bash -make build # bin/ax, bin/ax-controller, bin/ax-server +make build # bin/ax, bin/ax-server make install # install the ax CLI into $(go env GOPATH)/bin make build-task-runner # cross-compile the runner for linux/amd64 and build its image make push-task-runner # ...and push it (set TASK_RUNNER_REPO) diff --git a/docs/networking.md b/docs/networking.md index 2e57e6e2..a2753fe1 100644 --- a/docs/networking.md +++ b/docs/networking.md @@ -2,11 +2,11 @@ Tasks do not get a Kubernetes Service or Ingress of their own. Every request to a task goes through Agent Substrate's **atenet router**, the `atenet-router` Service in the `ate-system` namespace. The router reads a single header, `ate-target-actor`, resolves the actor to the worker it is running on, resumes it first if it was suspended, and proxies the request there. `Host` and `:authority` are left alone for your application; the header alone selects the target. -The header value is `/`. The controller always names a task's actor after the task, so `default/task123` reaches the task `task123` in the `default` atespace. +The header value is `/`. AX always names a task's actor after the task, so `default/task123` reaches the task `task123` in the `default` atespace. ## From inside the cluster -Use the Service DNS name and add the header. This is exactly how the controller polls a task's readiness. +Use the Service DNS name and add the header. This is exactly how AX polls a task's readiness. ```bash curl -H "ate-target-actor: default/task123" \ diff --git a/docs/roadmap.md b/docs/roadmap.md index 062237e5..1d1e70c6 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -16,7 +16,7 @@ Stabilize the `ax.io/v1alpha1` declarative schemas and lifecycle contracts acros Evolve how AX maps tasks onto Agent Substrate actors to improve security boundaries, cluster utilization, and stateful workflows: -- **Migration to the New Actor**: Migrate `ax-controller` and the Substrate integration layer (`internal/substrate`) to the new Agent Substrate Actor API and lifecycle model. +- **Migration to the New Actor**: Migrate the Substrate integration layer (`internal/substrate`) to the new Agent Substrate Actor API and lifecycle model. - **Splitting Task Workspace Setup into a Separate Actor**: Decouple maiden workspace initialization (Git repository cloning, MCP and skill materialization, and goal-driven bootstrap) from the primary task runtime by executing setup in a dedicated setup actor before handing off the prepared workspace state to the task actor. - **Minimally Privileged Policies**: Apply strict least-privilege policies tailored independently to the workspace setup actor and the task execution actor—scoping repository/registry credentials and setup egress exclusively to the initialization phase while enforcing minimal runtime permissions, network egress, and capabilities on the task actor. - **Idleness Detection and Automatic Suspension for Density**: Continuously monitor actor activity (process execution, I/O, network traffic, and active gRPC/SSH sessions) to detect idle tasks and automatically trigger `SuspendActor` checkpointing, reclaiming worker CPU and memory to maximize cluster density. diff --git a/docs/runner.md b/docs/runner.md index 30bc1653..16a03571 100644 --- a/docs/runner.md +++ b/docs/runner.md @@ -1,6 +1,6 @@ # Runners -A runner is the program that AX starts as PID 1 inside every task container. It is the bridge between the control plane and whatever your agent actually is: the controller hands it the `Task` and `Workspace` specs, and the runner turns them into a prepared workspace, a running command, and a small HTTP surface that the rest of AX uses to observe the sandbox. +A runner is the program that AX starts as PID 1 inside every task container. It is the bridge between the control plane and whatever your agent actually is: the control plane hands it the `Task` and `Workspace` specs, and the runner turns them into a prepared workspace, a running command, and a small HTTP surface that the rest of AX uses to observe the sandbox. AX ships a default runner, `ax-task-runner`, baked into the default task image. You do not have to use it. Any binary that honors the contract below can be packaged into a container image and named in `spec.image`, and the control plane will treat it exactly like the default. @@ -8,9 +8,9 @@ This page describes what a runner must do. For what the default runner exposes t ## How a runner is launched -The controller does not run `spec.command` as the container entrypoint. It always starts the container with a fixed command and lets the runner take it from there: +The control plane does not run `spec.command` as the container entrypoint. It always starts the container with a fixed command and lets the runner take it from there: -| What the controller sets | Value | +| What AX sets | Value | |---|---| | Container image | `spec.image`, or the default `ax-task-runner` image when unset | | Container command | `/usr/local/bin/ax-task-runner`, always | @@ -27,12 +27,12 @@ The `/workspace` volume is what survives suspend and resume. Agent Substrate sna ## What a runner must do -**Serve HTTP on port 80.** Both Agent Substrate and the AX controller probe the container on this port. The paths that matter: +**Serve HTTP on port 80.** Both Agent Substrate and the AX server probe the container on this port. The paths that matter: | Path | Behavior | |---|---| | `/healthz` | Return `200` as soon as the runner is alive. | -| `/readyz` | Return `503` until the workspace is prepared, then `200`. The controller polls this to set the task's `WorkspaceReady` condition, and `ax watch` shows the transition. | +| `/readyz` | Return `503` until the workspace is prepared, then `200`. AX polls this to set the task's `WorkspaceReady` condition, and `ax watch` shows the transition. | | `/metadata/v1alpha1/ax/task` | Return the `Task` as `application/yaml`. Optional, but your command and `ax` tooling may expect it. | | `/metadata/v1alpha1/ax/workspaces` | Return every bound `Workspace` as a multi-document YAML stream. Optional, as above. | @@ -148,7 +148,7 @@ spec: debug: true ``` -The controller provisions a dedicated Agent Substrate actor template for each distinct image and environment, so different tasks can run different runners side by side in the same atespace. +AX provisions a dedicated Agent Substrate actor template for each distinct image and environment, so different tasks can run different runners side by side in the same atespace. ## Testing a runner locally diff --git a/examples/simple.yaml b/examples/simple.yaml index bb0efbf1..7178b9bd 100644 --- a/examples/simple.yaml +++ b/examples/simple.yaml @@ -14,7 +14,7 @@ # The smallest useful Task: no Workspace or Model. # -# The controller runs it on the default task-runner image with an empty +# The server runs it on the default task-runner image with an empty # /workspace directory. # # ax apply -f examples/simple.yaml diff --git a/internal/controller/reconciler.go b/internal/controller/reconciler.go index 2b4a56ea..55814d7d 100644 --- a/internal/controller/reconciler.go +++ b/internal/controller/reconciler.go @@ -348,7 +348,7 @@ func (r *TaskReconciler) setCondition(task *v1alpha1.Task, condType, status, rea } // lookupGeminiKey resolves the Gemini API key for the task container, preferring the -// Kubernetes secret in the task's atespace and falling back to the controller's own +// Kubernetes secret in the task's atespace and falling back to the server's own // environment. It returns "" when neither source has a value. func (r *TaskReconciler) lookupGeminiKey(ctx context.Context, atespace string) string { if r.SecretResolver != nil { @@ -360,7 +360,7 @@ func (r *TaskReconciler) lookupGeminiKey(ctx context.Context, atespace string) s } } if key := os.Getenv(geminiSecretKey); key != "" { - slog.Info("resolved GEMINI_API_KEY from controller environment for actor template") + slog.Info("resolved GEMINI_API_KEY from server environment for actor template") return key } return "" diff --git a/internal/controller/reconciler_test.go b/internal/controller/reconciler_test.go index 71431095..6524324b 100644 --- a/internal/controller/reconciler_test.go +++ b/internal/controller/reconciler_test.go @@ -16,8 +16,10 @@ package controller_test import ( "context" + "errors" "net" "net/http" + "strings" "testing" "time" @@ -41,6 +43,7 @@ type mockControlServer struct { deletedActors []string actorTemplates map[string]bool deletedTemplates []string + getActorFunc func(ctx context.Context, req *ateapipb.GetActorRequest) (*ateapipb.Actor, error) } // noSecrets is a SecretResolver for tests: it never finds a key and never touches a cluster. @@ -129,6 +132,27 @@ func (m *mockControlServer) DeleteActor(ctx context.Context, req *ateapipb.Delet return &ateapipb.Actor{Metadata: &ateapipb.ResourceMetadata{Name: name}}, nil } +func (m *mockControlServer) GetActor(ctx context.Context, req *ateapipb.GetActorRequest) (*ateapipb.Actor, error) { + if m.getActorFunc != nil { + return m.getActorFunc(ctx, req) + } + name := req.GetActor().GetName() + for _, del := range m.deletedActors { + if del == name { + return nil, status.Errorf(codes.NotFound, "actor %q not found", name) + } + } + for _, a := range m.createdActors { + if a == name { + return &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Name: name}, + Status: &ateapipb.ActorStatus{State: ateapipb.ActorState_ACTOR_STATE_RUNNING}, + }, nil + } + } + return nil, status.Errorf(codes.NotFound, "actor %q not found", name) +} + func (m *mockControlServer) ListActorTemplates(ctx context.Context, req *ateapipb.ListActorTemplatesRequest) (*ateapipb.ListActorTemplatesResponse, error) { resp := &ateapipb.ListActorTemplatesResponse{} for name := range m.actorTemplates { @@ -459,3 +483,92 @@ func TestReconcileDelete_RemovesActorAndTemplates(t *testing.T) { } } } + +func TestReconcileDelete_BlocksUntilActorDeleted(t *testing.T) { + ctx := context.Background() + + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + defer lis.Close() + + getCalls := 0 + mockSrv := &mockControlServer{ + getActorFunc: func(ctx context.Context, req *ateapipb.GetActorRequest) (*ateapipb.Actor, error) { + getCalls++ + if getCalls < 3 { + // Simulate actor in deleting state during the first two poll checks + return &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Name: req.GetActor().GetName()}, + Status: &ateapipb.ActorStatus{State: ateapipb.ActorState_ACTOR_STATE_DELETING}, + }, nil + } + // Once actor is completely torn down + return nil, status.Errorf(codes.NotFound, "actor %q not found", req.GetActor().GetName()) + }, + } + grpcServer := grpc.NewServer() + ateapipb.RegisterControlServer(grpcServer, mockSrv) + go grpcServer.Serve(lis) + defer grpcServer.Stop() + + client, err := substrate.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("failed to create substrate client: %v", err) + } + defer client.Close() + + reconciler := controller.NewTaskReconciler(client, "test-template", "ax-system") + reconciler.SecretResolver = noSecrets + + if err := reconciler.ReconcileDelete(ctx, "default", "slow-delete-task"); err != nil { + t.Fatalf("ReconcileDelete failed: %v", err) + } + + if getCalls < 3 { + t.Errorf("expected at least 3 GetActor calls before deletion completes, got %d", getCalls) + } +} + +func TestReconcileDelete_ActorDeletionTimeout(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 250*time.Millisecond) + defer cancel() + + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + defer lis.Close() + + mockSrv := &mockControlServer{ + getActorFunc: func(ctx context.Context, req *ateapipb.GetActorRequest) (*ateapipb.Actor, error) { + // Actor remains in deleting state indefinitely + return &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Name: req.GetActor().GetName()}, + Status: &ateapipb.ActorStatus{State: ateapipb.ActorState_ACTOR_STATE_DELETING}, + }, nil + }, + } + grpcServer := grpc.NewServer() + ateapipb.RegisterControlServer(grpcServer, mockSrv) + go grpcServer.Serve(lis) + defer grpcServer.Stop() + + client, err := substrate.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("failed to create substrate client: %v", err) + } + defer client.Close() + + reconciler := controller.NewTaskReconciler(client, "test-template", "ax-system") + reconciler.SecretResolver = noSecrets + + err = reconciler.ReconcileDelete(ctx, "default", "stuck-task") + if err == nil { + t.Fatal("expected error due to timeout, got nil") + } + if !errors.Is(err, context.DeadlineExceeded) && !strings.Contains(err.Error(), "context deadline exceeded") { + t.Errorf("expected deadline exceeded error, got %v", err) + } +} diff --git a/internal/controller/worker.go b/internal/controller/worker.go deleted file mode 100644 index 3f36f411..00000000 --- a/internal/controller/worker.go +++ /dev/null @@ -1,163 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package controller - -import ( - "context" - "errors" - "fmt" - "log/slog" - "os" - "time" - - "github.com/google/ax/internal/store" - "github.com/google/ax/pkg/apis/v1alpha1" -) - -const ( - defaultWorkerGroup = "ax-controllers" - // readRetryDelay is how long the worker waits after a transient error from the - // event queue before trying again. - readRetryDelay = time.Second -) - -// Worker consumes task events from the store's event queue and reconciles each -// task against Substrate. Run several with the same group name to share the load; -// every event is handled by exactly one of them. -type Worker struct { - store store.Store - reconciler *TaskReconciler - group string - consumer string -} - -// NewWorker creates a worker that joins group as consumer. An empty group uses the -// default controller group; an empty consumer derives a unique name from the host. -func NewWorker(s store.Store, reconciler *TaskReconciler, group, consumer string) *Worker { - if group == "" { - group = defaultWorkerGroup - } - if consumer == "" { - hostname, _ := os.Hostname() - consumer = fmt.Sprintf("%s-%d", hostname, time.Now().UnixNano()%10000) - } - return &Worker{ - store: s, - reconciler: reconciler, - group: group, - consumer: consumer, - } -} - -// Run subscribes to task events and processes them until ctx is done. It returns -// ctx.Err() on shutdown; every event is acknowledged after processing, even when -// reconciliation fails, so a bad task cannot wedge the queue. -func (w *Worker) Run(ctx context.Context) error { - slog.Info("starting AX task worker", "group", w.group, "consumer", w.consumer) - - sub, err := w.store.Subscribe(ctx, w.group, w.consumer) - if err != nil { - return fmt.Errorf("subscribing to task events: %w", err) - } - defer sub.Close() - - for { - ev, err := sub.Next(ctx) - if err != nil { - if ctx.Err() != nil { - slog.Info("stopping AX task worker") - return ctx.Err() - } - slog.Error("error reading task events", "error", err) - select { - case <-ctx.Done(): - return ctx.Err() - case <-time.After(readRetryDelay): - } - continue - } - - if err := w.processEvent(ctx, ev); err != nil { - slog.Error("error processing task event", - "id", ev.ID, - "atespace", ev.Atespace, - "name", ev.Name, - "action", ev.Action, - "error", err, - ) - } - if err := sub.Ack(ctx, ev); err != nil { - slog.Warn("failed to acknowledge task event", "id", ev.ID, "error", err) - } - } -} - -func (w *Worker) processEvent(ctx context.Context, ev store.TaskEvent) error { - if ev.Action == "delete" { - slog.Info("handling task deletion event", "atespace", ev.Atespace, "name", ev.Name) - if err := w.reconciler.ReconcileDelete(ctx, ev.Atespace, ev.Name); err != nil { - // Leave the record in Terminating so the failure is visible; re-running - // `ax delete` republishes the event and retries the cleanup. - return fmt.Errorf("cleaning up task %s/%s: %w", ev.Atespace, ev.Name, err) - } - if err := w.store.DeleteTask(ctx, ev.Atespace, ev.Name); err != nil { - return fmt.Errorf("removing task record %s/%s: %w", ev.Atespace, ev.Name, err) - } - return nil - } - - task, err := w.store.GetTask(ctx, ev.Atespace, ev.Name) - if err != nil { - if errors.Is(err, store.ErrNotFound) { - slog.Info("task not found, skipping reconcile", "atespace", ev.Atespace, "name", ev.Name) - return nil - } - return fmt.Errorf("fetching task %s/%s: %w", ev.Atespace, ev.Name, err) - } - // A pending delete event owns this task now; reconciling would resume an actor - // that is about to be torn down and overwrite the Terminating phase. - if task.GetStatus().GetPhase() == v1alpha1.PhaseTerminating { - slog.Info("task is terminating, skipping reconcile", "atespace", ev.Atespace, "name", ev.Name) - return nil - } - - // Resolve every bound workspace. A missing one is skipped so the task still - // runs; the runner creates an empty directory at its path. - var workspaces []*v1alpha1.Workspace - for _, ref := range task.Spec.WorkspaceRefs() { - if ref.Name == "" { - continue - } - wsp, err := w.store.GetWorkspace(ctx, task.Metadata.Atespace, ref.Name) - if err == nil { - workspaces = append(workspaces, wsp) - } else if !errors.Is(err, store.ErrNotFound) { - slog.Warn("error fetching workspace", "name", ref.Name, "error", err) - } - } - - reconciled, err := w.reconciler.Reconcile(ctx, task, workspaces...) - if err != nil { - task.Status.Phase = "Failed" - _ = w.store.UpdateTaskStatus(ctx, task.Metadata.Atespace, task.Metadata.Name, task.Status) - return fmt.Errorf("reconciling task %s/%s: %w", task.Metadata.Atespace, task.Metadata.Name, err) - } - - if err := w.store.UpdateTaskStatus(ctx, task.Metadata.Atespace, task.Metadata.Name, reconciled.Status); err != nil { - return fmt.Errorf("updating task status %s/%s: %w", task.Metadata.Atespace, task.Metadata.Name, err) - } - - return nil -} diff --git a/internal/controller/worker_test.go b/internal/controller/worker_test.go deleted file mode 100644 index f152beff..00000000 --- a/internal/controller/worker_test.go +++ /dev/null @@ -1,233 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package controller_test - -import ( - "context" - "net" - "testing" - "time" - - "github.com/agent-substrate/substrate/pkg/proto/ateapipb" - "github.com/google/ax/internal/controller" - "github.com/google/ax/internal/store/memory" - "github.com/google/ax/internal/substrate" - "github.com/google/ax/pkg/apis/v1alpha1" - "google.golang.org/grpc" - "google.golang.org/grpc/credentials/insecure" -) - -func TestWorkerReconciliation(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - // 1. Start in-process mock Substrate server - lis, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - t.Fatalf("failed to listen: %v", err) - } - defer lis.Close() - - mockSrv := &mockControlServer{} - grpcServer := grpc.NewServer() - ateapipb.RegisterControlServer(grpcServer, mockSrv) - go grpcServer.Serve(lis) - defer grpcServer.Stop() - - // 2. Substrate client - subClient, err := substrate.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) - if err != nil { - t.Fatalf("failed to create substrate client: %v", err) - } - defer subClient.Close() - - reconciler := controller.NewTaskReconciler(subClient, "default-template", "ax-system") - reconciler.SecretResolver = noSecrets - reconciler.WorkspaceReadyTimeout = 200 * time.Millisecond - - // 3. In-memory store - memStore := memory.NewStore() - - // Save task - task := &v1alpha1.Task{ - Metadata: &v1alpha1.ObjectMeta{Name: "worker-task", Atespace: "default"}, - Spec: &v1alpha1.TaskSpec{ - Image: "ghrc.io/test/img", - }, - } - if err := memStore.SaveTask(ctx, task); err != nil { - t.Fatalf("failed to save task: %v", err) - } - - // 4. Start the worker in the background - worker := controller.NewWorker(memStore, reconciler, "test-group", "worker-1") - go func() { - _ = worker.Run(ctx) - }() - - // 5. Poll store until task reaches "Running" phase - deadline := time.Now().Add(3 * time.Second) - var finalTask *v1alpha1.Task - for time.Now().Before(deadline) { - tItem, err := memStore.GetTask(ctx, "default", "worker-task") - if err == nil && tItem.Status.Phase == "Running" { - finalTask = tItem - break - } - time.Sleep(50 * time.Millisecond) - } - - if finalTask == nil { - t.Fatalf("task did not transition to Running phase in time") - } - - if finalTask.Status.Actor != "worker-task" { - t.Errorf("expected actor 'worker-task', got %q", finalTask.Status.Actor) - } - if finalTask.Status.WorkerIp != "10.244.1.42" { - t.Errorf("expected worker IP '10.244.1.42', got %q", finalTask.Status.WorkerIp) - } -} - -func TestWorkerDeletion(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - lis, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - t.Fatalf("failed to listen: %v", err) - } - defer lis.Close() - - mockSrv := &mockControlServer{actorTemplates: map[string]bool{"doomed-tmpl-0a1b2c3d": true}} - grpcServer := grpc.NewServer() - ateapipb.RegisterControlServer(grpcServer, mockSrv) - go grpcServer.Serve(lis) - defer grpcServer.Stop() - - subClient, err := substrate.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) - if err != nil { - t.Fatalf("failed to create substrate client: %v", err) - } - defer subClient.Close() - - reconciler := controller.NewTaskReconciler(subClient, "default-template", "ax-system") - reconciler.SecretResolver = noSecrets - reconciler.WorkspaceReadyTimeout = 200 * time.Millisecond - - memStore := memory.NewStore() - task := &v1alpha1.Task{ - Metadata: &v1alpha1.ObjectMeta{Name: "doomed", Atespace: "default"}, - Spec: &v1alpha1.TaskSpec{Image: "ghcr.io/test/img"}, - Status: &v1alpha1.TaskStatus{Phase: "Running", Actor: "doomed"}, - } - if err := memStore.SaveTask(ctx, task); err != nil { - t.Fatalf("failed to save task: %v", err) - } - // Drain the reconcile event SaveTask published so only the delete is processed. - drain, _ := memStore.Subscribe(ctx, "drain", "drain") - drainCtx, drainCancel := context.WithTimeout(ctx, time.Second) - _, _ = drain.Next(drainCtx) - drainCancel() - - if err := memStore.MarkTaskDeleting(ctx, "default", "doomed"); err != nil { - t.Fatalf("MarkTaskDeleting failed: %v", err) - } - marked, err := memStore.GetTask(ctx, "default", "doomed") - if err != nil { - t.Fatalf("GetTask after mark failed: %v", err) - } - if marked.Status.Phase != v1alpha1.PhaseTerminating { - t.Fatalf("expected phase Terminating, got %q", marked.Status.Phase) - } - - worker := controller.NewWorker(memStore, reconciler, "test-group", "worker-1") - go func() { _ = worker.Run(ctx) }() - - deadline := time.Now().Add(3 * time.Second) - for time.Now().Before(deadline) { - if _, err := memStore.GetTask(ctx, "default", "doomed"); err != nil { - break - } - time.Sleep(50 * time.Millisecond) - } - if _, err := memStore.GetTask(ctx, "default", "doomed"); err == nil { - t.Fatalf("expected task record to be removed after cleanup") - } - if len(mockSrv.deletedActors) != 1 || mockSrv.deletedActors[0] != "doomed" { - t.Errorf("expected actor 'doomed' deleted, got %v", mockSrv.deletedActors) - } - if len(mockSrv.deletedTemplates) != 1 || mockSrv.deletedTemplates[0] != "doomed-tmpl-0a1b2c3d" { - t.Errorf("expected template deleted, got %v", mockSrv.deletedTemplates) - } -} - -func TestWorkerSkipsReconcileOfTerminatingTask(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - lis, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - t.Fatalf("failed to listen: %v", err) - } - defer lis.Close() - - mockSrv := &mockControlServer{} - grpcServer := grpc.NewServer() - ateapipb.RegisterControlServer(grpcServer, mockSrv) - go grpcServer.Serve(lis) - defer grpcServer.Stop() - - subClient, err := substrate.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) - if err != nil { - t.Fatalf("failed to create substrate client: %v", err) - } - defer subClient.Close() - - reconciler := controller.NewTaskReconciler(subClient, "default-template", "ax-system") - reconciler.SecretResolver = noSecrets - reconciler.WorkspaceReadyTimeout = 200 * time.Millisecond - - // Queue a reconcile and then a delete before the worker starts, as when a - // task is deleted while the controller is still busy with other events. - memStore := memory.NewStore() - task := &v1alpha1.Task{ - Metadata: &v1alpha1.ObjectMeta{Name: "doomed", Atespace: "default"}, - Spec: &v1alpha1.TaskSpec{Image: "ghcr.io/test/img"}, - } - if err := memStore.SaveTask(ctx, task); err != nil { - t.Fatalf("failed to save task: %v", err) - } - if err := memStore.MarkTaskDeleting(ctx, "default", "doomed"); err != nil { - t.Fatalf("MarkTaskDeleting failed: %v", err) - } - - worker := controller.NewWorker(memStore, reconciler, "test-group", "worker-1") - go func() { _ = worker.Run(ctx) }() - - deadline := time.Now().Add(3 * time.Second) - for time.Now().Before(deadline) { - if _, err := memStore.GetTask(ctx, "default", "doomed"); err != nil { - break - } - time.Sleep(50 * time.Millisecond) - } - if _, err := memStore.GetTask(ctx, "default", "doomed"); err == nil { - t.Fatalf("expected task record to be removed after cleanup") - } - if len(mockSrv.createdActors) != 0 || len(mockSrv.resumedActors) != 0 { - t.Errorf("terminating task was reconciled: created %v, resumed %v", mockSrv.createdActors, mockSrv.resumedActors) - } -} diff --git a/internal/lock/lock.go b/internal/lock/lock.go new file mode 100644 index 00000000..0b645498 --- /dev/null +++ b/internal/lock/lock.go @@ -0,0 +1,237 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package lock + +import ( + "context" + "crypto/rand" + "encoding/hex" + "errors" + "fmt" + "sync" + "time" + + "github.com/redis/go-redis/v9" +) + +var ( + ErrLockFailed = errors.New("failed to acquire lock") +) + +// Locker provides exclusive locking per resource kind, atespace, and name. +type Locker interface { + // Lock acquires an exclusive lock for the specified resource. + // It blocks until the lock is acquired or ctx is cancelled. + // The returned function releases the lock. + Lock(ctx context.Context, kind, atespace, name string) (unlock func(), err error) +} + +// Key formats the standard lock key. +func Key(kind, atespace, name string) string { + if atespace == "" { + atespace = "default" + } + return fmt.Sprintf("lock:%s:%s:%s", kind, atespace, name) +} + +// MemoryLocker is an in-memory keyed mutex implementation suitable for tests and local mode. +type MemoryLocker struct { + mu sync.Mutex + locks map[string]*entry +} + +type entry struct { + mu sync.Mutex + ref int +} + +// NewMemoryLocker creates an in-memory Locker. +func NewMemoryLocker() *MemoryLocker { + return &MemoryLocker{ + locks: make(map[string]*entry), + } +} + +func (m *MemoryLocker) Lock(ctx context.Context, kind, atespace, name string) (func(), error) { + key := Key(kind, atespace, name) + + m.mu.Lock() + e, ok := m.locks[key] + if !ok { + e = &entry{} + m.locks[key] = e + } + e.ref++ + m.mu.Unlock() + + // Acquire lock with ctx awareness + locked := make(chan struct{}) + go func() { + e.mu.Lock() + close(locked) + }() + + select { + case <-ctx.Done(): + // Clean up ref if timed out waiting + go func() { + <-locked + e.mu.Unlock() + m.mu.Lock() + e.ref-- + if e.ref == 0 { + delete(m.locks, key) + } + m.mu.Unlock() + }() + return nil, ctx.Err() + case <-locked: + } + + var once sync.Once + unlock := func() { + once.Do(func() { + e.mu.Unlock() + m.mu.Lock() + e.ref-- + if e.ref == 0 { + delete(m.locks, key) + } + m.mu.Unlock() + }) + } + return unlock, nil +} + +// RedisLocker implements distributed locking in Redis using SET NX PX with token validation +// and Redis Pub/Sub notification for instant contended lock hand-off. +type RedisLocker struct { + client *redis.Client + ttl time.Duration + fallbackInterval time.Duration +} + +type RedisLockerOptions struct { + TTL time.Duration + FallbackInterval time.Duration +} + +var releaseAndNotifyScript = redis.NewScript(` +if redis.call("get", KEYS[1]) == ARGV[1] then + local res = redis.call("del", KEYS[1]) + redis.call("publish", KEYS[2], "1") + return res +else + return 0 +end +`) + +// ChannelKey returns the Pub/Sub channel used to notify waiters when a lock is released. +func ChannelKey(kind, atespace, name string) string { + if atespace == "" { + atespace = "default" + } + return fmt.Sprintf("lock:chan:%s:%s:%s", kind, atespace, name) +} + +// NewRedisLocker creates a distributed Redis locker with Pub/Sub notification. +func NewRedisLocker(client *redis.Client, opts RedisLockerOptions) *RedisLocker { + ttl := opts.TTL + if ttl <= 0 { + ttl = 30 * time.Second + } + fallback := opts.FallbackInterval + if fallback <= 0 { + fallback = 1 * time.Second + } + return &RedisLocker{ + client: client, + ttl: ttl, + fallbackInterval: fallback, + } +} + +func (r *RedisLocker) Lock(ctx context.Context, kind, atespace, name string) (func(), error) { + key := Key(kind, atespace, name) + chanKey := ChannelKey(kind, atespace, name) + + tokenBytes := make([]byte, 16) + if _, err := rand.Read(tokenBytes); err != nil { + return nil, fmt.Errorf("generating lock token: %w", err) + } + token := hex.EncodeToString(tokenBytes) + + makeUnlock := func() func() { + var once sync.Once + return func() { + once.Do(func() { + releaseCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = releaseAndNotifyScript.Run(releaseCtx, r.client, []string{key, chanKey}, token).Err() + }) + } + } + + // 1. Fast Path: Try acquiring immediately without subscribing to Pub/Sub + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + + ok, err := r.client.SetNX(ctx, key, token, r.ttl).Result() + if err != nil && !errors.Is(err, context.Canceled) { + return nil, fmt.Errorf("acquiring lock %s: %w", key, err) + } + if ok { + return makeUnlock(), nil + } + + // 2. Slow Path: Subscribe to Pub/Sub channel for instant wakeup when released + pubsub := r.client.Subscribe(ctx, chanKey) + defer pubsub.Close() + + // Try acquiring again immediately in case it was released before subscribe completed + ok, err = r.client.SetNX(ctx, key, token, r.ttl).Result() + if err != nil && !errors.Is(err, context.Canceled) { + return nil, fmt.Errorf("acquiring lock %s: %w", key, err) + } + if ok { + return makeUnlock(), nil + } + + msgCh := pubsub.Channel() + ticker := time.NewTicker(r.fallbackInterval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-msgCh: + // Woken up instantly by lock release notification + case <-ticker.C: + // Fallback ticker in case of missed notification or TTL expiration without publish + } + + ok, err := r.client.SetNX(ctx, key, token, r.ttl).Result() + if err != nil && !errors.Is(err, context.Canceled) { + return nil, fmt.Errorf("acquiring lock %s: %w", key, err) + } + if ok { + return makeUnlock(), nil + } + } +} diff --git a/internal/lock/lock_test.go b/internal/lock/lock_test.go new file mode 100644 index 00000000..4128ce74 --- /dev/null +++ b/internal/lock/lock_test.go @@ -0,0 +1,96 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package lock_test + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/google/ax/internal/lock" +) + +func TestMemoryLocker_SerializesSameResource(t *testing.T) { + locker := lock.NewMemoryLocker() + ctx := context.Background() + + var counter int64 + var wg sync.WaitGroup + + for i := 0; i < 20; i++ { + wg.Add(1) + go func() { + defer wg.Done() + unlock, err := locker.Lock(ctx, "task", "default", "my-task") + if err != nil { + t.Errorf("unexpected lock error: %v", err) + return + } + defer unlock() + + current := atomic.AddInt64(&counter, 1) + time.Sleep(10 * time.Millisecond) + if atomic.LoadInt64(&counter) != current { + t.Errorf("race condition detected: counter changed while lock held") + } + }() + } + + wg.Wait() +} + +func TestMemoryLocker_ConcurrentDifferentResources(t *testing.T) { + locker := lock.NewMemoryLocker() + ctx := context.Background() + + unlock1, err := locker.Lock(ctx, "task", "default", "task-1") + if err != nil { + t.Fatalf("lock task-1: %v", err) + } + defer unlock1() + + // Different resource should acquire immediately + start := time.Now() + unlock2, err := locker.Lock(ctx, "task", "default", "task-2") + if err != nil { + t.Fatalf("lock task-2: %v", err) + } + defer unlock2() + + if time.Since(start) > 50*time.Millisecond { + t.Errorf("different resource lock was blocked") + } +} + +func TestMemoryLocker_Timeout(t *testing.T) { + locker := lock.NewMemoryLocker() + ctx := context.Background() + + unlock1, err := locker.Lock(ctx, "task", "default", "task-1") + if err != nil { + t.Fatalf("lock task-1: %v", err) + } + defer unlock1() + + timeoutCtx, cancel := context.WithTimeout(ctx, 20*time.Millisecond) + defer cancel() + + _, err = locker.Lock(timeoutCtx, "task", "default", "task-1") + if err == nil { + t.Fatalf("expected timeout error, got nil") + } +} diff --git a/internal/server/server.go b/internal/server/server.go index c9350509..2b54a30b 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -17,9 +17,11 @@ package server import ( "context" "errors" + "log/slog" "net/http" "strings" + "github.com/google/ax/internal/lock" "github.com/google/ax/internal/store" "github.com/google/ax/pkg/apis/v1alpha1" "google.golang.org/grpc" @@ -28,17 +30,42 @@ import ( "google.golang.org/protobuf/types/known/timestamppb" ) +// Reconciler coordinates sandbox/actor lifecycles on Agent Substrate directly. +type Reconciler interface { + Reconcile(ctx context.Context, task *v1alpha1.Task, workspaces ...*v1alpha1.Workspace) (*v1alpha1.Task, error) + ReconcileDelete(ctx context.Context, atespace, taskName string) error +} + +// Options configures the AX API Server. +type Options struct { + Locker lock.Locker + Reconciler Reconciler +} + // Server provides the gRPC API for AX. type Server struct { v1alpha1.UnimplementedAXServer store store.Store + locker lock.Locker + reconciler Reconciler grpcServer *grpc.Server } // NewServer creates a new AX API server. -func NewServer(s store.Store) *Server { +func NewServer(s store.Store, opts ...Options) *Server { + var opt Options + if len(opts) > 0 { + opt = opts[0] + } + locker := opt.Locker + if locker == nil { + locker = lock.NewMemoryLocker() + } + srv := &Server{ store: s, + locker: locker, + reconciler: opt.Reconciler, grpcServer: grpc.NewServer(), } v1alpha1.RegisterAXServer(srv.grpcServer, srv) @@ -124,9 +151,18 @@ func (s *Server) CreateTask(ctx context.Context, req *v1alpha1.CreateTaskRequest atespace = "default" task.Metadata.Atespace = atespace } - _, err := s.store.GetTask(ctx, atespace, task.Metadata.GetName()) + taskName := task.Metadata.GetName() + + // Acquire exclusive lock for this task + unlock, err := s.locker.Lock(ctx, "task", atespace, taskName) + if err != nil { + return nil, status.Errorf(codes.Aborted, "locking task %s/%s: %v", atespace, taskName, err) + } + defer unlock() + + _, err = s.store.GetTask(ctx, atespace, taskName) if err == nil { - return nil, status.Errorf(codes.FailedPrecondition, "task %s/%s already exists and is immutable", atespace, task.Metadata.GetName()) + return nil, status.Errorf(codes.FailedPrecondition, "task %s/%s already exists and is immutable", atespace, taskName) } if !errors.Is(err, store.ErrNotFound) { return nil, status.Errorf(codes.Internal, "checking existing task: %v", err) @@ -142,6 +178,23 @@ func (s *Server) CreateTask(ctx context.Context, req *v1alpha1.CreateTaskRequest if err := s.store.SaveTask(ctx, task); err != nil { return nil, status.Errorf(codes.Internal, "saving task: %v", err) } + + // Directly reconcile with Substrate + if s.reconciler != nil { + workspaces := s.fetchWorkspaces(ctx, atespace, task) + reconciled, err := s.reconciler.Reconcile(ctx, task, workspaces...) + if err != nil { + slog.Error("direct reconcile error on create task", "task", taskName, "error", err) + task.Status.Phase = "Failed" + _ = s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status) + return nil, status.Errorf(codes.Internal, "provisioning task on substrate: %v", err) + } + task.Status = reconciled.Status + if err := s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status); err != nil { + return nil, status.Errorf(codes.Internal, "updating task status: %v", err) + } + } + return task, nil } @@ -153,13 +206,50 @@ func (s *Server) DeleteTask(ctx context.Context, req *v1alpha1.DeleteTaskRequest if atespace == "" { atespace = "default" } - // Deletion is two-phase: mark the task Terminating and let the controller tear - // down the actor before the record is removed. Clients poll GetTask for NotFound. - if err := s.store.MarkTaskDeleting(ctx, atespace, req.Name); err != nil { + taskName := req.Name + + // Acquire exclusive lock for this task + unlock, err := s.locker.Lock(ctx, "task", atespace, taskName) + if err != nil { + return nil, status.Errorf(codes.Aborted, "locking task %s/%s: %v", atespace, taskName, err) + } + defer unlock() + + task, err := s.store.GetTask(ctx, atespace, taskName) + if err != nil { if errors.Is(err, store.ErrNotFound) { - return nil, status.Errorf(codes.NotFound, "task %q not found in atespace %q", req.Name, atespace) + return nil, status.Errorf(codes.NotFound, "task %q not found in atespace %q", taskName, atespace) + } + return nil, status.Errorf(codes.Internal, "getting task: %v", err) + } + + // Mark the task as terminating while actor cleanup is in progress. + if task.Status == nil { + task.Status = &v1alpha1.TaskStatus{} + } + task.Status.Phase = v1alpha1.PhaseTerminating + if err := s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status); err != nil { + return nil, status.Errorf(codes.Internal, "updating task status: %v", err) + } + + // Directly clean up Substrate actor and templates + if s.reconciler != nil { + if err := s.reconciler.ReconcileDelete(ctx, atespace, taskName); err != nil { + if errors.Is(err, context.Canceled) { + return nil, status.Error(codes.Canceled, err.Error()) + } + if errors.Is(err, context.DeadlineExceeded) { + return nil, status.Error(codes.DeadlineExceeded, err.Error()) + } + if st, ok := status.FromError(err); ok && st.Code() != codes.Unknown { + return nil, err + } + return nil, status.Errorf(codes.Internal, "cleaning up substrate resources for task %s/%s: %v", atespace, taskName, err) } - return nil, status.Errorf(codes.Internal, "deleting task: %v", err) + } + + if err := s.store.DeleteTask(ctx, atespace, taskName); err != nil { + return nil, status.Errorf(codes.Internal, "deleting task record: %v", err) } return &v1alpha1.DeleteTaskResponse{}, nil } @@ -172,10 +262,19 @@ func (s *Server) SuspendTask(ctx context.Context, req *v1alpha1.SuspendTaskReque if atespace == "" { atespace = "default" } - task, err := s.store.GetTask(ctx, atespace, req.Name) + taskName := req.Name + + // Acquire exclusive lock for this task + unlock, err := s.locker.Lock(ctx, "task", atespace, taskName) + if err != nil { + return nil, status.Errorf(codes.Aborted, "locking task %s/%s: %v", atespace, taskName, err) + } + defer unlock() + + task, err := s.store.GetTask(ctx, atespace, taskName) if err != nil { if errors.Is(err, store.ErrNotFound) { - return nil, status.Errorf(codes.NotFound, "task %q not found in atespace %q", req.Name, atespace) + return nil, status.Errorf(codes.NotFound, "task %q not found in atespace %q", taskName, atespace) } return nil, status.Errorf(codes.Internal, "getting task: %v", err) } @@ -183,9 +282,23 @@ func (s *Server) SuspendTask(ctx context.Context, req *v1alpha1.SuspendTaskReque task.Status = &v1alpha1.TaskStatus{} } task.Status.Phase = "Suspended" - if err := s.store.SaveTask(ctx, task); err != nil { - return nil, status.Errorf(codes.Internal, "suspending task: %v", err) + + if s.reconciler != nil { + workspaces := s.fetchWorkspaces(ctx, atespace, task) + reconciled, err := s.reconciler.Reconcile(ctx, task, workspaces...) + if err != nil { + return nil, status.Errorf(codes.Internal, "suspending task on substrate: %v", err) + } + task.Status = reconciled.Status + if err := s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status); err != nil { + return nil, status.Errorf(codes.Internal, "updating task status: %v", err) + } + } else { + if err := s.store.SaveTask(ctx, task); err != nil { + return nil, status.Errorf(codes.Internal, "suspending task: %v", err) + } } + return task, nil } @@ -197,10 +310,19 @@ func (s *Server) ResumeTask(ctx context.Context, req *v1alpha1.ResumeTaskRequest if atespace == "" { atespace = "default" } - task, err := s.store.GetTask(ctx, atespace, req.Name) + taskName := req.Name + + // Acquire exclusive lock for this task + unlock, err := s.locker.Lock(ctx, "task", atespace, taskName) + if err != nil { + return nil, status.Errorf(codes.Aborted, "locking task %s/%s: %v", atespace, taskName, err) + } + defer unlock() + + task, err := s.store.GetTask(ctx, atespace, taskName) if err != nil { if errors.Is(err, store.ErrNotFound) { - return nil, status.Errorf(codes.NotFound, "task %q not found in atespace %q", req.Name, atespace) + return nil, status.Errorf(codes.NotFound, "task %q not found in atespace %q", taskName, atespace) } return nil, status.Errorf(codes.Internal, "getting task: %v", err) } @@ -208,12 +330,44 @@ func (s *Server) ResumeTask(ctx context.Context, req *v1alpha1.ResumeTaskRequest task.Status = &v1alpha1.TaskStatus{} } task.Status.Phase = "Running" - if err := s.store.SaveTask(ctx, task); err != nil { - return nil, status.Errorf(codes.Internal, "resuming task: %v", err) + + if s.reconciler != nil { + workspaces := s.fetchWorkspaces(ctx, atespace, task) + reconciled, err := s.reconciler.Reconcile(ctx, task, workspaces...) + if err != nil { + task.Status.Phase = "Failed" + _ = s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status) + return nil, status.Errorf(codes.Internal, "resuming task on substrate: %v", err) + } + task.Status = reconciled.Status + if err := s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status); err != nil { + return nil, status.Errorf(codes.Internal, "updating task status: %v", err) + } + } else { + if err := s.store.SaveTask(ctx, task); err != nil { + return nil, status.Errorf(codes.Internal, "resuming task: %v", err) + } } + return task, nil } +func (s *Server) fetchWorkspaces(ctx context.Context, atespace string, task *v1alpha1.Task) []*v1alpha1.Workspace { + var workspaces []*v1alpha1.Workspace + if task.Spec == nil { + return workspaces + } + for _, ref := range task.Spec.WorkspaceRefs() { + if ref.Name == "" { + continue + } + if wsp, err := s.store.GetWorkspace(ctx, atespace, ref.Name); err == nil { + workspaces = append(workspaces, wsp) + } + } + return workspaces +} + func (s *Server) WatchTask(req *v1alpha1.WatchTaskRequest, stream grpc.ServerStreamingServer[v1alpha1.WatchTaskResponse]) error { if req == nil { return status.Error(codes.InvalidArgument, "missing request") @@ -246,14 +400,13 @@ func (s *Server) WatchTask(req *v1alpha1.WatchTaskRequest, stream grpc.ServerStr if err := stream.Send(&v1alpha1.WatchTaskResponse{Task: task, Action: "MODIFIED"}); err != nil { return err } - if task.Status != nil && (task.Status.Phase == "Running" || task.Status.Phase == "Failed" || task.Status.Phase == "Completed") { + if task.Status != nil && (task.Status.Phase == "Failed" || task.Status.Phase == "Completed") { return nil } } } } - // --- Workspaces --- func (s *Server) GetWorkspace(ctx context.Context, req *v1alpha1.GetWorkspaceRequest) (*v1alpha1.Workspace, error) { @@ -300,6 +453,17 @@ func (s *Server) UpdateWorkspace(ctx context.Context, req *v1alpha1.UpdateWorksp } return existing.GetMetadata() }) + + atespace := req.Workspace.Metadata.Atespace + wsName := req.Workspace.Metadata.Name + + // Acquire exclusive lock for this workspace + unlock, err := s.locker.Lock(ctx, "workspace", atespace, wsName) + if err != nil { + return nil, status.Errorf(codes.Aborted, "locking workspace %s/%s: %v", atespace, wsName, err) + } + defer unlock() + if err := s.store.SaveWorkspace(ctx, req.Workspace); err != nil { return nil, status.Errorf(codes.Internal, "saving workspace: %v", err) } @@ -314,7 +478,16 @@ func (s *Server) DeleteWorkspace(ctx context.Context, req *v1alpha1.DeleteWorksp if atespace == "" { atespace = "default" } - if err := s.store.DeleteWorkspace(ctx, atespace, req.Name); err != nil { + wsName := req.Name + + // Acquire exclusive lock for this workspace + unlock, err := s.locker.Lock(ctx, "workspace", atespace, wsName) + if err != nil { + return nil, status.Errorf(codes.Aborted, "locking workspace %s/%s: %v", atespace, wsName, err) + } + defer unlock() + + if err := s.store.DeleteWorkspace(ctx, atespace, wsName); err != nil { return nil, status.Errorf(codes.Internal, "deleting workspace: %v", err) } return &v1alpha1.DeleteWorkspaceResponse{}, nil @@ -366,15 +539,46 @@ func (s *Server) UpdateModel(ctx context.Context, req *v1alpha1.UpdateModelReque } return existing.GetMetadata() }) + + atespace := req.Model.Metadata.Atespace + modelName := req.Model.Metadata.Name + + // Acquire exclusive lock for this model + unlock, err := s.locker.Lock(ctx, "model", atespace, modelName) + if err != nil { + return nil, status.Errorf(codes.Aborted, "locking model %s/%s: %v", atespace, modelName, err) + } + defer unlock() + if err := s.store.SaveModel(ctx, req.Model); err != nil { return nil, status.Errorf(codes.Internal, "saving model: %v", err) } return req.Model, nil } -// defaultMetadata normalizes resource metadata before a save: a missing atespace -// becomes "default", and the creation timestamp is carried over from the existing -// resource (looked up via existing) or set to now for a new one. +func (s *Server) DeleteModel(ctx context.Context, req *v1alpha1.DeleteModelRequest) (*v1alpha1.DeleteModelResponse, error) { + if req == nil { + return nil, status.Error(codes.InvalidArgument, "missing request") + } + atespace := req.Atespace + if atespace == "" { + atespace = "default" + } + modelName := req.Name + + // Acquire exclusive lock for this model + unlock, err := s.locker.Lock(ctx, "model", atespace, modelName) + if err != nil { + return nil, status.Errorf(codes.Aborted, "locking model %s/%s: %v", atespace, modelName, err) + } + defer unlock() + + if err := s.store.DeleteModel(ctx, atespace, modelName); err != nil { + return nil, status.Errorf(codes.Internal, "deleting model: %v", err) + } + return &v1alpha1.DeleteModelResponse{}, nil +} + func defaultMetadata(meta *v1alpha1.ObjectMeta, existing func(atespace, name string) *v1alpha1.ObjectMeta) *v1alpha1.ObjectMeta { if meta == nil { meta = &v1alpha1.ObjectMeta{} @@ -396,16 +600,3 @@ func defaultMetadata(meta *v1alpha1.ObjectMeta, existing func(atespace, name str return meta } -func (s *Server) DeleteModel(ctx context.Context, req *v1alpha1.DeleteModelRequest) (*v1alpha1.DeleteModelResponse, error) { - if req == nil { - return nil, status.Error(codes.InvalidArgument, "missing request") - } - atespace := req.Atespace - if atespace == "" { - atespace = "default" - } - if err := s.store.DeleteModel(ctx, atespace, req.Name); err != nil { - return nil, status.Errorf(codes.Internal, "deleting model: %v", err) - } - return &v1alpha1.DeleteModelResponse{}, nil -} diff --git a/internal/server/server_test.go b/internal/server/server_test.go index c38af675..afb581cb 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -16,6 +16,7 @@ package server_test import ( "context" + "errors" "net" "net/http" "net/http/httptest" @@ -23,6 +24,7 @@ import ( "time" "github.com/google/ax/internal/server" + "github.com/google/ax/internal/store" "github.com/google/ax/internal/store/memory" "github.com/google/ax/pkg/apis/v1alpha1" "google.golang.org/grpc" @@ -219,25 +221,12 @@ func TestServerGRPC(t *testing.T) { } // 8. Delete operations - // Task deletion is two-phase: the RPC marks the task Terminating and the - // controller removes the record after tearing down the actor. if _, err := client.DeleteTask(ctx, &v1alpha1.DeleteTaskRequest{Atespace: "default", Name: "grpc-task"}); err != nil { t.Fatalf("DeleteTask failed: %v", err) } - terminating, err := client.GetTask(ctx, &v1alpha1.GetTaskRequest{Atespace: "default", Name: "grpc-task"}) - if err != nil { - t.Fatalf("GetTask after DeleteTask failed: %v", err) - } - if terminating.GetStatus().GetPhase() != v1alpha1.PhaseTerminating { - t.Errorf("expected phase %q after DeleteTask, got %q", v1alpha1.PhaseTerminating, terminating.GetStatus().GetPhase()) - } if _, err := client.DeleteTask(ctx, &v1alpha1.DeleteTaskRequest{Atespace: "default", Name: "no-such-task"}); status.Code(err) != codes.NotFound { t.Errorf("expected NotFound deleting a missing task, got %v", err) } - // Stand in for the controller finishing cleanup. - if err := memStore.DeleteTask(ctx, "default", "grpc-task"); err != nil { - t.Fatalf("removing task record failed: %v", err) - } if _, err := client.DeleteWorkspace(ctx, &v1alpha1.DeleteWorkspaceRequest{Atespace: "default", Name: "grpc-ws"}); err != nil { t.Fatalf("DeleteWorkspace failed: %v", err) } @@ -253,8 +242,8 @@ func TestServerGRPC(t *testing.T) { } // Names and atespaces become Substrate resource names, which must be RFC 1123 -// labels. The server rejects them up front instead of letting the controller -// fail asynchronously with ActorCreationFailed. +// labels. The server rejects them up front instead of failing during Substrate +// actor creation. func TestCreate_RejectsInvalidNames(t *testing.T) { srv := server.NewServer(memory.NewStore()) ctx := context.Background() @@ -318,3 +307,199 @@ func TestCreateTask_ValidatesWorkspaceBindings(t *testing.T) { t.Fatalf("expected a valid multi-workspace task to be accepted, got %v", err) } } + +type fakeReconciler struct { + reconcileCount int + deleteCount int + deleteErr error + onDelete func(ctx context.Context, atespace, taskName string) error +} + +func (f *fakeReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, workspaces ...*v1alpha1.Workspace) (*v1alpha1.Task, error) { + f.reconcileCount++ + task.Status = &v1alpha1.TaskStatus{ + Phase: task.GetStatus().GetPhase(), + } + return task, nil +} + +func (f *fakeReconciler) ReconcileDelete(ctx context.Context, atespace, taskName string) error { + f.deleteCount++ + if f.onDelete != nil { + return f.onDelete(ctx, atespace, taskName) + } + return f.deleteErr +} + +func TestServer_DirectReconcilerLifecycle(t *testing.T) { + rec := &fakeReconciler{} + srv := server.NewServer(memory.NewStore(), server.Options{Reconciler: rec}) + ctx := context.Background() + + // 1. CreateTask directly calls Reconciler + task, err := srv.CreateTask(ctx, &v1alpha1.CreateTaskRequest{ + Task: &v1alpha1.Task{ + Metadata: &v1alpha1.ObjectMeta{Name: "task-rec"}, + Spec: &v1alpha1.TaskSpec{Image: "alpine"}, + }, + }) + if err != nil { + t.Fatalf("CreateTask: %v", err) + } + if rec.reconcileCount != 1 { + t.Errorf("expected 1 reconcile call on CreateTask, got %d", rec.reconcileCount) + } + if task.GetStatus().GetPhase() != "Suspended" { + t.Errorf("expected phase Suspended, got %s", task.GetStatus().GetPhase()) + } + + // 2. ResumeTask directly calls Reconciler + task, err = srv.ResumeTask(ctx, &v1alpha1.ResumeTaskRequest{Name: "task-rec"}) + if err != nil { + t.Fatalf("ResumeTask: %v", err) + } + if rec.reconcileCount != 2 { + t.Errorf("expected 2 reconcile calls after ResumeTask, got %d", rec.reconcileCount) + } + if task.GetStatus().GetPhase() != "Running" { + t.Errorf("expected phase Running, got %s", task.GetStatus().GetPhase()) + } + + // 3. SuspendTask directly calls Reconciler + task, err = srv.SuspendTask(ctx, &v1alpha1.SuspendTaskRequest{Name: "task-rec"}) + if err != nil { + t.Fatalf("SuspendTask: %v", err) + } + if rec.reconcileCount != 3 { + t.Errorf("expected 3 reconcile calls after SuspendTask, got %d", rec.reconcileCount) + } + if task.GetStatus().GetPhase() != "Suspended" { + t.Errorf("expected phase Suspended, got %s", task.GetStatus().GetPhase()) + } + + // 4. DeleteTask directly calls ReconcileDelete + _, err = srv.DeleteTask(ctx, &v1alpha1.DeleteTaskRequest{Name: "task-rec"}) + if err != nil { + t.Fatalf("DeleteTask: %v", err) + } + if rec.deleteCount != 1 { + t.Errorf("expected 1 delete call, got %d", rec.deleteCount) + } + + // 5. Verify task is gone + _, err = srv.GetTask(ctx, &v1alpha1.GetTaskRequest{Name: "task-rec"}) + if status.Code(err) != codes.NotFound { + t.Errorf("expected NotFound after DeleteTask, got %v", err) + } +} + +func TestServer_DeleteTask_BlocksAndSetsPhaseTerminating(t *testing.T) { + st := memory.NewStore() + var observedPhaseDuringDelete string + rec := &fakeReconciler{ + onDelete: func(ctx context.Context, atespace, taskName string) error { + t, err := st.GetTask(ctx, atespace, taskName) + if err == nil && t.Status != nil { + observedPhaseDuringDelete = t.Status.Phase + } + return nil + }, + } + srv := server.NewServer(st, server.Options{Reconciler: rec}) + ctx := context.Background() + + _, err := srv.CreateTask(ctx, &v1alpha1.CreateTaskRequest{ + Task: &v1alpha1.Task{ + Metadata: &v1alpha1.ObjectMeta{Name: "task-term"}, + Spec: &v1alpha1.TaskSpec{Image: "alpine"}, + }, + }) + if err != nil { + t.Fatalf("CreateTask: %v", err) + } + + _, err = srv.DeleteTask(ctx, &v1alpha1.DeleteTaskRequest{Name: "task-term"}) + if err != nil { + t.Fatalf("DeleteTask: %v", err) + } + + if observedPhaseDuringDelete != v1alpha1.PhaseTerminating { + t.Errorf("expected phase %q during delete, got %q", v1alpha1.PhaseTerminating, observedPhaseDuringDelete) + } + + _, err = srv.GetTask(ctx, &v1alpha1.GetTaskRequest{Name: "task-term"}) + if status.Code(err) != codes.NotFound { + t.Errorf("expected NotFound after successful DeleteTask, got %v", err) + } +} + +func TestServer_DeleteTask_ReconcileErrorRetainsTask(t *testing.T) { + st := memory.NewStore() + rec := &fakeReconciler{ + deleteErr: errors.New("substrate timeout deleting actor"), + } + srv := server.NewServer(st, server.Options{Reconciler: rec}) + ctx := context.Background() + + _, err := srv.CreateTask(ctx, &v1alpha1.CreateTaskRequest{ + Task: &v1alpha1.Task{ + Metadata: &v1alpha1.ObjectMeta{Name: "task-err"}, + Spec: &v1alpha1.TaskSpec{Image: "alpine"}, + }, + }) + if err != nil { + t.Fatalf("CreateTask: %v", err) + } + + _, err = srv.DeleteTask(ctx, &v1alpha1.DeleteTaskRequest{Name: "task-err"}) + if err == nil { + t.Fatal("expected DeleteTask to fail when ReconcileDelete fails") + } + + task, err := srv.GetTask(ctx, &v1alpha1.GetTaskRequest{Name: "task-err"}) + if err != nil { + t.Fatalf("expected task to remain in store, got %v", err) + } + if task.GetStatus().GetPhase() != v1alpha1.PhaseTerminating { + t.Errorf("expected phase %q, got %q", v1alpha1.PhaseTerminating, task.GetStatus().GetPhase()) + } +} + +type failUpdateStatusStore struct { + store.Store + failUpdateStatus bool +} + +func (f *failUpdateStatusStore) UpdateTaskStatus(ctx context.Context, atespace, name string, status *v1alpha1.TaskStatus) error { + if f.failUpdateStatus { + return errors.New("simulated store failure") + } + return f.Store.UpdateTaskStatus(ctx, atespace, name, status) +} + +func TestServer_DeleteTask_UpdateStatusError(t *testing.T) { + base := memory.NewStore() + st := &failUpdateStatusStore{Store: base} + rec := &fakeReconciler{} + srv := server.NewServer(st, server.Options{Reconciler: rec}) + ctx := context.Background() + + _, err := srv.CreateTask(ctx, &v1alpha1.CreateTaskRequest{ + Task: &v1alpha1.Task{ + Metadata: &v1alpha1.ObjectMeta{Name: "task-status-fail"}, + Spec: &v1alpha1.TaskSpec{Image: "alpine"}, + }, + }) + if err != nil { + t.Fatalf("CreateTask: %v", err) + } + + st.failUpdateStatus = true + _, err = srv.DeleteTask(ctx, &v1alpha1.DeleteTaskRequest{Name: "task-status-fail"}) + if err == nil { + t.Fatal("expected DeleteTask to fail when UpdateTaskStatus fails") + } + if rec.deleteCount != 0 { + t.Errorf("expected ReconcileDelete not to be called if UpdateTaskStatus fails, got %d calls", rec.deleteCount) + } +} diff --git a/internal/store/memory/store.go b/internal/store/memory/store.go index c1788900..76adf03f 100644 --- a/internal/store/memory/store.go +++ b/internal/store/memory/store.go @@ -20,7 +20,6 @@ import ( "fmt" "io" "sync" - "time" "github.com/google/ax/internal/store" "github.com/google/ax/pkg/apis/v1alpha1" @@ -44,7 +43,6 @@ type MemoryStore struct { tasks map[string]*v1alpha1.Task models map[string]*v1alpha1.Model workspaces map[string]*v1alpha1.Workspace - events chan store.TaskEvent watchers map[string][]chan *v1alpha1.Task } @@ -54,7 +52,6 @@ func NewStore() *MemoryStore { tasks: make(map[string]*v1alpha1.Task), models: make(map[string]*v1alpha1.Model), workspaces: make(map[string]*v1alpha1.Workspace), - events: make(chan store.TaskEvent, 1000), watchers: make(map[string][]chan *v1alpha1.Task), } } @@ -89,13 +86,6 @@ func (s *MemoryStore) SaveTask(ctx context.Context, task *v1alpha1.Task) error { cp := clone(task) s.tasks[key] = cp - event := store.TaskEvent{ - ID: fmt.Sprintf("%d", time.Now().UnixNano()), - Atespace: task.Metadata.Atespace, - Name: task.Metadata.Name, - Action: "reconcile", - } - // Notify watchers if chs, ok := s.watchers[key]; ok { for _, ch := range chs { @@ -107,11 +97,6 @@ func (s *MemoryStore) SaveTask(ctx context.Context, task *v1alpha1.Task) error { } s.mu.Unlock() - select { - case s.events <- event: - default: - } - return nil } @@ -174,42 +159,6 @@ func (s *MemoryStore) UpdateTaskStatus(ctx context.Context, atespace, name strin return nil } -func (s *MemoryStore) MarkTaskDeleting(ctx context.Context, atespace, name string) error { - key := taskKey(atespace, name) - - s.mu.Lock() - t, ok := s.tasks[key] - if !ok { - s.mu.Unlock() - return store.ErrNotFound - } - if t.Status == nil { - t.Status = &v1alpha1.TaskStatus{} - } - t.Status.Phase = v1alpha1.PhaseTerminating - cp := clone(t) - if chs, ok := s.watchers[key]; ok { - for _, ch := range chs { - select { - case ch <- cp: - default: - } - } - } - s.mu.Unlock() - - select { - case s.events <- store.TaskEvent{ - ID: fmt.Sprintf("%d", time.Now().UnixNano()), - Atespace: atespace, - Name: name, - Action: "delete", - }: - default: - } - return nil -} - func (s *MemoryStore) DeleteTask(ctx context.Context, atespace, name string) error { s.mu.Lock() defer s.mu.Unlock() @@ -217,7 +166,6 @@ func (s *MemoryStore) DeleteTask(ctx context.Context, atespace, name string) err return nil } - func (s *MemoryStore) SaveModel(ctx context.Context, model *v1alpha1.Model) error { if model.Metadata.Name == "" { return errors.New("model name is required") @@ -314,36 +262,6 @@ func (s *MemoryStore) DeleteWorkspace(ctx context.Context, atespace, name string return nil } -// Subscribe returns a subscription over the store's single event channel. Every -// subscription shares that channel, so each event reaches exactly one subscriber, -// which matches the group semantics of the Redis implementation. Groups and -// consumers are accepted for interface parity but carry no meaning here. -func (s *MemoryStore) Subscribe(ctx context.Context, group, consumer string) (store.Subscription, error) { - return &subscription{events: s.events}, nil -} - -type subscription struct { - events <-chan store.TaskEvent -} - -func (sub *subscription) Next(ctx context.Context) (store.TaskEvent, error) { - select { - case ev := <-sub.events: - return ev, nil - case <-ctx.Done(): - return store.TaskEvent{}, ctx.Err() - } -} - -// Ack is a no-op: the channel hands each event out once, so there is nothing to retain. -func (sub *subscription) Ack(ctx context.Context, ev store.TaskEvent) error { - return nil -} - -func (sub *subscription) Close() error { - return nil -} - func (s *MemoryStore) WatchTask(ctx context.Context, atespace, name string) (<-chan *v1alpha1.Task, io.Closer, error) { key := taskKey(atespace, name) ch := make(chan *v1alpha1.Task, 10) diff --git a/internal/store/redis/store.go b/internal/store/redis/store.go index e96651d8..67a9cf81 100644 --- a/internal/store/redis/store.go +++ b/internal/store/redis/store.go @@ -28,12 +28,6 @@ import ( "google.golang.org/protobuf/encoding/protojson" ) -const ( - defaultStreamName = "ax:stream:tasks" - defaultReadBatchSize = 10 - defaultReadBlock = 2 * time.Second -) - var ( jsonMarshalOpts = protojson.MarshalOptions{UseProtoNames: false, EmitUnpopulated: false} jsonUnmarshalOpts = protojson.UnmarshalOptions{DiscardUnknown: true} @@ -41,15 +35,8 @@ var ( // Options contains configuration for the Redis store. type Options struct { - StreamName string - KeyPrefix string - TTL time.Duration // Optional TTL for task records - - // ReadBatchSize is how many events one XREADGROUP call may return. - ReadBatchSize int64 - // ReadBlock is how long one XREADGROUP call waits for events before returning - // empty. Shorter values make shutdown more responsive at the cost of more calls. - ReadBlock time.Duration + KeyPrefix string + TTL time.Duration // Optional TTL for task records } // Store is a Redis-backed implementation of store.Store. @@ -60,18 +47,9 @@ type Store struct { // NewStore creates a new Redis store. func NewStore(client *redis.Client, opts Options) *Store { - if opts.StreamName == "" { - opts.StreamName = defaultStreamName - } if opts.KeyPrefix == "" { opts.KeyPrefix = "ax" } - if opts.ReadBatchSize <= 0 { - opts.ReadBatchSize = defaultReadBatchSize - } - if opts.ReadBlock <= 0 { - opts.ReadBlock = defaultReadBlock - } return &Store{ client: client, opts: opts, @@ -157,14 +135,6 @@ func (s *Store) SaveTask(ctx context.Context, task *v1alpha1.Task) error { pipe.Set(ctx, s.taskKey(atespace, name), data, s.opts.TTL) pipe.ZAdd(ctx, s.taskIndexKey(), redis.Z{Score: score, Member: member}) pipe.ZAdd(ctx, s.taskAtespaceIndexKey(atespace), redis.Z{Score: score, Member: name}) - pipe.XAdd(ctx, &redis.XAddArgs{ - Stream: s.opts.StreamName, - Values: map[string]interface{}{ - "action": "reconcile", - "atespace": atespace, - "name": name, - }, - }) pipe.Publish(ctx, s.taskPubSubChannel(atespace, name), data) _, err = pipe.Exec(ctx) @@ -279,43 +249,7 @@ func (s *Store) UpdateTaskStatus(ctx context.Context, atespace, name string, sta return nil } -// MarkTaskDeleting flips the task to the Terminating phase, notifies watchers, and -// publishes a delete event for the controller. The record stays until DeleteTask. -func (s *Store) MarkTaskDeleting(ctx context.Context, atespace, name string) error { - if atespace == "" { - atespace = "default" - } - task, err := s.GetTask(ctx, atespace, name) - if err != nil { - return err - } - if task.Status == nil { - task.Status = &v1alpha1.TaskStatus{} - } - task.Status.Phase = v1alpha1.PhaseTerminating - data, err := protojson.Marshal(task) - if err != nil { - return fmt.Errorf("marshaling task: %w", err) - } - - pipe := s.client.TxPipeline() - pipe.Set(ctx, s.taskKey(atespace, name), data, s.opts.TTL) - pipe.Publish(ctx, s.taskPubSubChannel(atespace, name), data) - pipe.XAdd(ctx, &redis.XAddArgs{ - Stream: s.opts.StreamName, - Values: map[string]interface{}{ - "action": "delete", - "atespace": atespace, - "name": name, - }, - }) - if _, err := pipe.Exec(ctx); err != nil { - return fmt.Errorf("marking task deleting in redis: %w", err) - } - return nil -} - -// DeleteTask removes the task record and its index entries. No event is published. +// DeleteTask removes the task record and its index entries. func (s *Store) DeleteTask(ctx context.Context, atespace, name string) error { if atespace == "" { atespace = "default" @@ -587,98 +521,6 @@ func (s *Store) GetWorkspace(ctx context.Context, atespace, name string) (*v1alp return &w, nil } -// Subscribe joins a Redis Streams consumer group, creating the group (and the -// stream) if needed. A new group starts at the tail of the stream, so it only -// sees events published after it was created; an existing group keeps its -// position and any pending entries. -func (s *Store) Subscribe(ctx context.Context, group, consumer string) (store.Subscription, error) { - err := s.client.XGroupCreateMkStream(ctx, s.opts.StreamName, group, "$").Err() - if err != nil && !strings.Contains(err.Error(), "BUSYGROUP") { - return nil, fmt.Errorf("creating consumer group %q: %w", group, err) - } - return &subscription{store: s, group: group, consumer: consumer}, nil -} - -// subscription reads from a consumer group in batches and hands events out one -// at a time. Delivery is at-least-once: an event stays in the group's pending -// list until Ack is called for it. -type subscription struct { - store *Store - group string - consumer string - pending []store.TaskEvent -} - -func (sub *subscription) Next(ctx context.Context) (store.TaskEvent, error) { - for len(sub.pending) == 0 { - if err := ctx.Err(); err != nil { - return store.TaskEvent{}, err - } - events, err := sub.read(ctx) - if err != nil { - return store.TaskEvent{}, err - } - sub.pending = events - } - ev := sub.pending[0] - sub.pending = sub.pending[1:] - return ev, nil -} - -// read performs one blocking XREADGROUP call. It returns an empty slice, not an -// error, when the block time elapses without events. -func (sub *subscription) read(ctx context.Context) ([]store.TaskEvent, error) { - streams, err := sub.store.client.XReadGroup(ctx, &redis.XReadGroupArgs{ - Group: sub.group, - Consumer: sub.consumer, - Streams: []string{sub.store.opts.StreamName, ">"}, - Count: sub.store.opts.ReadBatchSize, - Block: sub.store.opts.ReadBlock, - }).Result() - if err != nil { - if errors.Is(err, redis.Nil) { - return nil, nil - } - return nil, fmt.Errorf("reading task events: %w", err) - } - - var events []store.TaskEvent - for _, stream := range streams { - for _, msg := range stream.Messages { - events = append(events, eventFromMessage(msg)) - } - } - return events, nil -} - -func (sub *subscription) Ack(ctx context.Context, ev store.TaskEvent) error { - return sub.store.client.XAck(ctx, sub.store.opts.StreamName, sub.group, ev.ID).Err() -} - -// Close releases the subscription. The consumer is deliberately left registered -// so that any events it had claimed but not acknowledged remain claimable. -func (sub *subscription) Close() error { - return nil -} - -// eventFromMessage decodes a stream entry. "namespace" is accepted as a legacy -// alias for "atespace". -func eventFromMessage(msg redis.XMessage) store.TaskEvent { - ev := store.TaskEvent{ID: msg.ID} - if atespace, ok := msg.Values["atespace"].(string); ok { - ev.Atespace = atespace - } else if ns, ok := msg.Values["namespace"].(string); ok { - ev.Atespace = ns - } - if name, ok := msg.Values["name"].(string); ok { - ev.Name = name - } - if action, ok := msg.Values["action"].(string); ok { - ev.Action = action - } - return ev -} - // WatchTask subscribes to status change notifications for a specific task. func (s *Store) WatchTask(ctx context.Context, atespace, name string) (<-chan *v1alpha1.Task, io.Closer, error) { if atespace == "" { diff --git a/internal/store/store.go b/internal/store/store.go index 41f3a89e..5d14d61a 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -26,48 +26,12 @@ var ( ErrNotFound = errors.New("resource not found") ) -// TaskEvent represents an event published to the task event stream. -type TaskEvent struct { - ID string - Atespace string - Name string - Action string // "reconcile", "delete" -} - -// EventQueue delivers task events to groups of cooperating workers. Every event -// is delivered to exactly one member of a group, and stays pending until that -// member acknowledges it, so a crashed worker's events can be picked up again. -type EventQueue interface { - // Subscribe joins group as consumer, creating the group if it does not exist. - // All members of a group share one stream of events. - Subscribe(ctx context.Context, group, consumer string) (Subscription, error) -} - -// Subscription is one consumer's view of an EventQueue group. -type Subscription interface { - // Next blocks until an event is available or ctx is done. - Next(ctx context.Context) (TaskEvent, error) - // Ack marks an event as processed so it is not delivered again. - Ack(ctx context.Context, ev TaskEvent) error - // Close releases the subscription. Unacknowledged events stay pending for the group. - Close() error -} - -// Store defines the storage and event streaming interface for AX resources. +// Store defines the storage interface for AX resources. type Store interface { - EventQueue - SaveTask(ctx context.Context, task *v1alpha1.Task) error GetTask(ctx context.Context, atespace, name string) (*v1alpha1.Task, error) ListTasks(ctx context.Context, atespace string, limit, offset int64) ([]*v1alpha1.Task, error) UpdateTaskStatus(ctx context.Context, atespace, name string, status *v1alpha1.TaskStatus) error - // MarkTaskDeleting begins a two-phase delete: the task's phase becomes - // "Terminating" and a delete event is published for the controller, which - // removes the actor and then calls DeleteTask. Returns ErrNotFound if the - // task does not exist. - MarkTaskDeleting(ctx context.Context, atespace, name string) error - // DeleteTask removes the task record. It publishes no event; callers are - // expected to have cleaned up the task's actor first. DeleteTask(ctx context.Context, atespace, name string) error SaveWorkspace(ctx context.Context, ws *v1alpha1.Workspace) error diff --git a/internal/substrate/client.go b/internal/substrate/client.go index b0273e87..de8ed991 100644 --- a/internal/substrate/client.go +++ b/internal/substrate/client.go @@ -205,7 +205,7 @@ func (c *Client) GetActorTemplate(ctx context.Context, atespace, templateName st const ( DefaultGuestCommand = "/usr/local/bin/ax-task-runner" - DefaultSnapshotsBucket = "gs://snapshot-substrate-test-ax-substrate/ate-env/" + DefaultSnapshotsBucket = "gs://dberkov-gke-dev3/ate-env/" ) // BuildActorTemplate constructs a Substrate ActorTemplate based on the standard ate-env specification. @@ -320,24 +320,19 @@ func (c *Client) EnsureActor(ctx context.Context, atespace, actorName, templateA actor, err := c.control.CreateActor(ctx, req) if err != nil { if status.Code(err) == codes.AlreadyExists { - existing, getErr := c.control.GetActor(ctx, &ateapipb.GetActorRequest{ - Actor: &ateapipb.ObjectRef{ - Atespace: atespace, - Name: actorName, - }, - }) + existing, getErr := c.GetActor(ctx, atespace, actorName) if getErr == nil && existing != nil { state := existing.GetStatus().GetState() if state == ateapipb.ActorState_ACTOR_STATE_CRASHED { slog.Warn("existing actor is crashed, deleting and recreating", "actor", actorName) - _, _ = c.control.DeleteActor(ctx, &ateapipb.DeleteActorRequest{ - Actor: &ateapipb.ObjectRef{ - Atespace: atespace, - Name: actorName, - }, - AnyState: true, - }) - state = ateapipb.ActorState_ACTOR_STATE_DELETING + if err := c.DeleteActor(ctx, atespace, actorName); err != nil { + return nil, err + } + actor, err = c.control.CreateActor(ctx, req) + if err == nil { + return actor, nil + } + return nil, fmt.Errorf("recreating actor %s/%s after crash: %w", atespace, actorName, err) } if state == ateapipb.ActorState_ACTOR_STATE_DELETING { // Wait briefly for previous actor deletion to finalize before recreating @@ -402,7 +397,22 @@ func (c *Client) SuspendActor(ctx context.Context, atespace, actorName string) e return nil } -// DeleteActor deletes the specified actor from Substrate. +// GetActor fetches an Actor by name. +func (c *Client) GetActor(ctx context.Context, atespace, actorName string) (*ateapipb.Actor, error) { + req := &ateapipb.GetActorRequest{ + Actor: &ateapipb.ObjectRef{ + Atespace: atespace, + Name: actorName, + }, + } + return c.control.GetActor(ctx, req) +} + +// actorDeletionPollInterval is the interval between checks when waiting for an actor to be deleted. +var actorDeletionPollInterval = 200 * time.Millisecond + +// DeleteActor deletes the specified actor from Substrate and blocks until the actor +// is fully deleted. func (c *Client) DeleteActor(ctx context.Context, atespace, actorName string) error { req := &ateapipb.DeleteActorRequest{ Actor: &ateapipb.ObjectRef{ @@ -412,10 +422,50 @@ func (c *Client) DeleteActor(ctx context.Context, atespace, actorName string) er AnyState: true, } _, err := c.control.DeleteActor(ctx, req) - if err != nil && status.Code(err) != codes.NotFound { + if err != nil { + if status.Code(err) == codes.NotFound { + return nil + } return fmt.Errorf("deleting actor %s/%s: %w", atespace, actorName, err) } - return nil + + getReq := &ateapipb.GetActorRequest{ + Actor: &ateapipb.ObjectRef{ + Atespace: atespace, + Name: actorName, + }, + } + + // Fast path: check if the actor was deleted synchronously. + _, err = c.control.GetActor(ctx, getReq) + if err != nil { + if status.Code(err) == codes.NotFound { + return nil + } + if ctx.Err() != nil { + return fmt.Errorf("waiting for actor %s/%s deletion: %w", atespace, actorName, ctx.Err()) + } + } + + ticker := time.NewTicker(actorDeletionPollInterval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return fmt.Errorf("waiting for actor %s/%s deletion: %w", atespace, actorName, ctx.Err()) + case <-ticker.C: + _, err := c.control.GetActor(ctx, getReq) + if err != nil { + if status.Code(err) == codes.NotFound { + return nil + } + if ctx.Err() != nil { + return fmt.Errorf("waiting for actor %s/%s deletion: %w", atespace, actorName, ctx.Err()) + } + } + } + } } // ListActorTemplates returns all ActorTemplates in the given atespace. diff --git a/internal/substrate/client_test.go b/internal/substrate/client_test.go new file mode 100644 index 00000000..2d889492 --- /dev/null +++ b/internal/substrate/client_test.go @@ -0,0 +1,142 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package substrate + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +type mockControlClient struct { + ateapipb.ControlClient + + mu sync.Mutex + deleteCalls []*ateapipb.DeleteActorRequest + getCalls []*ateapipb.GetActorRequest + deleteErr error + actorRemaining int // Number of GetActor calls before returning NotFound +} + +func (m *mockControlClient) DeleteActor(ctx context.Context, req *ateapipb.DeleteActorRequest, opts ...grpc.CallOption) (*ateapipb.Actor, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.deleteCalls = append(m.deleteCalls, req) + if m.deleteErr != nil { + return nil, m.deleteErr + } + return &ateapipb.Actor{}, nil +} + +func (m *mockControlClient) GetActor(ctx context.Context, req *ateapipb.GetActorRequest, opts ...grpc.CallOption) (*ateapipb.Actor, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.getCalls = append(m.getCalls, req) + if m.actorRemaining > 0 { + m.actorRemaining-- + return &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{ + Atespace: req.GetActor().GetAtespace(), + Name: req.GetActor().GetName(), + }, + Status: &ateapipb.ActorStatus{State: ateapipb.ActorState_ACTOR_STATE_DELETING}, + }, nil + } + return nil, status.Error(codes.NotFound, "actor not found") +} + +func TestClient_DeleteActor_AnyStateTrue(t *testing.T) { + mock := &mockControlClient{ + actorRemaining: 0, + } + c := &Client{control: mock} + ctx := context.Background() + + err := c.DeleteActor(ctx, "test-ns", "test-actor") + if err != nil { + t.Fatalf("DeleteActor failed: %v", err) + } + + mock.mu.Lock() + defer mock.mu.Unlock() + if len(mock.deleteCalls) != 1 { + t.Fatalf("expected 1 DeleteActor call, got %d", len(mock.deleteCalls)) + } + req := mock.deleteCalls[0] + if !req.GetAnyState() { + t.Errorf("expected AnyState to be true, got %v", req.GetAnyState()) + } + if req.GetActor().GetAtespace() != "test-ns" || req.GetActor().GetName() != "test-actor" { + t.Errorf("unexpected actor ref in request: %v", req.GetActor()) + } +} + +func TestClient_DeleteActor_PollsUntilNotFound(t *testing.T) { + oldInterval := actorDeletionPollInterval + actorDeletionPollInterval = 10 * time.Millisecond + defer func() { actorDeletionPollInterval = oldInterval }() + + mock := &mockControlClient{ + actorRemaining: 3, + } + c := &Client{control: mock} + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + err := c.DeleteActor(ctx, "test-ns", "test-actor") + if err != nil { + t.Fatalf("DeleteActor failed: %v", err) + } + + mock.mu.Lock() + defer mock.mu.Unlock() + // 1 fast-path GetActor call + 3 polling calls (last one returns NotFound) + if len(mock.getCalls) != 4 { + t.Errorf("expected 4 GetActor calls before NotFound, got %d", len(mock.getCalls)) + } +} + +func TestClient_DeleteActor_NotFoundIsIgnored(t *testing.T) { + mock := &mockControlClient{ + deleteErr: status.Error(codes.NotFound, "actor already gone"), + } + c := &Client{control: mock} + ctx := context.Background() + + err := c.DeleteActor(ctx, "test-ns", "test-actor") + if err != nil { + t.Fatalf("expected nil error when actor is NotFound, got: %v", err) + } +} + +func TestClient_DeleteActor_PropagatesRPCError(t *testing.T) { + mock := &mockControlClient{ + deleteErr: errors.New("internal atelet failure"), + } + c := &Client{control: mock} + ctx := context.Background() + + err := c.DeleteActor(ctx, "test-ns", "test-actor") + if err == nil { + t.Fatal("expected error, got nil") + } +}