From 1614478b209504a1e1e0a6eb86320e3ab2824bc5 Mon Sep 17 00:00:00 2001 From: Jaana Dogan Date: Sat, 24 Oct 2026 18:19:23 -0700 Subject: [PATCH 1/3] server: persist Redis and read actor state from Substrate Redis ran as a Deployment with no volume, so a pod restart dropped every Task, Workspace and Model. It now runs as a StatefulSet with a persistent volume, an append-only file synced every second, periodic RDB snapshots, and noeviction so a full instance rejects writes instead of dropping records. `make deploy-redis` deletes the old Deployment first, because its pods share the StatefulSet's labels. The e2e setup also deletes the Redis PVC so each run still starts empty. Task status also kept a copy of the actor's state (phase, actor, workerIP, Ready). It was written only when AX called Substrate, so it went stale as soon as Substrate changed the actor on its own, for example after a crash. The store now holds only what AX owns: spec, metadata, task ID, WorkspaceReady, and the Pending, Failed and Terminating phases. The rest is read from Substrate whenever a task is fetched, listed or watched, through the new TaskReconciler.Observe. That costs one GetActor call for a single task and one ListActors call per atespace for a list. A task whose actor is gone reports Failed with reason ActorNotFound. If Substrate can't be reached, the task reports Unknown. WatchTask now also polls Substrate every 2s, so changes made only on Substrate reach `ax watch`. CreateTask stores a task as Pending until its actor exists. Pending and Failed are now constants next to PhaseTerminating. --- DESIGN.md | 2 + Makefile | 3 + cmd/ax-server/main.go | 5 +- cmd/ax/main.go | 4 +- deploy/redis.yaml | 38 ++++- hack/setup-e2e-kind.sh | 8 +- internal/controller/observe.go | 149 +++++++++++++++++++ internal/controller/observe_test.go | 197 +++++++++++++++++++++++++ internal/controller/reconciler.go | 40 ++--- internal/controller/reconciler_test.go | 18 +-- internal/server/export_test.go | 25 ++++ internal/server/server.go | 133 +++++++++++++---- internal/server/server_test.go | 167 +++++++++++++++++++-- internal/store/memory/store.go | 2 +- internal/store/redis/store.go | 2 +- internal/substrate/client.go | 36 ++++- pkg/apis/v1alpha1/types.go | 19 +++ test/e2e/manifests/redis.yaml | 38 ++++- 18 files changed, 803 insertions(+), 83 deletions(-) create mode 100644 internal/controller/observe.go create mode 100644 internal/controller/observe_test.go create mode 100644 internal/server/export_test.go diff --git a/DESIGN.md b/DESIGN.md index 84f93c87..d1dc0f58 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -4,6 +4,8 @@ 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. +Redis runs as a StatefulSet with a persistent volume and append-only file, so resources survive restarts. It 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. + ``` ax apply -f task.yaml │ diff --git a/Makefile b/Makefile index 2d918f50..6e37dc32 100644 --- a/Makefile +++ b/Makefile @@ -68,6 +68,9 @@ deploy: deploy-redis deploy-server deploy-redis: @echo "==> Deploying Redis to ax-system namespace..." + @# Redis used to run as a Deployment without storage. Its pods share the + @# StatefulSet's labels, so remove it or the Service would route to both. + kubectl -n ax-system delete deployment ax-redis --ignore-not-found kubectl apply -f deploy/redis.yaml deploy-server: diff --git a/cmd/ax-server/main.go b/cmd/ax-server/main.go index 967b00b1..65a987a1 100644 --- a/cmd/ax-server/main.go +++ b/cmd/ax-server/main.go @@ -97,6 +97,7 @@ func main() { rLocker := lock.NewRedisLocker(rClient, lock.RedisLockerOptions{}) var reconciler server.Reconciler + var observer server.Observer subClient, err := substrate.NewClientWithOptions(substrate.ClientOptions{ Target: substrateEndpoint, Authority: substrateAuthority, @@ -109,12 +110,14 @@ func main() { slog.Warn("could not initialize substrate client; running without substrate reconciliation", "error", err) } else { defer subClient.Close() - reconciler = controller.NewTaskReconciler(subClient, defaultTemplate, defaultTemplateAtespace, defaultTaskImage) + r := controller.NewTaskReconciler(subClient, defaultTemplate, defaultTemplateAtespace, defaultTaskImage) + reconciler, observer = r, r } srv := server.NewServer(rStore, server.Options{ Locker: rLocker, Reconciler: reconciler, + Observer: observer, }) httpServer := &http.Server{ diff --git a/cmd/ax/main.go b/cmd/ax/main.go index 447a9284..62d1a3e3 100644 --- a/cmd/ax/main.go +++ b/cmd/ax/main.go @@ -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 == "Completed" || phase == v1alpha1.PhaseFailed { fmt.Printf("Task reached terminal phase %q.\n", phase) break } diff --git a/deploy/redis.yaml b/deploy/redis.yaml index 6a199b96..1c6f0920 100644 --- a/deploy/redis.yaml +++ b/deploy/redis.yaml @@ -18,7 +18,7 @@ metadata: name: ax-system --- apiVersion: apps/v1 -kind: Deployment +kind: StatefulSet metadata: name: ax-redis namespace: ax-system @@ -27,6 +27,7 @@ metadata: app.kubernetes.io/part-of: ax spec: replicas: 1 + serviceName: ax-redis selector: matchLabels: app.kubernetes.io/name: ax-redis @@ -35,12 +36,33 @@ spec: labels: app.kubernetes.io/name: ax-redis spec: + # Give Redis time to flush the append-only file on shutdown. + terminationGracePeriodSeconds: 60 containers: - name: redis image: redis:7-alpine + # Redis is AX's only store, so it persists to disk: every write is + # appended to an AOF that is fsynced once a second, with periodic RDB + # snapshots on top for faster restarts. noeviction makes a full + # instance reject writes instead of silently dropping records. + args: + - --dir + - /data + - --appendonly + - "yes" + - --appendfsync + - everysec + - --save + - "900 1 300 10 60 10000" + - --maxmemory-policy + - noeviction ports: - containerPort: 6379 name: redis + readinessProbe: + exec: + command: ["redis-cli", "ping"] + periodSeconds: 5 resources: requests: cpu: 100m @@ -48,6 +70,20 @@ spec: limits: cpu: 1000m memory: 1Gi + volumeMounts: + - name: data + mountPath: /data + volumeClaimTemplates: + - metadata: + name: data + labels: + app.kubernetes.io/name: ax-redis + app.kubernetes.io/part-of: ax + spec: + accessModes: ["ReadWriteOnce"] + resources: + requests: + storage: 10Gi --- apiVersion: v1 kind: Service diff --git a/hack/setup-e2e-kind.sh b/hack/setup-e2e-kind.sh index e9f7fbcc..91ea6853 100755 --- a/hack/setup-e2e-kind.sh +++ b/hack/setup-e2e-kind.sh @@ -261,13 +261,15 @@ 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. +# 5. Deploy Redis and ax-server. Delete existing workloads and Redis's volume +# 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 +k -n "${AX_SYSTEM_NAMESPACE}" delete statefulset ax-redis --ignore-not-found --wait +k -n "${AX_SYSTEM_NAMESPACE}" delete pvc -l app.kubernetes.io/name=ax-redis --ignore-not-found --wait render_manifest redis.yaml | k apply -f - 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 statefulset/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/observe.go b/internal/controller/observe.go new file mode 100644 index 00000000..2f38fdbd --- /dev/null +++ b/internal/controller/observe.go @@ -0,0 +1,149 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controller + +import ( + "context" + "log/slog" + "time" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/google/ax/internal/substrate" + "github.com/google/ax/pkg/apis/v1alpha1" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// Task status ownership. +// +// A task's actor is the source of truth for whether the task is running, where, +// and whether it is healthy. AX does not keep a copy of that: the stored status +// holds only what AX itself knows (the task ID, workspace setup, and the phases +// v1alpha1.IsAXPhase accepts), and Observe fills in the rest from Substrate +// each time a task is read. + +// Observe fills in the Substrate-owned parts of each task's status (actor, +// workerIP, phase and the Ready condition) from the task's actor. It costs one +// GetActor call for a single task and one ListActors call per atespace +// otherwise. If Substrate cannot be reached, the affected tasks are reported +// with phase "Unknown". +func (r *TaskReconciler) Observe(ctx context.Context, tasks []*v1alpha1.Task) { + if len(tasks) == 1 { + t := tasks[0] + actor, err := r.client.GetActor(ctx, atespaceOf(t), t.GetMetadata().GetName()) + if err != nil && status.Code(err) == codes.NotFound { + actor, err = nil, nil + } + applyActor(t, actor, err) + return + } + + byAtespace := make(map[string][]*v1alpha1.Task) + for _, t := range tasks { + byAtespace[atespaceOf(t)] = append(byAtespace[atespaceOf(t)], t) + } + for atespace, group := range byAtespace { + actors, err := r.client.ListActors(ctx, atespace) + byName := make(map[string]*ateapipb.Actor, len(actors)) + for _, a := range actors { + byName[a.GetMetadata().GetName()] = a + } + for _, t := range group { + applyActor(t, byName[t.GetMetadata().GetName()], err) + } + } +} + +// applyActor overlays the actor's state on the task's status. actor is nil when +// the task has no actor; err is set when the actor could not be looked up. +func applyActor(task *v1alpha1.Task, actor *ateapipb.Actor, err error) { + if task.Status == nil { + task.Status = &v1alpha1.TaskStatus{} + } + st := task.Status + st.Actor = "" + st.WorkerIp = "" + now := time.Now() + + switch { + case err != nil: + slog.Warn("could not observe task actor", "atespace", atespaceOf(task), "task", task.GetMetadata().GetName(), "error", err) + if !v1alpha1.IsAXPhase(st.Phase) { + st.Phase = "Unknown" + setConditionIfChanged(task, condReady, "Unknown", "SubstrateUnavailable", err.Error(), now) + } + 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 + setConditionIfChanged(task, condReady, "False", "ActorNotFound", "The task's actor no longer exists on Substrate", now) + } + default: + st.Actor = actor.GetMetadata().GetName() + st.WorkerIp = substrate.WorkerIP(actor) + if v1alpha1.IsAXPhase(st.Phase) { + return + } + phase, ready, reason, message := actorPhase(task, actor) + st.Phase = phase + setConditionIfChanged(task, condReady, ready, reason, message, now) + } +} + +// actorPhase maps an actor's state to the task's phase and Ready condition. +func actorPhase(task *v1alpha1.Task, actor *ateapipb.Actor) (phase, ready, reason, message string) { + switch state := actor.GetStatus().GetState(); state { + case ateapipb.ActorState_ACTOR_STATE_RUNNING: + if task.GetStatus() != nil && conditionTrue(task, condWorkspaceReady) { + return "Running", "True", "TaskRunning", "Task is running and its workspace is ready" + } + return "Running", "False", "WorkspaceInitializing", "Waiting for workspace setup to complete" + case ateapipb.ActorState_ACTOR_STATE_RESUMING, ateapipb.ActorState_ACTOR_STATE_REVERTING: + return "Resuming", "False", "ActorResuming", "Actor is being resumed" + case ateapipb.ActorState_ACTOR_STATE_SUSPENDING: + return "Suspending", "False", "ActorSuspending", "Actor is being suspended" + case ateapipb.ActorState_ACTOR_STATE_SUSPENDED: + return "Suspended", "False", "TaskSuspended", "Task is suspended" + case ateapipb.ActorState_ACTOR_STATE_CRASHED: + msg := actor.GetStatus().GetCrash().GetMessage() + if msg == "" { + msg = "Actor crashed" + } + return v1alpha1.PhaseFailed, "False", "ActorCrashed", msg + case ateapipb.ActorState_ACTOR_STATE_DELETING: + return v1alpha1.PhaseTerminating, "False", "ActorDeleting", "Actor is being deleted" + default: + return "Unknown", "Unknown", "UnknownActorState", state.String() + } +} + +// setConditionIfChanged sets a condition, keeping its update time when nothing +// about it changed so that repeated reads report a stable timestamp. +func setConditionIfChanged(task *v1alpha1.Task, condType, condStatus, reason, message string, t time.Time) { + for _, c := range task.Status.Conditions { + if c.Type == condType && c.Status == condStatus && c.Reason == reason && c.Message == message { + return + } + } + setCondition(task, condType, condStatus, reason, message, t) +} + +func atespaceOf(task *v1alpha1.Task) string { + if a := task.GetMetadata().GetAtespace(); a != "" { + return a + } + return "default" +} diff --git a/internal/controller/observe_test.go b/internal/controller/observe_test.go new file mode 100644 index 00000000..cd911167 --- /dev/null +++ b/internal/controller/observe_test.go @@ -0,0 +1,197 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controller_test + +import ( + "context" + "net" + "testing" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/google/ax/internal/controller" + "github.com/google/ax/internal/substrate" + "github.com/google/ax/pkg/apis/v1alpha1" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/status" +) + +// actorsServer serves a fixed set of actors, keyed by "atespace/name". +type actorsServer struct { + ateapipb.UnimplementedControlServer + actors map[string]*ateapipb.Actor + err error + getCalls int + listCalls int +} + +func (a *actorsServer) GetActor(_ context.Context, req *ateapipb.GetActorRequest) (*ateapipb.Actor, error) { + a.getCalls++ + if a.err != nil { + return nil, a.err + } + if actor, ok := a.actors[req.GetActor().GetAtespace()+"/"+req.GetActor().GetName()]; ok { + return actor, nil + } + return nil, status.Error(codes.NotFound, "actor not found") +} + +func (a *actorsServer) ListActors(_ context.Context, req *ateapipb.ListActorsRequest) (*ateapipb.ListActorsResponse, error) { + a.listCalls++ + if a.err != nil { + return nil, a.err + } + resp := &ateapipb.ListActorsResponse{} + for _, actor := range a.actors { + if actor.GetMetadata().GetAtespace() == req.GetAtespace() { + resp.Actors = append(resp.Actors, actor) + } + } + return resp, nil +} + +func newObserver(t *testing.T, srv *actorsServer) *controller.TaskReconciler { + 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, srv) + go g.Serve(lis) + t.Cleanup(g.Stop) + + client, err := substrate.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("substrate client: %v", err) + } + t.Cleanup(func() { client.Close() }) + return controller.NewTaskReconciler(client, "", "", "") +} + +func actor(atespace, name string, state ateapipb.ActorState, workerIP string) *ateapipb.Actor { + a := &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: atespace, 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 TestObserve(t *testing.T) { + workspaceReady := &v1alpha1.Condition{Type: "WorkspaceReady", Status: "True"} + crashed := actor("default", "crashed", ateapipb.ActorState_ACTOR_STATE_CRASHED, "") + crashed.Status.Crash = &ateapipb.ActorCrash{Message: "OOMKilled"} + srv := &actorsServer{actors: map[string]*ateapipb.Actor{ + "default/running": actor("default", "running", ateapipb.ActorState_ACTOR_STATE_RUNNING, "10.0.0.5"), + "default/setting-up": actor("default", "setting-up", ateapipb.ActorState_ACTOR_STATE_RUNNING, "10.0.0.6"), + "default/suspended": actor("default", "suspended", ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ""), + "default/crashed": crashed, + "default/failed": actor("default", "failed", ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ""), + }} + r := newObserver(t, srv) + + tests := []struct { + task *v1alpha1.Task + wantPhase, wantIP, wantActor string + wantReady, wantReadyReason string + }{ + {storedTask("running", "", workspaceReady), "Running", "10.0.0.5", "running", "True", "TaskRunning"}, + {storedTask("setting-up", ""), "Running", "10.0.0.6", "setting-up", "False", "WorkspaceInitializing"}, + {storedTask("suspended", "", workspaceReady), "Suspended", "", "suspended", "False", "TaskSuspended"}, + {storedTask("crashed", ""), v1alpha1.PhaseFailed, "", "crashed", "False", "ActorCrashed"}, + // The actor was removed outside AX. + {storedTask("gone", ""), v1alpha1.PhaseFailed, "", "", "False", "ActorNotFound"}, + // 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: "ActorResumeFailed"}), v1alpha1.PhaseFailed, "", "failed", "False", "ActorResumeFailed"}, + {storedTask("not-yet", v1alpha1.PhasePending), v1alpha1.PhasePending, "", "", "", ""}, + } + for _, tt := range tests { + t.Run(tt.task.GetMetadata().GetName(), func(t *testing.T) { + r.Observe(context.Background(), []*v1alpha1.Task{tt.task}) + 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 + } + assertCondition(t, tt.task, "Ready", tt.wantReady, tt.wantReadyReason) + }) + } + + t.Run("crash message", func(t *testing.T) { + task := storedTask("crashed", "") + r.Observe(context.Background(), []*v1alpha1.Task{task}) + for _, c := range task.GetStatus().GetConditions() { + if c.GetType() == "Ready" && c.GetMessage() != "OOMKilled" { + t.Errorf("Ready message = %q, want the actor's crash message", c.GetMessage()) + } + } + }) +} + +func TestObserve_ListUsesOneCallPerAtespace(t *testing.T) { + srv := &actorsServer{actors: map[string]*ateapipb.Actor{ + "default/a": actor("default", "a", ateapipb.ActorState_ACTOR_STATE_RUNNING, "10.0.0.1"), + "default/b": actor("default", "b", ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ""), + "team-x/c": actor("team-x", "c", ateapipb.ActorState_ACTOR_STATE_RUNNING, "10.0.0.3"), + }} + r := newObserver(t, srv) + + tasks := []*v1alpha1.Task{storedTask("a", ""), storedTask("b", ""), storedTask("c", "")} + tasks[2].Metadata.Atespace = "team-x" + r.Observe(context.Background(), tasks) + + if srv.listCalls != 2 || srv.getCalls != 0 { + t.Errorf("ListActors calls = %d, GetActor calls = %d, want 2 and 0", srv.listCalls, srv.getCalls) + } + for i, want := range []string{"Running", "Suspended", "Running"} { + if got := tasks[i].GetStatus().GetPhase(); got != want { + t.Errorf("task %s phase = %q, want %q", tasks[i].GetMetadata().GetName(), got, want) + } + } + if ip := tasks[2].GetStatus().GetWorkerIp(); ip != "10.0.0.3" { + t.Errorf("task c workerIP = %q, want 10.0.0.3", ip) + } +} + +func TestObserve_SubstrateErrorReportsUnknown(t *testing.T) { + srv := &actorsServer{err: status.Error(codes.PermissionDenied, "denied")} + r := newObserver(t, srv) + + task := storedTask("a", "") + terminating := storedTask("b", v1alpha1.PhaseTerminating) + r.Observe(context.Background(), []*v1alpha1.Task{task}) + r.Observe(context.Background(), []*v1alpha1.Task{terminating}) + + if got := task.GetStatus().GetPhase(); got != "Unknown" { + t.Errorf("phase = %q, want Unknown", got) + } + assertCondition(t, task, "Ready", "Unknown", "SubstrateUnavailable") + if got := terminating.GetStatus().GetPhase(); got != v1alpha1.PhaseTerminating { + t.Errorf("phase = %q, want %q to be kept", got, v1alpha1.PhaseTerminating) + } +} diff --git a/internal/controller/reconciler.go b/internal/controller/reconciler.go index 074a4810..15c6376a 100644 --- a/internal/controller/reconciler.go +++ b/internal/controller/reconciler.go @@ -128,11 +128,11 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws if err := r.client.EnsureAtespace(ctx, atespace); err != nil { if substrate.IsTransientError(err) { r.setNotReady(task, "SubstrateUnavailable", err.Error(), now) - task.Status.Phase = "Pending" + 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" + setCondition(task, "Ready", "False", "AtespaceCreationFailed", err.Error(), now) + task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("ensuring atespace: %w", err) } @@ -186,7 +186,7 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws // 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" + task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("validating resources: %w", err) } resources := substrate.ResourceLimits(task.Spec.Resources) @@ -202,7 +202,7 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws if err != nil { if substrate.IsTransientError(err) { r.setNotReady(task, "SubstrateUnavailable", err.Error(), now) - task.Status.Phase = "Pending" + 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 @@ -210,7 +210,7 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws // 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" + task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("ensuring actor template: %w", err) } if tmpl != nil && tmpl.Metadata != nil { @@ -224,11 +224,11 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws if err != nil { if substrate.IsTransientError(err) { r.setNotReady(task, "SubstrateUnavailable", err.Error(), now) - task.Status.Phase = "Pending" + task.Status.Phase = v1alpha1.PhasePending return task, fmt.Errorf("ensuring actor: %w", err) } r.setNotReady(task, "ActorCreationFailed", err.Error(), now) - task.Status.Phase = "Failed" + task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("ensuring actor: %w", err) } @@ -241,7 +241,7 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws } task.Status.WorkerIp = "" task.Status.Phase = "Suspended" - r.setCondition(task, condReady, "False", "TaskSuspended", "Task is suspended", now) + setCondition(task, condReady, "False", "TaskSuspended", "Task is suspended", now) slog.Info("task successfully suspended", "name", task.Metadata.Name, "actor", actorName) return task, nil } @@ -252,11 +252,11 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws if err != nil { if substrate.IsTransientError(err) { r.setNotReady(task, "SubstrateUnavailable", err.Error(), now) - task.Status.Phase = "Pending" + task.Status.Phase = v1alpha1.PhasePending return task, fmt.Errorf("resuming actor: %w", err) } r.setNotReady(task, "ActorResumeFailed", err.Error(), now) - task.Status.Phase = "Failed" + task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("resuming actor: %w", err) } @@ -273,7 +273,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, condWorkspaceReady) if workerIP != "" && !workspaceReady { // Poll briefly for workspace setup completion pollCtx, cancel := context.WithTimeout(ctx, r.WorkspaceReadyTimeout) @@ -327,13 +327,13 @@ DonePolling: // The task is Ready only once its actor is running and the workspace inside it is set up. 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, condWorkspaceReady) { + setCondition(task, condWorkspaceReady, "True", "SetupComplete", 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()) + 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, condWorkspaceReady, "False", "Initializing", fmt.Sprintf("Workspace is initializing at %s", workerIP), time.Now()) + setCondition(task, condReady, "False", "WorkspaceInitializing", "Waiting for workspace setup to complete", time.Now()) } slog.Info("task successfully reconciled and running", @@ -359,11 +359,11 @@ const ( // 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, condReady, "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 +392,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..55451feb 100644 --- a/internal/controller/reconciler_test.go +++ b/internal/controller/reconciler_test.go @@ -508,7 +508,7 @@ 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") @@ -579,7 +579,7 @@ 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") @@ -845,7 +845,7 @@ 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") @@ -868,7 +868,7 @@ 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") @@ -889,7 +889,7 @@ 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") @@ -1125,7 +1125,7 @@ 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") @@ -1141,7 +1141,7 @@ 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") @@ -1159,7 +1159,7 @@ 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") @@ -1185,7 +1185,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/server/export_test.go b/internal/server/export_test.go new file mode 100644 index 00000000..4e5b364a --- /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 polls Substrate 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..99351d61 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -20,6 +20,7 @@ import ( "log/slog" "net/http" "strings" + "time" "github.com/google/ax/internal/lock" "github.com/google/ax/internal/store" @@ -38,18 +39,31 @@ type Reconciler interface { ReconcileDelete(ctx context.Context, atespace, taskName string) error } +// Observer fills in the parts of a task's status that are owned by its actor +// on Agent Substrate. The store never holds them, so every task the API returns +// passes through the Observer first. +type Observer interface { + Observe(ctx context.Context, tasks []*v1alpha1.Task) +} + // Options configures the AX API Server. type Options struct { Locker lock.Locker Reconciler Reconciler + Observer Observer } +// watchPollInterval is how often WatchTask re-reads a task's actor, so changes +// made on Substrate (for example a crash) reach watchers without an AX write. +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 + observer Observer grpcServer *grpc.Server } @@ -68,6 +82,7 @@ func NewServer(s store.Store, opts ...Options) *Server { store: s, locker: locker, reconciler: opt.Reconciler, + observer: opt.Observer, grpcServer: grpc.NewServer(), } v1alpha1.RegisterAXServer(srv.grpcServer, srv) @@ -114,6 +129,7 @@ func (s *Server) GetTask(ctx context.Context, req *v1alpha1.GetTaskRequest) (*v1 } return nil, status.Errorf(codes.Internal, "getting task: %v", err) } + s.observe(ctx, task) return task, nil } @@ -134,6 +150,7 @@ func (s *Server) ListTasks(ctx context.Context, req *v1alpha1.ListTasksRequest) if err != nil { return nil, status.Errorf(codes.Internal, "listing tasks: %v", err) } + s.observe(ctx, tasks...) return &v1alpha1.ListTasksResponse{Tasks: tasks}, nil } @@ -164,7 +181,7 @@ func (s *Server) CreateTask(ctx context.Context, req *v1alpha1.CreateTaskRequest existing, err := s.store.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 +196,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.store.SaveTask(ctx, withStoredStatus(task)); err != nil { return nil, status.Errorf(codes.Internal, "saving task: %v", err) } + task.Status.Phase = "Suspended" // Directly reconcile with Substrate if s.reconciler != nil { @@ -193,20 +212,20 @@ 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.store.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(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.store.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(task.Status)); err != nil { + return nil, status.Errorf(codes.Internal, "updating task status: %v", err) } return task, nil @@ -242,7 +261,7 @@ 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.store.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(task.Status)); err != nil { return nil, status.Errorf(codes.Internal, "updating task status: %v", err) } @@ -306,11 +325,11 @@ 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.store.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(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.store.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(task.Status)); err != nil { return nil, status.Errorf(codes.Internal, "suspending task: %v", err) } } @@ -354,22 +373,22 @@ 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.store.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(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.store.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(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.store.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(task.Status)); err != nil { return nil, status.Errorf(codes.Internal, "resuming task: %v", err) } } @@ -377,6 +396,45 @@ func (s *Server) ResumeTask(ctx context.Context, req *v1alpha1.ResumeTaskRequest return task, nil } +// observe fills in the actor-owned parts of the tasks' status, when the server +// is connected to Substrate. +func (s *Server) observe(ctx context.Context, tasks ...*v1alpha1.Task) { + if s.observer != nil && len(tasks) > 0 { + s.observer.Observe(ctx, tasks) + } +} + +// 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: they are read from Substrate instead, so the store +// can never hold a stale copy of them. +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() != "Ready" { + conds = append(conds, c) + } + } + out.Conditions = conds + } + return out +} + +// withStoredStatus returns a copy of task whose status is storedStatus. +func withStoredStatus(task *v1alpha1.Task) *v1alpha1.Task { + out := proto.Clone(task).(*v1alpha1.Task) + out.Status = storedStatus(task.Status) + 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 { @@ -406,26 +464,49 @@ func (s *Server) WatchTask(req *v1alpha1.WatchTaskRequest, stream grpc.ServerStr } defer closer.Close() + var last *v1alpha1.Task + send := func(task *v1alpha1.Task, action string) (done bool, err error) { + s.observe(ctx, task) + 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 || phase == "Completed", nil + } + 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 done, err := send(initial, "INITIAL"); done || err != nil { return err } } + // Store events cover AX's own writes. Polling picks up changes that happen + // on Substrate alone, such as an actor crashing or being suspended. + ticker := time.NewTicker(watchPollInterval) + defer ticker.Stop() for { + var task *v1alpha1.Task select { case <-ctx.Done(): return ctx.Err() - case task, ok := <-ch: + case t, 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 + task = t + case <-ticker.C: + t, err := s.store.GetTask(ctx, atespace, req.Name) + if err != nil { + continue } + task = t + } + 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..b7c20e4f 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -20,6 +20,7 @@ import ( "net" "net/http" "net/http/httptest" + "sync" "testing" "time" @@ -589,18 +590,107 @@ 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) + } } - if stored.GetStatus().GetWorkerIp() != "" { - t.Errorf("stored worker IP = %q, want empty", stored.GetStatus().GetWorkerIp()) +} + +// fakeObserver stands in for Substrate: it reports every task's actor, named +// after the task, in the given phase. +type fakeObserver struct { + mu sync.Mutex + phase string + workerIP string +} + +func (f *fakeObserver) setPhase(phase string) { + f.mu.Lock() + defer f.mu.Unlock() + f.phase = phase +} + +func (f *fakeObserver) Observe(_ context.Context, tasks []*v1alpha1.Task) { + f.mu.Lock() + defer f.mu.Unlock() + for _, task := range tasks { + if task.Status == nil { + task.Status = &v1alpha1.TaskStatus{} + } + task.Status.Actor = task.GetMetadata().GetName() + task.Status.WorkerIp = f.workerIP + if !v1alpha1.IsAXPhase(task.Status.Phase) { + task.Status.Phase = f.phase + } } } +func TestServer_ReadsActorStateFromObserver(t *testing.T) { + st := memory.NewStore() + obs := &fakeObserver{phase: "Running", workerIP: "10.244.1.7"} + srv := server.NewServer(st, server.Options{Reconciler: &fakeReconciler{}, Observer: obs}) + 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 observed actor 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 observed actor state", task.GetMetadata().GetName(), task.GetStatus()) + } + } + + // A change on Substrate alone is visible on the next read. + obs.setPhase("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()) + } + + 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) { st := memory.NewStore() rec := &fakeReconciler{} @@ -620,13 +710,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 +740,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,7 +760,7 @@ 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"}, }, @@ -696,7 +784,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()) } @@ -770,7 +858,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 +867,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() + obs := &fakeObserver{phase: "Suspended"} + srv := server.NewServer(st, server.Options{Reconciler: &fakeReconciler{}, Observer: obs}) + 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. + obs.setPhase(v1alpha1.PhaseFailed) + 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/store/memory/store.go b/internal/store/memory/store.go index 76adf03f..de6814be 100644 --- a/internal/store/memory/store.go +++ b/internal/store/memory/store.go @@ -77,7 +77,7 @@ func (s *MemoryStore) SaveTask(ctx context.Context, task *v1alpha1.Task) error { 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) diff --git a/internal/store/redis/store.go b/internal/store/redis/store.go index 67a9cf81..467a636c 100644 --- a/internal/store/redis/store.go +++ b/internal/store/redis/store.go @@ -118,7 +118,7 @@ func (s *Store) SaveTask(ctx context.Context, task *v1alpha1.Task) error { task.Status = &v1alpha1.TaskStatus{} } if task.Status.Phase == "" { - task.Status.Phase = "Pending" + task.Status.Phase = v1alpha1.PhasePending } data, err := protojson.Marshal(task) 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..0ec510a9 100644 --- a/pkg/apis/v1alpha1/types.go +++ b/pkg/apis/v1alpha1/types.go @@ -38,11 +38,30 @@ 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" ) +// 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/manifests/redis.yaml b/test/e2e/manifests/redis.yaml index ddb31929..eb8de76b 100644 --- a/test/e2e/manifests/redis.yaml +++ b/test/e2e/manifests/redis.yaml @@ -1,5 +1,5 @@ apiVersion: apps/v1 -kind: Deployment +kind: StatefulSet metadata: name: ax-redis namespace: ${AX_SYSTEM_NAMESPACE} @@ -8,6 +8,7 @@ metadata: app.kubernetes.io/part-of: ax spec: replicas: 1 + serviceName: ax-redis selector: matchLabels: app.kubernetes.io/name: ax-redis @@ -16,12 +17,33 @@ spec: labels: app.kubernetes.io/name: ax-redis spec: + # Give Redis time to flush the append-only file on shutdown. + terminationGracePeriodSeconds: 60 containers: - name: redis image: redis:7-alpine + # Redis is AX's only store, so it persists to disk: every write is + # appended to an AOF that is fsynced once a second, with periodic RDB + # snapshots on top for faster restarts. noeviction makes a full + # instance reject writes instead of silently dropping records. + args: + - --dir + - /data + - --appendonly + - "yes" + - --appendfsync + - everysec + - --save + - "900 1 300 10 60 10000" + - --maxmemory-policy + - noeviction ports: - containerPort: 6379 name: redis + readinessProbe: + exec: + command: ["redis-cli", "ping"] + periodSeconds: 5 resources: requests: cpu: 100m @@ -29,6 +51,20 @@ spec: limits: cpu: 1000m memory: 1Gi + volumeMounts: + - name: data + mountPath: /data + volumeClaimTemplates: + - metadata: + name: data + labels: + app.kubernetes.io/name: ax-redis + app.kubernetes.io/part-of: ax + spec: + accessModes: ["ReadWriteOnce"] + resources: + requests: + storage: 1Gi --- apiVersion: v1 kind: Service From 79d64093d79d5b4884e1655ffb7a98f5d7cd98f5 Mon Sep 17 00:00:00 2001 From: Jaana Dogan Date: Sat, 24 Oct 2026 19:38:16 -0700 Subject: [PATCH 2/3] server, store: move to Postgres and read actor state from Substrate Replace the Redis store with PostgreSQL, and stop storing the parts of a task's status that belong to its actor on Agent Substrate. Storage: - internal/store/postgres holds resources as protojson JSONB documents. Open creates any missing tables. Locks are leases in a locks table, renewed while held. - Watches and lock waiters poll. There is no LISTEN/NOTIFY, so any connection pooler works. WatchTask is removed from store.Store. - The memory store is now test-only and no longer tracks save order. Newest-first ordering is covered by a Postgres-specific test. Task status: - The store keeps only what AX owns: spec, metadata, the Pending, Failed and Terminating phases, and WorkspaceReady. The actor, worker IP, live phase and Ready condition are read from Substrate on every get, list and watch. - Only a few helpers in server.go (getTask, listTasks, saveTask, updateTask, deleteTask) touch task storage, and a test enforces it. - Phases, condition types and condition reasons are now constants in pkg/apis/v1alpha1. Ready is derived from one rule in the server instead of being written in several places. - Ready reasons visible to clients change: TaskSuspended, ActorResuming, ActorSuspending, ActorCrashed, ActorDeleting and UnknownActorState become the phase name (Suspended, Resuming, Suspending, Failed, Terminating, Unknown). This work revealed major problems in the server, storage and controller implementations. Among them: - Task status was copied from Substrate into storage and could drift from the actor it described. - The controller wrote Ready conditions and phases that the server then discarded, or that disagreed with what clients saw. - Handlers pass the desired action to the reconciler by overwriting status.phase, mixing intent with observed state. - Phase and reason strings were spread across packages as literals. The server and CLI checked for a "Completed" phase that nothing produced. - The shared store tests encoded one backend's list ordering as a contract. This change fixes some of these, but the server, store and controller need further stabilization. @rakyll is taking ownership of that work. --- .github/CODEOWNERS | 6 + .github/workflows/go.yml | 16 + CONTRIBUTING.md | 2 +- DESIGN.md | 10 +- Makefile | 24 +- README.md | 2 +- cmd/ax-server/main.go | 58 +-- cmd/ax/main.go | 10 +- deploy/ax-server.yaml | 7 +- deploy/{redis.yaml => postgres.yaml} | 73 ++-- docs/concepts.md | 2 +- go.mod | 9 +- go.sum | 31 +- hack/setup-e2e-kind.sh | 18 +- internal/controller/observe.go | 149 ------- internal/controller/observe_test.go | 197 --------- internal/controller/reconciler.go | 50 +-- internal/controller/reconciler_test.go | 42 +- internal/lock/lock.go | 237 +---------- internal/lock/lock_test.go | 52 --- internal/server/actor_status_test.go | 128 ++++++ internal/server/export_test.go | 2 +- internal/server/server.go | 329 +++++++++++---- internal/server/server_test.go | 100 +++-- internal/server/status_test.go | 94 +++++ internal/store/memory/store.go | 104 +++-- internal/store/postgres/lock.go | 196 +++++++++ internal/store/postgres/schema.go | 52 +++ internal/store/postgres/schema.sql | 40 ++ internal/store/postgres/store.go | 355 ++++++++++++++++ internal/store/postgres/store_test.go | 183 ++++++++ internal/store/redis/store.go | 553 ------------------------- internal/store/store.go | 4 +- pkg/apis/v1alpha1/types.go | 67 +++ test/e2e/README.md | 6 +- test/e2e/e2e_test.go | 8 +- test/e2e/manifests/ax-server.yaml | 7 +- test/e2e/manifests/postgres.yaml | 97 +++++ test/e2e/manifests/redis.yaml | 83 ---- 39 files changed, 1790 insertions(+), 1613 deletions(-) create mode 100644 .github/CODEOWNERS rename deploy/{redis.yaml => postgres.yaml} (50%) delete mode 100644 internal/controller/observe.go delete mode 100644 internal/controller/observe_test.go create mode 100644 internal/server/actor_status_test.go create mode 100644 internal/server/status_test.go create mode 100644 internal/store/postgres/lock.go create mode 100644 internal/store/postgres/schema.go create mode 100644 internal/store/postgres/schema.sql create mode 100644 internal/store/postgres/store.go create mode 100644 internal/store/postgres/store_test.go delete mode 100644 internal/store/redis/store.go create mode 100644 test/e2e/manifests/postgres.yaml delete mode 100644 test/e2e/manifests/redis.yaml 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 d1dc0f58..17fec348 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -2,9 +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. -Redis runs as a StatefulSet with a persistent volume and append-only file, so resources survive restarts. It 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. +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 @@ -15,9 +15,9 @@ Redis runs as a StatefulSet with a persistent volume and append-only file, so re │ ┌──────────────────────┴──────────────────────┐ ▼ ▼ - Redis Agent Substrate + PostgreSQL Agent Substrate (Resource Store, ┌───────────────────────────────┐ - Locks, PubSub) │ • Atespace Provisioning │ + Locks) │ • Atespace Provisioning │ │ • Actor Creation & Activation │ │ • Worker Assignment │ └───────────────────────────────┘ @@ -28,7 +28,7 @@ Redis runs as a StatefulSet with a persistent volume and append-only file, so re | 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 6e37dc32..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,15 +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..." - @# Redis used to run as a Deployment without storage. Its pods share the - @# StatefulSet's labels, so remove it or the Service would route to both. - kubectl -n ax-system delete deployment ax-redis --ignore-not-found - 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 65a987a1..23ab5eed 100644 --- a/cmd/ax-server/main.go +++ b/cmd/ax-server/main.go @@ -19,25 +19,23 @@ import ( "flag" "log/slog" "net/http" + "net/url" "os" "os/signal" "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" ) func main() { var ( listenAddr string - redisAddr string - redisPassword string + storeURL string substrateEndpoint string substrateAuthority string substrateTokenFile string @@ -50,8 +48,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 +63,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,23 +78,22 @@ func main() { slog.Info("starting ax-server", "listenAddr", listenAddr, - "redisAddr", redisAddr, + "store", redactURL(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", redactURL(storeURL), "error", err) + os.Exit(1) + } + defer backend.Close() var reconciler server.Reconciler - var observer server.Observer subClient, err := substrate.NewClientWithOptions(substrate.ClientOptions{ Target: substrateEndpoint, Authority: substrateAuthority, @@ -110,14 +106,13 @@ func main() { slog.Warn("could not initialize substrate client; running without substrate reconciliation", "error", err) } else { defer subClient.Close() - r := controller.NewTaskReconciler(subClient, defaultTemplate, defaultTemplateAtespace, defaultTaskImage) - reconciler, observer = r, r + 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, - Observer: observer, + Substrate: subClient, }) httpServer := &http.Server{ @@ -145,3 +140,12 @@ func main() { defer shutdownCancel() _ = httpServer.Shutdown(shutdownCtx) } + +// redactURL hides any password in a store URL so that it can be logged. +func redactURL(raw string) string { + u, err := url.Parse(raw) + if err != nil { + return "" + } + return u.Redacted() +} diff --git a/cmd/ax/main.go b/cmd/ax/main.go index 62d1a3e3..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() @@ -769,7 +769,7 @@ func runWatch(serverURL, atespace string, args []string) error { actor, workerIP, ) - if phase == "Completed" || phase == v1alpha1.PhaseFailed { + 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/redis.yaml b/deploy/postgres.yaml similarity index 50% rename from deploy/redis.yaml rename to deploy/postgres.yaml index 1c6f0920..8e7685af 100644 --- a/deploy/redis.yaml +++ b/deploy/postgres.yaml @@ -12,6 +12,11 @@ # 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: @@ -20,64 +25,62 @@ metadata: apiVersion: apps/v1 kind: StatefulSet metadata: - name: ax-redis + name: ax-postgres namespace: ax-system labels: - app.kubernetes.io/name: ax-redis + app.kubernetes.io/name: ax-postgres app.kubernetes.io/part-of: ax spec: replicas: 1 - serviceName: ax-redis + serviceName: ax-postgres selector: matchLabels: - app.kubernetes.io/name: ax-redis + app.kubernetes.io/name: ax-postgres template: metadata: labels: - app.kubernetes.io/name: ax-redis + app.kubernetes.io/name: ax-postgres spec: - # Give Redis time to flush the append-only file on shutdown. + # Give Postgres time for a clean shutdown checkpoint. terminationGracePeriodSeconds: 60 containers: - - name: redis - image: redis:7-alpine - # Redis is AX's only store, so it persists to disk: every write is - # appended to an AOF that is fsynced once a second, with periodic RDB - # snapshots on top for faster restarts. noeviction makes a full - # instance reject writes instead of silently dropping records. - args: - - --dir - - /data - - --appendonly - - "yes" - - --appendfsync - - everysec - - --save - - "900 1 300 10 60 10000" - - --maxmemory-policy - - noeviction + - 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: 6379 - name: redis + - containerPort: 5432 + name: postgres readinessProbe: exec: - command: ["redis-cli", "ping"] + command: ["pg_isready", "-U", "ax", "-d", "ax"] periodSeconds: 5 resources: requests: cpu: 100m - memory: 128Mi + memory: 256Mi limits: cpu: 1000m memory: 1Gi volumeMounts: - name: data - mountPath: /data + mountPath: /var/lib/postgresql/data volumeClaimTemplates: - metadata: name: data labels: - app.kubernetes.io/name: ax-redis + app.kubernetes.io/name: ax-postgres app.kubernetes.io/part-of: ax spec: accessModes: ["ReadWriteOnce"] @@ -88,15 +91,15 @@ spec: apiVersion: v1 kind: Service metadata: - name: ax-redis + name: ax-postgres namespace: ax-system labels: - app.kubernetes.io/name: ax-redis + app.kubernetes.io/name: ax-postgres app.kubernetes.io/part-of: ax spec: ports: - - port: 6379 - targetPort: 6379 - name: redis + - port: 5432 + targetPort: 5432 + name: postgres selector: - app.kubernetes.io/name: ax-redis + app.kubernetes.io/name: ax-postgres 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 91ea6853..22cfbc90 100755 --- a/hack/setup-e2e-kind.sh +++ b/hack/setup-e2e-kind.sh @@ -261,15 +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 workloads and Redis's volume -# 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 -k -n "${AX_SYSTEM_NAMESPACE}" delete statefulset ax-redis --ignore-not-found --wait -k -n "${AX_SYSTEM_NAMESPACE}" delete pvc -l app.kubernetes.io/name=ax-redis --ignore-not-found --wait -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 statefulset/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/observe.go b/internal/controller/observe.go deleted file mode 100644 index 2f38fdbd..00000000 --- a/internal/controller/observe.go +++ /dev/null @@ -1,149 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package controller - -import ( - "context" - "log/slog" - "time" - - "github.com/agent-substrate/substrate/pkg/proto/ateapipb" - "github.com/google/ax/internal/substrate" - "github.com/google/ax/pkg/apis/v1alpha1" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/status" -) - -// Task status ownership. -// -// A task's actor is the source of truth for whether the task is running, where, -// and whether it is healthy. AX does not keep a copy of that: the stored status -// holds only what AX itself knows (the task ID, workspace setup, and the phases -// v1alpha1.IsAXPhase accepts), and Observe fills in the rest from Substrate -// each time a task is read. - -// Observe fills in the Substrate-owned parts of each task's status (actor, -// workerIP, phase and the Ready condition) from the task's actor. It costs one -// GetActor call for a single task and one ListActors call per atespace -// otherwise. If Substrate cannot be reached, the affected tasks are reported -// with phase "Unknown". -func (r *TaskReconciler) Observe(ctx context.Context, tasks []*v1alpha1.Task) { - if len(tasks) == 1 { - t := tasks[0] - actor, err := r.client.GetActor(ctx, atespaceOf(t), t.GetMetadata().GetName()) - if err != nil && status.Code(err) == codes.NotFound { - actor, err = nil, nil - } - applyActor(t, actor, err) - return - } - - byAtespace := make(map[string][]*v1alpha1.Task) - for _, t := range tasks { - byAtespace[atespaceOf(t)] = append(byAtespace[atespaceOf(t)], t) - } - for atespace, group := range byAtespace { - actors, err := r.client.ListActors(ctx, atespace) - byName := make(map[string]*ateapipb.Actor, len(actors)) - for _, a := range actors { - byName[a.GetMetadata().GetName()] = a - } - for _, t := range group { - applyActor(t, byName[t.GetMetadata().GetName()], err) - } - } -} - -// applyActor overlays the actor's state on the task's status. actor is nil when -// the task has no actor; err is set when the actor could not be looked up. -func applyActor(task *v1alpha1.Task, actor *ateapipb.Actor, err error) { - if task.Status == nil { - task.Status = &v1alpha1.TaskStatus{} - } - st := task.Status - st.Actor = "" - st.WorkerIp = "" - now := time.Now() - - switch { - case err != nil: - slog.Warn("could not observe task actor", "atespace", atespaceOf(task), "task", task.GetMetadata().GetName(), "error", err) - if !v1alpha1.IsAXPhase(st.Phase) { - st.Phase = "Unknown" - setConditionIfChanged(task, condReady, "Unknown", "SubstrateUnavailable", err.Error(), now) - } - 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 - setConditionIfChanged(task, condReady, "False", "ActorNotFound", "The task's actor no longer exists on Substrate", now) - } - default: - st.Actor = actor.GetMetadata().GetName() - st.WorkerIp = substrate.WorkerIP(actor) - if v1alpha1.IsAXPhase(st.Phase) { - return - } - phase, ready, reason, message := actorPhase(task, actor) - st.Phase = phase - setConditionIfChanged(task, condReady, ready, reason, message, now) - } -} - -// actorPhase maps an actor's state to the task's phase and Ready condition. -func actorPhase(task *v1alpha1.Task, actor *ateapipb.Actor) (phase, ready, reason, message string) { - switch state := actor.GetStatus().GetState(); state { - case ateapipb.ActorState_ACTOR_STATE_RUNNING: - if task.GetStatus() != nil && conditionTrue(task, condWorkspaceReady) { - return "Running", "True", "TaskRunning", "Task is running and its workspace is ready" - } - return "Running", "False", "WorkspaceInitializing", "Waiting for workspace setup to complete" - case ateapipb.ActorState_ACTOR_STATE_RESUMING, ateapipb.ActorState_ACTOR_STATE_REVERTING: - return "Resuming", "False", "ActorResuming", "Actor is being resumed" - case ateapipb.ActorState_ACTOR_STATE_SUSPENDING: - return "Suspending", "False", "ActorSuspending", "Actor is being suspended" - case ateapipb.ActorState_ACTOR_STATE_SUSPENDED: - return "Suspended", "False", "TaskSuspended", "Task is suspended" - case ateapipb.ActorState_ACTOR_STATE_CRASHED: - msg := actor.GetStatus().GetCrash().GetMessage() - if msg == "" { - msg = "Actor crashed" - } - return v1alpha1.PhaseFailed, "False", "ActorCrashed", msg - case ateapipb.ActorState_ACTOR_STATE_DELETING: - return v1alpha1.PhaseTerminating, "False", "ActorDeleting", "Actor is being deleted" - default: - return "Unknown", "Unknown", "UnknownActorState", state.String() - } -} - -// setConditionIfChanged sets a condition, keeping its update time when nothing -// about it changed so that repeated reads report a stable timestamp. -func setConditionIfChanged(task *v1alpha1.Task, condType, condStatus, reason, message string, t time.Time) { - for _, c := range task.Status.Conditions { - if c.Type == condType && c.Status == condStatus && c.Reason == reason && c.Message == message { - return - } - } - setCondition(task, condType, condStatus, reason, message, t) -} - -func atespaceOf(task *v1alpha1.Task) string { - if a := task.GetMetadata().GetAtespace(); a != "" { - return a - } - return "default" -} diff --git a/internal/controller/observe_test.go b/internal/controller/observe_test.go deleted file mode 100644 index cd911167..00000000 --- a/internal/controller/observe_test.go +++ /dev/null @@ -1,197 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package controller_test - -import ( - "context" - "net" - "testing" - - "github.com/agent-substrate/substrate/pkg/proto/ateapipb" - "github.com/google/ax/internal/controller" - "github.com/google/ax/internal/substrate" - "github.com/google/ax/pkg/apis/v1alpha1" - "google.golang.org/grpc" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/credentials/insecure" - "google.golang.org/grpc/status" -) - -// actorsServer serves a fixed set of actors, keyed by "atespace/name". -type actorsServer struct { - ateapipb.UnimplementedControlServer - actors map[string]*ateapipb.Actor - err error - getCalls int - listCalls int -} - -func (a *actorsServer) GetActor(_ context.Context, req *ateapipb.GetActorRequest) (*ateapipb.Actor, error) { - a.getCalls++ - if a.err != nil { - return nil, a.err - } - if actor, ok := a.actors[req.GetActor().GetAtespace()+"/"+req.GetActor().GetName()]; ok { - return actor, nil - } - return nil, status.Error(codes.NotFound, "actor not found") -} - -func (a *actorsServer) ListActors(_ context.Context, req *ateapipb.ListActorsRequest) (*ateapipb.ListActorsResponse, error) { - a.listCalls++ - if a.err != nil { - return nil, a.err - } - resp := &ateapipb.ListActorsResponse{} - for _, actor := range a.actors { - if actor.GetMetadata().GetAtespace() == req.GetAtespace() { - resp.Actors = append(resp.Actors, actor) - } - } - return resp, nil -} - -func newObserver(t *testing.T, srv *actorsServer) *controller.TaskReconciler { - 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, srv) - go g.Serve(lis) - t.Cleanup(g.Stop) - - client, err := substrate.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) - if err != nil { - t.Fatalf("substrate client: %v", err) - } - t.Cleanup(func() { client.Close() }) - return controller.NewTaskReconciler(client, "", "", "") -} - -func actor(atespace, name string, state ateapipb.ActorState, workerIP string) *ateapipb.Actor { - a := &ateapipb.Actor{ - Metadata: &ateapipb.ResourceMetadata{Atespace: atespace, 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 TestObserve(t *testing.T) { - workspaceReady := &v1alpha1.Condition{Type: "WorkspaceReady", Status: "True"} - crashed := actor("default", "crashed", ateapipb.ActorState_ACTOR_STATE_CRASHED, "") - crashed.Status.Crash = &ateapipb.ActorCrash{Message: "OOMKilled"} - srv := &actorsServer{actors: map[string]*ateapipb.Actor{ - "default/running": actor("default", "running", ateapipb.ActorState_ACTOR_STATE_RUNNING, "10.0.0.5"), - "default/setting-up": actor("default", "setting-up", ateapipb.ActorState_ACTOR_STATE_RUNNING, "10.0.0.6"), - "default/suspended": actor("default", "suspended", ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ""), - "default/crashed": crashed, - "default/failed": actor("default", "failed", ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ""), - }} - r := newObserver(t, srv) - - tests := []struct { - task *v1alpha1.Task - wantPhase, wantIP, wantActor string - wantReady, wantReadyReason string - }{ - {storedTask("running", "", workspaceReady), "Running", "10.0.0.5", "running", "True", "TaskRunning"}, - {storedTask("setting-up", ""), "Running", "10.0.0.6", "setting-up", "False", "WorkspaceInitializing"}, - {storedTask("suspended", "", workspaceReady), "Suspended", "", "suspended", "False", "TaskSuspended"}, - {storedTask("crashed", ""), v1alpha1.PhaseFailed, "", "crashed", "False", "ActorCrashed"}, - // The actor was removed outside AX. - {storedTask("gone", ""), v1alpha1.PhaseFailed, "", "", "False", "ActorNotFound"}, - // 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: "ActorResumeFailed"}), v1alpha1.PhaseFailed, "", "failed", "False", "ActorResumeFailed"}, - {storedTask("not-yet", v1alpha1.PhasePending), v1alpha1.PhasePending, "", "", "", ""}, - } - for _, tt := range tests { - t.Run(tt.task.GetMetadata().GetName(), func(t *testing.T) { - r.Observe(context.Background(), []*v1alpha1.Task{tt.task}) - 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 - } - assertCondition(t, tt.task, "Ready", tt.wantReady, tt.wantReadyReason) - }) - } - - t.Run("crash message", func(t *testing.T) { - task := storedTask("crashed", "") - r.Observe(context.Background(), []*v1alpha1.Task{task}) - for _, c := range task.GetStatus().GetConditions() { - if c.GetType() == "Ready" && c.GetMessage() != "OOMKilled" { - t.Errorf("Ready message = %q, want the actor's crash message", c.GetMessage()) - } - } - }) -} - -func TestObserve_ListUsesOneCallPerAtespace(t *testing.T) { - srv := &actorsServer{actors: map[string]*ateapipb.Actor{ - "default/a": actor("default", "a", ateapipb.ActorState_ACTOR_STATE_RUNNING, "10.0.0.1"), - "default/b": actor("default", "b", ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ""), - "team-x/c": actor("team-x", "c", ateapipb.ActorState_ACTOR_STATE_RUNNING, "10.0.0.3"), - }} - r := newObserver(t, srv) - - tasks := []*v1alpha1.Task{storedTask("a", ""), storedTask("b", ""), storedTask("c", "")} - tasks[2].Metadata.Atespace = "team-x" - r.Observe(context.Background(), tasks) - - if srv.listCalls != 2 || srv.getCalls != 0 { - t.Errorf("ListActors calls = %d, GetActor calls = %d, want 2 and 0", srv.listCalls, srv.getCalls) - } - for i, want := range []string{"Running", "Suspended", "Running"} { - if got := tasks[i].GetStatus().GetPhase(); got != want { - t.Errorf("task %s phase = %q, want %q", tasks[i].GetMetadata().GetName(), got, want) - } - } - if ip := tasks[2].GetStatus().GetWorkerIp(); ip != "10.0.0.3" { - t.Errorf("task c workerIP = %q, want 10.0.0.3", ip) - } -} - -func TestObserve_SubstrateErrorReportsUnknown(t *testing.T) { - srv := &actorsServer{err: status.Error(codes.PermissionDenied, "denied")} - r := newObserver(t, srv) - - task := storedTask("a", "") - terminating := storedTask("b", v1alpha1.PhaseTerminating) - r.Observe(context.Background(), []*v1alpha1.Task{task}) - r.Observe(context.Background(), []*v1alpha1.Task{terminating}) - - if got := task.GetStatus().GetPhase(); got != "Unknown" { - t.Errorf("phase = %q, want Unknown", got) - } - assertCondition(t, task, "Ready", "Unknown", "SubstrateUnavailable") - if got := terminating.GetStatus().GetPhase(); got != v1alpha1.PhaseTerminating { - t.Errorf("phase = %q, want %q to be kept", got, v1alpha1.PhaseTerminating) - } -} diff --git a/internal/controller/reconciler.go b/internal/controller/reconciler.go index 15c6376a..aa261e3b 100644 --- a/internal/controller/reconciler.go +++ b/internal/controller/reconciler.go @@ -127,11 +127,11 @@ 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) + r.setNotReady(task, v1alpha1.ReasonSubstrateUnavailable, err.Error(), now) task.Status.Phase = v1alpha1.PhasePending return task, fmt.Errorf("ensuring atespace: %w", err) } - setCondition(task, "Ready", "False", "AtespaceCreationFailed", err.Error(), now) + r.setNotReady(task, v1alpha1.ReasonAtespaceCreationFailed, err.Error(), now) task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("ensuring atespace: %w", err) } @@ -185,7 +185,7 @@ 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) + r.setNotReady(task, v1alpha1.ReasonInvalidResources, err.Error(), now) task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("validating resources: %w", err) } @@ -201,7 +201,7 @@ 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) + r.setNotReady(task, v1alpha1.ReasonSubstrateUnavailable, err.Error(), now) task.Status.Phase = v1alpha1.PhasePending return task, fmt.Errorf("ensuring actor template: %w", err) } @@ -209,7 +209,7 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, ws // 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) + r.setNotReady(task, v1alpha1.ReasonTemplateCreationFailed, substrateErrorMessage(err), now) task.Status.Phase = v1alpha1.PhaseFailed return task, fmt.Errorf("ensuring actor template: %w", err) } @@ -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) + 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) + 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" - 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) + 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) + 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 := 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 !conditionTrue(task, condWorkspaceReady) { - 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()) } - setCondition(task, condReady, "True", "TaskRunning", "Task is running and its workspace is ready", time.Now()) } else { - setCondition(task, condWorkspaceReady, "False", "Initializing", fmt.Sprintf("Workspace is initializing at %s", workerIP), time.Now()) - 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,19 +344,10 @@ 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) { - 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. diff --git a/internal/controller/reconciler_test.go b/internal/controller/reconciler_test.go index 55451feb..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 @@ -511,7 +511,7 @@ func TestTaskReconciler_TemplateCreationFailed(t *testing.T) { 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 @@ -582,7 +582,7 @@ func TestTaskReconciler_TemplateCreationTransient(t *testing.T) { 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) } @@ -848,7 +844,7 @@ func TestTaskReconciler_ResourceLimitsFailurePaths(t *testing.T) { 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", @@ -871,7 +867,7 @@ func TestTaskReconciler_ResourceLimitsFailurePaths(t *testing.T) { 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) } @@ -892,7 +888,7 @@ func TestTaskReconciler_ResourceLimitsFailurePaths(t *testing.T) { 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) } @@ -1128,7 +1124,7 @@ func TestTaskReconciler_TransientSubstrateUnavailable(t *testing.T) { 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". @@ -1144,7 +1140,7 @@ func TestTaskReconciler_TransientSubstrateUnavailable(t *testing.T) { 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", @@ -1162,7 +1158,7 @@ func TestTaskReconciler_TransientSubstrateUnavailable(t *testing.T) { 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 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 index 4e5b364a..dcbd11b5 100644 --- a/internal/server/export_test.go +++ b/internal/server/export_test.go @@ -16,7 +16,7 @@ package server import "time" -// SetWatchPollInterval overrides how often WatchTask polls Substrate and +// 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 diff --git a/internal/server/server.go b/internal/server/server.go index 99351d61..38a2c729 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -22,6 +22,7 @@ import ( "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" @@ -39,22 +40,19 @@ type Reconciler interface { ReconcileDelete(ctx context.Context, atespace, taskName string) error } -// Observer fills in the parts of a task's status that are owned by its actor -// on Agent Substrate. The store never holds them, so every task the API returns -// passes through the Observer first. -type Observer interface { - Observe(ctx context.Context, tasks []*v1alpha1.Task) -} - // 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 - Observer Observer + // 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's actor, so changes -// made on Substrate (for example a crash) reach watchers without an AX write. +// 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. @@ -63,7 +61,7 @@ type Server struct { store store.Store locker lock.Locker reconciler Reconciler - observer Observer + substrate *substrate.Client grpcServer *grpc.Server } @@ -82,7 +80,7 @@ func NewServer(s store.Store, opts ...Options) *Server { store: s, locker: locker, reconciler: opt.Reconciler, - observer: opt.Observer, + substrate: opt.Substrate, grpcServer: grpc.NewServer(), } v1alpha1.RegisterAXServer(srv.grpcServer, srv) @@ -122,14 +120,13 @@ 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) } return nil, status.Errorf(codes.Internal, "getting task: %v", err) } - s.observe(ctx, task) return task, nil } @@ -146,11 +143,10 @@ 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) } - s.observe(ctx, tasks...) return &v1alpha1.ListTasksResponse{Tasks: tasks}, nil } @@ -179,7 +175,7 @@ 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() != v1alpha1.PhasePending || !proto.Equal(existing.GetSpec(), task.GetSpec()) { return nil, status.Errorf(codes.FailedPrecondition, "task %s/%s already exists and is immutable", atespace, taskName) @@ -198,10 +194,10 @@ func (s *Server) CreateTask(ctx context.Context, req *v1alpha1.CreateTaskRequest } // The task is Pending until its actor exists. New tasks start suspended. task.Status.Phase = v1alpha1.PhasePending - if err := s.store.SaveTask(ctx, withStoredStatus(task)); err != nil { + if err := s.saveTask(ctx, task); err != nil { return nil, status.Errorf(codes.Internal, "saving task: %v", err) } - task.Status.Phase = "Suspended" + task.Status.Phase = v1alpha1.PhaseSuspended // Directly reconcile with Substrate if s.reconciler != nil { @@ -216,7 +212,7 @@ func (s *Server) CreateTask(ctx context.Context, req *v1alpha1.CreateTaskRequest } else if task.Status.Phase != v1alpha1.PhasePending { task.Status.Phase = v1alpha1.PhaseFailed } - _ = s.store.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(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) } @@ -224,10 +220,12 @@ func (s *Server) CreateTask(ctx context.Context, req *v1alpha1.CreateTaskRequest } task.Status = reconciled.Status } - if err := s.store.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(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) } + // Report the actor as it is now, as GetTask would. + s.applyActor(ctx, task) return task, nil } @@ -248,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) @@ -261,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, storedStatus(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 } @@ -304,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) @@ -314,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)) @@ -325,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, storedStatus(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.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(task.Status)); 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 } @@ -354,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) @@ -364,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)) @@ -377,37 +361,232 @@ func (s *Server) ResumeTask(ctx context.Context, req *v1alpha1.ResumeTaskRequest } else if task.Status.Phase != v1alpha1.PhasePending { task.Status.Phase = v1alpha1.PhaseFailed } - _ = s.store.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(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, storedStatus(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.UpdateTaskStatus(ctx, atespace, taskName, storedStatus(task.Status)); 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 } -// observe fills in the actor-owned parts of the tasks' status, when the server -// is connected to Substrate. -func (s *Server) observe(ctx context.Context, tasks ...*v1alpha1.Task) { - if s.observer != nil && len(tasks) > 0 { - s.observer.Observe(ctx, tasks) +// 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: they are read from Substrate instead, so the store -// can never hold a stale copy of them. +// actor's state are dropped. st is not modified. func storedStatus(st *v1alpha1.TaskStatus) *v1alpha1.TaskStatus { out := &v1alpha1.TaskStatus{} if st != nil { @@ -419,7 +598,7 @@ func storedStatus(st *v1alpha1.TaskStatus) *v1alpha1.TaskStatus { out.Phase = "" conds := out.Conditions[:0] for _, c := range out.Conditions { - if c.GetType() != "Ready" { + if c.GetType() != v1alpha1.ConditionReady { conds = append(conds, c) } } @@ -428,13 +607,6 @@ func storedStatus(st *v1alpha1.TaskStatus) *v1alpha1.TaskStatus { return out } -// withStoredStatus returns a copy of task whose status is storedStatus. -func withStoredStatus(task *v1alpha1.Task) *v1alpha1.Task { - out := proto.Clone(task).(*v1alpha1.Task) - out.Status = storedStatus(task.Status) - 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 { @@ -458,15 +630,9 @@ 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) - } - defer closer.Close() var last *v1alpha1.Task send := func(task *v1alpha1.Task, action string) (done bool, err error) { - s.observe(ctx, task) if last != nil && proto.Equal(last, task) { return false, nil } @@ -475,35 +641,26 @@ func (s *Server) WatchTask(req *v1alpha1.WatchTaskRequest, stream grpc.ServerStr return true, err } phase := task.GetStatus().GetPhase() - return phase == v1alpha1.PhaseFailed || phase == "Completed", nil + return phase == v1alpha1.PhaseFailed, nil } - if initial, err := s.store.GetTask(ctx, atespace, req.Name); err == nil { + if initial, err := s.getTask(ctx, atespace, req.Name); err == nil { if done, err := send(initial, "INITIAL"); done || err != nil { return err } } - // Store events cover AX's own writes. Polling picks up changes that happen - // on Substrate alone, such as an actor crashing or being suspended. ticker := time.NewTicker(watchPollInterval) defer ticker.Stop() for { - var task *v1alpha1.Task select { case <-ctx.Done(): return ctx.Err() - case t, ok := <-ch: - if !ok { - return nil - } - task = t case <-ticker.C: - t, err := s.store.GetTask(ctx, atespace, req.Name) - if err != nil { - continue - } - task = t + } + 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 b7c20e4f..e8166663 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -24,9 +24,11 @@ import ( "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" @@ -537,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) @@ -558,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()) } } @@ -576,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) @@ -612,39 +614,79 @@ func assertNoActorState(t *testing.T, stored *v1alpha1.Task) { } } -// fakeObserver stands in for Substrate: it reports every task's actor, named -// after the task, in the given phase. -type fakeObserver struct { +// 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 - phase string + names []string + state ateapipb.ActorState workerIP string } -func (f *fakeObserver) setPhase(phase string) { +func (f *fakeSubstrate) setState(state ateapipb.ActorState) { f.mu.Lock() defer f.mu.Unlock() - f.phase = phase + f.state = state } -func (f *fakeObserver) Observe(_ context.Context, tasks []*v1alpha1.Task) { +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 _, task := range tasks { - if task.Status == nil { - task.Status = &v1alpha1.TaskStatus{} + for _, name := range f.names { + if req.GetActor().GetAtespace() == "default" && req.GetActor().GetName() == name { + return f.actor(name), nil } - task.Status.Actor = task.GetMetadata().GetName() - task.Status.WorkerIp = f.workerIP - if !v1alpha1.IsAXPhase(task.Status.Phase) { - task.Status.Phase = f.phase + } + 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_ReadsActorStateFromObserver(t *testing.T) { +func TestServer_ReadsActorStateFromSubstrate(t *testing.T) { st := memory.NewStore() - obs := &fakeObserver{phase: "Running", workerIP: "10.244.1.7"} - srv := server.NewServer(st, server.Options{Reconciler: &fakeReconciler{}, Observer: obs}) + 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"} { @@ -661,7 +703,7 @@ func TestServer_ReadsActorStateFromObserver(t *testing.T) { 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 observed actor state", got.GetStatus()) + t.Errorf("GetTask status = %v, want the actor's state", got.GetStatus()) } list, err := srv.ListTasks(ctx, &v1alpha1.ListTasksRequest{Atespace: "default"}) @@ -670,12 +712,12 @@ func TestServer_ReadsActorStateFromObserver(t *testing.T) { } 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 observed actor state", task.GetMetadata().GetName(), task.GetStatus()) + 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. - obs.setPhase("Suspended") + 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) @@ -762,7 +804,7 @@ func TestServer_TransientSubstrateUnavailableRecovery(t *testing.T) { reconcileErr: status.Error(codes.Unavailable, "connection refused"), 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}) @@ -811,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) } @@ -832,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) } @@ -876,8 +918,8 @@ func TestServer_WatchTaskSeesSubstrateOnlyChanges(t *testing.T) { defer server.SetWatchPollInterval(10 * time.Millisecond)() st := memory.NewStore() - obs := &fakeObserver{phase: "Suspended"} - srv := server.NewServer(st, server.Options{Reconciler: &fakeReconciler{}, Observer: obs}) + 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() @@ -913,7 +955,7 @@ func TestServer_WatchTaskSeesSubstrateOnlyChanges(t *testing.T) { } // The actor crashes on Substrate. AX writes nothing, but the watcher hears about it. - obs.setPhase(v1alpha1.PhaseFailed) + sub.setState(ateapipb.ActorState_ACTOR_STATE_CRASHED) evt, err = stream.Recv() if err != nil { t.Fatalf("Recv: %v", err) 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 de6814be..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,6 +79,12 @@ 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{} } @@ -83,18 +95,7 @@ func (s *MemoryStore) SaveTask(ctx context.Context, task *v1alpha1.Task) error { 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 467a636c..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 = v1alpha1.PhasePending - } - - 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/pkg/apis/v1alpha1/types.go b/pkg/apis/v1alpha1/types.go index 0ec510a9..2ffe025e 100644 --- a/pkg/apis/v1alpha1/types.go +++ b/pkg/apis/v1alpha1/types.go @@ -48,6 +48,73 @@ const ( // 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 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 eb8de76b..00000000 --- a/test/e2e/manifests/redis.yaml +++ /dev/null @@ -1,83 +0,0 @@ -apiVersion: apps/v1 -kind: StatefulSet -metadata: - name: ax-redis - namespace: ${AX_SYSTEM_NAMESPACE} - labels: - app.kubernetes.io/name: ax-redis - app.kubernetes.io/part-of: ax -spec: - replicas: 1 - serviceName: ax-redis - selector: - matchLabels: - app.kubernetes.io/name: ax-redis - template: - metadata: - labels: - app.kubernetes.io/name: ax-redis - spec: - # Give Redis time to flush the append-only file on shutdown. - terminationGracePeriodSeconds: 60 - containers: - - name: redis - image: redis:7-alpine - # Redis is AX's only store, so it persists to disk: every write is - # appended to an AOF that is fsynced once a second, with periodic RDB - # snapshots on top for faster restarts. noeviction makes a full - # instance reject writes instead of silently dropping records. - args: - - --dir - - /data - - --appendonly - - "yes" - - --appendfsync - - everysec - - --save - - "900 1 300 10 60 10000" - - --maxmemory-policy - - noeviction - ports: - - containerPort: 6379 - name: redis - readinessProbe: - exec: - command: ["redis-cli", "ping"] - periodSeconds: 5 - resources: - requests: - cpu: 100m - memory: 128Mi - limits: - cpu: 1000m - memory: 1Gi - volumeMounts: - - name: data - mountPath: /data - volumeClaimTemplates: - - metadata: - name: data - labels: - app.kubernetes.io/name: ax-redis - app.kubernetes.io/part-of: ax - spec: - accessModes: ["ReadWriteOnce"] - resources: - requests: - storage: 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 From 0eeb8b71df49a7c11477c19009f1a3784b4ae7fe Mon Sep 17 00:00:00 2001 From: Jaana Dogan Date: Mon, 26 Oct 2026 18:51:52 -0700 Subject: [PATCH 3/3] ax-server: never log passwords from the store connection string redactURL relied on url.Redacted, which only masks a password in the URL's userinfo. A password passed as a query parameter was logged as-is, and a keyword/value DSN parsed as a bare path and was logged in full. Replace it with describeStore, which parses the string with pgconn.ParseConfig, the same parser the store uses, and logs only user@host:port/database. The password is never formatted, so the result is safe whatever form the connection string takes. --- cmd/ax-server/main.go | 20 ++++++++----- cmd/ax-server/main_test.go | 60 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 7 deletions(-) create mode 100644 cmd/ax-server/main_test.go diff --git a/cmd/ax-server/main.go b/cmd/ax-server/main.go index 23ab5eed..b955a51d 100644 --- a/cmd/ax-server/main.go +++ b/cmd/ax-server/main.go @@ -17,11 +17,13 @@ package main import ( "context" "flag" + "fmt" "log/slog" + "net" "net/http" - "net/url" "os" "os/signal" + "strconv" "syscall" "time" @@ -30,6 +32,7 @@ import ( "github.com/google/ax/internal/store/postgres" "github.com/google/ax/internal/substrate" "github.com/google/ax/pkg/apis/v1alpha1" + "github.com/jackc/pgx/v5/pgconn" ) func main() { @@ -78,7 +81,7 @@ func main() { slog.Info("starting ax-server", "listenAddr", listenAddr, - "store", redactURL(storeURL), + "store", describeStore(storeURL), "substrateEndpoint", substrateEndpoint, "template", defaultTemplate, "defaultTaskImage", defaultTaskImage, @@ -88,7 +91,7 @@ func main() { backend, err := postgres.Open(openCtx, storeURL, postgres.Options{}) openCancel() if err != nil { - slog.Error("opening store", "store", redactURL(storeURL), "error", err) + slog.Error("opening store", "store", describeStore(storeURL), "error", err) os.Exit(1) } defer backend.Close() @@ -141,11 +144,14 @@ func main() { _ = httpServer.Shutdown(shutdownCtx) } -// redactURL hides any password in a store URL so that it can be logged. -func redactURL(raw string) string { - u, err := url.Parse(raw) +// 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 u.Redacted() + 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) + } + }) + } +}