diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 00000000..abab283b --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,6 @@ +# Changes under these paths need a review from their owners. +# https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners + +/internal/server/ @rakyll @mayawang +/internal/store/ @rakyll @mayawang +/internal/controller/ @rakyll @mayawang diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index c2b65d83..79d1738f 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -26,6 +26,20 @@ permissions: jobs: build-and-test: runs-on: ubuntu-latest + services: + postgres: + image: postgres:17-alpine + env: + POSTGRES_USER: ax + POSTGRES_PASSWORD: ax + POSTGRES_DB: ax + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U ax -d ax" + --health-interval 5s + --health-timeout 5s + --health-retries 10 steps: - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 @@ -46,4 +60,6 @@ jobs: go build -v ./cmd/ax-task-runner - name: Test + env: + AX_TEST_POSTGRES_URL: postgres://ax:ax@localhost:5432/ax?sslmode=disable run: go test -v ./... diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 313c4cd7..46937744 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -90,7 +90,7 @@ make push-task-runner TASK_RUNNER_REPO=gcr.io//ax-task-runner ### Deploying to Kubernetes -Deploy Redis and server components to your cluster in the `ax-system` namespace: +Deploy PostgreSQL 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 84f93c87..17fec348 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -2,7 +2,9 @@ ## 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 stores its state in Redis and reconciles directly with Agent Substrate under fine-grained distributed locks. +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 PostgreSQL and reconciles directly with Agent Substrate under fine-grained distributed locks. + +PostgreSQL holds resources as protojson documents alongside the leases that back AX's distributed locks. Watches and lock waiters poll; nothing depends on `LISTEN/NOTIFY`, so any connection pooler works. `ax-server` creates any missing tables when it starts; there are no versioned migrations yet. The bundled manifest runs a single-instance StatefulSet with a persistent volume; any reachable PostgreSQL works. Storage sits behind the `store.Store` and `lock.Locker` interfaces. The store holds only what AX owns: specs, metadata, and the parts of a task's status that describe AX's own work (`Pending`, `Failed`, `Terminating`, the task ID, workspace setup). Everything that belongs to the task's actor (its phase while provisioned, `actor`, `workerIP`, and the `Ready` condition) is read from Agent Substrate whenever a task is fetched, listed, or watched, so AX never keeps a copy that can drift from the actor. The split lives in a few helpers in `internal/server/server.go` (`getTask`, `listTasks`, `saveTask`, `updateTask`, `deleteTask`), so storage implementations never need to know about Substrate. ``` ax apply -f task.yaml @@ -13,9 +15,9 @@ Storing millions of short-lived tasks as Kubernetes CRDs pushes etcd past its co │ ┌──────────────────────┴──────────────────────┐ ▼ ▼ - Redis Agent Substrate + PostgreSQL Agent Substrate (Resource Store, ┌───────────────────────────────┐ - Locks, PubSub) │ • Atespace Provisioning │ + Locks) │ • Atespace Provisioning │ │ • Actor Creation & Activation │ │ • Worker Assignment │ └───────────────────────────────┘ @@ -26,7 +28,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` | Direct-execution gRPC API on port 8080. Validates manifests, manages distributed locks, reconciles directly with Agent Substrate, and persists state to Redis. | +| `ax-server` | Direct-execution gRPC API on port 8080. Validates manifests, manages distributed locks, reconciles directly with Agent Substrate, and persists state to PostgreSQL. | | `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 2d918f50..3eb2d3a5 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-server deploy-redis apply-example test test-e2e test-e2e-head e2e-setup e2e-run e2e-cleanup clean +.PHONY: all build build-binaries build-task-runner install push push-task-runner deploy deploy-server deploy-postgres apply-example test test-e2e test-e2e-head e2e-setup e2e-run e2e-cleanup clean all: build @@ -63,12 +63,19 @@ push: push-task-runner ## Deployment Targets ## -------------------------------------- -# 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 all AX components to Kubernetes (Postgres, ax-server) +deploy: deploy-postgres deploy-server + +deploy-postgres: + @echo "==> Deploying Postgres to ax-system namespace..." + kubectl create namespace ax-system --dry-run=client -o yaml | kubectl apply -f - + @# Create the password once. Postgres only reads it when initializing an + @# empty volume, so it must not change on later deploys. + @kubectl -n ax-system get secret ax-postgres > /dev/null 2>&1 || \ + kubectl -n ax-system create secret generic ax-postgres \ + --from-literal=password="$$(openssl rand -hex 24)" + kubectl apply -f deploy/postgres.yaml + kubectl -n ax-system rollout status statefulset/ax-postgres --timeout=180s deploy-server: @echo "==> Building and deploying ax-server using ko..." diff --git a/README.md b/README.md index bc8a01d7..5b392c6c 100644 --- a/README.md +++ b/README.md @@ -88,7 +88,7 @@ With the [prerequisites](#prerequisites) in place — most importantly a reachab make deploy AX_IMAGE_REPO= ``` -This deploys Redis, then builds and deploys the control plane images with `ko`. Everything lands in the `ax-system` namespace. +This deploys PostgreSQL, then builds and deploys the control plane images with `ko`. Everything lands in the `ax-system` namespace. ### 3. Run your first task diff --git a/cmd/ax-server/main.go b/cmd/ax-server/main.go index 967b00b1..b955a51d 100644 --- a/cmd/ax-server/main.go +++ b/cmd/ax-server/main.go @@ -17,27 +17,28 @@ package main import ( "context" "flag" + "fmt" "log/slog" + "net" "net/http" "os" "os/signal" + "strconv" "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/store/postgres" "github.com/google/ax/internal/substrate" "github.com/google/ax/pkg/apis/v1alpha1" - goredis "github.com/redis/go-redis/v9" + "github.com/jackc/pgx/v5/pgconn" ) func main() { var ( listenAddr string - redisAddr string - redisPassword string + storeURL string substrateEndpoint string substrateAuthority string substrateTokenFile string @@ -50,8 +51,7 @@ func main() { ) 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(&storeURL, "store", "postgres://ax@localhost:5432/ax?sslmode=disable", "PostgreSQL connection URL (env: AX_STORE). libpq variables such as PGPASSWORD fill in what it leaves out") 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") @@ -66,11 +66,11 @@ func main() { if envAddr := os.Getenv("ADDR"); envAddr != "" { listenAddr = envAddr } - if envRedis := os.Getenv("REDIS_ADDR"); envRedis != "" { - redisAddr = envRedis - } - if envPass := os.Getenv("REDIS_PASSWORD"); envPass != "" { - redisPassword = envPass + // TODO: Reject or warn on passwords embedded in the store URL. They end up + // in process args and pod specs; credentials should come from PGPASSWORD, + // a password file, or a Secret instead. + if envStore := os.Getenv("AX_STORE"); envStore != "" { + storeURL = envStore } if envImage := os.Getenv("AX_DEFAULT_TASK_IMAGE"); envImage != "" { defaultTaskImage = envImage @@ -81,20 +81,20 @@ func main() { slog.Info("starting ax-server", "listenAddr", listenAddr, - "redisAddr", redisAddr, + "store", describeStore(storeURL), "substrateEndpoint", substrateEndpoint, "template", defaultTemplate, "defaultTaskImage", defaultTaskImage, ) - rClient := goredis.NewClient(&goredis.Options{ - Addr: redisAddr, - Password: redisPassword, - }) - defer rClient.Close() - - rStore := redis.NewStore(rClient, redis.Options{}) - rLocker := lock.NewRedisLocker(rClient, lock.RedisLockerOptions{}) + openCtx, openCancel := context.WithTimeout(context.Background(), time.Minute) + backend, err := postgres.Open(openCtx, storeURL, postgres.Options{}) + openCancel() + if err != nil { + slog.Error("opening store", "store", describeStore(storeURL), "error", err) + os.Exit(1) + } + defer backend.Close() var reconciler server.Reconciler subClient, err := substrate.NewClientWithOptions(substrate.ClientOptions{ @@ -112,9 +112,10 @@ func main() { reconciler = controller.NewTaskReconciler(subClient, defaultTemplate, defaultTemplateAtespace, defaultTaskImage) } - srv := server.NewServer(rStore, server.Options{ - Locker: rLocker, + srv := server.NewServer(backend, server.Options{ + Locker: backend, Reconciler: reconciler, + Substrate: subClient, }) httpServer := &http.Server{ @@ -142,3 +143,15 @@ func main() { defer shutdownCancel() _ = httpServer.Shutdown(shutdownCtx) } + +// describeStore summarizes a store connection string for logs. It parses the +// string the way the store does, so URLs, query parameters and keyword/value +// DSNs are all handled, and reports only the user, host, port and database. +// Passwords and other connection parameters are never included. +func describeStore(connString string) string { + cfg, err := pgconn.ParseConfig(connString) + if err != nil { + return "" + } + return fmt.Sprintf("%s@%s/%s", cfg.User, net.JoinHostPort(cfg.Host, strconv.Itoa(int(cfg.Port))), cfg.Database) +} diff --git a/cmd/ax-server/main_test.go b/cmd/ax-server/main_test.go new file mode 100644 index 00000000..48c1cc05 --- /dev/null +++ b/cmd/ax-server/main_test.go @@ -0,0 +1,60 @@ +// 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 ( + "strings" + "testing" +) + +func TestDescribeStore(t *testing.T) { + tests := []struct { + name string + conn string + want string + }{ + { + name: "password in userinfo", + conn: "postgres://ax:secret@db:5432/ax?sslmode=disable", + want: "ax@db:5432/ax", + }, + { + name: "password in query", + conn: "postgres://ax@db:5432/ax?sslmode=disable&password=secret", + want: "ax@db:5432/ax", + }, + { + name: "keyword DSN", + conn: "host=db port=5432 user=ax password=secret dbname=ax sslmode=disable", + want: "ax@db:5432/ax", + }, + { + name: "unparseable", + conn: "postgres://ax:secret@db:notaport/ax", + want: "", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := describeStore(tt.conn) + if strings.Contains(got, "secret") { + t.Fatalf("describeStore(%q) = %q, leaks the password", tt.conn, got) + } + if got != tt.want { + t.Errorf("describeStore(%q) = %q, want %q", tt.conn, got, tt.want) + } + }) + } +} diff --git a/cmd/ax/main.go b/cmd/ax/main.go index 447a9284..601b5a1e 100644 --- a/cmd/ax/main.go +++ b/cmd/ax/main.go @@ -64,7 +64,7 @@ func main() { // Commands that don't require an AX server connection switch cmd { case "version": - fmt.Println("ax version v1alpha1 (standalone redis engine)") + fmt.Println("ax version v1alpha1") return case "help", "-h", "--help": printUsage() @@ -417,7 +417,7 @@ func runGet(serverURL, atespace string, args []string) error { for _, t := range tasks { workerIP := "" actor := "" - phase := "Pending" + phase := v1alpha1.PhasePending if t.Status != nil { workerIP = t.Status.WorkerIp actor = t.Status.Actor @@ -769,7 +769,7 @@ func runWatch(serverURL, atespace string, args []string) error { actor, workerIP, ) - if phase == "Completed" || phase == "Failed" { + if phase == v1alpha1.PhaseFailed { fmt.Printf("Task reached terminal phase %q.\n", phase) break } @@ -871,9 +871,9 @@ func manifestFromArgs(args []string) (data []byte, ok bool, err error) { // lifecycleRPCTimeout bounds suspend/resume RPCs, which wait for Substrate to // snapshot or restore the task's sandbox. MicroVM-backed sandboxes can take // well over a minute to snapshot/restore, so this is deliberately generous. -// The server renews its per-task Redis lock for as long as the RPC is in +// The server renews its per-task lock lease for as long as the RPC is in // flight, so this deadline is not constrained by the lock TTL (see -// lock.RedisLocker). +// postgres.Store.Lock). const lifecycleRPCTimeout = 5 * time.Minute func runSuspend(serverURL, atespace string, args []string) error { @@ -1091,7 +1091,7 @@ func runSSH(serverURL, atespace, kubeContext string, args []string) error { return fmt.Errorf("task %q has no status available", taskName) } - if task.Status.Phase != "Running" { + if task.Status.Phase != v1alpha1.PhaseRunning { return fmt.Errorf("task %q is in phase %q (must be Running to ssh)", taskName, task.Status.Phase) } diff --git a/deploy/ax-server.yaml b/deploy/ax-server.yaml index 6f31f4ae..a9d3c52e 100644 --- a/deploy/ax-server.yaml +++ b/deploy/ax-server.yaml @@ -72,7 +72,7 @@ spec: imagePullPolicy: IfNotPresent args: - "--addr=:8080" - - "--redis-addr=ax-redis.ax-system.svc.cluster.local:6379" + - "--store=postgres://ax@ax-postgres.ax-system.svc.cluster.local:5432/ax?sslmode=disable" - "--substrate-endpoint=api.ate-system.svc.cluster.local:443" - "--substrate-authority=api.ate-system.svc" - "--substrate-token-file=/var/run/secrets/ateapi/token" @@ -80,6 +80,11 @@ spec: - "--template=default-template" - "--template-atespace=ax-system" env: + - name: PGPASSWORD + valueFrom: + secretKeyRef: + name: ax-postgres + key: password - name: ATENET_ROUTER_ADDR value: "atenet-router.ate-system.svc.cluster.local:80" - name: AX_SNAPSHOTS_BUCKET diff --git a/deploy/postgres.yaml b/deploy/postgres.yaml new file mode 100644 index 00000000..8e7685af --- /dev/null +++ b/deploy/postgres.yaml @@ -0,0 +1,105 @@ +# 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. + +# A single-instance PostgreSQL for AX. It reads its password from the +# ax-postgres Secret (key: password), which `make deploy-postgres` creates with +# a random value if it does not exist yet. For production, consider a managed +# database (e.g. Cloud SQL) or an operator with backups and failover, and point +# ax-server's --store at it instead. +apiVersion: v1 +kind: Namespace +metadata: + name: ax-system +--- +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: ax-postgres + namespace: ax-system + labels: + app.kubernetes.io/name: ax-postgres + app.kubernetes.io/part-of: ax +spec: + replicas: 1 + serviceName: ax-postgres + selector: + matchLabels: + app.kubernetes.io/name: ax-postgres + template: + metadata: + labels: + app.kubernetes.io/name: ax-postgres + spec: + # Give Postgres time for a clean shutdown checkpoint. + terminationGracePeriodSeconds: 60 + containers: + - name: postgres + image: postgres:17-alpine + env: + - name: POSTGRES_USER + value: ax + - name: POSTGRES_DB + value: ax + - name: POSTGRES_PASSWORD + valueFrom: + secretKeyRef: + name: ax-postgres + key: password + # A subdirectory, because the volume root may contain lost+found. + - name: PGDATA + value: /var/lib/postgresql/data/pgdata + ports: + - containerPort: 5432 + name: postgres + readinessProbe: + exec: + command: ["pg_isready", "-U", "ax", "-d", "ax"] + periodSeconds: 5 + resources: + requests: + cpu: 100m + memory: 256Mi + limits: + cpu: 1000m + memory: 1Gi + volumeMounts: + - name: data + mountPath: /var/lib/postgresql/data + volumeClaimTemplates: + - metadata: + name: data + labels: + app.kubernetes.io/name: ax-postgres + app.kubernetes.io/part-of: ax + spec: + accessModes: ["ReadWriteOnce"] + resources: + requests: + storage: 10Gi +--- +apiVersion: v1 +kind: Service +metadata: + name: ax-postgres + namespace: ax-system + labels: + app.kubernetes.io/name: ax-postgres + app.kubernetes.io/part-of: ax +spec: + ports: + - port: 5432 + targetPort: 5432 + name: postgres + selector: + app.kubernetes.io/name: ax-postgres diff --git a/deploy/redis.yaml b/deploy/redis.yaml deleted file mode 100644 index 6a199b96..00000000 --- a/deploy/redis.yaml +++ /dev/null @@ -1,66 +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: apps/v1 -kind: Deployment -metadata: - name: ax-redis - namespace: ax-system - labels: - app.kubernetes.io/name: ax-redis - app.kubernetes.io/part-of: ax -spec: - replicas: 1 - selector: - matchLabels: - app.kubernetes.io/name: ax-redis - template: - metadata: - labels: - app.kubernetes.io/name: ax-redis - spec: - containers: - - name: redis - image: redis:7-alpine - ports: - - containerPort: 6379 - name: redis - resources: - requests: - cpu: 100m - memory: 128Mi - limits: - cpu: 1000m - memory: 1Gi ---- -apiVersion: v1 -kind: Service -metadata: - name: ax-redis - namespace: ax-system - labels: - app.kubernetes.io/name: ax-redis - app.kubernetes.io/part-of: ax -spec: - ports: - - port: 6379 - targetPort: 6379 - name: redis - selector: - app.kubernetes.io/name: ax-redis diff --git a/docs/concepts.md b/docs/concepts.md index 881c6ba6..27975a1b 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` | The 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 tears down the sandbox on Agent Substrate and removes the record. `ax delete` blocks until that has completed. +Two transitions are worth knowing. Suspending a task sets `Ready` to False with reason `Suspended`; 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/go.mod b/go.mod index 508533b4..6b8e4283 100644 --- a/go.mod +++ b/go.mod @@ -6,20 +6,21 @@ require ( github.com/agent-substrate/env v0.0.11-0.20260912052224-4468a200b170 github.com/agent-substrate/substrate v0.5.0 github.com/cenkalti/backoff/v5 v5.0.3 - github.com/redis/go-redis/v9 v9.22.0 + github.com/jackc/pgx/v5 v5.11.0 google.golang.org/grpc v1.83.2 google.golang.org/protobuf v1.36.12 gopkg.in/yaml.v3 v3.0.1 ) require ( - github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/kr/pretty v0.3.1 // indirect github.com/rogpeppe/go-internal v1.14.1 // indirect - go.uber.org/atomic v1.11.0 // indirect golang.org/x/net v0.60.0 // indirect + golang.org/x/sync v0.23.0 // indirect golang.org/x/sys v0.48.0 // indirect golang.org/x/text v0.42.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260803160001-6ac0973c030d // indirect - gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect ) diff --git a/go.sum b/go.sum index b6a0f9a5..0cd02838 100644 --- a/go.sum +++ b/go.sum @@ -2,15 +2,12 @@ github.com/agent-substrate/env v0.0.11-0.20260912052224-4468a200b170 h1:8AxhEx3B github.com/agent-substrate/env v0.0.11-0.20260912052224-4468a200b170/go.mod h1:RZ4ACmDqWQx7fgX9b1EaFJskeYTZh1Y7uZinDtp7U2o= github.com/agent-substrate/substrate v0.5.0 h1:R80f/jItcprNoe7VpV2ZJAcY4Lk7WuhCVgSEwn0Sfnk= github.com/agent-substrate/substrate v0.5.0/go.mod h1:2Pm1s1pBoB1u4OycVtezqda2BVCjnwcmRQXobWTV3eQ= -github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= -github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= -github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= -github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM= github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-logr/logr v1.4.4 h1:tG4xh9yMsRCAiodLVTxyrkzSZ9+o0L1Kg/+cPVcbP/8= github.com/go-logr/logr v1.4.4/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= @@ -21,25 +18,28 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE= -github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= -github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.11.0 h1:IzBBtyK9AHqf98cctWFifYSci2hgQR/cd56wB4p+ogg= +github.com/jackc/pgx/v5 v5.11.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= -github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= -github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= -github.com/redis/go-redis/v9 v9.22.0 h1:laDvpYXTJtZLloinw1fA5Kqd6HAEH2XKxOkG/PDq2F0= -github.com/redis/go-redis/v9 v9.22.0/go.mod h1:y2g0Wj8rQvuK0ELM+oxSudcLtC09JScs98I/X9gRWY4= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= -github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= -github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= go.opentelemetry.io/otel v1.45.0 h1:pdrWmLHofpubmArBv1LgFSv1Z0Ie/ppdZzu+kUN5EeU= @@ -52,12 +52,12 @@ go.opentelemetry.io/otel/sdk/metric v1.45.0 h1:oVFszMfyj1Am6s24Vtc7wBb8BKLcwepJj go.opentelemetry.io/otel/sdk/metric v1.45.0/go.mod h1:vUWUxDZvu1WVRj8JA8S0AdhsPrZoDpA2DdZauIh4mDA= go.opentelemetry.io/otel/trace v1.45.0 h1:l/mP6Uv7oNO7/TblbhpbgMidxhq1uO/rPsikOyVhxag= go.opentelemetry.io/otel/trace v1.45.0/go.mod h1:qoJJA2xNMnxRrdISU/kLtfUH2wNeQbiv+jhs/CxI8bc= -go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= -go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= golang.org/x/net v0.60.0 h1:79p50tfZlm0J9YfoDsSi639qSXNGVwEzOPLCxM2FsYU= golang.org/x/net v0.60.0/go.mod h1:2DA/G1UfVbCpQPeWTmMPGY7Cs2PkBkwu743bVX5PIVg= +golang.org/x/sync v0.23.0 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk= +golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0= golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo= golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og= golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI= @@ -73,5 +73,6 @@ google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/hack/setup-e2e-kind.sh b/hack/setup-e2e-kind.sh index e9f7fbcc..22cfbc90 100755 --- a/hack/setup-e2e-kind.sh +++ b/hack/setup-e2e-kind.sh @@ -261,13 +261,17 @@ AX_DEFAULT_TASK_IMAGE="$(build_and_push "${REGISTRY_ADDR}/ax-task-runner:e2e" \ export AX_SERVER_IMAGE AX_DEFAULT_TASK_IMAGE printf ' AX_SERVER_IMAGE=%s\n AX_DEFAULT_TASK_IMAGE=%s\n' "${AX_SERVER_IMAGE}" "${AX_DEFAULT_TASK_IMAGE}" -# 5. Deploy Redis and ax-server. Delete existing deployments first so every -# run starts from empty Redis state. -log "Deploying ax-redis and ax-server in namespace '${AX_SYSTEM_NAMESPACE}'..." -k -n "${AX_SYSTEM_NAMESPACE}" delete deployment ax-redis ax-server --ignore-not-found -render_manifest redis.yaml | k apply -f - +# 5. Deploy Postgres and ax-server. Delete existing workloads and Postgres's +# volume first so every run starts from an empty database. +log "Deploying ax-postgres and ax-server in namespace '${AX_SYSTEM_NAMESPACE}'..." +k -n "${AX_SYSTEM_NAMESPACE}" delete deployment ax-server --ignore-not-found +k -n "${AX_SYSTEM_NAMESPACE}" delete statefulset ax-postgres --ignore-not-found --wait +k -n "${AX_SYSTEM_NAMESPACE}" delete pvc -l app.kubernetes.io/name=ax-postgres --ignore-not-found --wait +k -n "${AX_SYSTEM_NAMESPACE}" create secret generic ax-postgres \ + --from-literal=password="$(openssl rand -hex 24)" --dry-run=client -o yaml | k apply -f - +render_manifest postgres.yaml | k apply -f - +k -n "${AX_SYSTEM_NAMESPACE}" rollout status statefulset/ax-postgres --timeout=180s render_manifest ax-server.yaml | k apply -f - -k -n "${AX_SYSTEM_NAMESPACE}" rollout status deployment/ax-redis --timeout=120s k -n "${AX_SYSTEM_NAMESPACE}" rollout status deployment/ax-server --timeout=120s # 6. Persist the environment for hack/run-e2e.sh and the Go test harness. diff --git a/internal/controller/reconciler.go b/internal/controller/reconciler.go index 074a4810..aa261e3b 100644 --- a/internal/controller/reconciler.go +++ b/internal/controller/reconciler.go @@ -127,12 +127,12 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws // 1. Ensure the Atespace exists in Substrate if err := r.client.EnsureAtespace(ctx, atespace); err != nil { if substrate.IsTransientError(err) { - r.setNotReady(task, "SubstrateUnavailable", err.Error(), now) - task.Status.Phase = "Pending" + r.setNotReady(task, v1alpha1.ReasonSubstrateUnavailable, err.Error(), now) + task.Status.Phase = v1alpha1.PhasePending return task, fmt.Errorf("ensuring atespace: %w", err) } - r.setCondition(task, "Ready", "False", "AtespaceCreationFailed", err.Error(), now) - task.Status.Phase = "Failed" + r.setNotReady(task, v1alpha1.ReasonAtespaceCreationFailed, err.Error(), now) + task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("ensuring atespace: %w", err) } @@ -185,8 +185,8 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws // Resource limits are enforced by the per-task template, so they are checked // here as well as at apply time: a task must never run without limits it asked for. if err := v1alpha1.ValidateResources(task.Spec.Resources); err != nil { - r.setNotReady(task, "InvalidResources", err.Error(), now) - task.Status.Phase = "Failed" + r.setNotReady(task, v1alpha1.ReasonInvalidResources, err.Error(), now) + task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("validating resources: %w", err) } resources := substrate.ResourceLimits(task.Spec.Resources) @@ -201,16 +201,16 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws tmpl, err := r.client.EnsureActorTemplateWithImage(ctx, templateAtespace, templateName, atespace, customTemplateName, task.Spec.Image, resources, extraEnv) if err != nil { if substrate.IsTransientError(err) { - r.setNotReady(task, "SubstrateUnavailable", err.Error(), now) - task.Status.Phase = "Pending" + r.setNotReady(task, v1alpha1.ReasonSubstrateUnavailable, err.Error(), now) + task.Status.Phase = v1alpha1.PhasePending return task, fmt.Errorf("ensuring actor template: %w", err) } // Falling back to the default template would run a different image than // the user asked for, without the limits it asked for, or fail later with // a misleading "actor template not found". Surface the real error on the // task and stop here. - r.setNotReady(task, "TemplateCreationFailed", substrateErrorMessage(err), now) - task.Status.Phase = "Failed" + r.setNotReady(task, v1alpha1.ReasonTemplateCreationFailed, substrateErrorMessage(err), now) + task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("ensuring actor template: %w", err) } if tmpl != nil && tmpl.Metadata != nil { @@ -223,25 +223,24 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws _, err := r.client.EnsureActor(ctx, atespace, actorName, templateAtespace, templateName) if err != nil { if substrate.IsTransientError(err) { - r.setNotReady(task, "SubstrateUnavailable", err.Error(), now) - task.Status.Phase = "Pending" + r.setNotReady(task, v1alpha1.ReasonSubstrateUnavailable, err.Error(), now) + task.Status.Phase = v1alpha1.PhasePending return task, fmt.Errorf("ensuring actor: %w", err) } - r.setNotReady(task, "ActorCreationFailed", err.Error(), now) - task.Status.Phase = "Failed" + r.setNotReady(task, v1alpha1.ReasonActorCreationFailed, err.Error(), now) + task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("ensuring actor: %w", err) } // 4. Suspend or Resume the Actor - // Tasks are suspended by default upon creation until explicitly resumed to "Running". - if task.Status.Phase == "Suspended" || task.Status.Phase == "" { + // Tasks are suspended by default upon creation until explicitly resumed. + if task.Status.Phase == v1alpha1.PhaseSuspended || task.Status.Phase == "" { slog.Info("suspending actor on Substrate", "actor", actorName) if err := r.client.SuspendActor(ctx, atespace, actorName); err != nil { return task, fmt.Errorf("suspending actor: %w", err) } task.Status.WorkerIp = "" - task.Status.Phase = "Suspended" - r.setCondition(task, condReady, "False", "TaskSuspended", "Task is suspended", now) + task.Status.Phase = v1alpha1.PhaseSuspended slog.Info("task successfully suspended", "name", task.Metadata.Name, "actor", actorName) return task, nil } @@ -251,17 +250,17 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws _, workerIP, err := r.client.ResumeActor(ctx, atespace, actorName) if err != nil { if substrate.IsTransientError(err) { - r.setNotReady(task, "SubstrateUnavailable", err.Error(), now) - task.Status.Phase = "Pending" + r.setNotReady(task, v1alpha1.ReasonSubstrateUnavailable, err.Error(), now) + task.Status.Phase = v1alpha1.PhasePending return task, fmt.Errorf("resuming actor: %w", err) } - r.setNotReady(task, "ActorResumeFailed", err.Error(), now) - task.Status.Phase = "Failed" + r.setNotReady(task, v1alpha1.ReasonActorResumeFailed, err.Error(), now) + task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("resuming actor: %w", err) } task.Status.WorkerIp = workerIP - task.Status.Phase = "Running" + task.Status.Phase = v1alpha1.PhaseRunning // Check if workspace setup inside the actor has completed host := workerIP @@ -273,7 +272,7 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws readyURL := fmt.Sprintf("http://%s/readyz?check=workspace", net.JoinHostPort(host, port)) // Workspace setup happens once per task. After it has completed, WorkspaceReady stays // True across suspend/resume cycles, so only poll while it is still initializing. - workspaceReady := r.conditionTrue(task, condWorkspaceReady) + workspaceReady := conditionTrue(task, v1alpha1.ConditionWorkspaceReady) if workerIP != "" && !workspaceReady { // Poll briefly for workspace setup completion pollCtx, cancel := context.WithTimeout(ctx, r.WorkspaceReadyTimeout) @@ -325,15 +324,13 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws } DonePolling: - // The task is Ready only once its actor is running and the workspace inside it is set up. + // The server derives Ready from the actor's state and WorkspaceReady. if workspaceReady { - if !r.conditionTrue(task, condWorkspaceReady) { - r.setCondition(task, condWorkspaceReady, "True", "SetupComplete", fmt.Sprintf("Workspace setup completed at %s", workerIP), time.Now()) + if !conditionTrue(task, v1alpha1.ConditionWorkspaceReady) { + setCondition(task, v1alpha1.ConditionWorkspaceReady, "True", v1alpha1.ReasonSetupComplete, fmt.Sprintf("Workspace setup completed at %s", workerIP), time.Now()) } - r.setCondition(task, condReady, "True", "TaskRunning", "Task is running and its workspace is ready", time.Now()) } else { - r.setCondition(task, condWorkspaceReady, "False", "Initializing", fmt.Sprintf("Workspace is initializing at %s", workerIP), time.Now()) - r.setCondition(task, condReady, "False", "WorkspaceInitializing", "Waiting for workspace setup to complete", time.Now()) + setCondition(task, v1alpha1.ConditionWorkspaceReady, "False", v1alpha1.ReasonInitializing, fmt.Sprintf("Workspace is initializing at %s", workerIP), time.Now()) } slog.Info("task successfully reconciled and running", @@ -347,23 +344,14 @@ DonePolling: return task, nil } -// Condition types reported on Task status. -const ( - // condReady reports whether the task as a whole is ready to do work: its actor is - // running and the workspace inside it has finished setting up. - condReady = "Ready" - // condWorkspaceReady reports whether the workspace inside the actor has finished setting up. - condWorkspaceReady = "WorkspaceReady" -) - // setNotReady marks the task's Ready condition False. WorkspaceReady is left untouched: // workspace setup is a one-time step whose result outlives actor failures and suspends. func (r *TaskReconciler) setNotReady(task *v1alpha1.Task, reason, message string, t time.Time) { - r.setCondition(task, condReady, "False", reason, message, t) + setCondition(task, v1alpha1.ConditionReady, "False", reason, message, t) } // conditionTrue reports whether the task currently has the given condition with status True. -func (r *TaskReconciler) conditionTrue(task *v1alpha1.Task, condType string) bool { +func conditionTrue(task *v1alpha1.Task, condType string) bool { if task.Status == nil { return false } @@ -392,7 +380,7 @@ func substrateErrorMessage(err error) string { return err.Error() } -func (r *TaskReconciler) setCondition(task *v1alpha1.Task, condType, status, reason, message string, t time.Time) { +func setCondition(task *v1alpha1.Task, condType, status, reason, message string, t time.Time) { ts := timestamppb.New(t) for i, c := range task.Status.Conditions { if c.Type == condType { diff --git a/internal/controller/reconciler_test.go b/internal/controller/reconciler_test.go index 1ca00001..e7f39ff5 100644 --- a/internal/controller/reconciler_test.go +++ b/internal/controller/reconciler_test.go @@ -440,7 +440,7 @@ func TestTaskReconciler_SuspendError(t *testing.T) { Status: &v1alpha1.TaskStatus{ Phase: "Suspended", WorkerIp: workerIp, - Conditions: []*v1alpha1.Condition{{Type: "Ready", Status: "True", Reason: "TaskRunning"}}, + Conditions: []*v1alpha1.Condition{{Type: "Ready", Status: "True", Reason: v1alpha1.ReasonTaskRunning}}, }, } @@ -451,7 +451,7 @@ func TestTaskReconciler_SuspendError(t *testing.T) { if reconciled.GetStatus().GetWorkerIp() != workerIp { t.Errorf("expected worker IP %q to be preserved, got %q", workerIp, reconciled.GetStatus().GetWorkerIp()) } - assertCondition(t, reconciled, "Ready", "True", "TaskRunning") + assertCondition(t, reconciled, "Ready", "True", v1alpha1.ReasonTaskRunning) } // TestTaskReconciler_TemplateCreationFailed covers google/ax#366: when Substrate @@ -508,10 +508,10 @@ func TestTaskReconciler_TemplateCreationFailed(t *testing.T) { t.Errorf("expected error to carry the Substrate message, got %v", err) } - if reconciled.Status.Phase != "Failed" { + if reconciled.Status.Phase != v1alpha1.PhaseFailed { t.Errorf("expected phase 'Failed', got %q", reconciled.Status.Phase) } - assertCondition(t, reconciled, "Ready", "False", "TemplateCreationFailed") + assertCondition(t, reconciled, "Ready", "False", v1alpha1.ReasonTemplateCreationFailed) for _, c := range reconciled.Status.Conditions { if c.Type != "Ready" { continue @@ -579,10 +579,10 @@ func TestTaskReconciler_TemplateCreationTransient(t *testing.T) { if err == nil { t.Fatal("expected Reconcile to return an error on transient template failure") } - if reconciled.Status.Phase != "Pending" { + if reconciled.Status.Phase != v1alpha1.PhasePending { t.Errorf("expected phase 'Pending', got %q", reconciled.Status.Phase) } - assertCondition(t, reconciled, "Ready", "False", "SubstrateUnavailable") + assertCondition(t, reconciled, "Ready", "False", v1alpha1.ReasonSubstrateUnavailable) if len(mockSrv.createdActors) != 0 { t.Errorf("expected no actor to be created, got %v", mockSrv.createdActors) } @@ -650,25 +650,23 @@ func TestTaskReconciler_WorkspaceReady(t *testing.T) { }, } - // Case 1: Worker not responding on readyz -> WorkspaceReady=False and Ready=False. + // Case 1: Worker not responding on readyz -> WorkspaceReady=False. reconciled, err := reconciler.Reconcile(ctx, task, nil) if err != nil { t.Fatalf("Reconcile failed: %v", err) } - assertCondition(t, reconciled, "WorkspaceReady", "False", "Initializing") - assertCondition(t, reconciled, "Ready", "False", "WorkspaceInitializing") + assertCondition(t, reconciled, "WorkspaceReady", "False", v1alpha1.ReasonInitializing) - // Case 2: Worker readyz endpoint succeeds -> WorkspaceReady=True and Ready=True. + // Case 2: Worker readyz endpoint succeeds -> WorkspaceReady=True. mockSrv.workerIP = httpLis.Addr().String() reconciledReady, err := reconciler.Reconcile(ctx, task, nil) if err != nil { t.Fatalf("Reconcile with ready worker failed: %v", err) } - assertCondition(t, reconciledReady, "WorkspaceReady", "True", "SetupComplete") - assertCondition(t, reconciledReady, "Ready", "True", "TaskRunning") + assertCondition(t, reconciledReady, "WorkspaceReady", "True", v1alpha1.ReasonSetupComplete) - // Case 3: Suspending the task -> Ready=False (TaskSuspended), but the workspace was - // already initialized so WorkspaceReady stays True. + // Case 3: Suspending the task -> the workspace was already initialized, so + // WorkspaceReady stays True. task = reconciledReady task.Status.Phase = "Suspended" reconciledSuspended, err := reconciler.Reconcile(ctx, task, nil) @@ -678,11 +676,10 @@ func TestTaskReconciler_WorkspaceReady(t *testing.T) { if reconciledSuspended.Status.Phase != "Suspended" { t.Errorf("expected phase Suspended, got %s", reconciledSuspended.Status.Phase) } - assertCondition(t, reconciledSuspended, "Ready", "False", "TaskSuspended") - assertCondition(t, reconciledSuspended, "WorkspaceReady", "True", "SetupComplete") + assertCondition(t, reconciledSuspended, "WorkspaceReady", "True", v1alpha1.ReasonSetupComplete) // Case 4: Resuming with the worker unreachable -> the reconciler trusts the recorded - // WorkspaceReady instead of re-polling, so the task is Ready again immediately. + // WorkspaceReady instead of re-polling. mockSrv.workerIP = "127.0.0.1:1" task = reconciledSuspended task.Status.Phase = "Running" @@ -690,8 +687,7 @@ func TestTaskReconciler_WorkspaceReady(t *testing.T) { if err != nil { t.Fatalf("Reconcile with resume failed: %v", err) } - assertCondition(t, reconciledResumed, "WorkspaceReady", "True", "SetupComplete") - assertCondition(t, reconciledResumed, "Ready", "True", "TaskRunning") + assertCondition(t, reconciledResumed, "WorkspaceReady", "True", v1alpha1.ReasonSetupComplete) if got := len(mockSrv.actorTemplates); got != 1 { t.Fatalf("readiness updates and suspend/resume created %d templates, want 1", got) } @@ -845,10 +841,10 @@ func TestTaskReconciler_ResourceLimitsFailurePaths(t *testing.T) { if err == nil { t.Errorf("Reconcile(%v) succeeded, want error", reqs) } - if reconciled.Status.Phase != "Failed" { + if reconciled.Status.Phase != v1alpha1.PhaseFailed { t.Errorf("Reconcile(%v) phase = %q, want Failed", reqs, reconciled.Status.Phase) } - assertCondition(t, reconciled, "Ready", "False", "InvalidResources") + assertCondition(t, reconciled, "Ready", "False", v1alpha1.ReasonInvalidResources) } if len(mockSrv.createdTemplates) != 0 || len(mockSrv.createdActors) != 0 || len(mockSrv.resumedActors) != 0 { t.Errorf("invalid limits reached Substrate: templates=%d actors=%v resumed=%v", @@ -868,10 +864,10 @@ func TestTaskReconciler_ResourceLimitsFailurePaths(t *testing.T) { if err == nil { t.Fatal("Reconcile succeeded although the template with limits was rejected") } - if reconciled.Status.Phase != "Failed" { + if reconciled.Status.Phase != v1alpha1.PhaseFailed { t.Errorf("phase = %q, want Failed", reconciled.Status.Phase) } - assertCondition(t, reconciled, "Ready", "False", "TemplateCreationFailed") + assertCondition(t, reconciled, "Ready", "False", v1alpha1.ReasonTemplateCreationFailed) if len(mockSrv.createdActors) != 0 || len(mockSrv.resumedActors) != 0 { t.Errorf("task ran on a fallback template without its limits: actors=%v resumed=%v", mockSrv.createdActors, mockSrv.resumedActors) } @@ -889,10 +885,10 @@ func TestTaskReconciler_ResourceLimitsFailurePaths(t *testing.T) { } // Falling back would run a different image than the task asked for, so the // template error is terminal regardless of whether limits are set (#366). - if reconciled.Status.Phase != "Failed" { + if reconciled.Status.Phase != v1alpha1.PhaseFailed { t.Errorf("phase = %q, want Failed", reconciled.Status.Phase) } - assertCondition(t, reconciled, "Ready", "False", "TemplateCreationFailed") + assertCondition(t, reconciled, "Ready", "False", v1alpha1.ReasonTemplateCreationFailed) if len(mockSrv.createdActors) != 0 { t.Errorf("task ran on a fallback template: actors=%v", mockSrv.createdActors) } @@ -1125,10 +1121,10 @@ func TestTaskReconciler_TransientSubstrateUnavailable(t *testing.T) { if err == nil { t.Fatal("expected error from Reconcile when EnsureAtespace fails, got nil") } - if reconciled.GetStatus().GetPhase() != "Pending" { + if reconciled.GetStatus().GetPhase() != v1alpha1.PhasePending { t.Errorf("expected Phase='Pending' on transient EnsureAtespace error, got %q", reconciled.GetStatus().GetPhase()) } - assertCondition(t, reconciled, "Ready", "False", "SubstrateUnavailable") + assertCondition(t, reconciled, "Ready", "False", v1alpha1.ReasonSubstrateUnavailable) mockSrv.createAtespaceErr = nil // 2. Transient error during EnsureActor leaves Task in Phase="Pending" with Reason="SubstrateUnavailable". @@ -1141,10 +1137,10 @@ func TestTaskReconciler_TransientSubstrateUnavailable(t *testing.T) { if err == nil { t.Fatal("expected error from Reconcile when EnsureActor fails, got nil") } - if reconciled2.GetStatus().GetPhase() != "Pending" { + if reconciled2.GetStatus().GetPhase() != v1alpha1.PhasePending { t.Errorf("expected Phase='Pending' on transient EnsureActor error, got %q", reconciled2.GetStatus().GetPhase()) } - assertCondition(t, reconciled2, "Ready", "False", "SubstrateUnavailable") + assertCondition(t, reconciled2, "Ready", "False", v1alpha1.ReasonSubstrateUnavailable) mockSrv.createActorErr = nil // 3. Transient error during ResumeActor leaves Task in Phase="Pending" with Reason="SubstrateUnavailable", @@ -1159,10 +1155,10 @@ func TestTaskReconciler_TransientSubstrateUnavailable(t *testing.T) { if err == nil { t.Fatal("expected error from Reconcile when ResumeActor fails, got nil") } - if reconciled3.GetStatus().GetPhase() != "Pending" { + if reconciled3.GetStatus().GetPhase() != v1alpha1.PhasePending { t.Errorf("expected Phase='Pending' on transient ResumeActor error, got %q", reconciled3.GetStatus().GetPhase()) } - assertCondition(t, reconciled3, "Ready", "False", "SubstrateUnavailable") + assertCondition(t, reconciled3, "Ready", "False", v1alpha1.ReasonSubstrateUnavailable) // Recovery on retry: mockSrv.resumeActorErr = nil @@ -1185,7 +1181,7 @@ func TestTaskReconciler_TransientSubstrateUnavailable(t *testing.T) { if err == nil { t.Fatal("expected error on InvalidArgument, got nil") } - if reconciled4.GetStatus().GetPhase() != "Failed" { + if reconciled4.GetStatus().GetPhase() != v1alpha1.PhaseFailed { t.Errorf("expected Phase='Failed' on permanent error, got %q", reconciled4.GetStatus().GetPhase()) } mockSrv.createActorErr = nil diff --git a/internal/lock/lock.go b/internal/lock/lock.go index dd1fb215..5b986c3d 100644 --- a/internal/lock/lock.go +++ b/internal/lock/lock.go @@ -16,22 +16,24 @@ package lock import ( "context" - "crypto/rand" - "encoding/hex" "errors" "fmt" - "log/slog" "sync" - "time" - - "github.com/redis/go-redis/v9" ) +// TODO(jbd): Move Locker interface and implementations +// to storage. + var ( ErrLockFailed = errors.New("failed to acquire lock") ) // Locker provides exclusive locking per resource kind, atespace, and name. +// +// When ax-server runs with more than one replica, every replica must use a +// Locker backed by the same shared system. Distributed implementations are +// provided by storage packages (for example postgres.Store) so that locks live +// next to the data they protect. type Locker interface { // Lock acquires an exclusive lock for the specified resource. // It blocks until the lock is acquired or ctx is cancelled. @@ -115,226 +117,3 @@ func (m *MemoryLocker) Lock(ctx context.Context, kind, atespace, name string) (f } 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. -// -// # Lease renewal -// -// A held lock is kept alive by a background goroutine that extends its TTL -// every TTL/3 (only while the key still holds this holder's token) until the -// unlock function is called. The TTL therefore does not bound how long a lock -// may be held; it only bounds how long a lock can be orphaned after its holder -// crashes or loses connectivity to Redis. -// -// We deliberately renew a short lease instead of using a single long TTL: -// -// - The server does not control how long a critical section runs. Handlers -// hold the lock across Substrate calls (e.g. snapshot/restore of a microVM -// during SuspendTask/ResumeTask) that run under the caller's request -// context, so their duration is set by the client's deadline, not by -// ax-server. A fixed TTL would have to exceed the longest deadline any -// client might use, and would silently allow concurrent holders whenever -// an operation ran past it. -// - The TTL is the recovery time after a crash. A long TTL (minutes) would -// wedge the affected task/workspace/model for that long if ax-server died -// mid-operation; a short renewed lease (30s by default) keeps recovery -// fast. -// - One locker guards every resource kind and operation, so renewal avoids -// having to retune the TTL whenever a slower operation is added. -// -// Limitation: renewal is best-effort and there is no fencing token. If renewals -// keep failing for longer than the TTL (e.g. a prolonged Redis outage), the -// lease can expire while the holder is still working, and the holder is not -// interrupted; the loss is only logged. The alternative considered was a long -// TTL combined with a server-side deadline shorter than it on every locked -// operation, which avoids the goroutine but trades away fast crash recovery. -type RedisLocker struct { - client *redis.Client - ttl time.Duration - fallbackInterval time.Duration -} - -type RedisLockerOptions struct { - // TTL is the lease duration of a lock key in Redis (default 30s). Held - // locks are renewed every TTL/3, so this does not limit how long a lock - // may be held; it bounds how long a lock survives a crashed holder. See - // the RedisLocker doc for why a short renewed lease is preferred over a - // long TTL. - TTL time.Duration - // FallbackInterval is how often waiters retry acquisition in case a - // release notification is missed or a lease expires without a publish - // (default 1s). - 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 -`) - -// renewScript extends the lock TTL only if it is still held by the caller's token. -var renewScript = redis.NewScript(` -if redis.call("get", KEYS[1]) == ARGV[1] then - return redis.call("pexpire", KEYS[1], ARGV[2]) -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 is called once the lock is acquired. It starts the lease - // renewal goroutine (see "Lease renewal" on RedisLocker) and returns an - // unlock function that stops renewal before releasing the key. - makeUnlock := func() func() { - renewCtx, cancelRenew := context.WithCancel(context.Background()) - done := make(chan struct{}) - go func() { - defer close(done) - r.renew(renewCtx, key, token) - }() - - var once sync.Once - return func() { - once.Do(func() { - cancelRenew() - <-done - 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 - } - } -} - -// renew periodically extends the TTL of the lock identified by key/token until -// ctx is canceled or the lock is found to be no longer held by token. -func (r *RedisLocker) renew(ctx context.Context, key, token string) { - interval := r.ttl / 3 - if interval <= 0 { - interval = time.Millisecond - } - ticker := time.NewTicker(interval) - defer ticker.Stop() - - ttlMillis := r.ttl.Milliseconds() - if ttlMillis <= 0 { - ttlMillis = 1 - } - - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - } - - renewCtx, cancel := context.WithTimeout(ctx, interval) - res, err := renewScript.Run(renewCtx, r.client, []string{key}, token, ttlMillis).Int() - cancel() - if ctx.Err() != nil { - return - } - if err != nil { - // Transient Redis error: keep trying on the next tick while the - // remaining TTL still covers us. - slog.Warn("renewing lock lease", "key", key, "error", err) - continue - } - if res == 0 { - // The lock expired or was taken over; there is nothing left to renew. - slog.Error("lock lease lost before release", "key", key) - return - } - } -} diff --git a/internal/lock/lock_test.go b/internal/lock/lock_test.go index c887cd18..4128ce74 100644 --- a/internal/lock/lock_test.go +++ b/internal/lock/lock_test.go @@ -16,16 +16,12 @@ package lock_test import ( "context" - "errors" - "fmt" - "os" "sync" "sync/atomic" "testing" "time" "github.com/google/ax/internal/lock" - "github.com/redis/go-redis/v9" ) func TestMemoryLocker_SerializesSameResource(t *testing.T) { @@ -98,51 +94,3 @@ func TestMemoryLocker_Timeout(t *testing.T) { t.Fatalf("expected timeout error, got nil") } } - -// newTestRedisClient returns a client for the Redis at $REDIS_ADDR, skipping the -// test when it is unset. -func newTestRedisClient(t *testing.T) *redis.Client { - t.Helper() - addr := os.Getenv("REDIS_ADDR") - if addr == "" { - t.Skip("REDIS_ADDR not set; skipping Redis-backed lock test") - } - client := redis.NewClient(&redis.Options{Addr: addr}) - t.Cleanup(func() { _ = client.Close() }) - if err := client.Ping(context.Background()).Err(); err != nil { - t.Skipf("redis at %s unreachable: %v", addr, err) - } - return client -} - -// TestRedisLocker_RenewsLeaseWhileHeld verifies that a held lock outlives its -// TTL (because it is renewed) and that waiters acquire it only after release. -func TestRedisLocker_RenewsLeaseWhileHeld(t *testing.T) { - client := newTestRedisClient(t) - ttl := 300 * time.Millisecond - locker := lock.NewRedisLocker(client, lock.RedisLockerOptions{TTL: ttl, FallbackInterval: 20 * time.Millisecond}) - ctx := context.Background() - name := fmt.Sprintf("renew-%d", time.Now().UnixNano()) - - unlock, err := locker.Lock(ctx, "task", "default", name) - if err != nil { - t.Fatalf("lock: %v", err) - } - - // Hold the lock for several TTLs; a contender must not get it meanwhile. - waitCtx, cancel := context.WithTimeout(ctx, 4*ttl) - defer cancel() - if _, err := locker.Lock(waitCtx, "task", "default", name); !errors.Is(err, context.DeadlineExceeded) { - t.Fatalf("expected contender deadline while lease was held, got %v", err) - } - - unlock() - - acquireCtx, cancel2 := context.WithTimeout(ctx, ttl) - defer cancel2() - unlock2, err := locker.Lock(acquireCtx, "task", "default", name) - if err != nil { - t.Fatalf("lock after release: %v", err) - } - unlock2() -} diff --git a/internal/server/actor_status_test.go b/internal/server/actor_status_test.go new file mode 100644 index 00000000..c7159c2c --- /dev/null +++ b/internal/server/actor_status_test.go @@ -0,0 +1,128 @@ +// 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 server + +import ( + "errors" + "testing" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/google/ax/pkg/apis/v1alpha1" +) + +func testActor(name string, state ateapipb.ActorState, workerIP string) *ateapipb.Actor { + a := &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "default", Name: name}, + Status: &ateapipb.ActorStatus{State: state}, + } + if workerIP != "" { + a.Status.WorkerAssignment = &ateapipb.WorkerAssignment{WorkerPodIps: []string{workerIP}} + } + return a +} + +func storedTask(name, phase string, conds ...*v1alpha1.Condition) *v1alpha1.Task { + return &v1alpha1.Task{ + Metadata: &v1alpha1.ObjectMeta{Name: name, Atespace: "default"}, + Status: &v1alpha1.TaskStatus{Phase: phase, Conditions: conds}, + } +} + +func readyCondition(task *v1alpha1.Task) *v1alpha1.Condition { + for _, c := range task.GetStatus().GetConditions() { + if c.GetType() == "Ready" { + return c + } + } + return nil +} + +func TestSetActorStatus(t *testing.T) { + workspaceReady := &v1alpha1.Condition{Type: "WorkspaceReady", Status: "True"} + crashed := testActor("crashed", ateapipb.ActorState_ACTOR_STATE_CRASHED, "") + crashed.Status.Crash = &ateapipb.ActorCrash{Message: "OOMKilled"} + + tests := []struct { + task *v1alpha1.Task + actor *ateapipb.Actor + wantPhase, wantIP, wantActor string + wantReady, wantReadyReason string + }{ + {storedTask("running", "", workspaceReady), testActor("running", ateapipb.ActorState_ACTOR_STATE_RUNNING, "10.0.0.5"), "Running", "10.0.0.5", "running", "True", v1alpha1.ReasonTaskRunning}, + {storedTask("setting-up", ""), testActor("setting-up", ateapipb.ActorState_ACTOR_STATE_RUNNING, "10.0.0.6"), "Running", "10.0.0.6", "setting-up", "False", v1alpha1.ReasonWorkspaceInitializing}, + {storedTask("suspended", "", workspaceReady), testActor("suspended", ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ""), "Suspended", "", "suspended", "False", v1alpha1.PhaseSuspended}, + {storedTask("crashed", ""), crashed, v1alpha1.PhaseFailed, "", "crashed", "False", v1alpha1.PhaseFailed}, + {storedTask("reverting", ""), testActor("reverting", ateapipb.ActorState_ACTOR_STATE_REVERTING, ""), v1alpha1.PhaseResuming, "", "reverting", "False", v1alpha1.PhaseResuming}, + {storedTask("suspending", ""), testActor("suspending", ateapipb.ActorState_ACTOR_STATE_SUSPENDING, ""), v1alpha1.PhaseSuspending, "", "suspending", "False", v1alpha1.PhaseSuspending}, + {storedTask("deleting", ""), testActor("deleting", ateapipb.ActorState_ACTOR_STATE_DELETING, ""), v1alpha1.PhaseTerminating, "", "deleting", "False", v1alpha1.PhaseTerminating}, + {storedTask("odd", ""), testActor("odd", ateapipb.ActorState_ACTOR_STATE_UNSPECIFIED, ""), v1alpha1.PhaseUnknown, "", "odd", "Unknown", v1alpha1.PhaseUnknown}, + // The actor was removed outside AX. + {storedTask("gone", ""), nil, v1alpha1.PhaseFailed, "", "", "False", v1alpha1.ReasonActorNotFound}, + // Phases that describe AX's own work are kept, and so is their Ready reason. + {storedTask("failed", v1alpha1.PhaseFailed, &v1alpha1.Condition{Type: "Ready", Status: "False", Reason: v1alpha1.ReasonActorResumeFailed}), testActor("failed", ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ""), v1alpha1.PhaseFailed, "", "failed", "False", v1alpha1.ReasonActorResumeFailed}, + {storedTask("not-yet", v1alpha1.PhasePending), nil, v1alpha1.PhasePending, "", "", "", ""}, + } + for _, tt := range tests { + t.Run(tt.task.GetMetadata().GetName(), func(t *testing.T) { + setActorStatus(tt.task, tt.actor, nil) + st := tt.task.GetStatus() + if st.GetPhase() != tt.wantPhase || st.GetWorkerIp() != tt.wantIP || st.GetActor() != tt.wantActor { + t.Errorf("phase=%q workerIP=%q actor=%q, want %q %q %q", st.GetPhase(), st.GetWorkerIp(), st.GetActor(), tt.wantPhase, tt.wantIP, tt.wantActor) + } + if tt.wantReady == "" { + return + } + if c := readyCondition(tt.task); c.GetStatus() != tt.wantReady || c.GetReason() != tt.wantReadyReason { + t.Errorf("Ready = %v, want %s/%s", c, tt.wantReady, tt.wantReadyReason) + } + }) + } + + t.Run("crash message", func(t *testing.T) { + task := storedTask("crashed", "") + setActorStatus(task, crashed, nil) + if msg := readyCondition(task).GetMessage(); msg != "OOMKilled" { + t.Errorf("Ready message = %q, want the actor's crash message", msg) + } + }) + + t.Run("stable update time", func(t *testing.T) { + task := storedTask("suspended", "") + setActorStatus(task, testActor("suspended", ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ""), nil) + first := readyCondition(task).GetUpdateTime() + setActorStatus(task, testActor("suspended", ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ""), nil) + if got := readyCondition(task).GetUpdateTime(); got != first { + t.Errorf("Ready update time changed on an unchanged read") + } + }) +} + +func TestSetActorStatus_SubstrateErrorReportsUnknown(t *testing.T) { + err := errors.New("denied") + task := storedTask("a", "") + terminating := storedTask("b", v1alpha1.PhaseTerminating) + setActorStatus(task, nil, err) + setActorStatus(terminating, nil, err) + + if got := task.GetStatus().GetPhase(); got != "Unknown" { + t.Errorf("phase = %q, want Unknown", got) + } + if c := readyCondition(task); c.GetStatus() != "Unknown" || c.GetReason() != v1alpha1.ReasonSubstrateUnavailable { + t.Errorf("Ready = %v, want Unknown/SubstrateUnavailable", c) + } + if got := terminating.GetStatus().GetPhase(); got != v1alpha1.PhaseTerminating { + t.Errorf("phase = %q, want %q to be kept", got, v1alpha1.PhaseTerminating) + } +} diff --git a/internal/server/export_test.go b/internal/server/export_test.go new file mode 100644 index 00000000..dcbd11b5 --- /dev/null +++ b/internal/server/export_test.go @@ -0,0 +1,25 @@ +// 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 server + +import "time" + +// SetWatchPollInterval overrides how often WatchTask re-reads a task and +// returns a function that restores the previous value. +func SetWatchPollInterval(d time.Duration) func() { + prev := watchPollInterval + watchPollInterval = d + return func() { watchPollInterval = prev } +} diff --git a/internal/server/server.go b/internal/server/server.go index ac59496f..38a2c729 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -20,7 +20,9 @@ import ( "log/slog" "net/http" "strings" + "time" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "github.com/google/ax/internal/lock" "github.com/google/ax/internal/store" "github.com/google/ax/internal/substrate" @@ -40,16 +42,26 @@ type Reconciler interface { // Options configures the AX API Server. type Options struct { + // Locker serializes changes to a resource. It must be shared by all + // ax-server replicas; it defaults to an in-process lock.MemoryLocker. Locker lock.Locker Reconciler Reconciler + // Substrate is read for the parts of a task's status owned by its actor, + // which the store never holds. Without it, tasks are returned as stored. + Substrate *substrate.Client } +// watchPollInterval is how often WatchTask re-reads a task, picking up both +// AX's own writes and changes made on Substrate alone (for example a crash). +var watchPollInterval = 2 * time.Second + // Server provides the gRPC API for AX. type Server struct { v1alpha1.UnimplementedAXServer store store.Store locker lock.Locker reconciler Reconciler + substrate *substrate.Client grpcServer *grpc.Server } @@ -68,6 +80,7 @@ func NewServer(s store.Store, opts ...Options) *Server { store: s, locker: locker, reconciler: opt.Reconciler, + substrate: opt.Substrate, grpcServer: grpc.NewServer(), } v1alpha1.RegisterAXServer(srv.grpcServer, srv) @@ -107,7 +120,7 @@ func (s *Server) GetTask(ctx context.Context, req *v1alpha1.GetTaskRequest) (*v1 if atespace == "" { atespace = "default" } - task, err := s.store.GetTask(ctx, atespace, req.Name) + task, err := s.getTask(ctx, atespace, req.Name) 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) @@ -130,7 +143,7 @@ func (s *Server) ListTasks(ctx context.Context, req *v1alpha1.ListTasksRequest) offset = req.Offset } } - tasks, err := s.store.ListTasks(ctx, atespace, limit, offset) + tasks, err := s.listTasks(ctx, atespace, limit, offset) if err != nil { return nil, status.Errorf(codes.Internal, "listing tasks: %v", err) } @@ -162,9 +175,9 @@ func (s *Server) CreateTask(ctx context.Context, req *v1alpha1.CreateTaskRequest } defer unlock() - existing, err := s.store.GetTask(ctx, atespace, taskName) + existing, err := s.getTask(ctx, atespace, taskName) if err == nil { - if existing.GetStatus().GetPhase() != "Pending" || !proto.Equal(existing.GetSpec(), task.GetSpec()) { + if existing.GetStatus().GetPhase() != v1alpha1.PhasePending || !proto.Equal(existing.GetSpec(), task.GetSpec()) { return nil, status.Errorf(codes.FailedPrecondition, "task %s/%s already exists and is immutable", atespace, taskName) } task = existing @@ -179,10 +192,12 @@ func (s *Server) CreateTask(ctx context.Context, req *v1alpha1.CreateTaskRequest if task.Status == nil { task.Status = &v1alpha1.TaskStatus{} } - task.Status.Phase = "Suspended" - if err := s.store.SaveTask(ctx, task); err != nil { + // The task is Pending until its actor exists. New tasks start suspended. + task.Status.Phase = v1alpha1.PhasePending + if err := s.saveTask(ctx, task); err != nil { return nil, status.Errorf(codes.Internal, "saving task: %v", err) } + task.Status.Phase = v1alpha1.PhaseSuspended // Directly reconcile with Substrate if s.reconciler != nil { @@ -193,22 +208,24 @@ func (s *Server) CreateTask(ctx context.Context, req *v1alpha1.CreateTaskRequest task.Status = reconciled.Status } if substrate.IsTransientError(err) { - task.Status.Phase = "Pending" - } else if task.Status.Phase != "Pending" { - task.Status.Phase = "Failed" + task.Status.Phase = v1alpha1.PhasePending + } else if task.Status.Phase != v1alpha1.PhasePending { + task.Status.Phase = v1alpha1.PhaseFailed } - _ = s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status) + _ = s.updateTask(ctx, atespace, taskName, task.Status) if substrate.IsTransientError(err) { return nil, status.Errorf(codes.Unavailable, "provisioning task on substrate: %v", err) } 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) - } + } + if err := s.updateTask(ctx, atespace, taskName, task.Status); err != nil { + return nil, status.Errorf(codes.Internal, "updating task status: %v", err) } + // Report the actor as it is now, as GetTask would. + s.applyActor(ctx, task) return task, nil } @@ -229,7 +246,7 @@ func (s *Server) DeleteTask(ctx context.Context, req *v1alpha1.DeleteTaskRequest } defer unlock() - task, err := s.store.GetTask(ctx, atespace, taskName) + task, err := s.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", taskName, atespace) @@ -242,28 +259,12 @@ func (s *Server) DeleteTask(ctx context.Context, req *v1alpha1.DeleteTaskRequest task.Status = &v1alpha1.TaskStatus{} } task.Status.Phase = v1alpha1.PhaseTerminating - if err := s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status); err != nil { + if err := s.updateTask(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) - } - } - - if err := s.store.DeleteTask(ctx, atespace, taskName); err != nil { - return nil, status.Errorf(codes.Internal, "deleting task record: %v", err) + if err := s.deleteTask(ctx, atespace, taskName); err != nil { + return nil, err } return &v1alpha1.DeleteTaskResponse{}, nil } @@ -285,7 +286,7 @@ func (s *Server) SuspendTask(ctx context.Context, req *v1alpha1.SuspendTaskReque } defer unlock() - task, err := s.store.GetTask(ctx, atespace, taskName) + task, err := s.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", taskName, atespace) @@ -295,7 +296,7 @@ func (s *Server) SuspendTask(ctx context.Context, req *v1alpha1.SuspendTaskReque if task.Status == nil { task.Status = &v1alpha1.TaskStatus{} } - task.Status.Phase = "Suspended" + task.Status.Phase = v1alpha1.PhaseSuspended if s.reconciler != nil { reconciled, err := s.reconciler.Reconcile(ctx, task, s.fetchWorkspace(ctx, atespace, task)) @@ -306,15 +307,17 @@ func (s *Server) SuspendTask(ctx context.Context, req *v1alpha1.SuspendTaskReque 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 { + if err := s.updateTask(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 { + if err := s.updateTask(ctx, atespace, taskName, task.Status); err != nil { return nil, status.Errorf(codes.Internal, "suspending task: %v", err) } } + // Report the actor as it is now, as GetTask would. + s.applyActor(ctx, task) return task, nil } @@ -335,7 +338,7 @@ func (s *Server) ResumeTask(ctx context.Context, req *v1alpha1.ResumeTaskRequest } defer unlock() - task, err := s.store.GetTask(ctx, atespace, taskName) + task, err := s.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", taskName, atespace) @@ -345,7 +348,7 @@ func (s *Server) ResumeTask(ctx context.Context, req *v1alpha1.ResumeTaskRequest if task.Status == nil { task.Status = &v1alpha1.TaskStatus{} } - task.Status.Phase = "Running" + task.Status.Phase = v1alpha1.PhaseRunning if s.reconciler != nil { reconciled, err := s.reconciler.Reconcile(ctx, task, s.fetchWorkspace(ctx, atespace, task)) @@ -354,29 +357,256 @@ func (s *Server) ResumeTask(ctx context.Context, req *v1alpha1.ResumeTaskRequest task.Status = reconciled.Status } if substrate.IsTransientError(err) { - task.Status.Phase = "Pending" - } else if task.Status.Phase != "Pending" { - task.Status.Phase = "Failed" + task.Status.Phase = v1alpha1.PhasePending + } else if task.Status.Phase != v1alpha1.PhasePending { + task.Status.Phase = v1alpha1.PhaseFailed } - _ = s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status) + _ = s.updateTask(ctx, atespace, taskName, task.Status) if substrate.IsTransientError(err) { return nil, status.Errorf(codes.Unavailable, "resuming task on substrate: %v", err) } 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 { + if err := s.updateTask(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 { + if err := s.updateTask(ctx, atespace, taskName, task.Status); err != nil { return nil, status.Errorf(codes.Internal, "resuming task: %v", err) } } + // Report the actor as it is now, as GetTask would. + s.applyActor(ctx, task) + return task, nil +} + +// Task storage and Substrate. +// +// A task's actor on Agent Substrate is the source of truth for whether the task +// is running, where, and whether it is healthy. The store keeps only what AX +// itself knows. +// +// Handlers never call s.store for tasks; they use the methods below. saveTask +// and updateTask drop the actor-owned status before writing. getTask and +// listTasks fill it in from Substrate. + +// getTask returns a task with its Substrate-owned status filled in. +func (s *Server) getTask(ctx context.Context, atespace, name string) (*v1alpha1.Task, error) { + task, err := s.store.GetTask(ctx, atespace, name) + if err != nil { + return nil, err + } + s.applyActor(ctx, task) return task, nil } +// listTasks returns a page of tasks with their Substrate-owned status filled +// in, using one ListActors call per atespace. +func (s *Server) listTasks(ctx context.Context, atespace string, limit, offset int64) ([]*v1alpha1.Task, error) { + tasks, err := s.store.ListTasks(ctx, atespace, limit, offset) + if err != nil || s.substrate == nil { + return tasks, err + } + byAtespace := make(map[string][]*v1alpha1.Task) + for _, t := range tasks { + ns := t.GetMetadata().GetAtespace() + byAtespace[ns] = append(byAtespace[ns], t) + } + for ns, group := range byAtespace { + actors, err := s.substrate.ListActors(ctx, ns) + byName := make(map[string]*ateapipb.Actor, len(actors)) + for _, a := range actors { + byName[a.GetMetadata().GetName()] = a + } + for _, t := range group { + setActorStatus(t, byName[t.GetMetadata().GetName()], err) + } + } + return tasks, nil +} + +// saveTask stores task without its Substrate-owned status. task is not modified. +func (s *Server) saveTask(ctx context.Context, task *v1alpha1.Task) error { + stored := proto.Clone(task).(*v1alpha1.Task) + stored.Status = storedStatus(task.Status) + return s.store.SaveTask(ctx, stored) +} + +// updateTask stores a task's status without its Substrate-owned parts. st is +// not modified. +func (s *Server) updateTask(ctx context.Context, atespace, name string, st *v1alpha1.TaskStatus) error { + return s.store.UpdateTaskStatus(ctx, atespace, name, storedStatus(st)) +} + +// deleteTask removes a task's actor and templates from Substrate, then its +// record from the store. It returns a gRPC status error. +func (s *Server) deleteTask(ctx context.Context, atespace, name string) error { + if s.reconciler != nil { + if err := s.reconciler.ReconcileDelete(ctx, atespace, name); err != nil { + if errors.Is(err, context.Canceled) { + return status.Error(codes.Canceled, err.Error()) + } + if errors.Is(err, context.DeadlineExceeded) { + return status.Error(codes.DeadlineExceeded, err.Error()) + } + if st, ok := status.FromError(err); ok && st.Code() != codes.Unknown { + return err + } + return status.Errorf(codes.Internal, "cleaning up substrate resources for task %s/%s: %v", atespace, name, err) + } + } + if err := s.store.DeleteTask(ctx, atespace, name); err != nil { + return status.Errorf(codes.Internal, "deleting task record: %v", err) + } + return nil +} + +// applyActor fills in the Substrate-owned status of task from its actor, when +// the server is connected to Substrate. +func (s *Server) applyActor(ctx context.Context, task *v1alpha1.Task) { + if s.substrate == nil { + return + } + actor, err := s.substrate.GetActor(ctx, task.GetMetadata().GetAtespace(), task.GetMetadata().GetName()) + if status.Code(err) == codes.NotFound { + actor, err = nil, nil + } + setActorStatus(task, actor, err) +} + +// setActorStatus overlays the actor's state on the task's status: actor, +// workerIP, and, unless the phase is one of AX's own, the phase and Ready +// condition. actor is nil when the task has no actor; err is set when the actor +// could not be looked up, in which case the task is reported with phase +// Unknown. +func setActorStatus(task *v1alpha1.Task, actor *ateapipb.Actor, err error) { + if task.Status == nil { + task.Status = &v1alpha1.TaskStatus{} + } + st := task.Status + st.Actor = "" + st.WorkerIp = "" + + switch { + case err != nil: + slog.Warn("could not read task actor", "atespace", task.GetMetadata().GetAtespace(), "task", task.GetMetadata().GetName(), "error", err) + if !v1alpha1.IsAXPhase(st.Phase) { + st.Phase = v1alpha1.PhaseUnknown + setReady(task, "Unknown", v1alpha1.ReasonSubstrateUnavailable, err.Error()) + } + case actor == nil: + // A task that finished provisioning always has an actor, so a missing one + // was removed outside AX. + if !v1alpha1.IsAXPhase(st.Phase) { + st.Phase = v1alpha1.PhaseFailed + setReady(task, "False", v1alpha1.ReasonActorNotFound, "The task's actor no longer exists on Substrate") + } + default: + st.Actor = actor.GetMetadata().GetName() + st.WorkerIp = substrate.WorkerIP(actor) + if v1alpha1.IsAXPhase(st.Phase) { + return + } + state := actor.GetStatus().GetState() + st.Phase = actorPhase(state) + switch st.Phase { + case v1alpha1.PhaseRunning: + if conditionTrue(task, v1alpha1.ConditionWorkspaceReady) { + setReady(task, "True", v1alpha1.ReasonTaskRunning, "Task is running and its workspace is ready") + } else { + setReady(task, "False", v1alpha1.ReasonWorkspaceInitializing, "Waiting for workspace setup to complete") + } + case v1alpha1.PhaseUnknown: + setReady(task, "Unknown", v1alpha1.PhaseUnknown, "Actor is in state "+state.String()) + default: + // The task is not ready, and its phase says why. + message := actor.GetStatus().GetCrash().GetMessage() + if message == "" { + message = "Actor is " + strings.ToLower(strings.TrimPrefix(state.String(), "ACTOR_STATE_")) + } + setReady(task, "False", st.Phase, message) + } + } +} + +// actorPhase maps an actor's state to the task's phase. +func actorPhase(state ateapipb.ActorState) string { + switch state { + case ateapipb.ActorState_ACTOR_STATE_RUNNING: + return v1alpha1.PhaseRunning + case ateapipb.ActorState_ACTOR_STATE_RESUMING, ateapipb.ActorState_ACTOR_STATE_REVERTING: + return v1alpha1.PhaseResuming + case ateapipb.ActorState_ACTOR_STATE_SUSPENDING: + return v1alpha1.PhaseSuspending + case ateapipb.ActorState_ACTOR_STATE_SUSPENDED: + return v1alpha1.PhaseSuspended + case ateapipb.ActorState_ACTOR_STATE_CRASHED: + return v1alpha1.PhaseFailed + case ateapipb.ActorState_ACTOR_STATE_DELETING: + return v1alpha1.PhaseTerminating + default: + return v1alpha1.PhaseUnknown + } +} + +// conditionTrue reports whether the task has the given condition with status +// True. +func conditionTrue(task *v1alpha1.Task, condType string) bool { + for _, c := range task.GetStatus().GetConditions() { + if c.GetType() == condType { + return c.GetStatus() == "True" + } + } + return false +} + +// setReady sets the task's Ready condition. Its update time changes only when +// the condition does, so that repeated reads report a stable timestamp. +func setReady(task *v1alpha1.Task, condStatus, reason, message string) { + for _, c := range task.Status.Conditions { + if c.GetType() != v1alpha1.ConditionReady { + continue + } + if c.Status != condStatus || c.Reason != reason || c.Message != message { + c.Status, c.Reason, c.Message = condStatus, reason, message + c.UpdateTime = timestamppb.Now() + } + return + } + task.Status.Conditions = append(task.Status.Conditions, &v1alpha1.Condition{ + Type: v1alpha1.ConditionReady, + Status: condStatus, + Reason: reason, + Message: message, + UpdateTime: timestamppb.Now(), + }) +} + +// storedStatus returns the part of a task's status that AX owns and stores. The +// actor, its worker, and the phase and Ready condition that follow from the +// actor's state are dropped. st is not modified. +func storedStatus(st *v1alpha1.TaskStatus) *v1alpha1.TaskStatus { + out := &v1alpha1.TaskStatus{} + if st != nil { + out = proto.Clone(st).(*v1alpha1.TaskStatus) + } + out.Actor = "" + out.WorkerIp = "" + if !v1alpha1.IsAXPhase(out.Phase) { + out.Phase = "" + conds := out.Conditions[:0] + for _, c := range out.Conditions { + if c.GetType() != v1alpha1.ConditionReady { + conds = append(conds, c) + } + } + out.Conditions = conds + } + return out +} + // fetchWorkspace loads the Workspace the task binds, or nil when the task // binds none or the Workspace does not exist. func (s *Server) fetchWorkspace(ctx context.Context, atespace string, task *v1alpha1.Task) *v1alpha1.Workspace { @@ -400,32 +630,40 @@ func (s *Server) WatchTask(req *v1alpha1.WatchTaskRequest, stream grpc.ServerStr atespace = "default" } ctx := stream.Context() - ch, closer, err := s.store.WatchTask(ctx, atespace, req.Name) - if err != nil { - return status.Errorf(codes.Internal, "watching task: %v", err) + + var last *v1alpha1.Task + send := func(task *v1alpha1.Task, action string) (done bool, err error) { + if last != nil && proto.Equal(last, task) { + return false, nil + } + last = task + if err := stream.Send(&v1alpha1.WatchTaskResponse{Task: task, Action: action}); err != nil { + return true, err + } + phase := task.GetStatus().GetPhase() + return phase == v1alpha1.PhaseFailed, nil } - defer closer.Close() - if initial, err := s.store.GetTask(ctx, atespace, req.Name); err == nil { - if err := stream.Send(&v1alpha1.WatchTaskResponse{Task: initial, Action: "INITIAL"}); err != nil { + if initial, err := s.getTask(ctx, atespace, req.Name); err == nil { + if done, err := send(initial, "INITIAL"); done || err != nil { return err } } + ticker := time.NewTicker(watchPollInterval) + defer ticker.Stop() for { select { case <-ctx.Done(): return ctx.Err() - case task, ok := <-ch: - if !ok { - return nil - } - if err := stream.Send(&v1alpha1.WatchTaskResponse{Task: task, Action: "MODIFIED"}); err != nil { - return err - } - if task.Status != nil && (task.Status.Phase == "Failed" || task.Status.Phase == "Completed") { - return nil - } + case <-ticker.C: + } + task, err := s.getTask(ctx, atespace, req.Name) + if err != nil { + continue + } + if done, err := send(task, "MODIFIED"); done || err != nil { + return err } } } diff --git a/internal/server/server_test.go b/internal/server/server_test.go index dc4be1cd..e8166663 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -20,12 +20,15 @@ import ( "net" "net/http" "net/http/httptest" + "sync" "testing" "time" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "github.com/google/ax/internal/server" "github.com/google/ax/internal/store" "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/codes" @@ -536,7 +539,7 @@ func TestServer_SuspendTask_ReconcileErrorRetainsRunningState(t *testing.T) { Phase: "Running", Actor: "task-suspend-err", WorkerIp: "10.244.1.42", - Conditions: []*v1alpha1.Condition{{Type: "Ready", Status: "True", Reason: "TaskRunning"}}, + Conditions: []*v1alpha1.Condition{{Type: "Ready", Status: "True", Reason: v1alpha1.ReasonTaskRunning}}, }, }); err != nil { t.Fatalf("SaveTask: %v", err) @@ -557,7 +560,7 @@ func TestServer_SuspendTask_ReconcileErrorRetainsRunningState(t *testing.T) { if stored.GetStatus().GetWorkerIp() != "10.244.1.42" { t.Errorf("stored worker IP = %q, want 10.244.1.42", stored.GetStatus().GetWorkerIp()) } - if len(stored.GetStatus().GetConditions()) != 1 || stored.GetStatus().GetConditions()[0].GetReason() != "TaskRunning" { + if len(stored.GetStatus().GetConditions()) != 1 || stored.GetStatus().GetConditions()[0].GetReason() != v1alpha1.ReasonTaskRunning { t.Errorf("stored conditions = %v, want Ready=True (TaskRunning)", stored.GetStatus().GetConditions()) } } @@ -575,7 +578,7 @@ func TestServer_SuspendTask_SuccessUpdatesStoredState(t *testing.T) { Phase: "Running", Actor: "task-suspend-ok", WorkerIp: "10.244.1.42", - Conditions: []*v1alpha1.Condition{{Type: "Ready", Status: "True", Reason: "TaskRunning"}}, + Conditions: []*v1alpha1.Condition{{Type: "Ready", Status: "True", Reason: v1alpha1.ReasonTaskRunning}}, }, }); err != nil { t.Fatalf("SaveTask: %v", err) @@ -589,16 +592,145 @@ func TestServer_SuspendTask_SuccessUpdatesStoredState(t *testing.T) { t.Errorf("response phase = %q, want Suspended", resp.GetStatus().GetPhase()) } - stored, err := srv.GetTask(ctx, &v1alpha1.GetTaskRequest{Name: "task-suspend-ok"}) + stored, err := st.GetTask(ctx, "default", "task-suspend-ok") if err != nil { t.Fatalf("GetTask: %v", err) } - if stored.GetStatus().GetPhase() != "Suspended" { - t.Errorf("stored phase = %q, want Suspended", stored.GetStatus().GetPhase()) + assertNoActorState(t, stored) +} + +// assertNoActorState checks that a stored task holds none of the status that +// belongs to its actor on Substrate. +func assertNoActorState(t *testing.T, stored *v1alpha1.Task) { + t.Helper() + st := stored.GetStatus() + if st.GetPhase() != "" || st.GetActor() != "" || st.GetWorkerIp() != "" { + t.Errorf("stored status has phase=%q actor=%q workerIP=%q, want all empty", st.GetPhase(), st.GetActor(), st.GetWorkerIp()) + } + for _, c := range st.GetConditions() { + if c.GetType() == "Ready" { + t.Errorf("stored status has a Ready condition %v, want none", c) + } + } +} + +// fakeSubstrate stands in for Substrate's control API: every actor named in +// names exists in the default atespace, in the given state. +type fakeSubstrate struct { + ateapipb.UnimplementedControlServer + mu sync.Mutex + names []string + state ateapipb.ActorState + workerIP string +} + +func (f *fakeSubstrate) setState(state ateapipb.ActorState) { + f.mu.Lock() + defer f.mu.Unlock() + f.state = state +} + +func (f *fakeSubstrate) actor(name string) *ateapipb.Actor { + a := &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "default", Name: name}, + Status: &ateapipb.ActorStatus{State: f.state}, + } + if f.workerIP != "" { + a.Status.WorkerAssignment = &ateapipb.WorkerAssignment{WorkerPodIps: []string{f.workerIP}} + } + return a +} + +func (f *fakeSubstrate) GetActor(_ context.Context, req *ateapipb.GetActorRequest) (*ateapipb.Actor, error) { + f.mu.Lock() + defer f.mu.Unlock() + for _, name := range f.names { + if req.GetActor().GetAtespace() == "default" && req.GetActor().GetName() == name { + return f.actor(name), nil + } + } + return nil, status.Error(codes.NotFound, "actor not found") +} + +func (f *fakeSubstrate) ListActors(_ context.Context, req *ateapipb.ListActorsRequest) (*ateapipb.ListActorsResponse, error) { + f.mu.Lock() + defer f.mu.Unlock() + resp := &ateapipb.ListActorsResponse{} + if req.GetAtespace() == "default" { + for _, name := range f.names { + resp.Actors = append(resp.Actors, f.actor(name)) + } + } + return resp, nil +} + +// client serves f and returns a Substrate client connected to it. +func (f *fakeSubstrate) client(t *testing.T) *substrate.Client { + t.Helper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + g := grpc.NewServer() + ateapipb.RegisterControlServer(g, f) + go g.Serve(lis) + t.Cleanup(g.Stop) + c, err := substrate.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("substrate client: %v", err) + } + t.Cleanup(func() { c.Close() }) + return c +} + +func TestServer_ReadsActorStateFromSubstrate(t *testing.T) { + st := memory.NewStore() + sub := &fakeSubstrate{names: []string{"observed-a", "observed-b"}, state: ateapipb.ActorState_ACTOR_STATE_RUNNING, workerIP: "10.244.1.7"} + srv := server.NewServer(st, server.Options{Reconciler: &fakeReconciler{}, Substrate: sub.client(t)}) + ctx := context.Background() + + for _, name := range []string{"observed-a", "observed-b"} { + if _, err := srv.CreateTask(ctx, &v1alpha1.CreateTaskRequest{Task: &v1alpha1.Task{ + Metadata: &v1alpha1.ObjectMeta{Name: name}, + Spec: &v1alpha1.TaskSpec{Image: "alpine"}, + }}); err != nil { + t.Fatalf("CreateTask(%s): %v", name, err) + } + } + + got, err := srv.GetTask(ctx, &v1alpha1.GetTaskRequest{Name: "observed-a"}) + if err != nil { + t.Fatalf("GetTask: %v", err) + } + if got.GetStatus().GetPhase() != "Running" || got.GetStatus().GetWorkerIp() != "10.244.1.7" || got.GetStatus().GetActor() != "observed-a" { + t.Errorf("GetTask status = %v, want the actor's state", got.GetStatus()) + } + + list, err := srv.ListTasks(ctx, &v1alpha1.ListTasksRequest{Atespace: "default"}) + if err != nil { + t.Fatalf("ListTasks: %v", err) + } + for _, task := range list.GetTasks() { + if task.GetStatus().GetPhase() != "Running" || task.GetStatus().GetActor() != task.GetMetadata().GetName() { + t.Errorf("ListTasks status for %s = %v, want the actor's state", task.GetMetadata().GetName(), task.GetStatus()) + } + } + + // A change on Substrate alone is visible on the next read. + sub.setState(ateapipb.ActorState_ACTOR_STATE_SUSPENDED) + got, err = srv.GetTask(ctx, &v1alpha1.GetTaskRequest{Name: "observed-a"}) + if err != nil { + t.Fatalf("GetTask: %v", err) + } + if got.GetStatus().GetPhase() != "Suspended" { + t.Errorf("phase after the actor was suspended on Substrate = %q, want Suspended", got.GetStatus().GetPhase()) } - if stored.GetStatus().GetWorkerIp() != "" { - t.Errorf("stored worker IP = %q, want empty", stored.GetStatus().GetWorkerIp()) + + stored, err := st.GetTask(ctx, "default", "observed-a") + if err != nil { + t.Fatalf("store GetTask: %v", err) } + assertNoActorState(t, stored) } func TestServer_CreateTask_InitialPhaseSuspended(t *testing.T) { @@ -620,13 +752,11 @@ func TestServer_CreateTask_InitialPhaseSuspended(t *testing.T) { t.Errorf("returned task phase = %q, want Suspended", task.GetStatus().GetPhase()) } - stored, err := srv.GetTask(ctx, &v1alpha1.GetTaskRequest{Name: "task-create-suspended"}) + stored, err := st.GetTask(ctx, "default", "task-create-suspended") if err != nil { t.Fatalf("GetTask: %v", err) } - if stored.GetStatus().GetPhase() != "Suspended" { - t.Errorf("stored phase = %q, want Suspended", stored.GetStatus().GetPhase()) - } + assertNoActorState(t, stored) } func TestServer_CreateTask_ReconcileErrorMarksFailed(t *testing.T) { @@ -652,7 +782,7 @@ func TestServer_CreateTask_ReconcileErrorMarksFailed(t *testing.T) { if err != nil { t.Fatalf("expected task to exist in store after reconcile error, got %v", err) } - if stored.GetStatus().GetPhase() != "Failed" { + if stored.GetStatus().GetPhase() != v1alpha1.PhaseFailed { t.Errorf("stored phase = %q, want Failed", stored.GetStatus().GetPhase()) } @@ -672,9 +802,9 @@ func TestServer_TransientSubstrateUnavailableRecovery(t *testing.T) { st := memory.NewStore() rec := &fakeReconciler{ reconcileErr: status.Error(codes.Unavailable, "connection refused"), - reconcilePhase: "Pending", + reconcilePhase: v1alpha1.PhasePending, reconcileConditions: []*v1alpha1.Condition{ - {Type: "Ready", Status: "False", Reason: "SubstrateUnavailable"}, + {Type: "Ready", Status: "False", Reason: v1alpha1.ReasonSubstrateUnavailable}, }, } srv := server.NewServer(st, server.Options{Reconciler: rec}) @@ -696,7 +826,7 @@ func TestServer_TransientSubstrateUnavailableRecovery(t *testing.T) { if err != nil { t.Fatalf("GetTask failed: %v", err) } - if stored.GetStatus().GetPhase() != "Pending" { + if stored.GetStatus().GetPhase() != v1alpha1.PhasePending { t.Errorf("expected stored phase 'Pending' after transient Substrate error, got %q", stored.GetStatus().GetPhase()) } @@ -723,7 +853,7 @@ func TestServer_TransientSubstrateUnavailableRecovery(t *testing.T) { if status.Code(err) != codes.Unavailable { t.Fatalf("expected Unavailable on retry while Substrate is still down, got %v", err) } - if conds := rec.observedTask.GetStatus().GetConditions(); len(conds) != 1 || conds[0].GetReason() != "SubstrateUnavailable" { + if conds := rec.observedTask.GetStatus().GetConditions(); len(conds) != 1 || conds[0].GetReason() != v1alpha1.ReasonSubstrateUnavailable { t.Errorf("expected stored condition SubstrateUnavailable to be preserved on omitted-status retry, got %v", conds) } @@ -744,7 +874,7 @@ func TestServer_TransientSubstrateUnavailableRecovery(t *testing.T) { if status.Code(err) != codes.Unavailable { t.Fatalf("expected Unavailable on retry while Substrate is still down, got %v", err) } - if conds := rec.observedTask.GetStatus().GetConditions(); len(conds) != 1 || conds[0].GetReason() != "SubstrateUnavailable" { + if conds := rec.observedTask.GetStatus().GetConditions(); len(conds) != 1 || conds[0].GetReason() != v1alpha1.ReasonSubstrateUnavailable { t.Errorf("expected client-supplied conditions to be ignored on Pending retry, got %v", conds) } @@ -770,7 +900,7 @@ func TestServer_TransientSubstrateUnavailableRecovery(t *testing.T) { // 6. ResumeTask during a brief Substrate outage keeps the task Pending instead of Failed. rec.reconcileErr = status.Error(codes.Unavailable, "substrate restarting") - rec.reconcilePhase = "Pending" + rec.reconcilePhase = v1alpha1.PhasePending _, err = srv.ResumeTask(ctx, &v1alpha1.ResumeTaskRequest{Name: "simple-task"}) if status.Code(err) != codes.Unavailable { t.Fatalf("expected Unavailable from ResumeTask when Substrate is unreachable, got %v", err) @@ -779,7 +909,58 @@ func TestServer_TransientSubstrateUnavailableRecovery(t *testing.T) { if err != nil { t.Fatalf("GetTask failed: %v", err) } - if storedAfterResumeErr.GetStatus().GetPhase() != "Pending" { + if storedAfterResumeErr.GetStatus().GetPhase() != v1alpha1.PhasePending { t.Errorf("expected stored phase 'Pending' after transient ResumeTask error, got %q", storedAfterResumeErr.GetStatus().GetPhase()) } } + +func TestServer_WatchTaskSeesSubstrateOnlyChanges(t *testing.T) { + defer server.SetWatchPollInterval(10 * time.Millisecond)() + + st := memory.NewStore() + sub := &fakeSubstrate{names: []string{"watched"}, state: ateapipb.ActorState_ACTOR_STATE_SUSPENDED} + srv := server.NewServer(st, server.Options{Reconciler: &fakeReconciler{}, Substrate: sub.client(t)}) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + if _, err := srv.CreateTask(ctx, &v1alpha1.CreateTaskRequest{Task: &v1alpha1.Task{ + Metadata: &v1alpha1.ObjectMeta{Name: "watched"}, + Spec: &v1alpha1.TaskSpec{Image: "alpine"}, + }}); err != nil { + t.Fatalf("CreateTask: %v", err) + } + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + go srv.GRPCServer().Serve(ln) + defer srv.GRPCServer().Stop() + conn, err := grpc.NewClient(ln.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + stream, err := v1alpha1.NewAXClient(conn).WatchTask(ctx, &v1alpha1.WatchTaskRequest{Name: "watched"}) + if err != nil { + t.Fatalf("WatchTask: %v", err) + } + evt, err := stream.Recv() + if err != nil { + t.Fatalf("Recv: %v", err) + } + if evt.GetAction() != "INITIAL" || evt.GetTask().GetStatus().GetPhase() != "Suspended" { + t.Fatalf("first event = %s/%s, want INITIAL/Suspended", evt.GetAction(), evt.GetTask().GetStatus().GetPhase()) + } + + // The actor crashes on Substrate. AX writes nothing, but the watcher hears about it. + sub.setState(ateapipb.ActorState_ACTOR_STATE_CRASHED) + evt, err = stream.Recv() + if err != nil { + t.Fatalf("Recv: %v", err) + } + if evt.GetAction() != "MODIFIED" || evt.GetTask().GetStatus().GetPhase() != v1alpha1.PhaseFailed { + t.Fatalf("second event = %s/%s, want MODIFIED/Failed", evt.GetAction(), evt.GetTask().GetStatus().GetPhase()) + } +} diff --git a/internal/server/status_test.go b/internal/server/status_test.go new file mode 100644 index 00000000..3a16628a --- /dev/null +++ b/internal/server/status_test.go @@ -0,0 +1,94 @@ +// 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 server + +import ( + "go/ast" + "go/parser" + "go/token" + "path/filepath" + "strings" + "testing" + + "github.com/google/ax/pkg/apis/v1alpha1" +) + +func TestStoredStatus(t *testing.T) { + in := &v1alpha1.TaskStatus{ + Id: "id", + Phase: "Running", + Actor: "actor", + WorkerIp: "10.0.0.1", + Conditions: []*v1alpha1.Condition{{Type: "Ready"}, {Type: "WorkspaceReady"}}, + } + got := storedStatus(in) + if got.Id != "id" || got.Phase != "" || got.Actor != "" || got.WorkerIp != "" { + t.Errorf("storedStatus = %v, want only Id kept", got) + } + if len(got.Conditions) != 1 || got.Conditions[0].Type != "WorkspaceReady" { + t.Errorf("conditions = %v, want only WorkspaceReady", got.Conditions) + } + if in.Actor != "actor" || len(in.Conditions) != 2 { + t.Errorf("storedStatus modified its input: %v", in) + } + + in = &v1alpha1.TaskStatus{Phase: v1alpha1.PhaseFailed, Conditions: []*v1alpha1.Condition{{Type: "Ready"}}} + if got := storedStatus(in); got.Phase != v1alpha1.PhaseFailed || len(got.Conditions) != 1 { + t.Errorf("storedStatus(%v) = %v, want AX phase and its conditions kept", in, got) + } + if got := storedStatus(nil); got == nil { + t.Error("storedStatus(nil) = nil, want empty status") + } +} + +// TestTasksGoThroughHelpers fails if any code in the package other than the +// task helpers calls a task method on s.store, which would skip the split +// between stored and Substrate-owned status. +func TestTasksGoThroughHelpers(t *testing.T) { + helpers := map[string]bool{ + "getTask": true, "listTasks": true, "saveTask": true, "updateTask": true, + "deleteTask": true, + } + files, err := filepath.Glob("*.go") + if err != nil { + t.Fatal(err) + } + fset := token.NewFileSet() + for _, name := range files { + if strings.HasSuffix(name, "_test.go") { + continue + } + f, err := parser.ParseFile(fset, name, nil, 0) + if err != nil { + t.Fatal(err) + } + for _, decl := range f.Decls { + fn, ok := decl.(*ast.FuncDecl) + if !ok || fn.Body == nil || (fn.Recv != nil && helpers[fn.Name.Name]) { + continue + } + ast.Inspect(fn.Body, func(n ast.Node) bool { + sel, ok := n.(*ast.SelectorExpr) + if !ok || !strings.Contains(sel.Sel.Name, "Task") { + return true + } + if recv, ok := sel.X.(*ast.SelectorExpr); ok && recv.Sel.Name == "store" { + t.Errorf("%s: %s calls store.%s directly; use a task helper", fset.Position(sel.Pos()), fn.Name.Name, sel.Sel.Name) + } + return true + }) + } + } +} diff --git a/internal/store/memory/store.go b/internal/store/memory/store.go index 76adf03f..6235eecc 100644 --- a/internal/store/memory/store.go +++ b/internal/store/memory/store.go @@ -18,7 +18,7 @@ import ( "context" "errors" "fmt" - "io" + "sort" "sync" "github.com/google/ax/internal/store" @@ -26,10 +26,6 @@ import ( "google.golang.org/protobuf/proto" ) -type closerFunc func() error - -func (f closerFunc) Close() error { return f() } - // clone returns a deep copy of a protobuf message so stored resources are never // shared with callers. Copying the struct by value is unsafe: generated messages // embed internal state, including a mutex. @@ -43,7 +39,6 @@ type MemoryStore struct { tasks map[string]*v1alpha1.Task models map[string]*v1alpha1.Model workspaces map[string]*v1alpha1.Workspace - watchers map[string][]chan *v1alpha1.Task } // NewStore creates a new in-memory Store. @@ -52,10 +47,21 @@ func NewStore() *MemoryStore { tasks: make(map[string]*v1alpha1.Task), models: make(map[string]*v1alpha1.Model), workspaces: make(map[string]*v1alpha1.Workspace), - watchers: make(map[string][]chan *v1alpha1.Task), } } +// sortByKey orders resources by atespace and name, so that lists, and pages +// of them, come back in a stable order. +func sortByKey[T interface{ GetMetadata() *v1alpha1.ObjectMeta }](items []T) { + sort.Slice(items, func(i, j int) bool { + return metaKey(items[i].GetMetadata()) < metaKey(items[j].GetMetadata()) + }) +} + +func metaKey(m *v1alpha1.ObjectMeta) string { + return taskKey(m.GetAtespace(), m.GetName()) +} + func taskKey(atespace, name string) string { if atespace == "" { atespace = "default" @@ -73,28 +79,23 @@ func (s *MemoryStore) SaveTask(ctx context.Context, task *v1alpha1.Task) error { if task.Metadata.Atespace == "" { task.Metadata.Atespace = "default" } + if task.ApiVersion == "" { + task.ApiVersion = v1alpha1.APIVersion + } + if task.Kind == "" { + task.Kind = v1alpha1.KindTask + } if task.Status == nil { task.Status = &v1alpha1.TaskStatus{} } if task.Status.Phase == "" { - task.Status.Phase = "Pending" + task.Status.Phase = v1alpha1.PhasePending } key := taskKey(task.Metadata.Atespace, task.Metadata.Name) s.mu.Lock() - cp := clone(task) - s.tasks[key] = cp - - // Notify watchers - if chs, ok := s.watchers[key]; ok { - for _, ch := range chs { - select { - case ch <- cp: - default: - } - } - } + s.tasks[key] = clone(task) s.mu.Unlock() return nil @@ -124,12 +125,16 @@ func (s *MemoryStore) ListTasks(ctx context.Context, atespace string, limit, off result = append(result, cp) } } + sortByKey(result) if offset >= int64(len(result)) { return []*v1alpha1.Task{}, nil } end := offset + limit - if limit <= 0 || end > int64(len(result)) { + if limit <= 0 { + end = offset + 50 + } + if end > int64(len(result)) { end = int64(len(result)) } return result[offset:end], nil @@ -146,16 +151,6 @@ func (s *MemoryStore) UpdateTaskStatus(ctx context.Context, atespace, name strin return store.ErrNotFound } t.Status = status - cp := clone(t) - - if chs, ok := s.watchers[key]; ok { - for _, ch := range chs { - select { - case ch <- cp: - default: - } - } - } return nil } @@ -167,16 +162,25 @@ func (s *MemoryStore) DeleteTask(ctx context.Context, atespace, name string) err } func (s *MemoryStore) SaveModel(ctx context.Context, model *v1alpha1.Model) error { + if model.Metadata == nil { + model.Metadata = &v1alpha1.ObjectMeta{} + } if model.Metadata.Name == "" { return errors.New("model name is required") } if model.Metadata.Atespace == "" { model.Metadata.Atespace = "default" } + if model.ApiVersion == "" { + model.ApiVersion = v1alpha1.APIVersion + } + if model.Kind == "" { + model.Kind = v1alpha1.KindModel + } s.mu.Lock() defer s.mu.Unlock() cp := clone(model) - s.models[taskKey(model.Metadata.Atespace, model.Metadata.Name)] = cp + s.models[metaKey(model.Metadata)] = cp return nil } @@ -203,6 +207,7 @@ func (s *MemoryStore) ListModels(ctx context.Context, atespace string) ([]*v1alp result = append(result, cp) } } + sortByKey(result) return result, nil } @@ -215,16 +220,25 @@ func (s *MemoryStore) DeleteModel(ctx context.Context, atespace, name string) er } func (s *MemoryStore) SaveWorkspace(ctx context.Context, ws *v1alpha1.Workspace) error { + if ws.Metadata == nil { + ws.Metadata = &v1alpha1.ObjectMeta{} + } if ws.Metadata.Name == "" { return errors.New("workspace name is required") } if ws.Metadata.Atespace == "" { ws.Metadata.Atespace = "default" } + if ws.ApiVersion == "" { + ws.ApiVersion = v1alpha1.APIVersion + } + if ws.Kind == "" { + ws.Kind = v1alpha1.KindWorkspace + } s.mu.Lock() defer s.mu.Unlock() cp := clone(ws) - s.workspaces[taskKey(ws.Metadata.Atespace, ws.Metadata.Name)] = cp + s.workspaces[metaKey(ws.Metadata)] = cp return nil } @@ -251,6 +265,7 @@ func (s *MemoryStore) ListWorkspaces(ctx context.Context, atespace string) ([]*v result = append(result, cp) } } + sortByKey(result) return result, nil } @@ -262,31 +277,6 @@ func (s *MemoryStore) DeleteWorkspace(ctx context.Context, atespace, name string 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) - - s.mu.Lock() - s.watchers[key] = append(s.watchers[key], ch) - s.mu.Unlock() - - closer := closerFunc(func() error { - s.mu.Lock() - defer s.mu.Unlock() - chs := s.watchers[key] - for i, c := range chs { - if c == ch { - s.watchers[key] = append(chs[:i], chs[i+1:]...) - close(ch) - break - } - } - return nil - }) - - return ch, closer, nil -} - func (s *MemoryStore) Close() error { return nil } diff --git a/internal/store/postgres/lock.go b/internal/store/postgres/lock.go new file mode 100644 index 00000000..0515b445 --- /dev/null +++ b/internal/store/postgres/lock.go @@ -0,0 +1,196 @@ +// 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 postgres + +import ( + "context" + "crypto/rand" + "encoding/hex" + "fmt" + "log/slog" + "sync" + "time" + + "github.com/google/ax/internal/lock" +) + +// Lock implements lock.Locker with leases stored in the locks table. All +// lease times are taken from the database clock, so ax-server replicas do not +// need synchronized clocks. +// +// A lock is a row holding a random token and an expiry. Acquiring inserts the +// row, or takes it over if its lease has expired; releasing deletes it (only +// if it still holds the caller's token). Waiters retry every +// LockRetryInterval. +// +// Advisory locks were not used because they belong to a database session: a +// lock held across a long Substrate call would pin a pooled connection for +// that long, and would be released silently if that connection dropped. +// +// # Lease renewal +// +// A held lock is kept alive by a background goroutine that extends its lease +// every LockTTL/3 (only while the row still holds this holder's token) until +// the unlock function is called. The TTL therefore does not bound how long a +// lock may be held; it only bounds how long a lock can be orphaned after its +// holder crashes or loses connectivity to Postgres. +// +// We deliberately renew a short lease instead of using a single long TTL: +// +// - The server does not control how long a critical section runs. Handlers +// hold the lock across Substrate calls (e.g. snapshot/restore of a microVM +// during SuspendTask/ResumeTask) that run under the caller's request +// context, so their duration is set by the client's deadline, not by +// ax-server. A fixed TTL would have to exceed the longest deadline any +// client might use, and would silently allow concurrent holders whenever +// an operation ran past it. +// - The TTL is the recovery time after a crash. A long TTL (minutes) would +// wedge the affected task/workspace/model for that long if ax-server died +// mid-operation; a short renewed lease (30s by default) keeps recovery +// fast. +// - One locker guards every resource kind and operation, so renewal avoids +// having to retune the TTL whenever a slower operation is added. +// +// Limitation: renewal is best-effort and there is no fencing token. If renewals +// keep failing for longer than the TTL (e.g. a prolonged database outage), the +// lease can expire while the holder is still working, and the holder is not +// interrupted; the loss is only logged. +func (s *Store) Lock(ctx context.Context, kind, atespace, name string) (func(), error) { + key := lock.Key(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) + + if err := ctx.Err(); err != nil { + return nil, err + } + + // Fast path: the lock is usually free. + ok, err := s.tryLock(ctx, key, token) + if err != nil { + return nil, err + } + if ok { + return s.makeUnlock(key, token), nil + } + + // Slow path: retry until the holder releases the lock or its lease + // expires. + ticker := time.NewTicker(s.opts.LockRetryInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-ticker.C: + } + ok, err := s.tryLock(ctx, key, token) + if err != nil { + return nil, err + } + if ok { + return s.makeUnlock(key, token), nil + } + } +} + +// tryLock attempts to take the lease for key, succeeding if no lease exists or +// the existing one has expired. +func (s *Store) tryLock(ctx context.Context, key, token string) (bool, error) { + tag, err := s.pool.Exec(ctx, ` + INSERT INTO locks (key, token, expires_at) + VALUES ($1, $2, clock_timestamp() + $3 * interval '1 millisecond') + ON CONFLICT (key) DO UPDATE + SET token = EXCLUDED.token, expires_at = EXCLUDED.expires_at + WHERE locks.expires_at <= clock_timestamp()`, + key, token, s.ttlMillis()) + if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return false, ctxErr + } + return false, fmt.Errorf("acquiring lock %s: %w", key, err) + } + return tag.RowsAffected() == 1, nil +} + +func (s *Store) ttlMillis() int64 { + return max(s.opts.LockTTL.Milliseconds(), 1) +} + +// makeUnlock starts lease renewal for an acquired lock and returns an +// idempotent function that stops renewal and releases the lock. +func (s *Store) makeUnlock(key, token string) func() { + renewCtx, cancelRenew := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + s.renew(renewCtx, key, token) + }() + + var once sync.Once + return func() { + once.Do(func() { + cancelRenew() + <-done + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _, err := s.pool.Exec(ctx, `DELETE FROM locks WHERE key = $1 AND token = $2`, key, token) + if err != nil { + // The lease will expire on its own after LockTTL. + slog.Warn("releasing lock", "key", key, "error", err) + } + }) + } +} + +// renew periodically extends the lease of the lock identified by key/token +// until ctx is canceled or the lock is found to be no longer held by token. +func (s *Store) renew(ctx context.Context, key, token string) { + interval := max(s.opts.LockTTL/3, time.Millisecond) + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + + renewCtx, cancel := context.WithTimeout(ctx, interval) + tag, err := s.pool.Exec(renewCtx, ` + UPDATE locks SET expires_at = clock_timestamp() + $3 * interval '1 millisecond' + WHERE key = $1 AND token = $2`, + key, token, s.ttlMillis()) + cancel() + if ctx.Err() != nil { + return + } + if err != nil { + // Transient error: keep trying on the next tick while the + // remaining lease still covers us. + slog.Warn("renewing lock lease", "key", key, "error", err) + continue + } + if tag.RowsAffected() == 0 { + // The lease expired and was taken over; there is nothing left to renew. + slog.Error("lock lease lost before release", "key", key) + return + } + } +} diff --git a/internal/store/postgres/schema.go b/internal/store/postgres/schema.go new file mode 100644 index 00000000..67512bf9 --- /dev/null +++ b/internal/store/postgres/schema.go @@ -0,0 +1,52 @@ +// 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 postgres + +import ( + "context" + _ "embed" + "fmt" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +// schema creates every table and index the store needs. It only adds what is +// missing, so it is safe to run on every start. +// +// TODO: Add versioned migrations before the schema first changes after release. +// +//go:embed schema.sql +var schema string + +// schemaLockID is the advisory lock key that serializes schema creation when +// several ax-server replicas start at once; concurrent CREATE ... IF NOT EXISTS +// statements can otherwise fail. +const schemaLockID = 0x6178_7363_6865_6d61 // "axschema" + +// createSchema creates any missing tables and indexes. +func createSchema(ctx context.Context, pool *pgxpool.Pool) error { + err := pgx.BeginFunc(ctx, pool, func(tx pgx.Tx) error { + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1)`, int64(schemaLockID)); err != nil { + return err + } + _, err := tx.Exec(ctx, schema) + return err + }) + if err != nil { + return fmt.Errorf("creating postgres schema: %w", err) + } + return nil +} diff --git a/internal/store/postgres/schema.sql b/internal/store/postgres/schema.sql new file mode 100644 index 00000000..e4a33c6e --- /dev/null +++ b/internal/store/postgres/schema.sql @@ -0,0 +1,40 @@ +-- Resources are stored as protojson documents. Only the columns needed to +-- address and order them are broken out, so adding fields to the API does not +-- require a schema change. saved_at changes on every Save but not on status +-- updates; lists are ordered by it, newest first. + +CREATE TABLE IF NOT EXISTS tasks ( + atespace text NOT NULL, + name text NOT NULL, + data jsonb NOT NULL, + saved_at timestamptz NOT NULL, + PRIMARY KEY (atespace, name) +); +CREATE INDEX IF NOT EXISTS tasks_saved_at_idx ON tasks (saved_at DESC, atespace, name); +CREATE INDEX IF NOT EXISTS tasks_atespace_saved_at_idx ON tasks (atespace, saved_at DESC, name); + +CREATE TABLE IF NOT EXISTS workspaces ( + atespace text NOT NULL, + name text NOT NULL, + data jsonb NOT NULL, + saved_at timestamptz NOT NULL, + PRIMARY KEY (atespace, name) +); +CREATE INDEX IF NOT EXISTS workspaces_atespace_saved_at_idx ON workspaces (atespace, saved_at DESC, name); + +CREATE TABLE IF NOT EXISTS models ( + atespace text NOT NULL, + name text NOT NULL, + data jsonb NOT NULL, + saved_at timestamptz NOT NULL, + PRIMARY KEY (atespace, name) +); +CREATE INDEX IF NOT EXISTS models_atespace_saved_at_idx ON models (atespace, saved_at DESC, name); + +-- Leases backing the distributed Locker. A row whose expires_at has passed is +-- free and may be taken over by the next caller. +CREATE TABLE IF NOT EXISTS locks ( + key text PRIMARY KEY, + token text NOT NULL, + expires_at timestamptz NOT NULL +); diff --git a/internal/store/postgres/store.go b/internal/store/postgres/store.go new file mode 100644 index 00000000..f85be678 --- /dev/null +++ b/internal/store/postgres/store.go @@ -0,0 +1,355 @@ +// 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 postgres implements store.Store and lock.Locker on PostgreSQL. +// +// Open takes a libpq-style connection URL; standard libpq environment +// variables such as PGPASSWORD fill in what it leaves out. +// Open creates any missing tables and indexes. +package postgres + +import ( + "context" + "errors" + "fmt" + "log/slog" + "time" + + "github.com/google/ax/internal/lock" + "github.com/google/ax/internal/store" + "github.com/google/ax/pkg/apis/v1alpha1" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/proto" +) + +var unmarshalOpts = protojson.UnmarshalOptions{DiscardUnknown: true} + +// Options configures the Postgres backend. +type Options struct { + // LockTTL is the lease duration of a lock (default 30s). Held locks are + // renewed every LockTTL/3, so this does not limit how long a lock may be + // held; it bounds how long a lock survives a crashed holder. See Lock. + LockTTL time.Duration + // LockRetryInterval is how often lock waiters retry acquisition + // (default 1s). + LockRetryInterval time.Duration +} + +// Store is a PostgreSQL implementation of store.Store. It is also a +// lock.Locker (see Lock) whose locks are shared by every ax-server replica +// using the same database. +type Store struct { + pool *pgxpool.Pool + opts Options +} + +var ( + _ store.Store = (*Store)(nil) + _ lock.Locker = (*Store)(nil) +) + +// Open connects to the database at connURL and creates any missing tables. +func Open(ctx context.Context, connURL string, opts Options) (*Store, error) { + if opts.LockTTL <= 0 { + opts.LockTTL = 30 * time.Second + } + if opts.LockRetryInterval <= 0 { + opts.LockRetryInterval = time.Second + } + + cfg, err := pgxpool.ParseConfig(connURL) + if err != nil { + return nil, fmt.Errorf("parsing postgres URL: %w", err) + } + pool, err := pgxpool.NewWithConfig(ctx, cfg) + if err != nil { + return nil, fmt.Errorf("connecting to postgres: %w", err) + } + if err := pool.Ping(ctx); err != nil { + pool.Close() + return nil, fmt.Errorf("connecting to postgres: %w", err) + } + if err := createSchema(ctx, pool); err != nil { + pool.Close() + return nil, err + } + return &Store{pool: pool, opts: opts}, nil +} + +// Close closes all connections. +func (s *Store) Close() error { + s.pool.Close() + return nil +} + +func defaultAtespace(atespace string) string { + if atespace == "" { + return "default" + } + return atespace +} + +// table describes how a resource kind is stored. Table names are constants, +// never user input. +type table string + +const ( + tasksTable table = "tasks" + workspacesTable table = "workspaces" + modelsTable table = "models" +) + +// save upserts a resource. +func (s *Store) save(ctx context.Context, t table, meta *v1alpha1.ObjectMeta, m proto.Message) error { + data, err := protojson.Marshal(m) + if err != nil { + return fmt.Errorf("marshaling %s: %w", t, err) + } + _, err = s.pool.Exec(ctx, ` + INSERT INTO `+string(t)+` (atespace, name, data, saved_at) + VALUES ($1, $2, $3, clock_timestamp()) + ON CONFLICT (atespace, name) DO UPDATE + SET data = EXCLUDED.data, saved_at = EXCLUDED.saved_at`, + meta.Atespace, meta.Name, string(data)) + if err != nil { + return fmt.Errorf("saving %s: %w", t, err) + } + return nil +} + +// get loads a single resource into m. +func (s *Store) get(ctx context.Context, t table, atespace, name string, m proto.Message) error { + var data []byte + err := s.pool.QueryRow(ctx, `SELECT data FROM `+string(t)+` WHERE atespace = $1 AND name = $2`, + defaultAtespace(atespace), name).Scan(&data) + if errors.Is(err, pgx.ErrNoRows) { + return store.ErrNotFound + } + if err != nil { + return fmt.Errorf("getting %s: %w", t, err) + } + if err := unmarshalOpts.Unmarshal(data, m); err != nil { + return fmt.Errorf("unmarshaling %s: %w", t, err) + } + return nil +} + +// list returns resources of one kind, most recently saved first, in atespace +// or across all atespaces when atespace is "" or "*". A limit <= 0 means no +// limit. Records that no longer parse are skipped and logged. +func list[T proto.Message](ctx context.Context, s *Store, t table, atespace string, limit, offset int64, newT func() T) ([]T, error) { + query := `SELECT data FROM ` + string(t) + args := []any{} + if atespace != "" && atespace != "*" { + query += ` WHERE atespace = $1` + args = append(args, atespace) + } + query += ` ORDER BY saved_at DESC, atespace, name` + if limit > 0 { + args = append(args, limit) + query += fmt.Sprintf(` LIMIT $%d`, len(args)) + } + if offset > 0 { + args = append(args, offset) + query += fmt.Sprintf(` OFFSET $%d`, len(args)) + } + + rows, err := s.pool.Query(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("listing %s: %w", t, err) + } + datas, err := pgx.CollectRows(rows, pgx.RowTo[[]byte]) + if err != nil { + return nil, fmt.Errorf("listing %s: %w", t, err) + } + out := make([]T, 0, len(datas)) + for _, data := range datas { + m := newT() + if err := unmarshalOpts.Unmarshal(data, m); err != nil { + slog.Warn("skipping unreadable record", "table", t, "error", err) + continue + } + out = append(out, m) + } + return out, nil +} + +func (s *Store) delete(ctx context.Context, t table, atespace, name string) error { + if _, err := s.pool.Exec(ctx, `DELETE FROM `+string(t)+` WHERE atespace = $1 AND name = $2`, + defaultAtespace(atespace), name); err != nil { + return fmt.Errorf("deleting %s: %w", t, err) + } + return nil +} + +// --- Tasks --- + +// SaveTask stores or replaces a task. +func (s *Store) SaveTask(ctx context.Context, task *v1alpha1.Task) error { + if task.Metadata == nil { + task.Metadata = &v1alpha1.ObjectMeta{} + } + if task.Metadata.Name == "" { + return errors.New("task name is required") + } + task.Metadata.Atespace = defaultAtespace(task.Metadata.Atespace) + if task.ApiVersion == "" { + task.ApiVersion = v1alpha1.APIVersion + } + if task.Kind == "" { + task.Kind = v1alpha1.KindTask + } + if task.Status == nil { + task.Status = &v1alpha1.TaskStatus{} + } + if task.Status.Phase == "" { + task.Status.Phase = v1alpha1.PhasePending + } + return s.save(ctx, tasksTable, task.Metadata, task) +} + +// GetTask retrieves a task by atespace and name. +func (s *Store) GetTask(ctx context.Context, atespace, name string) (*v1alpha1.Task, error) { + var task v1alpha1.Task + if err := s.get(ctx, tasksTable, atespace, name, &task); err != nil { + return nil, err + } + return &task, nil +} + +// ListTasks lists tasks ordered by most recently saved first. A limit <= 0 +// defaults to 50. +func (s *Store) ListTasks(ctx context.Context, atespace string, limit, offset int64) ([]*v1alpha1.Task, error) { + if limit <= 0 { + limit = 50 + } + return list(ctx, s, tasksTable, atespace, limit, offset, func() *v1alpha1.Task { return &v1alpha1.Task{} }) +} + +// UpdateTaskStatus replaces only the status of a task. It does not change the +// task's position in lists. +func (s *Store) UpdateTaskStatus(ctx context.Context, atespace, name string, status *v1alpha1.TaskStatus) error { + atespace = defaultAtespace(atespace) + return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { + var data []byte + err := tx.QueryRow(ctx, `SELECT data FROM tasks WHERE atespace = $1 AND name = $2 FOR UPDATE`, + atespace, name).Scan(&data) + if errors.Is(err, pgx.ErrNoRows) { + return store.ErrNotFound + } + if err != nil { + return fmt.Errorf("getting task: %w", err) + } + var task v1alpha1.Task + if err := unmarshalOpts.Unmarshal(data, &task); err != nil { + return fmt.Errorf("unmarshaling task: %w", err) + } + task.Status = status + if data, err = protojson.Marshal(&task); err != nil { + return fmt.Errorf("marshaling task status: %w", err) + } + if _, err := tx.Exec(ctx, `UPDATE tasks SET data = $3 WHERE atespace = $1 AND name = $2`, + atespace, name, string(data)); err != nil { + return fmt.Errorf("updating task status: %w", err) + } + return nil + }) +} + +// DeleteTask removes a task. Deleting a missing task is not an error. +func (s *Store) DeleteTask(ctx context.Context, atespace, name string) error { + return s.delete(ctx, tasksTable, atespace, name) +} + +// --- Workspaces --- + +// SaveWorkspace stores or replaces a workspace. +func (s *Store) SaveWorkspace(ctx context.Context, ws *v1alpha1.Workspace) error { + if ws.Metadata == nil { + ws.Metadata = &v1alpha1.ObjectMeta{} + } + if ws.Metadata.Name == "" { + return errors.New("workspace name is required") + } + ws.Metadata.Atespace = defaultAtespace(ws.Metadata.Atespace) + if ws.ApiVersion == "" { + ws.ApiVersion = v1alpha1.APIVersion + } + if ws.Kind == "" { + ws.Kind = v1alpha1.KindWorkspace + } + return s.save(ctx, workspacesTable, ws.Metadata, ws) +} + +// GetWorkspace retrieves a workspace by atespace and name. +func (s *Store) GetWorkspace(ctx context.Context, atespace, name string) (*v1alpha1.Workspace, error) { + var ws v1alpha1.Workspace + if err := s.get(ctx, workspacesTable, atespace, name, &ws); err != nil { + return nil, err + } + return &ws, nil +} + +// ListWorkspaces lists workspaces in an atespace, or across all atespaces when +// atespace is "" or "*", most recently saved first. +func (s *Store) ListWorkspaces(ctx context.Context, atespace string) ([]*v1alpha1.Workspace, error) { + return list(ctx, s, workspacesTable, atespace, 0, 0, func() *v1alpha1.Workspace { return &v1alpha1.Workspace{} }) +} + +// DeleteWorkspace removes a workspace. Deleting a missing workspace is not an error. +func (s *Store) DeleteWorkspace(ctx context.Context, atespace, name string) error { + return s.delete(ctx, workspacesTable, atespace, name) +} + +// --- Models --- + +// SaveModel stores or replaces a model. +func (s *Store) SaveModel(ctx context.Context, model *v1alpha1.Model) error { + if model.Metadata == nil { + model.Metadata = &v1alpha1.ObjectMeta{} + } + if model.Metadata.Name == "" { + return errors.New("model name is required") + } + model.Metadata.Atespace = defaultAtespace(model.Metadata.Atespace) + if model.ApiVersion == "" { + model.ApiVersion = v1alpha1.APIVersion + } + if model.Kind == "" { + model.Kind = v1alpha1.KindModel + } + return s.save(ctx, modelsTable, model.Metadata, model) +} + +// GetModel retrieves a model by atespace and name. +func (s *Store) GetModel(ctx context.Context, atespace, name string) (*v1alpha1.Model, error) { + var m v1alpha1.Model + if err := s.get(ctx, modelsTable, atespace, name, &m); err != nil { + return nil, err + } + return &m, nil +} + +// ListModels lists models in an atespace, or across all atespaces when +// atespace is "" or "*", most recently saved first. +func (s *Store) ListModels(ctx context.Context, atespace string) ([]*v1alpha1.Model, error) { + return list(ctx, s, modelsTable, atespace, 0, 0, func() *v1alpha1.Model { return &v1alpha1.Model{} }) +} + +// DeleteModel removes a model. Deleting a missing model is not an error. +func (s *Store) DeleteModel(ctx context.Context, atespace, name string) error { + return s.delete(ctx, modelsTable, atespace, name) +} diff --git a/internal/store/postgres/store_test.go b/internal/store/postgres/store_test.go new file mode 100644 index 00000000..64a5ba9c --- /dev/null +++ b/internal/store/postgres/store_test.go @@ -0,0 +1,183 @@ +// 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 postgres + +import ( + "context" + "errors" + "fmt" + "os" + "testing" + "time" + + "github.com/google/ax/pkg/apis/v1alpha1" +) + +// testURL returns the database URL from $AX_TEST_POSTGRES_URL, skipping the +// test when it is unset. Tests drop and recreate all AX tables in it. +func testURL(t *testing.T) string { + t.Helper() + u := os.Getenv("AX_TEST_POSTGRES_URL") + if u == "" { + t.Skip("AX_TEST_POSTGRES_URL not set; skipping Postgres tests") + } + return u +} + +// openTest opens a Store on an empty schema. +func openTest(t *testing.T, opts Options) *Store { + t.Helper() + ctx := context.Background() + s, err := Open(ctx, testURL(t), opts) + if err != nil { + t.Fatalf("Open: %v", err) + } + if _, err := s.pool.Exec(ctx, `TRUNCATE tasks, workspaces, models, locks`); err != nil { + t.Fatalf("truncating tables: %v", err) + } + t.Cleanup(func() { _ = s.Close() }) + return s +} + +func TestCreateSchemaConcurrent(t *testing.T) { + s := openTest(t, Options{}) + ctx := context.Background() + if _, err := s.pool.Exec(ctx, `DROP TABLE tasks, workspaces, models, locks`); err != nil { + t.Fatal(err) + } + // Simulate several replicas starting at once. + errs := make(chan error, 4) + for range 4 { + go func() { errs <- createSchema(ctx, s.pool) }() + } + for range 4 { + if err := <-errs; err != nil { + t.Fatalf("createSchema: %v", err) + } + } + if err := s.SaveTask(ctx, &v1alpha1.Task{Metadata: &v1alpha1.ObjectMeta{Name: "after-create"}}); err != nil { + t.Fatalf("SaveTask after createSchema: %v", err) + } +} + +// TestListNewestFirst checks that lists are ordered by most recent save, as +// they were with Redis: re-saving a resource moves it to the front, and a +// status update does not. +func TestListNewestFirst(t *testing.T) { + s := openTest(t, Options{}) + ctx := context.Background() + save := func(atespace, name string) { + t.Helper() + task := &v1alpha1.Task{Metadata: &v1alpha1.ObjectMeta{Atespace: atespace, Name: name}} + if err := s.SaveTask(ctx, task); err != nil { + t.Fatalf("SaveTask: %v", err) + } + } + save("a", "a1") + save("a", "a2") + save("a", "a3") + save("b", "b1") + save("a", "a1") + if err := s.UpdateTaskStatus(ctx, "a", "a2", &v1alpha1.TaskStatus{Phase: v1alpha1.PhaseFailed}); err != nil { + t.Fatalf("UpdateTaskStatus: %v", err) + } + + check := func(atespace string, limit, offset int64, want ...string) { + t.Helper() + tasks, err := s.ListTasks(ctx, atespace, limit, offset) + if err != nil { + t.Fatalf("ListTasks(%q, %d, %d): %v", atespace, limit, offset, err) + } + var got []string + for _, task := range tasks { + got = append(got, task.GetMetadata().GetAtespace()+"/"+task.GetMetadata().GetName()) + } + if fmt.Sprint(got) != fmt.Sprint(want) { + t.Errorf("ListTasks(%q, %d, %d) = %v, want %v", atespace, limit, offset, got, want) + } + } + check("a", 10, 0, "a/a1", "a/a3", "a/a2") + check("*", 10, 0, "a/a1", "b/b1", "a/a3", "a/a2") + check("*", 2, 2, "a/a3", "a/a2") + check("a", 1, 1, "a/a3") + + for _, w := range []string{"w1", "w2"} { + if err := s.SaveWorkspace(ctx, &v1alpha1.Workspace{Metadata: &v1alpha1.ObjectMeta{Name: w}}); err != nil { + t.Fatalf("SaveWorkspace: %v", err) + } + } + if list, err := s.ListWorkspaces(ctx, "*"); err != nil || len(list) != 2 || list[0].GetMetadata().GetName() != "w2" { + t.Errorf("ListWorkspaces = %v, %v; want w2 first", list, err) + } + for _, m := range []string{"m1", "m2"} { + if err := s.SaveModel(ctx, &v1alpha1.Model{Metadata: &v1alpha1.ObjectMeta{Name: m}}); err != nil { + t.Fatalf("SaveModel: %v", err) + } + } + if list, err := s.ListModels(ctx, "*"); err != nil || len(list) != 2 || list[0].GetMetadata().GetName() != "m2" { + t.Errorf("ListModels = %v, %v; want m2 first", list, err) + } +} + +// TestLockRenewsLeaseWhileHeld verifies that a held lock outlives its TTL +// (because it is renewed) and that waiters acquire it only after release. +func TestLockRenewsLeaseWhileHeld(t *testing.T) { + ttl := 300 * time.Millisecond + s := openTest(t, Options{LockTTL: ttl, LockRetryInterval: 20 * time.Millisecond}) + ctx := context.Background() + + unlock, err := s.Lock(ctx, "task", "default", "renew") + if err != nil { + t.Fatalf("Lock: %v", err) + } + waitCtx, cancel := context.WithTimeout(ctx, 4*ttl) + defer cancel() + if _, err := s.Lock(waitCtx, "task", "default", "renew"); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected contender deadline while lease was held, got %v", err) + } + unlock() + + acquireCtx, cancel2 := context.WithTimeout(ctx, ttl) + defer cancel2() + unlock2, err := s.Lock(acquireCtx, "task", "default", "renew") + if err != nil { + t.Fatalf("Lock after release: %v", err) + } + unlock2() +} + +// TestLockTakesOverExpiredLease verifies that a lock orphaned by a crashed +// holder becomes available once its lease expires. +func TestLockTakesOverExpiredLease(t *testing.T) { + ttl := 200 * time.Millisecond + s := openTest(t, Options{LockTTL: ttl, LockRetryInterval: 20 * time.Millisecond}) + ctx := context.Background() + + if _, err := s.pool.Exec(ctx, `INSERT INTO locks (key, token, expires_at) VALUES ($1, 'crashed', clock_timestamp() + interval '200 milliseconds')`, + "lock:task:default:orphan"); err != nil { + t.Fatal(err) + } + start := time.Now() + acquireCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + unlock, err := s.Lock(acquireCtx, "task", "default", "orphan") + if err != nil { + t.Fatalf("Lock: %v", err) + } + defer unlock() + if elapsed := time.Since(start); elapsed < 150*time.Millisecond { + t.Fatalf("acquired orphaned lock after %v, before its lease expired", elapsed) + } +} diff --git a/internal/store/redis/store.go b/internal/store/redis/store.go deleted file mode 100644 index 67a9cf81..00000000 --- a/internal/store/redis/store.go +++ /dev/null @@ -1,553 +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 redis - -import ( - "context" - "errors" - "fmt" - "io" - "strings" - "time" - - "github.com/google/ax/internal/store" - "github.com/google/ax/pkg/apis/v1alpha1" - "github.com/redis/go-redis/v9" - "google.golang.org/protobuf/encoding/protojson" -) - -var ( - jsonMarshalOpts = protojson.MarshalOptions{UseProtoNames: false, EmitUnpopulated: false} - jsonUnmarshalOpts = protojson.UnmarshalOptions{DiscardUnknown: true} -) - -// Options contains configuration for the Redis store. -type Options struct { - KeyPrefix string - TTL time.Duration // Optional TTL for task records -} - -// Store is a Redis-backed implementation of store.Store. -type Store struct { - client *redis.Client - opts Options -} - -// NewStore creates a new Redis store. -func NewStore(client *redis.Client, opts Options) *Store { - if opts.KeyPrefix == "" { - opts.KeyPrefix = "ax" - } - return &Store{ - client: client, - opts: opts, - } -} - -func (s *Store) taskKey(atespace, name string) string { - return fmt.Sprintf("%s:task:%s:%s", s.opts.KeyPrefix, atespace, name) -} - - -func (s *Store) modelKey(atespace, name string) string { - return fmt.Sprintf("%s:model:%s:%s", s.opts.KeyPrefix, atespace, name) -} - -func (s *Store) modelIndexKey() string { - return fmt.Sprintf("%s:models:index", s.opts.KeyPrefix) -} - -func (s *Store) modelAtespaceIndexKey(atespace string) string { - return fmt.Sprintf("%s:models:atespace:%s", s.opts.KeyPrefix, atespace) -} - -func (s *Store) wsKey(atespace, name string) string { - return fmt.Sprintf("%s:workspace:%s:%s", s.opts.KeyPrefix, atespace, name) -} - -func (s *Store) wsIndexKey() string { - return fmt.Sprintf("%s:workspaces:index", s.opts.KeyPrefix) -} - -func (s *Store) wsAtespaceIndexKey(atespace string) string { - return fmt.Sprintf("%s:workspaces:atespace:%s", s.opts.KeyPrefix, atespace) -} - -func (s *Store) taskIndexKey() string { - return fmt.Sprintf("%s:tasks:index", s.opts.KeyPrefix) -} - -func (s *Store) taskAtespaceIndexKey(atespace string) string { - return fmt.Sprintf("%s:tasks:atespace:%s", s.opts.KeyPrefix, atespace) -} - -func (s *Store) taskPubSubChannel(atespace, name string) string { - return fmt.Sprintf("%s:pubsub:task:%s:%s", s.opts.KeyPrefix, atespace, name) -} - -// SaveTask stores or updates a task and publishes a reconcile event to the stream. -func (s *Store) SaveTask(ctx context.Context, task *v1alpha1.Task) error { - if task.Metadata == nil { - task.Metadata = &v1alpha1.ObjectMeta{} - } - if task.Metadata.Name == "" { - return errors.New("task name is required") - } - if task.Metadata.Atespace == "" { - task.Metadata.Atespace = "default" - } - if task.ApiVersion == "" { - task.ApiVersion = v1alpha1.APIVersion - } - if task.Kind == "" { - task.Kind = v1alpha1.KindTask - } - if task.Status == nil { - task.Status = &v1alpha1.TaskStatus{} - } - if task.Status.Phase == "" { - task.Status.Phase = "Pending" - } - - data, err := protojson.Marshal(task) - if err != nil { - return fmt.Errorf("marshaling task: %w", err) - } - - atespace := task.Metadata.Atespace - name := task.Metadata.Name - score := float64(time.Now().UnixNano()) - member := fmt.Sprintf("%s:%s", atespace, name) - - pipe := s.client.TxPipeline() - 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.Publish(ctx, s.taskPubSubChannel(atespace, name), data) - - _, err = pipe.Exec(ctx) - if err != nil { - return fmt.Errorf("saving task to redis: %w", err) - } - return nil -} - -// GetTask retrieves a task by atespace and name. -func (s *Store) GetTask(ctx context.Context, atespace, name string) (*v1alpha1.Task, error) { - if atespace == "" { - atespace = "default" - } - val, err := s.client.Get(ctx, s.taskKey(atespace, name)).Result() - if err != nil { - if errors.Is(err, redis.Nil) { - return nil, store.ErrNotFound - } - return nil, fmt.Errorf("getting task from redis: %w", err) - } - - var task v1alpha1.Task - if err := jsonUnmarshalOpts.Unmarshal([]byte(val), &task); err != nil { - return nil, fmt.Errorf("unmarshaling task: %w", err) - } - return &task, nil -} - -// ListTasks lists tasks ordered by newest first. -func (s *Store) ListTasks(ctx context.Context, atespace string, limit, offset int64) ([]*v1alpha1.Task, error) { - if limit <= 0 { - limit = 50 - } - start := offset - stop := offset + limit - 1 - - var members []string - var err error - - if atespace == "" || atespace == "*" { - members, err = s.client.ZRevRange(ctx, s.taskIndexKey(), start, stop).Result() - } else { - names, nErr := s.client.ZRevRange(ctx, s.taskAtespaceIndexKey(atespace), start, stop).Result() - if nErr == nil { - for _, n := range names { - members = append(members, fmt.Sprintf("%s:%s", atespace, n)) - } - } - err = nErr - } - - if err != nil { - return nil, fmt.Errorf("listing task index: %w", err) - } - if len(members) == 0 { - return []*v1alpha1.Task{}, nil - } - - keys := make([]string, len(members)) - for i, m := range members { - parts := strings.SplitN(m, ":", 2) - if len(parts) == 2 { - keys[i] = s.taskKey(parts[0], parts[1]) - } else { - keys[i] = s.taskKey("default", m) - } - } - - vals, err := s.client.MGet(ctx, keys...).Result() - if err != nil { - return nil, fmt.Errorf("batch fetching tasks: %w", err) - } - - tasks := make([]*v1alpha1.Task, 0, len(vals)) - for _, v := range vals { - if v == nil { - continue - } - str, ok := v.(string) - if !ok { - continue - } - var t v1alpha1.Task - if err := jsonUnmarshalOpts.Unmarshal([]byte(str), &t); err == nil { - tasks = append(tasks, &t) - } - } - return tasks, nil -} - -// UpdateTaskStatus updates only the status portion of a task. -func (s *Store) UpdateTaskStatus(ctx context.Context, atespace, name string, status *v1alpha1.TaskStatus) error { - task, err := s.GetTask(ctx, atespace, name) - if err != nil { - return err - } - - task.Status = status - data, err := jsonMarshalOpts.Marshal(task) - if err != nil { - return fmt.Errorf("marshaling task status: %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) - _, err = pipe.Exec(ctx) - if err != nil { - return fmt.Errorf("updating task status in redis: %w", err) - } - return nil -} - -// DeleteTask removes the task record and its index entries. -func (s *Store) DeleteTask(ctx context.Context, atespace, name string) error { - if atespace == "" { - atespace = "default" - } - member := fmt.Sprintf("%s:%s", atespace, name) - - pipe := s.client.TxPipeline() - pipe.Del(ctx, s.taskKey(atespace, name)) - pipe.ZRem(ctx, s.taskIndexKey(), member) - pipe.ZRem(ctx, s.taskAtespaceIndexKey(atespace), name) - if _, err := pipe.Exec(ctx); err != nil { - return fmt.Errorf("deleting task from redis: %w", err) - } - return nil -} - - -// SaveModel stores a model. -func (s *Store) SaveModel(ctx context.Context, model *v1alpha1.Model) error { - if model.Metadata == nil { - model.Metadata = &v1alpha1.ObjectMeta{} - } - if model.Metadata.Name == "" { - return errors.New("model name is required") - } - if model.Metadata.Atespace == "" { - model.Metadata.Atespace = "default" - } - if model.ApiVersion == "" { - model.ApiVersion = v1alpha1.APIVersion - } - if model.Kind == "" { - model.Kind = v1alpha1.KindModel - } - - data, err := protojson.Marshal(model) - if err != nil { - return fmt.Errorf("marshaling model: %w", err) - } - - atespace := model.Metadata.Atespace - name := model.Metadata.Name - score := float64(time.Now().UnixNano()) - member := fmt.Sprintf("%s:%s", atespace, name) - - pipe := s.client.TxPipeline() - pipe.Set(ctx, s.modelKey(atespace, name), data, 0) - pipe.ZAdd(ctx, s.modelIndexKey(), redis.Z{Score: score, Member: member}) - pipe.ZAdd(ctx, s.modelAtespaceIndexKey(atespace), redis.Z{Score: score, Member: name}) - _, err = pipe.Exec(ctx) - return err -} - -// DeleteModel removes a model. -func (s *Store) DeleteModel(ctx context.Context, atespace, name string) error { - if atespace == "" { - atespace = "default" - } - member := fmt.Sprintf("%s:%s", atespace, name) - - pipe := s.client.TxPipeline() - pipe.Del(ctx, s.modelKey(atespace, name)) - pipe.ZRem(ctx, s.modelIndexKey(), member) - pipe.ZRem(ctx, s.modelAtespaceIndexKey(atespace), name) - _, err := pipe.Exec(ctx) - return err -} - -// ListModels lists models for an atespace or across all atespaces. -func (s *Store) ListModels(ctx context.Context, atespace string) ([]*v1alpha1.Model, error) { - var members []string - var err error - - if atespace == "" || atespace == "*" { - members, err = s.client.ZRevRange(ctx, s.modelIndexKey(), 0, -1).Result() - } else { - names, nErr := s.client.ZRevRange(ctx, s.modelAtespaceIndexKey(atespace), 0, -1).Result() - if nErr == nil { - for _, n := range names { - members = append(members, fmt.Sprintf("%s:%s", atespace, n)) - } - } - err = nErr - } - - if err != nil { - return nil, fmt.Errorf("listing model index: %w", err) - } - if len(members) == 0 { - return []*v1alpha1.Model{}, nil - } - - keys := make([]string, len(members)) - for i, m := range members { - parts := strings.SplitN(m, ":", 2) - if len(parts) == 2 { - keys[i] = s.modelKey(parts[0], parts[1]) - } else { - keys[i] = s.modelKey("default", m) - } - } - - vals, err := s.client.MGet(ctx, keys...).Result() - if err != nil { - return nil, fmt.Errorf("batch fetching models: %w", err) - } - - models := make([]*v1alpha1.Model, 0, len(vals)) - for _, v := range vals { - if v == nil { - continue - } - str, ok := v.(string) - if !ok { - continue - } - var m v1alpha1.Model - if err := jsonUnmarshalOpts.Unmarshal([]byte(str), &m); err == nil { - models = append(models, &m) - } - } - return models, nil -} - -// GetModel retrieves a model by atespace and name. -func (s *Store) GetModel(ctx context.Context, atespace, name string) (*v1alpha1.Model, error) { - if atespace == "" { - atespace = "default" - } - val, err := s.client.Get(ctx, s.modelKey(atespace, name)).Result() - if err != nil { - if errors.Is(err, redis.Nil) { - return nil, store.ErrNotFound - } - return nil, fmt.Errorf("getting model from redis: %w", err) - } - - var m v1alpha1.Model - if err := jsonUnmarshalOpts.Unmarshal([]byte(val), &m); err != nil { - return nil, fmt.Errorf("unmarshaling model: %w", err) - } - return &m, nil -} - -// SaveWorkspace stores a workspace. -func (s *Store) SaveWorkspace(ctx context.Context, ws *v1alpha1.Workspace) error { - if ws.Metadata == nil { - ws.Metadata = &v1alpha1.ObjectMeta{} - } - if ws.Metadata.Name == "" { - return errors.New("workspace name is required") - } - if ws.Metadata.Atespace == "" { - ws.Metadata.Atespace = "default" - } - if ws.ApiVersion == "" { - ws.ApiVersion = v1alpha1.APIVersion - } - if ws.Kind == "" { - ws.Kind = v1alpha1.KindWorkspace - } - - data, err := protojson.Marshal(ws) - if err != nil { - return fmt.Errorf("marshaling workspace: %w", err) - } - - atespace := ws.Metadata.Atespace - name := ws.Metadata.Name - score := float64(time.Now().UnixNano()) - member := fmt.Sprintf("%s:%s", atespace, name) - - pipe := s.client.TxPipeline() - pipe.Set(ctx, s.wsKey(atespace, name), data, 0) - pipe.ZAdd(ctx, s.wsIndexKey(), redis.Z{Score: score, Member: member}) - pipe.ZAdd(ctx, s.wsAtespaceIndexKey(atespace), redis.Z{Score: score, Member: name}) - _, err = pipe.Exec(ctx) - return err -} - -// DeleteWorkspace removes a workspace. -func (s *Store) DeleteWorkspace(ctx context.Context, atespace, name string) error { - if atespace == "" { - atespace = "default" - } - member := fmt.Sprintf("%s:%s", atespace, name) - - pipe := s.client.TxPipeline() - pipe.Del(ctx, s.wsKey(atespace, name)) - pipe.ZRem(ctx, s.wsIndexKey(), member) - pipe.ZRem(ctx, s.wsAtespaceIndexKey(atespace), name) - _, err := pipe.Exec(ctx) - return err -} - -// ListWorkspaces lists workspaces for an atespace or across all atespaces. -func (s *Store) ListWorkspaces(ctx context.Context, atespace string) ([]*v1alpha1.Workspace, error) { - var members []string - var err error - - if atespace == "" || atespace == "*" { - members, err = s.client.ZRevRange(ctx, s.wsIndexKey(), 0, -1).Result() - } else { - names, nErr := s.client.ZRevRange(ctx, s.wsAtespaceIndexKey(atespace), 0, -1).Result() - if nErr == nil { - for _, n := range names { - members = append(members, fmt.Sprintf("%s:%s", atespace, n)) - } - } - err = nErr - } - - if err != nil { - return nil, fmt.Errorf("listing workspace index: %w", err) - } - if len(members) == 0 { - return []*v1alpha1.Workspace{}, nil - } - - keys := make([]string, len(members)) - for i, m := range members { - parts := strings.SplitN(m, ":", 2) - if len(parts) == 2 { - keys[i] = s.wsKey(parts[0], parts[1]) - } else { - keys[i] = s.wsKey("default", m) - } - } - - vals, err := s.client.MGet(ctx, keys...).Result() - if err != nil { - return nil, fmt.Errorf("batch fetching workspaces: %w", err) - } - - workspaces := make([]*v1alpha1.Workspace, 0, len(vals)) - for _, v := range vals { - if v == nil { - continue - } - str, ok := v.(string) - if !ok { - continue - } - var w v1alpha1.Workspace - if err := jsonUnmarshalOpts.Unmarshal([]byte(str), &w); err == nil { - workspaces = append(workspaces, &w) - } - } - return workspaces, nil -} - -// GetWorkspace retrieves a workspace by atespace and name. -func (s *Store) GetWorkspace(ctx context.Context, atespace, name string) (*v1alpha1.Workspace, error) { - if atespace == "" { - atespace = "default" - } - val, err := s.client.Get(ctx, s.wsKey(atespace, name)).Result() - if err != nil { - if errors.Is(err, redis.Nil) { - return nil, store.ErrNotFound - } - return nil, fmt.Errorf("getting workspace from redis: %w", err) - } - - var w v1alpha1.Workspace - if err := jsonUnmarshalOpts.Unmarshal([]byte(val), &w); err != nil { - return nil, fmt.Errorf("unmarshaling workspace: %w", err) - } - return &w, nil -} - -// 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 == "" { - atespace = "default" - } - pubsub := s.client.Subscribe(ctx, s.taskPubSubChannel(atespace, name)) - _, err := pubsub.Receive(ctx) - if err != nil { - return nil, nil, fmt.Errorf("subscribing to task watch: %w", err) - } - - ch := make(chan *v1alpha1.Task, 10) - go func() { - defer close(ch) - msgCh := pubsub.Channel() - for msg := range msgCh { - var t v1alpha1.Task - if err := jsonUnmarshalOpts.Unmarshal([]byte(msg.Payload), &t); err == nil { - ch <- &t - } - } - }() - - return ch, pubsub, nil -} - -// Close closes the underlying Redis client connection. -func (s *Store) Close() error { - return s.client.Close() -} diff --git a/internal/store/store.go b/internal/store/store.go index 5d14d61a..9f4360e1 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -17,7 +17,6 @@ package store import ( "context" "errors" - "io" "github.com/google/ax/pkg/apis/v1alpha1" ) @@ -26,6 +25,8 @@ var ( ErrNotFound = errors.New("resource not found") ) +// TODO(jbd): Better pagination. + // Store defines the storage interface for AX resources. type Store interface { SaveTask(ctx context.Context, task *v1alpha1.Task) error @@ -44,6 +45,5 @@ type Store interface { ListModels(ctx context.Context, atespace string) ([]*v1alpha1.Model, error) DeleteModel(ctx context.Context, atespace, name string) error - WatchTask(ctx context.Context, atespace, name string) (<-chan *v1alpha1.Task, io.Closer, error) Close() error } diff --git a/internal/substrate/client.go b/internal/substrate/client.go index abd06ffe..35274c6b 100644 --- a/internal/substrate/client.go +++ b/internal/substrate/client.go @@ -511,11 +511,16 @@ func (c *Client) ResumeActor(ctx context.Context, atespace, actorName string) (* if resp.GetActor() == nil { return nil, "", fmt.Errorf("nil actor returned when resuming %s/%s", atespace, actorName) } - var workerIP string - if ips := resp.GetActor().GetStatus().GetWorkerAssignment().GetWorkerPodIps(); len(ips) > 0 { - workerIP = ips[0] + return resp.GetActor(), WorkerIP(resp.GetActor()), nil +} + +// WorkerIP returns the pod IP of the worker hosting the actor, or "" when the +// actor has no worker (for example while it is suspended). +func WorkerIP(actor *ateapipb.Actor) string { + if ips := actor.GetStatus().GetWorkerAssignment().GetWorkerPodIps(); len(ips) > 0 { + return ips[0] } - return resp.GetActor(), workerIP, nil + return "" } // SuspendActor suspends the specified actor, triggering state checkpointing. @@ -553,6 +558,29 @@ func (c *Client) GetActor(ctx context.Context, atespace, actorName string) (*ate return actor, err } +// ListActors returns every Actor in the given atespace, or in all atespaces +// when atespace is empty. +func (c *Client) ListActors(ctx context.Context, atespace string) ([]*ateapipb.Actor, error) { + var actors []*ateapipb.Actor + req := &ateapipb.ListActorsRequest{Atespace: atespace, Options: &ateapipb.ListOptions{}} + for { + var resp *ateapipb.ListActorsResponse + err := retryOnCodes(ctx, func() error { + var err error + resp, err = c.control.ListActors(ctx, req) + return err + }, codes.Unavailable) + if err != nil { + return nil, fmt.Errorf("listing actors in %q: %w", atespace, err) + } + actors = append(actors, resp.GetActors()...) + if resp.GetNextPageToken() == "" { + return actors, nil + } + req.Options.PageToken = resp.GetNextPageToken() + } +} + // 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 { diff --git a/pkg/apis/v1alpha1/types.go b/pkg/apis/v1alpha1/types.go index 21382230..2ffe025e 100644 --- a/pkg/apis/v1alpha1/types.go +++ b/pkg/apis/v1alpha1/types.go @@ -38,11 +38,97 @@ const ( DefaultTaskImage = "gcr.io/ax-substrate/ate-images/ax-task-runner" DefaultSystemNamespace = "ax-system" + // PhasePending marks a task whose actor has not been provisioned yet, for + // example because Substrate was unavailable. Retrying the operation can + // still succeed. + PhasePending = "Pending" + // PhaseFailed marks a task that AX could not provision or operate, or whose + // actor crashed or disappeared. + PhaseFailed = "Failed" // PhaseTerminating marks a task whose deletion has been requested and whose // actor is being torn down. The record disappears once cleanup completes. PhaseTerminating = "Terminating" + + // The phases below follow from the state of the task's actor on Agent + // Substrate. They are never stored. + + // PhaseRunning marks a task whose actor is running. + PhaseRunning = "Running" + // PhaseResuming marks a task whose actor is being resumed. + PhaseResuming = "Resuming" + // PhaseSuspending marks a task whose actor is being suspended. + PhaseSuspending = "Suspending" + // PhaseSuspended marks a task whose actor is suspended. + PhaseSuspended = "Suspended" + // PhaseUnknown marks a task whose actor could not be read, or is in a state + // AX does not recognize. + PhaseUnknown = "Unknown" +) + +// Condition types reported on Task status. +const ( + // ConditionReady reports whether the task as a whole is ready to do work: + // its actor is running and the workspace inside it has finished setting up. + ConditionReady = "Ready" + // ConditionWorkspaceReady reports whether the workspace inside the actor has + // finished setting up. + ConditionWorkspaceReady = "WorkspaceReady" ) +// TODO(jbd): ReasonAtespaceCreationFailed, ReasonSubstrateUnavailable, +// ReasonActorNotFound, ReasonActorCreationFailed, ReasonActorResumeFailed +// are invalid states. These should be returned as errors, shouldn't be +// leaking as state. + +// Condition reasons. A Ready condition that is False because the actor is +// not running (for example while it is suspended) uses the task's phase as its +// reason instead. +const ( + // ReasonTaskRunning: Ready is True; the actor is running and the workspace + // is set up. + ReasonTaskRunning = "TaskRunning" + // ReasonWorkspaceInitializing: Ready is False; the actor is running but + // the workspace is still being set up. + ReasonWorkspaceInitializing = "WorkspaceInitializing" + // ReasonSubstrateUnavailable: Agent Substrate could not be reached. Ready + // is False with phase Pending while provisioning, or Unknown with phase + // Unknown while reading the actor. + ReasonSubstrateUnavailable = "SubstrateUnavailable" + // ReasonActorNotFound: Ready is False; the task's actor was removed + // outside AX. + ReasonActorNotFound = "ActorNotFound" + // ReasonAtespaceCreationFailed: Ready is False; the atespace could not be + // created on Agent Substrate. + ReasonAtespaceCreationFailed = "AtespaceCreationFailed" + // ReasonInvalidResources: Ready is False; spec.resources is invalid. + ReasonInvalidResources = "InvalidResources" + // ReasonTemplateCreationFailed: Ready is False; Agent Substrate rejected + // the task's ActorTemplate, for example because of its image. + ReasonTemplateCreationFailed = "TemplateCreationFailed" + // ReasonActorCreationFailed: Ready is False; the actor could not be + // created. + ReasonActorCreationFailed = "ActorCreationFailed" + // ReasonActorResumeFailed: Ready is False; the actor could not be resumed. + ReasonActorResumeFailed = "ActorResumeFailed" + + // ReasonSetupComplete: WorkspaceReady is True. + ReasonSetupComplete = "SetupComplete" + // ReasonInitializing: WorkspaceReady is False; setup is still running. + ReasonInitializing = "Initializing" +) + +// IsAXPhase reports whether phase describes the control plane's own work on a +// task (provisioning it, failing to, or tearing it down) rather than the state +// of the task's actor. Only these phases are stored; every other phase is read +// from the actor on Agent Substrate when the task is fetched. +func IsAXPhase(phase string) bool { + switch phase { + case PhasePending, PhaseFailed, PhaseTerminating: + return true + } + return false +} + // YAML encoding. // // The API kinds are protobuf messages, so protojson is the single source of diff --git a/test/e2e/README.md b/test/e2e/README.md index 5754129a..2fd69a0b 100644 --- a/test/e2e/README.md +++ b/test/e2e/README.md @@ -5,14 +5,14 @@ This directory contains the end-to-end test suite for **AX** running on the manifests and Dockerfiles used to deploy AX for the tests. The suite exercises real gVisor sandboxes, Substrate's `ate-api-server`, -in-cluster RustFS snapshot storage, Redis, `ax-server`, `ax-task-runner`, and +in-cluster RustFS snapshot storage, PostgreSQL, `ax-server`, `ax-task-runner`, and the `ax` CLI, against a configurable Substrate branch, tag, or commit. | Path | Purpose | |---|---| | `harness_test.go` | `TestMain`, cluster connection setup, and shared helpers | | `e2e_test.go` | The test cases | -| `manifests/` | Namespace, WorkerPool, Redis, and ax-server manifests (rendered with `envsubst`) | +| `manifests/` | Namespace, WorkerPool, PostgreSQL, and ax-server manifests (rendered with `envsubst`) | | `Dockerfile.server` | ax-server image | | `Dockerfile.task-runner.e2e` | Extends `Dockerfile.task-runner` with test fixtures (a sample Git repo) | @@ -59,7 +59,7 @@ resolved values the tests need to `.tmp/e2e.env`. | `SUBSTRATE_REPO` | `https://github.com/agent-substrate/substrate.git` | Cloned into `.tmp/substrate` when no local checkout exists | | `KIND_CLUSTER_NAME` | `ax-e2e` | Kind cluster name | | `KUBECONFIG` | `.tmp/kubeconfig-${KIND_CLUSTER_NAME}` | Kubeconfig for the Kind cluster | -| `AX_SYSTEM_NAMESPACE` | `ax-system` | Namespace for `ax-redis`, `ax-server`, and the `ax-workers` WorkerPool | +| `AX_SYSTEM_NAMESPACE` | `ax-system` | Namespace for `ax-postgres`, `ax-server`, and the `ax-workers` WorkerPool | | `AX_WORKER_REPLICAS` | `4` | gVisor worker replicas | | `AX_SNAPSHOTS_BUCKET` | `gs://ate-snapshots/ax/` | Snapshot location in Substrate's in-cluster RustFS | | `RECREATE_KIND` | `false` | Delete and recreate the Kind cluster | diff --git a/test/e2e/e2e_test.go b/test/e2e/e2e_test.go index fc9fffd7..07990047 100644 --- a/test/e2e/e2e_test.go +++ b/test/e2e/e2e_test.go @@ -253,8 +253,8 @@ func TestE2E_TaskLifecycleAndSubstrateSync(t *testing.T) { assertContains(t, "'ax resume task'", env.mustRunAX(t, ns, "resume", "task", name), "resumed") task = env.getTask(t, ns, name) assertPhase(t, task, phaseRunning) - assertCondition(t, task, "WorkspaceReady", "True", "SetupComplete") - assertCondition(t, task, "Ready", "True", "TaskRunning") + assertCondition(t, task, "WorkspaceReady", "True", v1alpha1.ReasonSetupComplete) + assertCondition(t, task, "Ready", "True", v1alpha1.ReasonTaskRunning) workerIP := task.GetStatus().GetWorkerIp() if workerIP == "" { t.Fatal("expected non-empty WorkerIp on running task") @@ -481,7 +481,7 @@ func TestE2E_TaskSuspendAndResumePersistence(t *testing.T) { assertContains(t, "'ax suspend task'", env.mustRunAX(t, ns, "suspend", "task", task), "suspended") got := env.getTask(t, ns, task) assertPhase(t, got, phaseSuspended) - assertCondition(t, got, "Ready", "False", "TaskSuspended") + assertCondition(t, got, "Ready", "False", v1alpha1.PhaseSuspended) if ip := got.GetStatus().GetWorkerIp(); ip != "" { t.Fatalf("expected empty WorkerIp when suspended, got %q", ip) } @@ -560,5 +560,5 @@ func TestE2E_WatchDescribeAndAtespaceIsolation(t *testing.T) { // 3. `ax describe task` reflects the running task and its conditions. assertContains(t, "'ax describe task'", env.mustRunAX(t, ns, "describe", "task", task), - task, ns, phaseRunning, "WorkspaceReady", "TaskRunning") + task, ns, phaseRunning, "WorkspaceReady", v1alpha1.ReasonTaskRunning) } diff --git a/test/e2e/manifests/ax-server.yaml b/test/e2e/manifests/ax-server.yaml index c0127f61..3cd88429 100644 --- a/test/e2e/manifests/ax-server.yaml +++ b/test/e2e/manifests/ax-server.yaml @@ -51,7 +51,7 @@ spec: imagePullPolicy: IfNotPresent args: - "--addr=:8080" - - "--redis-addr=ax-redis.${AX_SYSTEM_NAMESPACE}.svc.cluster.local:6379" + - "--store=postgres://ax@ax-postgres.${AX_SYSTEM_NAMESPACE}.svc.cluster.local:5432/ax?sslmode=disable" - "--substrate-endpoint=api.ate-system.svc.cluster.local:443" - "--substrate-authority=api.ate-system.svc" - "--substrate-token-file=/var/run/secrets/ateapi/token" @@ -59,6 +59,11 @@ spec: - "--template=default-template" - "--template-atespace=${AX_SYSTEM_NAMESPACE}" env: + - name: PGPASSWORD + valueFrom: + secretKeyRef: + name: ax-postgres + key: password - name: ATENET_ROUTER_ADDR value: "atenet-router.ate-system.svc.cluster.local:80" - name: AX_SNAPSHOTS_BUCKET diff --git a/test/e2e/manifests/postgres.yaml b/test/e2e/manifests/postgres.yaml new file mode 100644 index 00000000..835f5e62 --- /dev/null +++ b/test/e2e/manifests/postgres.yaml @@ -0,0 +1,97 @@ +# 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. + +# PostgreSQL for the E2E suite. hack/setup-e2e-kind.sh creates the +# ax-postgres Secret and wipes the volume before every run. +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: ax-postgres + namespace: ${AX_SYSTEM_NAMESPACE} + labels: + app.kubernetes.io/name: ax-postgres + app.kubernetes.io/part-of: ax +spec: + replicas: 1 + serviceName: ax-postgres + selector: + matchLabels: + app.kubernetes.io/name: ax-postgres + template: + metadata: + labels: + app.kubernetes.io/name: ax-postgres + spec: + # Give Postgres time for a clean shutdown checkpoint. + terminationGracePeriodSeconds: 60 + containers: + - name: postgres + image: postgres:17-alpine + env: + - name: POSTGRES_USER + value: ax + - name: POSTGRES_DB + value: ax + - name: POSTGRES_PASSWORD + valueFrom: + secretKeyRef: + name: ax-postgres + key: password + # A subdirectory, because the volume root may contain lost+found. + - name: PGDATA + value: /var/lib/postgresql/data/pgdata + ports: + - containerPort: 5432 + name: postgres + readinessProbe: + exec: + command: ["pg_isready", "-U", "ax", "-d", "ax"] + periodSeconds: 5 + resources: + requests: + cpu: 100m + memory: 256Mi + limits: + cpu: 1000m + memory: 1Gi + volumeMounts: + - name: data + mountPath: /var/lib/postgresql/data + volumeClaimTemplates: + - metadata: + name: data + labels: + app.kubernetes.io/name: ax-postgres + app.kubernetes.io/part-of: ax + spec: + accessModes: ["ReadWriteOnce"] + resources: + requests: + storage: 1Gi +--- +apiVersion: v1 +kind: Service +metadata: + name: ax-postgres + namespace: ${AX_SYSTEM_NAMESPACE} + labels: + app.kubernetes.io/name: ax-postgres + app.kubernetes.io/part-of: ax +spec: + ports: + - port: 5432 + targetPort: 5432 + name: postgres + selector: + app.kubernetes.io/name: ax-postgres diff --git a/test/e2e/manifests/redis.yaml b/test/e2e/manifests/redis.yaml deleted file mode 100644 index ddb31929..00000000 --- a/test/e2e/manifests/redis.yaml +++ /dev/null @@ -1,47 +0,0 @@ -apiVersion: apps/v1 -kind: Deployment -metadata: - name: ax-redis - namespace: ${AX_SYSTEM_NAMESPACE} - labels: - app.kubernetes.io/name: ax-redis - app.kubernetes.io/part-of: ax -spec: - replicas: 1 - selector: - matchLabels: - app.kubernetes.io/name: ax-redis - template: - metadata: - labels: - app.kubernetes.io/name: ax-redis - spec: - containers: - - name: redis - image: redis:7-alpine - ports: - - containerPort: 6379 - name: redis - resources: - requests: - cpu: 100m - memory: 128Mi - limits: - cpu: 1000m - memory: 1Gi ---- -apiVersion: v1 -kind: Service -metadata: - name: ax-redis - namespace: ${AX_SYSTEM_NAMESPACE} - labels: - app.kubernetes.io/name: ax-redis - app.kubernetes.io/part-of: ax -spec: - ports: - - port: 6379 - targetPort: 6379 - name: redis - selector: - app.kubernetes.io/name: ax-redis