From d6294613a5fdef6c401e377e34a83b418bdd258d Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Mon, 17 Aug 2026 16:19:03 -0400 Subject: [PATCH 01/14] postgres retry Signed-off-by: Jet Chiang --- cmd/ateapi/main.go | 34 ++++++++++++++++++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/cmd/ateapi/main.go b/cmd/ateapi/main.go index 5d1d3e0d5a..345ea631f0 100644 --- a/cmd/ateapi/main.go +++ b/cmd/ateapi/main.go @@ -43,6 +43,7 @@ import ( "github.com/agent-substrate/substrate/pkg/client/clientset/versioned" "github.com/agent-substrate/substrate/pkg/client/informers/externalversions" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/jackc/pgx/v5/pgxpool" "github.com/redis/go-redis/v9" "github.com/spf13/pflag" "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" @@ -133,7 +134,7 @@ func main() { serverboot.Fatal(ctx, "Failed to initialize JWT providers", err) } - persistence, err := connectStore(ctx) + persistence, err := connectStore(shutdownCtx) if err != nil { serverboot.Fatal(ctx, "Failed to set up persistence backend", err) } @@ -331,7 +332,10 @@ func connectStore(ctx context.Context) (store.Interface, error) { if *postgresConnectionString == "" { return nil, fmt.Errorf("--store-backend=postgres requires --postgres-connection-string") } - persistence, err := atepg.Connect(ctx, *postgresConnectionString) + if _, err := pgxpool.ParseConfig(*postgresConnectionString); err != nil { + return nil, fmt.Errorf("parsing PostgreSQL connection string: %w", err) + } + persistence, err := connectPostgresWithRetries(ctx) if err != nil { return nil, fmt.Errorf("setting up PostgreSQL: %w", err) } @@ -341,6 +345,32 @@ func connectStore(ctx context.Context) (store.Interface, error) { } } +var ( + postgresConnectTries = 30 + postgresConnectPeriod = 2 * time.Second +) + +func connectPostgresWithRetries(ctx context.Context) (*atepg.Persistence, error) { + var connectErr error + for attempt := 1; attempt <= postgresConnectTries; attempt++ { + persistence, err := atepg.Connect(ctx, *postgresConnectionString) + if err == nil { + return persistence, nil + } + connectErr = err + slog.WarnContext(ctx, "Failed to connect to PostgreSQL, retrying...", slog.Int("attempt", attempt), slog.Any("err", err)) + if attempt == postgresConnectTries { + break + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(postgresConnectPeriod): + } + } + return nil, fmt.Errorf("connect to PostgreSQL after %d attempts: %w", postgresConnectTries, connectErr) +} + // connectRedis builds the Redis/Valkey TLS config, plumbs IAM auth if // requested, opens the cluster client, and pings with retries. func connectRedis(ctx context.Context) (*redis.ClusterClient, error) { From 0c6613fcbf6bf055b2fc8e6850858dbcd3b48252 Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Mon, 17 Aug 2026 16:22:05 -0400 Subject: [PATCH 02/14] fix: fk constraints on atv reference to at, improve pagetoken errors in list ops, and replace snapshot tag creation precheck with fk Signed-off-by: Jet Chiang --- cmd/ateapi/internal/store/atepg/atepg.go | 53 +++++++++++++++--- cmd/ateapi/internal/store/atepg/atepg_test.go | 55 +++++++++++++++++++ cmd/ateapi/internal/store/atepg/schema.go | 10 +++- 3 files changed, 108 insertions(+), 10 deletions(-) diff --git a/cmd/ateapi/internal/store/atepg/atepg.go b/cmd/ateapi/internal/store/atepg/atepg.go index 4583839de4..c29117dbd7 100644 --- a/cmd/ateapi/internal/store/atepg/atepg.go +++ b/cmd/ateapi/internal/store/atepg/atepg.go @@ -128,6 +128,14 @@ func pgErrCode(err error) string { return "" } +func pgErrConstraint(err error) string { + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) { + return pgErr.ConstraintName + } + return "" +} + // --- Atespaces --- func (p *Persistence) CreateAtespace(ctx context.Context, atespace *ateapipb.Atespace) (*ateapipb.Atespace, error) { @@ -183,6 +191,10 @@ func (p *Persistence) AtespaceExists(ctx context.Context, name string) (bool, er } func (p *Persistence) ListAtespaces(ctx context.Context, opts store.ListOptions) (store.ListResponse[*ateapipb.Atespace], error) { + opts, err := store.NormalizeListOptions(opts) + if err != nil { + return store.ListResponse[*ateapipb.Atespace]{}, err + } pageSize, pageTokenStr := opts.PageSize, opts.PageToken token, err := decodePageToken(pageTokenStr, kindAtespace, "", 1) if err != nil { @@ -360,6 +372,10 @@ func (p *Persistence) UpdateActorTemplate(ctx context.Context, templateRef resou } func (p *Persistence) ListActorTemplates(ctx context.Context, atespace string, opts store.ListOptions) (store.ListResponse[*ateapipb.ActorTemplate], error) { + opts, err := store.NormalizeListOptions(opts) + if err != nil { + return store.ListResponse[*ateapipb.ActorTemplate]{}, err + } pageSize, pageTokenStr := opts.PageSize, opts.PageToken keyParts := 2 if atespace != "" { @@ -620,9 +636,12 @@ func (p *Persistence) DeleteActor(ctx context.Context, actorRef resources.ActorR } func (p *Persistence) ListActors(ctx context.Context, atespace string, opts store.ListOptions) (store.ListResponse[*ateapipb.Actor], error) { + opts, err := store.NormalizeListOptions(opts) + if err != nil { + return store.ListResponse[*ateapipb.Actor]{}, err + } var items []*ateapipb.Actor var nextToken string - var err error if atespace != "" { items, nextToken, err = p.listActorsScoped(ctx, atespace, opts.PageSize, opts.PageToken) } else { @@ -793,9 +812,12 @@ func (p *Persistence) GetActorSnapshotTag(ctx context.Context, atespace, name st } func (p *Persistence) ListActorSnapshots(ctx context.Context, atespace string, opts store.ListOptions) (store.ListResponse[*ateapipb.ActorSnapshot], error) { + opts, err := store.NormalizeListOptions(opts) + if err != nil { + return store.ListResponse[*ateapipb.ActorSnapshot]{}, err + } var items []*ateapipb.ActorSnapshot var nextToken string - var err error if atespace != "" { items, nextToken, err = p.listActorSnapshotsScoped(ctx, atespace, opts.PageSize, opts.PageToken) } else { @@ -915,10 +937,6 @@ func (p *Persistence) CreateActorSnapshotTag(ctx context.Context, snapshotAtespa return nil, fmt.Errorf("beginning actor snapshot tag create: %w", err) } defer tx.Rollback(ctx) //nolint:errcheck // no-op once committed - if _, err := getActorSnapshotRow(ctx, tx, snapshotAtespace, snapshotName); err != nil { - return nil, err - } - var inserted []byte err = tx.QueryRow(ctx, ` INSERT INTO actor_snapshot_tags @@ -934,7 +952,14 @@ func (p *Persistence) CreateActorSnapshotTag(ctx context.Context, snapshotAtespa return dbTag, nil } if isForeignKeyViolation(err) { - return nil, store.ErrFailedPrecondition + switch pgErrConstraint(err) { + case "actor_snapshot_tags_snapshot_fk": + return nil, store.ErrNotFound + case "actor_snapshot_tags_atespace_fk": + return nil, store.ErrFailedPrecondition + default: + return nil, fmt.Errorf("inserting actor snapshot tag %s/%s violated unknown foreign key %q: %w", tagAtespace, tagName, pgErrConstraint(err), err) + } } if !errors.Is(err, pgx.ErrNoRows) { return nil, fmt.Errorf("inserting actor snapshot tag %s/%s: %w", tagAtespace, tagName, err) @@ -1228,6 +1253,10 @@ func (p *Persistence) DeleteWorker(ctx context.Context, namespace, poolName, pod } func (p *Persistence) ListWorkers(ctx context.Context, opts store.ListOptions) (store.ListResponse[*ateapipb.Worker], error) { + opts, err := store.NormalizeListOptions(opts) + if err != nil { + return store.ListResponse[*ateapipb.Worker]{}, err + } pageSize, pageTokenStr := opts.PageSize, opts.PageToken token, err := decodePageToken(pageTokenStr, kindWorker, "", 3) if err != nil { @@ -1333,6 +1362,9 @@ const defaultLockTTL = 30 * time.Second func (p *Persistence) AcquireLock(ctx context.Context, key string) (*store.Lock, error) { ttl := p.lockTTL token := uuid.NewString() + if err := p.cleanupExpiredLeases(ctx); err != nil { + slog.WarnContext(ctx, "failed to clean up expired PostgreSQL leases", "error", err) + } acquired, err := p.acquireLease(ctx, key, token, ttl) if err != nil { @@ -1363,6 +1395,13 @@ func (p *Persistence) AcquireLock(ctx context.Context, key string) (*store.Lock, return store.NewLock(leaseCtx, closeFn), nil } +func (p *Persistence) cleanupExpiredLeases(ctx context.Context) error { + if _, err := p.pool.Exec(ctx, `DELETE FROM leases WHERE expires_at <= clock_timestamp()`); err != nil { + return fmt.Errorf("deleting expired leases: %w", err) + } + return nil +} + func (p *Persistence) acquireLease(ctx context.Context, key, token string, ttl time.Duration) (bool, error) { var returnedKey string err := p.pool.QueryRow(ctx, ` diff --git a/cmd/ateapi/internal/store/atepg/atepg_test.go b/cmd/ateapi/internal/store/atepg/atepg_test.go index f19663fc99..a0b2f81798 100644 --- a/cmd/ateapi/internal/store/atepg/atepg_test.go +++ b/cmd/ateapi/internal/store/atepg/atepg_test.go @@ -129,6 +129,61 @@ func newTestAtespace(name string) *ateapipb.Atespace { return &ateapipb.Atespace{Metadata: &ateapipb.ResourceMetadata{Name: name}} } +func createTestAtespace(t *testing.T, s *Persistence, name string) { + t.Helper() + if _, err := s.CreateAtespace(context.Background(), newTestAtespace(name)); err != nil { + t.Fatalf("CreateAtespace(%q) failed: %v", name, err) + } +} + +func TestCreateActorSnapshotTag_ForeignKeyErrors(t *testing.T) { + s := setupPostgresPersistence(t) + ctx := context.Background() + createTestAtespace(t, s, "team-a") + tag := func() *ateapipb.ActorSnapshotTag { + return &ateapipb.ActorSnapshotTag{Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "latest"}} + } + + if _, err := s.CreateActorSnapshotTag(ctx, "team-a", "missing", tag()); !errors.Is(err, store.ErrNotFound) { + t.Errorf("missing snapshot error = %v, want ErrNotFound", err) + } + if _, err := s.CreateActorSnapshot(ctx, &ateapipb.ActorSnapshot{Metadata: &ateapipb.ResourceMetadata{Atespace: "gone", Name: "snapshot"}}); err != nil { + t.Fatalf("CreateActorSnapshot: %v", err) + } + tagWithoutAtespace := tag() + tagWithoutAtespace.Metadata.Atespace = "gone" + if _, err := s.CreateActorSnapshotTag(ctx, "gone", "snapshot", tagWithoutAtespace); !errors.Is(err, store.ErrFailedPrecondition) { + t.Errorf("missing tag atespace error = %v, want ErrFailedPrecondition", err) + } +} + +func TestAcquireLock_CleansExpiredLeases(t *testing.T) { + s := setupPostgresPersistence(t) + ctx := context.Background() + if _, err := s.pool.Exec(ctx, ` + INSERT INTO leases (key, token, expires_at) VALUES + ('expired', 'old', clock_timestamp() - interval '1 minute'), + ('active', 'live', clock_timestamp() + interval '1 hour')`); err != nil { + t.Fatalf("seeding leases: %v", err) + } + lock, err := s.AcquireLock(ctx, "new") + if err != nil { + t.Fatalf("AcquireLock: %v", err) + } + defer lock.Close() + + var expired, active int + if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM leases WHERE key = 'expired'`).Scan(&expired); err != nil { + t.Fatalf("counting expired lease: %v", err) + } + if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM leases WHERE key = 'active'`).Scan(&active); err != nil { + t.Fatalf("counting active lease: %v", err) + } + if expired != 0 || active != 1 { + t.Errorf("lease counts = expired:%d active:%d, want 0 and 1", expired, active) + } +} + // TestCreateActor_MissingAtespace_FailedPrecondition exercises the // foreign-key race the doc calls out: CreateActor rejects an actor whose // atespace doesn't exist (including a concurrently-deleted one), closing the diff --git a/cmd/ateapi/internal/store/atepg/schema.go b/cmd/ateapi/internal/store/atepg/schema.go index 078a1115ad..63fb62c096 100644 --- a/cmd/ateapi/internal/store/atepg/schema.go +++ b/cmd/ateapi/internal/store/atepg/schema.go @@ -60,15 +60,17 @@ CREATE TABLE IF NOT EXISTS actor_snapshots ( ); CREATE TABLE IF NOT EXISTS actor_snapshot_tags ( - atespace text NOT NULL - REFERENCES atespaces(name) ON DELETE RESTRICT, + atespace text NOT NULL, name text NOT NULL, snapshot_atespace text NOT NULL, snapshot_name text NOT NULL, version bigint NOT NULL, proto bytea NOT NULL, PRIMARY KEY (atespace, name), - FOREIGN KEY (snapshot_atespace, snapshot_name) + CONSTRAINT actor_snapshot_tags_atespace_fk + FOREIGN KEY (atespace) REFERENCES atespaces(name) ON DELETE RESTRICT, + CONSTRAINT actor_snapshot_tags_snapshot_fk + FOREIGN KEY (snapshot_atespace, snapshot_name) REFERENCES actor_snapshots(atespace, name) ON DELETE RESTRICT ); @@ -86,6 +88,8 @@ CREATE TABLE IF NOT EXISTS leases ( token text NOT NULL, expires_at timestamptz NOT NULL ); + +CREATE INDEX IF NOT EXISTS leases_expires_at_idx ON leases (expires_at); ` // applySchema idempotently creates atepg's tables. From 728e58f93fbbc9e105a6029f910efd469bb6c45c Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Mon, 17 Aug 2026 16:24:01 -0400 Subject: [PATCH 03/14] pagination errors and handlers Signed-off-by: Jet Chiang --- .../internal/controlapi/actor_snapshot.go | 2 +- cmd/ateapi/internal/controlapi/atespace.go | 2 +- cmd/ateapi/internal/controlapi/pagination.go | 18 +++++++++++ cmd/ateapi/internal/controlapi/worker.go | 2 +- cmd/ateapi/internal/store/store.go | 30 +++++++++++++++++-- 5 files changed, 49 insertions(+), 5 deletions(-) diff --git a/cmd/ateapi/internal/controlapi/actor_snapshot.go b/cmd/ateapi/internal/controlapi/actor_snapshot.go index 7e34d8de06..9496b668df 100644 --- a/cmd/ateapi/internal/controlapi/actor_snapshot.go +++ b/cmd/ateapi/internal/controlapi/actor_snapshot.go @@ -107,7 +107,7 @@ func (s *Service) ListActorSnapshots(ctx context.Context, req *ateapipb.ListActo } page, err := s.persistence.ListActorSnapshots(ctx, req.GetAtespace(), store.ListOptions{PageSize: effectivePageSize(req.GetPageSize()), PageToken: req.GetPageToken()}) if err != nil { - return nil, fmt.Errorf("while listing actor snapshots: %w", err) + return nil, mapListError(fmt.Errorf("while listing actor snapshots: %w", err)) } return &ateapipb.ListActorSnapshotsResponse{Snapshots: page.Items, NextPageToken: page.NextPageToken}, nil } diff --git a/cmd/ateapi/internal/controlapi/atespace.go b/cmd/ateapi/internal/controlapi/atespace.go index 220275791e..8d7fd18913 100644 --- a/cmd/ateapi/internal/controlapi/atespace.go +++ b/cmd/ateapi/internal/controlapi/atespace.go @@ -110,7 +110,7 @@ func (s *Service) ListAtespaces(ctx context.Context, req *ateapipb.ListAtespaces page, err := s.persistence.ListAtespaces(ctx, store.ListOptions{PageSize: effectivePageSize(req.GetPageSize()), PageToken: req.GetPageToken()}) if err != nil { - return nil, fmt.Errorf("while listing atespaces in db: %w", err) + return nil, mapListError(fmt.Errorf("while listing atespaces in db: %w", err)) } return &ateapipb.ListAtespacesResponse{ Atespaces: page.Items, diff --git a/cmd/ateapi/internal/controlapi/pagination.go b/cmd/ateapi/internal/controlapi/pagination.go index 1d9d591f2f..9a458be401 100644 --- a/cmd/ateapi/internal/controlapi/pagination.go +++ b/cmd/ateapi/internal/controlapi/pagination.go @@ -14,6 +14,14 @@ package controlapi +import ( + "errors" + + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + const maxPageSize = 1000 // effectivePageSize applies the server-chosen default for an unset page_size @@ -24,3 +32,13 @@ func effectivePageSize(requested int32) int32 { } return requested } + +func mapListError(err error) error { + if errors.Is(err, store.ErrInvalidPageToken) { + return status.Error(codes.InvalidArgument, "invalid page_token") + } + if errors.Is(err, store.ErrInvalidPageSize) { + return status.Error(codes.InvalidArgument, "invalid page_size") + } + return err +} diff --git a/cmd/ateapi/internal/controlapi/worker.go b/cmd/ateapi/internal/controlapi/worker.go index 5c394742bd..0e7cc3ef12 100644 --- a/cmd/ateapi/internal/controlapi/worker.go +++ b/cmd/ateapi/internal/controlapi/worker.go @@ -30,7 +30,7 @@ func (s *Service) ListWorkers(ctx context.Context, req *ateapipb.ListWorkersRequ page, err := s.persistence.ListWorkers(ctx, store.ListOptions{PageSize: effectivePageSize(req.GetPageSize()), PageToken: req.GetPageToken()}) if err != nil { - return nil, fmt.Errorf("while listing workers in db: %w", err) + return nil, mapListError(fmt.Errorf("while listing workers in db: %w", err)) } return &ateapipb.ListWorkersResponse{ Workers: page.Items, diff --git a/cmd/ateapi/internal/store/store.go b/cmd/ateapi/internal/store/store.go index 534740a7ac..0ae7b1b0f2 100644 --- a/cmd/ateapi/internal/store/store.go +++ b/cmd/ateapi/internal/store/store.go @@ -42,6 +42,13 @@ var ( // ErrLockConflict indicates that a distributed lock is already held by another client. ErrLockConflict = errors.New("persistence: lock conflict") + // ErrInvalidPageToken indicates that a list page token is malformed or was + // issued for a different list operation or scope. + ErrInvalidPageToken = errors.New("persistence: invalid page token") + + // ErrInvalidPageSize indicates that a negative page size was supplied. + ErrInvalidPageSize = errors.New("persistence: invalid page size") + // ErrUIDConflict indicates a precondition pinned a uid the stored object does // not carry, meaning the name now addresses a different incarnation. Retrying // can never resolve it. @@ -52,7 +59,8 @@ var ( type Interface interface { // Stores a new actor in suspended state and returns the stored resource with // server-assigned metadata (uid, version, timestamps). The input is not - // mutated. Returns ErrAlreadyExists if key is taken. + // mutated. Returns ErrAlreadyExists if key is taken, or + // ErrFailedPrecondition if the actor's atespace does not exist. CreateActor(ctx context.Context, actor *ateapipb.Actor) (*ateapipb.Actor, error) // Fetches an actor by reference. Returns ErrNotFound if missing. @@ -90,7 +98,9 @@ type Interface interface { // Lists ActorSnapshots in one atespace, or all atespaces when empty. ListActorSnapshots(ctx context.Context, atespace string, opts ListOptions) (ListResponse[*ateapipb.ActorSnapshot], error) - // Adds an immutable Atespace-owned tag to an ActorSnapshot. + // Adds an immutable Atespace-owned tag to an ActorSnapshot. Returns + // ErrNotFound if the snapshot does not exist, or ErrFailedPrecondition if + // the tag's atespace does not exist. CreateActorSnapshotTag(ctx context.Context, atespace, name string, tag *ateapipb.ActorSnapshotTag) (*ateapipb.ActorSnapshotTag, error) // Fetches an Atespace-owned tag. Returns ErrNotFound if missing. The tag's @@ -307,6 +317,22 @@ type ListOptions struct { PageToken string } +// DefaultPageSize is used by store implementations when PageSize is unset. +const DefaultPageSize int32 = 1000 + +// NormalizeListOptions applies the store default and rejects invalid sizes. +// RPC handlers validate user input separately, but store callers also need a +// safe contract because list implementations use PageSize in slice indexes. +func NormalizeListOptions(opts ListOptions) (ListOptions, error) { + if opts.PageSize < 0 { + return ListOptions{}, ErrInvalidPageSize + } + if opts.PageSize == 0 { + opts.PageSize = DefaultPageSize + } + return opts, nil +} + // ListResponse is the return value of a List method: the page of items it // addressed, plus the token to fetch the next page. NextPageToken is empty // once the listing has reached its last page. From 8c532050dbb630f26de94efd7a9ceff64e718642 Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Mon, 17 Aug 2026 16:28:52 -0400 Subject: [PATCH 04/14] fix: handle ErrFailedPrecondition when atespace not found in CreateActor RPC Signed-off-by: Jet Chiang --- cmd/ateapi/internal/controlapi/actor.go | 5 +++- cmd/ateapi/internal/controlapi/actor_test.go | 27 ++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/cmd/ateapi/internal/controlapi/actor.go b/cmd/ateapi/internal/controlapi/actor.go index c7100b63fa..a5c52881cf 100644 --- a/cmd/ateapi/internal/controlapi/actor.go +++ b/cmd/ateapi/internal/controlapi/actor.go @@ -108,6 +108,9 @@ func (s *Service) CreateActor(ctx context.Context, req *ateapipb.CreateActorRequ if errors.Is(err, store.ErrAlreadyExists) { return nil, status.Errorf(codes.AlreadyExists, "Actor %s already exists", name) } + if errors.Is(err, store.ErrFailedPrecondition) { + return nil, status.Errorf(codes.FailedPrecondition, "Atespace %s not found", atespace) + } return nil, fmt.Errorf("while recording actor: %w", err) } @@ -244,7 +247,7 @@ func (s *Service) ListActors(ctx context.Context, req *ateapipb.ListActorsReques page, err := s.persistence.ListActors(ctx, req.GetAtespace(), store.ListOptions{PageSize: effectivePageSize(req.GetPageSize()), PageToken: req.GetPageToken()}) if err != nil { - return nil, fmt.Errorf("while listing actors in db: %w", err) + return nil, mapListError(fmt.Errorf("while listing actors in db: %w", err)) } return &ateapipb.ListActorsResponse{ Actors: page.Items, diff --git a/cmd/ateapi/internal/controlapi/actor_test.go b/cmd/ateapi/internal/controlapi/actor_test.go index 333ff65df9..c0675f23cf 100644 --- a/cmd/ateapi/internal/controlapi/actor_test.go +++ b/cmd/ateapi/internal/controlapi/actor_test.go @@ -21,6 +21,7 @@ import ( "testing" "time" + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" "github.com/agent-substrate/substrate/cmd/ateapi/internal/store/storetest" "github.com/agent-substrate/substrate/internal/ateattr" "github.com/agent-substrate/substrate/internal/resources" @@ -39,6 +40,32 @@ import ( "k8s.io/apimachinery/pkg/util/wait" ) +type createActorErrorStore struct { + serviceStore + err error +} + +func (s *createActorErrorStore) CreateActor(context.Context, *ateapipb.Actor) (*ateapipb.Actor, error) { + return nil, s.err +} + +func TestCreateActor_AtespaceDeletedAfterPrecheck(t *testing.T) { + ns := namespaceForTest("ns-create-atespace-race") + tc := setupTest(t, ns) + defer tc.cleanup() + createTemplate(t, tc, ns) + tc.service.persistence = &createActorErrorStore{serviceStore: tc.service.persistence, err: store.ErrFailedPrecondition} + + _, err := tc.service.CreateActor(context.Background(), &ateapipb.CreateActorRequest{Actor: &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: "racing-create"}, + ActorTemplateNamespace: ns, + ActorTemplateName: "tmpl1", + }}) + if status.Code(err) != codes.FailedPrecondition { + t.Fatalf("CreateActor status = %v, want FailedPrecondition (error: %v)", status.Code(err), err) + } +} + // CreateActor is the only lifecycle op with the full identity (incl. version) // available in the request, so the whole ate.* set should land on its span. func TestCreateActor_StampsFullSpanIdentity(t *testing.T) { From 908e61e03f9c5270814fbe71d776ef9bc72a4c82 Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Mon, 17 Aug 2026 16:47:30 -0400 Subject: [PATCH 05/14] fix race conditions in worker resume flow Signed-off-by: Jet Chiang --- .../internal/controlapi/workflow_resume.go | 19 +++++ .../controlapi/workflow_resume_test.go | 82 +++++++++++++++++++ cmd/ateapi/internal/store/atepg/pagetoken.go | 14 ++-- .../internal/workercache/workercache.go | 10 +++ 4 files changed, 119 insertions(+), 6 deletions(-) diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index 5dfce905f1..b2a04ac978 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume.go @@ -83,6 +83,18 @@ func (w *ActorWorkflow) ResumeActor(ctx context.Context, actorRef resources.Acto lifecycleOpAttrs(actor, actorTemplate, tele.SnapshotKind, tele.WireSnapshotScope)...) }() + // Routed requests call ResumeActor even when the actor is already running. + // Read before taking the distributed lease so that hot-path checks do not + // upsert and delete a PostgreSQL lease row. Any state that needs work is read + // again under the lock below. + actor, err = w.store.GetActor(ctx, actorRef) + if err != nil { + return nil, false, err + } + if wasRunning = actor.GetStatus() == ateapipb.Actor_STATUS_RUNNING; wasRunning { + return actor, false, nil + } + lockCtx, lock, err := w.acquireActorLock(ctx, actorRef) if err != nil { return nil, false, err @@ -329,6 +341,9 @@ func (w *ActorWorkflow) ensureWorkerAssigned(ctx context.Context, actorRef resou return false, attemptErr }) if err != nil { + if wait.Interrupted(err) && !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + return nil, nil, store.ErrVersionConflict + } return nil, nil, err } return assignedActor, assignedWorker, nil @@ -505,6 +520,10 @@ func (w *ActorWorkflow) assignWorkerAttempt(ctx context.Context, actorRef resour } if err := w.store.UpdateWorker(ctx, assignedWorker, assignedWorker.Version); err != nil { + if errors.Is(err, store.ErrNotFound) { + w.workerCache.Forget(assignedWorker.GetWorkerNamespace(), assignedWorker.GetWorkerPod()) + return nil, nil, fmt.Errorf("selected worker disappeared before claim: %w", store.ErrVersionConflict) + } return nil, nil, err } diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index 0439571ee4..f3c1391e29 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume_test.go @@ -62,6 +62,88 @@ func TestSchedulerRecordable(t *testing.T) { } } +type lockCountingStore struct { + store.Interface + acquireCalls int +} + +func (s *lockCountingStore) AcquireLock(ctx context.Context, key string) (*store.Lock, error) { + s.acquireCalls++ + return s.Interface.AcquireLock(ctx, key) +} + +func TestResumeActor_RunningFastPathDoesNotAcquireLock(t *testing.T) { + ctx := context.Background() + persistence := newTestPersistence(t) + created, err := persistence.CreateActor(ctx, &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "id1"}, + Status: ateapipb.Actor_STATUS_RUNNING, + }) + if err != nil { + t.Fatalf("CreateActor: %v", err) + } + st := &lockCountingStore{Interface: persistence} + w := &ActorWorkflow{store: st} + + got, resumed, err := w.ResumeActor(ctx, resources.ActorRef{Atespace: "team-a", Name: "id1"}, false) + if err != nil { + t.Fatalf("ResumeActor: %v", err) + } + if resumed { + t.Error("ResumeActor resumed = true, want false") + } + if !proto.Equal(got, created) { + t.Errorf("ResumeActor actor = %v, want %v", got, created) + } + if st.acquireCalls != 0 { + t.Errorf("AcquireLock calls = %d, want 0", st.acquireCalls) + } +} + +type updateWorkerErrorStore struct { + store.Interface + err error +} + +func (s *updateWorkerErrorStore) UpdateWorker(context.Context, *ateapipb.Worker, int64) error { + return s.err +} + +func TestAssignWorkerAttempt_MissingSelectedWorkerIsRetried(t *testing.T) { + ctx := context.Background() + persistence := newTestPersistence(t) + actor, wc := seedAssignFixture(t, ctx, persistence) + st := &updateWorkerErrorStore{Interface: persistence, err: store.ErrNotFound} + w := &ActorWorkflow{store: st, workerCache: wc, scheduler: scheduling.New(wc)} + tmpl := &atev1alpha1.ActorTemplate{Spec: atev1alpha1.ActorTemplateSpec{SandboxClass: atev1alpha1.SandboxClassGvisor}} + + _, _, err := w.assignWorkerAttempt(ctx, resources.ActorRef{Atespace: "team-a", Name: "id1"}, actor, tmpl) + if !errors.Is(err, store.ErrVersionConflict) { + t.Fatalf("assignWorkerAttempt error = %v, want ErrVersionConflict", err) + } + workers, err := wc.Workers() + if err != nil { + t.Fatalf("Workers: %v", err) + } + if len(workers) != 0 { + t.Errorf("cached workers after missing claim = %d, want 0", len(workers)) + } +} + +func TestEnsureWorkerAssigned_ConflictExhaustionIsRetryable(t *testing.T) { + ctx := context.Background() + persistence := newTestPersistence(t) + actor, wc := seedAssignFixture(t, ctx, persistence) + st := &updateWorkerErrorStore{Interface: persistence, err: store.ErrVersionConflict} + w := &ActorWorkflow{store: st, workerCache: wc, scheduler: scheduling.New(wc)} + tmpl := &atev1alpha1.ActorTemplate{Spec: atev1alpha1.ActorTemplateSpec{SandboxClass: atev1alpha1.SandboxClassGvisor}} + + _, _, err := w.ensureWorkerAssigned(ctx, resources.ActorRef{Atespace: "team-a", Name: "id1"}, actor, tmpl) + if !errors.Is(err, store.ErrVersionConflict) { + t.Fatalf("ensureWorkerAssigned error = %v, want ErrVersionConflict", err) + } +} + func TestAssignWorkerAttempt_SkipsWorkerAssignedInOtherAtespace(t *testing.T) { ctx := context.Background() persistence := newTestPersistence(t) diff --git a/cmd/ateapi/internal/store/atepg/pagetoken.go b/cmd/ateapi/internal/store/atepg/pagetoken.go index 46e395e700..28e818a27e 100644 --- a/cmd/ateapi/internal/store/atepg/pagetoken.go +++ b/cmd/ateapi/internal/store/atepg/pagetoken.go @@ -18,6 +18,8 @@ import ( "encoding/base64" "encoding/json" "fmt" + + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" ) // pageTokenVersion guards against decoding a token produced by an incompatible @@ -61,23 +63,23 @@ func decodePageToken(tokenStr string, wantKind resourceKind, wantScope string, w } b, err := base64.StdEncoding.DecodeString(tokenStr) if err != nil { - return pageToken{}, fmt.Errorf("invalid page token: %w", err) + return pageToken{}, fmt.Errorf("%w: %v", store.ErrInvalidPageToken, err) } var token pageToken if err := json.Unmarshal(b, &token); err != nil { - return pageToken{}, fmt.Errorf("invalid page token: %w", err) + return pageToken{}, fmt.Errorf("%w: %v", store.ErrInvalidPageToken, err) } if token.Version != pageTokenVersion { - return pageToken{}, fmt.Errorf("invalid page token: unsupported version %d", token.Version) + return pageToken{}, fmt.Errorf("%w: unsupported version %d", store.ErrInvalidPageToken, token.Version) } if token.Kind != wantKind { - return pageToken{}, fmt.Errorf("invalid page token: for %q, used with %q", token.Kind, wantKind) + return pageToken{}, fmt.Errorf("%w: for %q, used with %q", store.ErrInvalidPageToken, token.Kind, wantKind) } if token.Scope != wantScope { - return pageToken{}, fmt.Errorf("invalid page token: for scope %q, used with scope %q", token.Scope, wantScope) + return pageToken{}, fmt.Errorf("%w: for scope %q, used with scope %q", store.ErrInvalidPageToken, token.Scope, wantScope) } if len(token.Last) != wantKeyParts { - return pageToken{}, fmt.Errorf("invalid page token: got %d key parts, want %d", len(token.Last), wantKeyParts) + return pageToken{}, fmt.Errorf("%w: got %d key parts, want %d", store.ErrInvalidPageToken, len(token.Last), wantKeyParts) } return token, nil } diff --git a/cmd/ateapi/internal/workercache/workercache.go b/cmd/ateapi/internal/workercache/workercache.go index 56b1583b47..1408f69b3a 100644 --- a/cmd/ateapi/internal/workercache/workercache.go +++ b/cmd/ateapi/internal/workercache/workercache.go @@ -110,6 +110,16 @@ func (c *Cache) Worker(namespace, pod string) (*ateapipb.Worker, error) { return worker, nil } +// Forget removes a worker that a store write proved no longer exists. The +// normal delete watch remains authoritative, but this closes the short race in +// which scheduling selected a worker just after its row was deleted and before +// the watch event reached the cache. +func (c *Cache) Forget(namespace, pod string) { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.workers, namespace+":"+pod) +} + func (c *Cache) sync(ctx context.Context) (*store.WorkerWatch, error) { watch, err := c.store.WatchWorkers(ctx) if err != nil { From b8f95f5e73ff2121f64b7a7877cd4fb2782e3abf Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Mon, 17 Aug 2026 16:48:56 -0400 Subject: [PATCH 06/14] add relevant changes to ateredis for store contract tests Signed-off-by: Jet Chiang --- cmd/ateapi/internal/store/ateredis/ateredis.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/cmd/ateapi/internal/store/ateredis/ateredis.go b/cmd/ateapi/internal/store/ateredis/ateredis.go index 3ed64a1745..06a657e762 100644 --- a/cmd/ateapi/internal/store/ateredis/ateredis.go +++ b/cmd/ateapi/internal/store/ateredis/ateredis.go @@ -1154,9 +1154,14 @@ func (s *Persistence) ListActors(ctx context.Context, atespace string, opts stor // listPage SCANs pattern across the redis masters from the page token, feeding key batches to collect and returns the next-page token. func (s *Persistence) listPage(ctx context.Context, pattern string, pageSize int32, pageTokenStr string, collect func(ctx context.Context, master *redis.Client, keys []string) (int, error)) (string, error) { + normalized, err := store.NormalizeListOptions(store.ListOptions{PageSize: pageSize}) + if err != nil { + return "", err + } + pageSize = normalized.PageSize token, err := decodePageToken(pageTokenStr) if err != nil { - return "", fmt.Errorf("invalid page token: %w", err) + return "", fmt.Errorf("%w: %v", store.ErrInvalidPageToken, err) } masters, err := s.getSortedMasters(ctx) From 5e5ca6ddb746f06415cf3eae4fc7b5f60cf01c8e Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Mon, 17 Aug 2026 16:49:13 -0400 Subject: [PATCH 07/14] update ateapi tests and store tests Signed-off-by: Jet Chiang --- .../controlapi/actor_snapshot_test.go | 17 ++++++ .../internal/controlapi/functional_test.go | 20 +++++++ cmd/ateapi/internal/store/store_test.go | 13 +++++ .../internal/store/storecontract/contract.go | 52 +++++++++++++++++++ 4 files changed, 102 insertions(+) diff --git a/cmd/ateapi/internal/controlapi/actor_snapshot_test.go b/cmd/ateapi/internal/controlapi/actor_snapshot_test.go index e7ef8c68bf..de8efb1b39 100644 --- a/cmd/ateapi/internal/controlapi/actor_snapshot_test.go +++ b/cmd/ateapi/internal/controlapi/actor_snapshot_test.go @@ -252,6 +252,23 @@ func TestValidateUpdateActorSnapshotTagRequest(t *testing.T) { } } +func TestCreateActorSnapshotTag_MissingSnapshotIsNotFound(t *testing.T) { + persistence, cleanup := storetest.SetupTestStore(t) + t.Cleanup(cleanup) + s := &Service{persistence: persistence} + + _, err := s.CreateActorSnapshotTag(context.Background(), &ateapipb.CreateActorSnapshotTagRequest{ + ActorSnapshotTag: &ateapipb.ActorSnapshotTag{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "latest"}, + Snapshot: &ateapipb.ObjectRef{Atespace: "team-a", Name: "missing"}, + Scope: ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_ATESPACE, + }, + }) + if status.Code(err) != codes.NotFound { + t.Fatalf("CreateActorSnapshotTag status = %v, want NotFound (error: %v)", status.Code(err), err) + } +} + func TestUpdateActorSnapshotTag_FieldMasks(t *testing.T) { tests := []struct { name string diff --git a/cmd/ateapi/internal/controlapi/functional_test.go b/cmd/ateapi/internal/controlapi/functional_test.go index 95756b7d06..aac17fb504 100644 --- a/cmd/ateapi/internal/controlapi/functional_test.go +++ b/cmd/ateapi/internal/controlapi/functional_test.go @@ -2979,16 +2979,36 @@ func TestValidation(t *testing.T) { assertGrpcErrorRegex(t, err, codes.InvalidArgument, "page_size: Invalid value") }) + t.Run("ListActors invalid token", func(t *testing.T) { + _, err := tc.client.ListActors(context.Background(), &ateapipb.ListActorsRequest{PageToken: "%%%"}) + assertGrpcError(t, err, codes.InvalidArgument, "invalid page_token") + }) + t.Run("ListWorkers", func(t *testing.T) { _, err := tc.client.ListWorkers(context.Background(), &ateapipb.ListWorkersRequest{PageSize: -1}) assertGrpcErrorRegex(t, err, codes.InvalidArgument, "page_size: Invalid value") }) + t.Run("ListWorkers invalid token", func(t *testing.T) { + _, err := tc.client.ListWorkers(context.Background(), &ateapipb.ListWorkersRequest{PageToken: "%%%"}) + assertGrpcError(t, err, codes.InvalidArgument, "invalid page_token") + }) + t.Run("ListAtespaces", func(t *testing.T) { _, err := tc.client.ListAtespaces(context.Background(), &ateapipb.ListAtespacesRequest{PageSize: -1}) assertGrpcErrorRegex(t, err, codes.InvalidArgument, "page_size: Invalid value") }) + t.Run("ListAtespaces invalid token", func(t *testing.T) { + _, err := tc.client.ListAtespaces(context.Background(), &ateapipb.ListAtespacesRequest{PageToken: "%%%"}) + assertGrpcError(t, err, codes.InvalidArgument, "invalid page_token") + }) + + t.Run("ListActorSnapshots invalid token", func(t *testing.T) { + _, err := tc.client.ListActorSnapshots(context.Background(), &ateapipb.ListActorSnapshotsRequest{PageToken: "%%%"}) + assertGrpcError(t, err, codes.InvalidArgument, "invalid page_token") + }) + t.Run("CreateAtespace", func(t *testing.T) { _, err := tc.client.CreateAtespace(context.Background(), &ateapipb.CreateAtespaceRequest{}) assertGrpcErrorRegex(t, err, codes.InvalidArgument, "atespace: Required value") diff --git a/cmd/ateapi/internal/store/store_test.go b/cmd/ateapi/internal/store/store_test.go index 38609b1275..7673362f5f 100644 --- a/cmd/ateapi/internal/store/store_test.go +++ b/cmd/ateapi/internal/store/store_test.go @@ -145,3 +145,16 @@ func TestWithPrecondition(t *testing.T) { }) } } + +func TestNormalizeListOptions(t *testing.T) { + got, err := NormalizeListOptions(ListOptions{}) + if err != nil { + t.Fatalf("NormalizeListOptions(zero) error = %v", err) + } + if got.PageSize != DefaultPageSize { + t.Errorf("NormalizeListOptions(zero).PageSize = %d, want %d", got.PageSize, DefaultPageSize) + } + if _, err := NormalizeListOptions(ListOptions{PageSize: -1}); !errors.Is(err, ErrInvalidPageSize) { + t.Errorf("NormalizeListOptions(negative) error = %v, want ErrInvalidPageSize", err) + } +} diff --git a/cmd/ateapi/internal/store/storecontract/contract.go b/cmd/ateapi/internal/store/storecontract/contract.go index b46a79ae05..3e9794b83e 100644 --- a/cmd/ateapi/internal/store/storecontract/contract.go +++ b/cmd/ateapi/internal/store/storecontract/contract.go @@ -115,9 +115,61 @@ func RunContractTests(t *testing.T, setup func(t *testing.T) store.Interface) { runActorTemplateContractTests(t, setup) runActorSnapshotContractTests(t, setup) runLockContractTests(t, setup) + runListOptionsContractTests(t, setup) runDebugContractTests(t, setup) } +func runListOptionsContractTests(t *testing.T, setup func(t *testing.T) store.Interface) { + t.Helper() + + t.Run("ListOptions_InvalidPageSize", func(t *testing.T) { + s := setup(t) + ctx := context.Background() + calls := []struct { + name string + call func(store.ListOptions) error + }{ + {"atespaces", func(opts store.ListOptions) error { _, err := s.ListAtespaces(ctx, opts); return err }}, + {"actors", func(opts store.ListOptions) error { _, err := s.ListActors(ctx, "", opts); return err }}, + {"actor templates", func(opts store.ListOptions) error { _, err := s.ListActorTemplates(ctx, "", opts); return err }}, + {"actor snapshots", func(opts store.ListOptions) error { _, err := s.ListActorSnapshots(ctx, "", opts); return err }}, + {"workers", func(opts store.ListOptions) error { _, err := s.ListWorkers(ctx, opts); return err }}, + } + for _, call := range calls { + t.Run(call.name, func(t *testing.T) { + if err := call.call(store.ListOptions{PageSize: -1}); !errors.Is(err, store.ErrInvalidPageSize) { + t.Errorf("negative PageSize error = %v, want ErrInvalidPageSize", err) + } + if err := call.call(store.ListOptions{}); err != nil { + t.Errorf("zero PageSize error = %v, want nil", err) + } + }) + } + }) + + t.Run("ListOptions_InvalidPageToken", func(t *testing.T) { + s := setup(t) + ctx := context.Background() + calls := []struct { + name string + call func(store.ListOptions) error + }{ + {"atespaces", func(opts store.ListOptions) error { _, err := s.ListAtespaces(ctx, opts); return err }}, + {"actors", func(opts store.ListOptions) error { _, err := s.ListActors(ctx, "", opts); return err }}, + {"actor templates", func(opts store.ListOptions) error { _, err := s.ListActorTemplates(ctx, "", opts); return err }}, + {"actor snapshots", func(opts store.ListOptions) error { _, err := s.ListActorSnapshots(ctx, "", opts); return err }}, + {"workers", func(opts store.ListOptions) error { _, err := s.ListWorkers(ctx, opts); return err }}, + } + for _, call := range calls { + t.Run(call.name, func(t *testing.T) { + if err := call.call(store.ListOptions{PageSize: 1, PageToken: "%%%"}); !errors.Is(err, store.ErrInvalidPageToken) { + t.Errorf("malformed PageToken error = %v, want ErrInvalidPageToken", err) + } + }) + } + }) +} + func runActorContractTests(t *testing.T, setup func(t *testing.T) store.Interface) { t.Helper() From 40f4b8e5d794e69065f5fe42ca5cb5bc45b41c0a Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Mon, 17 Aug 2026 17:28:47 -0400 Subject: [PATCH 08/14] remove unique on uid and add missing fk index Signed-off-by: Jet Chiang --- cmd/ateapi/internal/store/atepg/schema.go | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/cmd/ateapi/internal/store/atepg/schema.go b/cmd/ateapi/internal/store/atepg/schema.go index 63fb62c096..04ba10fb19 100644 --- a/cmd/ateapi/internal/store/atepg/schema.go +++ b/cmd/ateapi/internal/store/atepg/schema.go @@ -36,7 +36,7 @@ CREATE TABLE IF NOT EXISTS actors ( atespace text NOT NULL REFERENCES atespaces(name) ON DELETE RESTRICT, name text NOT NULL, - uid text NOT NULL UNIQUE, + uid text NOT NULL, version bigint NOT NULL, proto bytea NOT NULL, PRIMARY KEY (atespace, name) @@ -46,7 +46,7 @@ CREATE TABLE IF NOT EXISTS actor_templates ( atespace text NOT NULL REFERENCES atespaces(name) ON DELETE RESTRICT, name text NOT NULL, - uid text NOT NULL UNIQUE, + uid text NOT NULL, version bigint NOT NULL, proto bytea NOT NULL, PRIMARY KEY (atespace, name) @@ -64,6 +64,7 @@ CREATE TABLE IF NOT EXISTS actor_snapshot_tags ( name text NOT NULL, snapshot_atespace text NOT NULL, snapshot_name text NOT NULL, + uid text NOT NULL, version bigint NOT NULL, proto bytea NOT NULL, PRIMARY KEY (atespace, name), @@ -74,6 +75,9 @@ CREATE TABLE IF NOT EXISTS actor_snapshot_tags ( REFERENCES actor_snapshots(atespace, name) ON DELETE RESTRICT ); +CREATE INDEX IF NOT EXISTS actor_snapshot_tags_snapshot_idx + ON actor_snapshot_tags (snapshot_atespace, snapshot_name); + CREATE TABLE IF NOT EXISTS workers ( worker_namespace text NOT NULL, worker_pool text NOT NULL, From 5d92b83d5855006f053772c3f1b9527fb46618a0 Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Mon, 17 Aug 2026 17:42:20 -0400 Subject: [PATCH 09/14] change pessimistic locking to use optimistic concurrency control, also fixes potential data loss on corrupt worker notifications Signed-off-by: Jet Chiang --- cmd/ateapi/internal/store/atepg/atepg.go | 299 +++++++++--------- cmd/ateapi/internal/store/atepg/atepg_test.go | 198 ++++++++++++ cmd/ateapi/internal/store/atepg/pagetoken.go | 1 - .../internal/store/ateredis/ateredis.go | 6 +- cmd/ateapi/internal/store/store.go | 3 +- 5 files changed, 358 insertions(+), 149 deletions(-) diff --git a/cmd/ateapi/internal/store/atepg/atepg.go b/cmd/ateapi/internal/store/atepg/atepg.go index c29117dbd7..1c7bf7a35b 100644 --- a/cmd/ateapi/internal/store/atepg/atepg.go +++ b/cmd/ateapi/internal/store/atepg/atepg.go @@ -105,6 +105,20 @@ func newUpdateMetadata(current *ateapipb.ResourceMetadata) *ateapipb.ResourceMet return metadata } +// updateMaxAttempts bounds how many times a read-modify-write is retried after +// its optimistic uid/version check loses to a concurrent writer. +const updateMaxAttempts = 5 + +func validateMetadataProjection(resource string, metadata *ateapipb.ResourceMetadata, uid string, version int64) error { + if metadata.GetUid() != uid { + return fmt.Errorf("%s uid projection %q does not match proto metadata uid %q", resource, uid, metadata.GetUid()) + } + if metadata.GetVersion() != version { + return fmt.Errorf("%s version projection %d does not match proto metadata version %d", resource, version, metadata.GetVersion()) + } + return nil +} + func isUniqueViolation(err error) bool { return pgErrCode(err) == "23505" } // isForeignKeyViolation matches both the insert/update-side violation @@ -326,49 +340,53 @@ func validateUpdateActorTemplateMutation(storedTemplate, mutatedTemplate *ateapi } func (p *Persistence) UpdateActorTemplate(ctx context.Context, templateRef resources.ActorTemplateRef, mutate func(*ateapipb.ActorTemplate) error) (*ateapipb.ActorTemplate, error) { - tx, err := p.pool.Begin(ctx) - if err != nil { - return nil, fmt.Errorf("beginning actor template update: %w", err) - } - defer tx.Rollback(ctx) //nolint:errcheck // no-op once committed - - var currentBytes []byte - if err := tx.QueryRow(ctx, ` - SELECT proto FROM actor_templates - WHERE atespace = $1 AND name = $2 - FOR UPDATE`, templateRef.Atespace, templateRef.Name).Scan(¤tBytes); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return nil, store.ErrNotFound + for range updateMaxAttempts { + var currentUID string + var currentVersion int64 + var currentBytes []byte + if err := p.pool.QueryRow(ctx, ` + SELECT uid, version, proto FROM actor_templates + WHERE atespace = $1 AND name = $2`, templateRef.Atespace, templateRef.Name).Scan(¤tUID, ¤tVersion, ¤tBytes); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, store.ErrNotFound + } + return nil, fmt.Errorf("getting actor template %s for update: %w", templateRef, err) } - return nil, fmt.Errorf("locking actor template %s for update: %w", templateRef, err) - } - dbTemplate := &ateapipb.ActorTemplate{} - if err := proto.Unmarshal(currentBytes, dbTemplate); err != nil { - return nil, fmt.Errorf("unmarshaling actor template for update: %w", err) - } - templateBeforeMutation := proto.Clone(dbTemplate).(*ateapipb.ActorTemplate) - if err := mutate(dbTemplate); err != nil { - return nil, err - } - if err := validateUpdateActorTemplateMutation(templateBeforeMutation, dbTemplate); err != nil { - return nil, err - } - dbTemplate.Metadata = newUpdateMetadata(templateBeforeMutation.GetMetadata()) - updatedBytes, err := proto.Marshal(dbTemplate) - if err != nil { - return nil, fmt.Errorf("marshaling actor template: %w", err) - } - if _, err := tx.Exec(ctx, ` - UPDATE actor_templates SET version = $1, proto = $2 - WHERE atespace = $3 AND name = $4`, - dbTemplate.GetMetadata().GetVersion(), updatedBytes, templateRef.Atespace, templateRef.Name); err != nil { - return nil, fmt.Errorf("updating actor template %s: %w", templateRef, err) - } - if err := tx.Commit(ctx); err != nil { - return nil, fmt.Errorf("committing actor template update: %w", err) + dbTemplate := &ateapipb.ActorTemplate{} + if err := proto.Unmarshal(currentBytes, dbTemplate); err != nil { + return nil, fmt.Errorf("unmarshaling actor template for update: %w", err) + } + if err := validateMetadataProjection("actor template "+templateRef.String(), dbTemplate.GetMetadata(), currentUID, currentVersion); err != nil { + return nil, err + } + templateBeforeMutation := proto.Clone(dbTemplate).(*ateapipb.ActorTemplate) + if err := mutate(dbTemplate); err != nil { + return nil, err + } + if err := validateUpdateActorTemplateMutation(templateBeforeMutation, dbTemplate); err != nil { + return nil, err + } + dbTemplate.Metadata = newUpdateMetadata(templateBeforeMutation.GetMetadata()) + updatedBytes, err := proto.Marshal(dbTemplate) + if err != nil { + return nil, fmt.Errorf("marshaling actor template: %w", err) + } + commandTag, err := p.pool.Exec(ctx, ` + UPDATE actor_templates SET version = $1, proto = $2 + WHERE atespace = $3 AND name = $4 AND uid = $5 AND version = $6`, + dbTemplate.GetMetadata().GetVersion(), updatedBytes, templateRef.Atespace, templateRef.Name, currentUID, currentVersion) + if err != nil { + return nil, fmt.Errorf("updating actor template %s: %w", templateRef, err) + } + if commandTag.RowsAffected() == 1 { + return dbTemplate, nil + } + if commandTag.RowsAffected() != 0 { + return nil, fmt.Errorf("updating actor template %s affected %d rows, want at most 1", templateRef, commandTag.RowsAffected()) + } } - return dbTemplate, nil + return nil, store.ErrVersionConflict } func (p *Persistence) ListActorTemplates(ctx context.Context, atespace string, opts store.ListOptions) (store.ListResponse[*ateapipb.ActorTemplate], error) { @@ -541,60 +559,57 @@ func validateUpdateActorMutation(storedActor, mutatedActor *ateapipb.Actor) erro func (p *Persistence) UpdateActor(ctx context.Context, actorRef resources.ActorRef, mutate func(*ateapipb.Actor) error) (*ateapipb.Actor, error) { atespace, name := actorRef.Atespace, actorRef.Name - - tx, err := p.pool.Begin(ctx) - if err != nil { - return nil, fmt.Errorf("beginning actor update: %w", err) - } - defer tx.Rollback(ctx) //nolint:errcheck // no-op once committed - - var protoBytes []byte - if err := tx.QueryRow(ctx, ` - SELECT proto FROM actors - WHERE atespace = $1 AND name = $2 - FOR UPDATE`, atespace, name).Scan(&protoBytes); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return nil, store.ErrNotFound + for range updateMaxAttempts { + var currentUID string + var currentVersion int64 + var currentBytes []byte + if err := p.pool.QueryRow(ctx, ` + SELECT uid, version, proto FROM actors + WHERE atespace = $1 AND name = $2`, atespace, name).Scan(¤tUID, ¤tVersion, ¤tBytes); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, store.ErrNotFound + } + return nil, fmt.Errorf("getting actor %s/%s for update: %w", atespace, name, err) } - return nil, fmt.Errorf("locking actor %s/%s for update: %w", atespace, name, err) - } - - dbActor := &ateapipb.Actor{} - if err := proto.Unmarshal(protoBytes, dbActor); err != nil { - return nil, fmt.Errorf("unmarshaling actor for update: %w", err) - } - actorBeforeMutation := proto.Clone(dbActor).(*ateapipb.Actor) - if err := mutate(dbActor); err != nil { - return nil, err - } - if err := validateUpdateActorMutation(actorBeforeMutation, dbActor); err != nil { - return nil, err - } - // Stored metadata is authoritative; discard any metadata edits made by the - // closure and derive the next revision from the transactionally read actor. - dbActor.Metadata = newUpdateMetadata(actorBeforeMutation.GetMetadata()) - - protoBytes, err = proto.Marshal(dbActor) - if err != nil { - return nil, fmt.Errorf("marshaling actor: %w", err) - } - commandTag, err := tx.Exec(ctx, ` - UPDATE actors - SET version = $1, proto = $2 - WHERE atespace = $3 AND name = $4`, - dbActor.GetMetadata().GetVersion(), protoBytes, atespace, name) - if err != nil { - return nil, fmt.Errorf("updating actor %s/%s: %w", atespace, name, err) - } - if commandTag.RowsAffected() != 1 { - return nil, fmt.Errorf("updating actor %s/%s affected %d rows, want 1", atespace, name, commandTag.RowsAffected()) - } + dbActor := &ateapipb.Actor{} + if err := proto.Unmarshal(currentBytes, dbActor); err != nil { + return nil, fmt.Errorf("unmarshaling actor for update: %w", err) + } + if err := validateMetadataProjection("actor "+actorRef.String(), dbActor.GetMetadata(), currentUID, currentVersion); err != nil { + return nil, err + } + actorBeforeMutation := proto.Clone(dbActor).(*ateapipb.Actor) + if err := mutate(dbActor); err != nil { + return nil, err + } + if err := validateUpdateActorMutation(actorBeforeMutation, dbActor); err != nil { + return nil, err + } + // Stored metadata is authoritative; discard any metadata edits made by the + // closure and derive the next revision from the state this attempt read. + dbActor.Metadata = newUpdateMetadata(actorBeforeMutation.GetMetadata()) - if err := tx.Commit(ctx); err != nil { - return nil, fmt.Errorf("committing actor update: %w", err) + updatedBytes, err := proto.Marshal(dbActor) + if err != nil { + return nil, fmt.Errorf("marshaling actor: %w", err) + } + commandTag, err := p.pool.Exec(ctx, ` + UPDATE actors + SET version = $1, proto = $2 + WHERE atespace = $3 AND name = $4 AND uid = $5 AND version = $6`, + dbActor.GetMetadata().GetVersion(), updatedBytes, atespace, name, currentUID, currentVersion) + if err != nil { + return nil, fmt.Errorf("updating actor %s/%s: %w", atespace, name, err) + } + if commandTag.RowsAffected() == 1 { + return dbActor, nil + } + if commandTag.RowsAffected() != 0 { + return nil, fmt.Errorf("updating actor %s/%s affected %d rows, want at most 1", atespace, name, commandTag.RowsAffected()) + } } - return dbActor, nil + return nil, store.ErrVersionConflict } func (p *Persistence) DeleteActor(ctx context.Context, actorRef resources.ActorRef) (*ateapipb.Actor, error) { @@ -940,11 +955,11 @@ func (p *Persistence) CreateActorSnapshotTag(ctx context.Context, snapshotAtespa var inserted []byte err = tx.QueryRow(ctx, ` INSERT INTO actor_snapshot_tags - (atespace, name, snapshot_atespace, snapshot_name, version, proto) - VALUES ($1, $2, $3, $4, $5, $6) + (atespace, name, snapshot_atespace, snapshot_name, uid, version, proto) + VALUES ($1, $2, $3, $4, $5, $6, $7) ON CONFLICT (atespace, name) DO NOTHING RETURNING proto`, tagAtespace, tagName, snapshotAtespace, snapshotName, - dbTag.GetMetadata().GetVersion(), protoBytes).Scan(&inserted) + dbTag.GetMetadata().GetUid(), dbTag.GetMetadata().GetVersion(), protoBytes).Scan(&inserted) if err == nil { if err := tx.Commit(ctx); err != nil { return nil, fmt.Errorf("committing actor snapshot tag create: %w", err) @@ -1001,57 +1016,57 @@ func validateUpdateActorSnapshotTagMutation(storedTag, mutatedTag *ateapipb.Acto } func (p *Persistence) UpdateActorSnapshotTag(ctx context.Context, atespace, name string, mutate func(*ateapipb.ActorSnapshotTag) error) (*ateapipb.ActorSnapshotTag, error) { - tx, err := p.pool.Begin(ctx) - if err != nil { - return nil, fmt.Errorf("beginning actor snapshot tag update: %w", err) - } - defer tx.Rollback(ctx) //nolint:errcheck // no-op once committed - - var currentBytes []byte - if err := tx.QueryRow(ctx, ` - SELECT proto FROM actor_snapshot_tags - WHERE atespace = $1 AND name = $2 - FOR UPDATE`, atespace, name).Scan(¤tBytes); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return nil, store.ErrNotFound + for range updateMaxAttempts { + var currentUID string + var currentVersion int64 + var currentBytes []byte + if err := p.pool.QueryRow(ctx, ` + SELECT uid, version, proto FROM actor_snapshot_tags + WHERE atespace = $1 AND name = $2`, atespace, name).Scan(¤tUID, ¤tVersion, ¤tBytes); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, store.ErrNotFound + } + return nil, fmt.Errorf("getting actor snapshot tag %s/%s for update: %w", atespace, name, err) } - return nil, fmt.Errorf("locking actor snapshot tag %s/%s for update: %w", atespace, name, err) - } - dbTag := &ateapipb.ActorSnapshotTag{} - if err := proto.Unmarshal(currentBytes, dbTag); err != nil { - return nil, fmt.Errorf("unmarshaling actor snapshot tag: %w", err) - } - tagBeforeMutation := proto.Clone(dbTag).(*ateapipb.ActorSnapshotTag) - if err := mutate(dbTag); err != nil { - return nil, err - } - if err := validateUpdateActorSnapshotTagMutation(tagBeforeMutation, dbTag); err != nil { - return nil, err - } - // Stored metadata is authoritative; discard any metadata edits made by the - // closure and derive the next revision from the transactionally read tag. - dbTag.Metadata = newUpdateMetadata(tagBeforeMutation.GetMetadata()) + dbTag := &ateapipb.ActorSnapshotTag{} + if err := proto.Unmarshal(currentBytes, dbTag); err != nil { + return nil, fmt.Errorf("unmarshaling actor snapshot tag: %w", err) + } + if err := validateMetadataProjection(fmt.Sprintf("actor snapshot tag %s/%s", atespace, name), dbTag.GetMetadata(), currentUID, currentVersion); err != nil { + return nil, err + } + tagBeforeMutation := proto.Clone(dbTag).(*ateapipb.ActorSnapshotTag) + if err := mutate(dbTag); err != nil { + return nil, err + } + if err := validateUpdateActorSnapshotTagMutation(tagBeforeMutation, dbTag); err != nil { + return nil, err + } + // Stored metadata is authoritative; discard any metadata edits made by the + // closure and derive the next revision from the state this attempt read. + dbTag.Metadata = newUpdateMetadata(tagBeforeMutation.GetMetadata()) - updatedBytes, err := proto.Marshal(dbTag) - if err != nil { - return nil, fmt.Errorf("marshaling actor snapshot tag: %w", err) - } - commandTag, err := tx.Exec(ctx, ` - UPDATE actor_snapshot_tags - SET version = $1, proto = $2 - WHERE atespace = $3 AND name = $4`, - dbTag.GetMetadata().GetVersion(), updatedBytes, atespace, name) - if err != nil { - return nil, fmt.Errorf("updating actor snapshot tag %s/%s: %w", atespace, name, err) - } - if commandTag.RowsAffected() != 1 { - return nil, fmt.Errorf("updating actor snapshot tag %s/%s affected %d rows, want 1", atespace, name, commandTag.RowsAffected()) - } - if err := tx.Commit(ctx); err != nil { - return nil, fmt.Errorf("committing actor snapshot tag update: %w", err) + updatedBytes, err := proto.Marshal(dbTag) + if err != nil { + return nil, fmt.Errorf("marshaling actor snapshot tag: %w", err) + } + commandTag, err := p.pool.Exec(ctx, ` + UPDATE actor_snapshot_tags + SET version = $1, proto = $2 + WHERE atespace = $3 AND name = $4 AND uid = $5 AND version = $6`, + dbTag.GetMetadata().GetVersion(), updatedBytes, atespace, name, currentUID, currentVersion) + if err != nil { + return nil, fmt.Errorf("updating actor snapshot tag %s/%s: %w", atespace, name, err) + } + if commandTag.RowsAffected() == 1 { + return dbTag, nil + } + if commandTag.RowsAffected() != 0 { + return nil, fmt.Errorf("updating actor snapshot tag %s/%s affected %d rows, want at most 1", atespace, name, commandTag.RowsAffected()) + } } - return dbTag, nil + return nil, store.ErrVersionConflict } func (p *Persistence) DeleteActorSnapshotTag(ctx context.Context, atespace, name string) (*ateapipb.ActorSnapshotTag, error) { @@ -1340,8 +1355,8 @@ func (p *Persistence) WatchWorkers(ctx context.Context) (*store.WorkerWatch, err } event, err := unmarshalWorkerEvent(notification.Payload) if err != nil { - slog.ErrorContext(ctx, "worker event unmarshal failed", slog.Any("err", err)) - continue + slog.ErrorContext(ctx, "worker event unmarshal failed; closing watch", slog.Any("err", err)) + return } select { case ch <- event: diff --git a/cmd/ateapi/internal/store/atepg/atepg_test.go b/cmd/ateapi/internal/store/atepg/atepg_test.go index a0b2f81798..f1c77c7e40 100644 --- a/cmd/ateapi/internal/store/atepg/atepg_test.go +++ b/cmd/ateapi/internal/store/atepg/atepg_test.go @@ -136,6 +136,180 @@ func createTestAtespace(t *testing.T, s *Persistence, name string) { } } +func createTestActorTemplate(t *testing.T, s *Persistence, atespace, name string) { + t.Helper() + if _, err := s.CreateActorTemplate(context.Background(), &ateapipb.ActorTemplate{ + Metadata: &ateapipb.ResourceMetadata{Atespace: atespace, Name: name}, + }); err != nil { + t.Fatalf("CreateActorTemplate(%q/%q) failed: %v", atespace, name, err) + } +} + +func TestUpdateActor_RetriesConcurrentWrite(t *testing.T) { + s := setupPostgresPersistence(t) + ctx := context.Background() + createTestAtespace(t, s, "team-a") + created, err := s.CreateActor(ctx, &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "actor-a"}, + ActorTemplateNamespace: "default", + ActorTemplateName: "template-a", + Status: ateapipb.Actor_STATUS_SUSPENDED, + }) + if err != nil { + t.Fatalf("CreateActor failed: %v", err) + } + actorRef := resources.ActorRefFromActor(created) + + attempts := 0 + updated, err := s.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error { + attempts++ + if attempts == 1 { + if _, err := s.UpdateActor(ctx, actorRef, func(concurrent *ateapipb.Actor) error { + concurrent.WorkerSelector = &ateapipb.Selector{MatchLabels: map[string]string{"tier": "paid"}} + return nil + }); err != nil { + return fmt.Errorf("concurrent actor update: %w", err) + } + } + toUpdate.Status = ateapipb.Actor_STATUS_RUNNING + return nil + }) + if err != nil { + t.Fatalf("UpdateActor failed: %v", err) + } + if attempts != 2 { + t.Errorf("mutate ran %d times, want 2", attempts) + } + if updated.GetStatus() != ateapipb.Actor_STATUS_RUNNING { + t.Errorf("status = %v, want RUNNING", updated.GetStatus()) + } + if got := updated.GetWorkerSelector().GetMatchLabels()["tier"]; got != "paid" { + t.Errorf("worker selector tier = %q, want paid: concurrent update was lost", got) + } + if got, want := updated.GetMetadata().GetVersion(), created.GetMetadata().GetVersion()+2; got != want { + t.Errorf("version = %d, want %d", got, want) + } +} + +func TestUpdateActor_ExhaustsOptimisticRetries(t *testing.T) { + s := setupPostgresPersistence(t) + ctx := context.Background() + createTestAtespace(t, s, "team-a") + created, err := s.CreateActor(ctx, &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "actor-a"}, + ActorTemplateNamespace: "default", + ActorTemplateName: "template-a", + Status: ateapipb.Actor_STATUS_SUSPENDED, + }) + if err != nil { + t.Fatalf("CreateActor failed: %v", err) + } + actorRef := resources.ActorRefFromActor(created) + + attempts := 0 + _, err = s.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error { + attempts++ + _, err := s.UpdateActor(ctx, actorRef, func(concurrent *ateapipb.Actor) error { + concurrent.Status = ateapipb.Actor_STATUS_RUNNING + return nil + }) + return err + }) + if !errors.Is(err, store.ErrVersionConflict) { + t.Fatalf("UpdateActor error = %v, want ErrVersionConflict", err) + } + if attempts != updateMaxAttempts { + t.Errorf("mutate ran %d times, want %d", attempts, updateMaxAttempts) + } +} + +func TestUpdateActorTemplate_RetriesConcurrentWrite(t *testing.T) { + s := setupPostgresPersistence(t) + ctx := context.Background() + createTestAtespace(t, s, "team-a") + createTestActorTemplate(t, s, "team-a", "template-a") + templateRef := resources.ActorTemplateRef{Atespace: "team-a", Name: "template-a"} + + attempts := 0 + updated, err := s.UpdateActorTemplate(ctx, templateRef, func(toUpdate *ateapipb.ActorTemplate) error { + attempts++ + if attempts == 1 { + if _, err := s.UpdateActorTemplate(ctx, templateRef, func(concurrent *ateapipb.ActorTemplate) error { + concurrent.WorkerSelector = &ateapipb.Selector{MatchLabels: map[string]string{"tier": "paid"}} + return nil + }); err != nil { + return fmt.Errorf("concurrent actor template update: %w", err) + } + } + toUpdate.Phase = &ateapipb.ActorTemplatePhase{Message: "ready"} + return nil + }) + if err != nil { + t.Fatalf("UpdateActorTemplate failed: %v", err) + } + if attempts != 2 { + t.Errorf("mutate ran %d times, want 2", attempts) + } + if got := updated.GetWorkerSelector().GetMatchLabels()["tier"]; got != "paid" { + t.Errorf("worker selector tier = %q, want paid: concurrent update was lost", got) + } + if got := updated.GetPhase().GetMessage(); got != "ready" { + t.Errorf("phase message = %q, want ready", got) + } +} + +func TestUpdateActorSnapshotTag_UIDPreventsDeleteRecreateABA(t *testing.T) { + s := setupPostgresPersistence(t) + ctx := context.Background() + createTestAtespace(t, s, "team-a") + for _, name := range []string{"snapshot-a", "snapshot-b"} { + if _, err := s.CreateActorSnapshot(ctx, &ateapipb.ActorSnapshot{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: name}, + SnapshotUri: "gs://bucket/" + name, + }); err != nil { + t.Fatalf("CreateActorSnapshot(%q) failed: %v", name, err) + } + } + original, err := s.CreateActorSnapshotTag(ctx, "team-a", "snapshot-a", &ateapipb.ActorSnapshotTag{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "tag-a"}, + Scope: ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_ATESPACE, + }) + if err != nil { + t.Fatalf("CreateActorSnapshotTag failed: %v", err) + } + + mutations := 0 + var recreated *ateapipb.ActorSnapshotTag + _, err = s.UpdateActorSnapshotTag(ctx, "team-a", "tag-a", store.WithPrecondition(original, func(toUpdate *ateapipb.ActorSnapshotTag) error { + mutations++ + if _, err := s.DeleteActorSnapshotTag(ctx, "team-a", "tag-a"); err != nil { + return fmt.Errorf("deleting original tag: %w", err) + } + recreated, err = s.CreateActorSnapshotTag(ctx, "team-a", "snapshot-b", &ateapipb.ActorSnapshotTag{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "tag-a"}, + Scope: ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_ATESPACE, + }) + if err != nil { + return fmt.Errorf("recreating tag: %w", err) + } + toUpdate.Scope = ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED + return nil + })) + if !errors.Is(err, store.ErrUIDConflict) { + t.Fatalf("UpdateActorSnapshotTag error = %v, want ErrUIDConflict", err) + } + if mutations != 1 { + t.Errorf("guarded mutation ran %d times, want 1", mutations) + } + stored, err := s.GetActorSnapshotTag(ctx, "team-a", "tag-a") + if err != nil { + t.Fatalf("GetActorSnapshotTag failed: %v", err) + } + if diff := cmp.Diff(recreated, stored, protocmp.Transform()); diff != "" { + t.Errorf("recreated tag was overwritten (-want +got):\n%s", diff) + } +} + func TestCreateActorSnapshotTag_ForeignKeyErrors(t *testing.T) { s := setupPostgresPersistence(t) ctx := context.Background() @@ -267,6 +441,30 @@ func TestWorkerNotification_OnlyAfterCommit(t *testing.T) { } } +func TestWatchWorkers_MalformedNotificationClosesWatch(t *testing.T) { + s := setupPostgresStore(t).(*Persistence) + ctx := context.Background() + + watch, err := s.WatchWorkers(ctx) + if err != nil { + t.Fatalf("WatchWorkers failed: %v", err) + } + defer watch.Close() + + if _, err := s.pool.Exec(ctx, `SELECT pg_notify($1, $2)`, workerChangeChannel, "not-json"); err != nil { + t.Fatalf("pg_notify failed: %v", err) + } + + select { + case event, ok := <-watch.Events: + if ok { + t.Fatalf("received event %+v from malformed notification; want closed watch", event) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for malformed notification to close watch") + } +} + func TestListActors_InvalidPageToken(t *testing.T) { s := setupPostgresStore(t).(*Persistence) ctx := context.Background() diff --git a/cmd/ateapi/internal/store/atepg/pagetoken.go b/cmd/ateapi/internal/store/atepg/pagetoken.go index 28e818a27e..df3e14ee28 100644 --- a/cmd/ateapi/internal/store/atepg/pagetoken.go +++ b/cmd/ateapi/internal/store/atepg/pagetoken.go @@ -34,7 +34,6 @@ const ( kindAtespace resourceKind = "atespace" kindActor resourceKind = "actor" kindActorTemplate resourceKind = "actor-template" - kindActorTemplateVersion resourceKind = "actor-template-version" kindSnapshot resourceKind = "snapshot" kindWorker resourceKind = "worker" ) diff --git a/cmd/ateapi/internal/store/ateredis/ateredis.go b/cmd/ateapi/internal/store/ateredis/ateredis.go index 06a657e762..572b621ee7 100644 --- a/cmd/ateapi/internal/store/ateredis/ateredis.go +++ b/cmd/ateapi/internal/store/ateredis/ateredis.go @@ -430,10 +430,8 @@ func (s *Persistence) ListActorTemplates(ctx context.Context, atespace string, o return store.ListResponse[*ateapipb.ActorTemplate]{Items: result, NextPageToken: nextToken}, nil } -// DeleteActorTemplate deletes an ActorTemplate with no remaining versions. -// Returns store.ErrNotFound if the template does not exist, or -// store.ErrFailedPrecondition while any ActorTemplateVersion still names it -// as parent. +// DeleteActorTemplate deletes an ActorTemplate. +// Returns store.ErrNotFound if the template does not exist. func (s *Persistence) DeleteActorTemplate(ctx context.Context, templateRef resources.ActorTemplateRef) (*ateapipb.ActorTemplate, error) { dbKey := actorTemplateDBKey(templateRef) diff --git a/cmd/ateapi/internal/store/store.go b/cmd/ateapi/internal/store/store.go index 0ae7b1b0f2..9437f08fce 100644 --- a/cmd/ateapi/internal/store/store.go +++ b/cmd/ateapi/internal/store/store.go @@ -165,8 +165,7 @@ type Interface interface { UpdateActorTemplate(ctx context.Context, templateRef resources.ActorTemplateRef, mutate func(dbTemplate *ateapipb.ActorTemplate) error) (*ateapipb.ActorTemplate, error) // Removes an ActorTemplate and returns the deleted resource. Returns - // ErrNotFound if missing, or ErrFailedPrecondition while any - // ActorTemplateVersion still names it as parent. + // ErrNotFound if missing. DeleteActorTemplate(ctx context.Context, templateRef resources.ActorTemplateRef) (*ateapipb.ActorTemplate, error) // Registers a new idle worker. Returns ErrAlreadyExists if already registered. From 29b18b124e8ae61249e9350228f31dd5699fa38c Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Tue, 18 Aug 2026 11:57:51 -0400 Subject: [PATCH 10/14] review comments Signed-off-by: Jet Chiang --- cmd/ateapi/internal/store/atepg/atepg.go | 10 ++++++---- cmd/ateapi/internal/store/atepg/atepg_test.go | 1 + cmd/ateapi/internal/store/atepg/schema.go | 2 ++ 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/cmd/ateapi/internal/store/atepg/atepg.go b/cmd/ateapi/internal/store/atepg/atepg.go index 1c7bf7a35b..31a5dfa15b 100644 --- a/cmd/ateapi/internal/store/atepg/atepg.go +++ b/cmd/ateapi/internal/store/atepg/atepg.go @@ -109,7 +109,9 @@ func newUpdateMetadata(current *ateapipb.ResourceMetadata) *ateapipb.ResourceMet // its optimistic uid/version check loses to a concurrent writer. const updateMaxAttempts = 5 -func validateMetadataProjection(resource string, metadata *ateapipb.ResourceMetadata, uid string, version int64) error { +// validateProtoMetadataMatchesColumns verifies that the metadata in the database +// matches the metadata in the proto. +func validateProtoMetadataMatchesColumns(resource string, metadata *ateapipb.ResourceMetadata, uid string, version int64) error { if metadata.GetUid() != uid { return fmt.Errorf("%s uid projection %q does not match proto metadata uid %q", resource, uid, metadata.GetUid()) } @@ -357,7 +359,7 @@ func (p *Persistence) UpdateActorTemplate(ctx context.Context, templateRef resou if err := proto.Unmarshal(currentBytes, dbTemplate); err != nil { return nil, fmt.Errorf("unmarshaling actor template for update: %w", err) } - if err := validateMetadataProjection("actor template "+templateRef.String(), dbTemplate.GetMetadata(), currentUID, currentVersion); err != nil { + if err := validateProtoMetadataMatchesColumns("actor template "+templateRef.String(), dbTemplate.GetMetadata(), currentUID, currentVersion); err != nil { return nil, err } templateBeforeMutation := proto.Clone(dbTemplate).(*ateapipb.ActorTemplate) @@ -576,7 +578,7 @@ func (p *Persistence) UpdateActor(ctx context.Context, actorRef resources.ActorR if err := proto.Unmarshal(currentBytes, dbActor); err != nil { return nil, fmt.Errorf("unmarshaling actor for update: %w", err) } - if err := validateMetadataProjection("actor "+actorRef.String(), dbActor.GetMetadata(), currentUID, currentVersion); err != nil { + if err := validateProtoMetadataMatchesColumns("actor "+actorRef.String(), dbActor.GetMetadata(), currentUID, currentVersion); err != nil { return nil, err } actorBeforeMutation := proto.Clone(dbActor).(*ateapipb.Actor) @@ -1033,7 +1035,7 @@ func (p *Persistence) UpdateActorSnapshotTag(ctx context.Context, atespace, name if err := proto.Unmarshal(currentBytes, dbTag); err != nil { return nil, fmt.Errorf("unmarshaling actor snapshot tag: %w", err) } - if err := validateMetadataProjection(fmt.Sprintf("actor snapshot tag %s/%s", atespace, name), dbTag.GetMetadata(), currentUID, currentVersion); err != nil { + if err := validateProtoMetadataMatchesColumns(fmt.Sprintf("actor snapshot tag %s/%s", atespace, name), dbTag.GetMetadata(), currentUID, currentVersion); err != nil { return nil, err } tagBeforeMutation := proto.Clone(dbTag).(*ateapipb.ActorSnapshotTag) diff --git a/cmd/ateapi/internal/store/atepg/atepg_test.go b/cmd/ateapi/internal/store/atepg/atepg_test.go index f1c77c7e40..5d589dc361 100644 --- a/cmd/ateapi/internal/store/atepg/atepg_test.go +++ b/cmd/ateapi/internal/store/atepg/atepg_test.go @@ -30,6 +30,7 @@ import ( "google.golang.org/protobuf/testing/protocmp" "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" + "github.com/agent-substrate/substrate/internal/resources" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" ) diff --git a/cmd/ateapi/internal/store/atepg/schema.go b/cmd/ateapi/internal/store/atepg/schema.go index 04ba10fb19..3c7507a60c 100644 --- a/cmd/ateapi/internal/store/atepg/schema.go +++ b/cmd/ateapi/internal/store/atepg/schema.go @@ -29,6 +29,8 @@ import ( const schema = ` CREATE TABLE IF NOT EXISTS atespaces ( name text PRIMARY KEY, + uid text NOT NULL, + version bigint NOT NULL, proto bytea NOT NULL ); From a61a4bb9140a3fb98966b8825ed98a7a97811ac8 Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Tue, 18 Aug 2026 13:28:59 -0400 Subject: [PATCH 11/14] add uid and version for actor snapshot table Signed-off-by: Jet Chiang --- cmd/ateapi/internal/store/atepg/schema.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cmd/ateapi/internal/store/atepg/schema.go b/cmd/ateapi/internal/store/atepg/schema.go index 3c7507a60c..3d7ee1b790 100644 --- a/cmd/ateapi/internal/store/atepg/schema.go +++ b/cmd/ateapi/internal/store/atepg/schema.go @@ -57,6 +57,8 @@ CREATE TABLE IF NOT EXISTS actor_templates ( CREATE TABLE IF NOT EXISTS actor_snapshots ( atespace text NOT NULL, name text NOT NULL, + uid text NOT NULL, + version bigint NOT NULL, proto bytea NOT NULL, PRIMARY KEY (atespace, name) ); From 6a94e546b5e2de35dc1eb2cb9b13dd212cc9a30d Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Tue, 18 Aug 2026 14:15:01 -0400 Subject: [PATCH 12/14] fix create queries Signed-off-by: Jet Chiang --- cmd/ateapi/internal/store/atepg/atepg.go | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/cmd/ateapi/internal/store/atepg/atepg.go b/cmd/ateapi/internal/store/atepg/atepg.go index 31a5dfa15b..faab1c0b94 100644 --- a/cmd/ateapi/internal/store/atepg/atepg.go +++ b/cmd/ateapi/internal/store/atepg/atepg.go @@ -166,9 +166,9 @@ func (p *Persistence) CreateAtespace(ctx context.Context, atespace *ateapipb.Ate } _, err = p.pool.Exec(ctx, ` - INSERT INTO atespaces (name, proto) - VALUES ($1, $2)`, - name, protoBytes) + INSERT INTO atespaces (name, uid, version, proto) + VALUES ($1, $2, $3, $4)`, + name, dbAtespace.GetMetadata().GetUid(), dbAtespace.GetMetadata().GetVersion(), protoBytes) if err != nil { if isUniqueViolation(err) { return nil, store.ErrAlreadyExists @@ -779,9 +779,9 @@ func (p *Persistence) CreateActorSnapshot(ctx context.Context, snapshot *ateapip return nil, fmt.Errorf("marshaling actor snapshot: %w", err) } if _, err := p.pool.Exec(ctx, ` - INSERT INTO actor_snapshots (atespace, name, proto) - VALUES ($1, $2, $3)`, - atespace, name, protoBytes); err != nil { + INSERT INTO actor_snapshots (atespace, name, uid, version, proto) + VALUES ($1, $2, $3, $4, $5)`, + atespace, name, dbSnapshot.GetMetadata().GetUid(), dbSnapshot.GetMetadata().GetVersion(), protoBytes); err != nil { if isUniqueViolation(err) { return nil, store.ErrAlreadyExists } From c91e51bcca3b9d88928ceebfb8b8cd4c9a5399a7 Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Tue, 18 Aug 2026 15:12:02 -0400 Subject: [PATCH 13/14] fmt Signed-off-by: Jet Chiang --- cmd/ateapi/internal/store/atepg/pagetoken.go | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/cmd/ateapi/internal/store/atepg/pagetoken.go b/cmd/ateapi/internal/store/atepg/pagetoken.go index df3e14ee28..174c95c45d 100644 --- a/cmd/ateapi/internal/store/atepg/pagetoken.go +++ b/cmd/ateapi/internal/store/atepg/pagetoken.go @@ -31,11 +31,11 @@ const pageTokenVersion = 1 type resourceKind string const ( - kindAtespace resourceKind = "atespace" - kindActor resourceKind = "actor" - kindActorTemplate resourceKind = "actor-template" - kindSnapshot resourceKind = "snapshot" - kindWorker resourceKind = "worker" + kindAtespace resourceKind = "atespace" + kindActor resourceKind = "actor" + kindActorTemplate resourceKind = "actor-template" + kindSnapshot resourceKind = "snapshot" + kindWorker resourceKind = "worker" ) // pageToken is PostgreSQL's opaque keyset page token. Unlike ateredis's From 403cdb33e668cd7619856a0073e2a19e3ee70118 Mon Sep 17 00:00:00 2001 From: Jet Chiang Date: Tue, 18 Aug 2026 16:14:22 -0400 Subject: [PATCH 14/14] sync changes from status field PR Signed-off-by: Jet Chiang --- .../internal/controlapi/workflow_resume.go | 2 +- .../controlapi/workflow_resume_test.go | 2 +- cmd/ateapi/internal/store/atepg/atepg_test.go | 20 +++++++++---------- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index b2a04ac978..530ee5529b 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume.go @@ -91,7 +91,7 @@ func (w *ActorWorkflow) ResumeActor(ctx context.Context, actorRef resources.Acto if err != nil { return nil, false, err } - if wasRunning = actor.GetStatus() == ateapipb.Actor_STATUS_RUNNING; wasRunning { + if wasRunning = actor.GetStatus().GetState() == ateapipb.ActorState_ACTOR_STATE_RUNNING; wasRunning { return actor, false, nil } diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index f3c1391e29..ae9d614f32 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume_test.go @@ -77,7 +77,7 @@ func TestResumeActor_RunningFastPathDoesNotAcquireLock(t *testing.T) { persistence := newTestPersistence(t) created, err := persistence.CreateActor(ctx, &ateapipb.Actor{ Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "id1"}, - Status: ateapipb.Actor_STATUS_RUNNING, + Status: &ateapipb.ActorStatus{State: ateapipb.ActorState_ACTOR_STATE_RUNNING}, }) if err != nil { t.Fatalf("CreateActor: %v", err) diff --git a/cmd/ateapi/internal/store/atepg/atepg_test.go b/cmd/ateapi/internal/store/atepg/atepg_test.go index 5d589dc361..448b5b2d3f 100644 --- a/cmd/ateapi/internal/store/atepg/atepg_test.go +++ b/cmd/ateapi/internal/store/atepg/atepg_test.go @@ -154,7 +154,7 @@ func TestUpdateActor_RetriesConcurrentWrite(t *testing.T) { Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "actor-a"}, ActorTemplateNamespace: "default", ActorTemplateName: "template-a", - Status: ateapipb.Actor_STATUS_SUSPENDED, + Status: &ateapipb.ActorStatus{State: ateapipb.ActorState_ACTOR_STATE_SUSPENDED}, }) if err != nil { t.Fatalf("CreateActor failed: %v", err) @@ -172,7 +172,7 @@ func TestUpdateActor_RetriesConcurrentWrite(t *testing.T) { return fmt.Errorf("concurrent actor update: %w", err) } } - toUpdate.Status = ateapipb.Actor_STATUS_RUNNING + toUpdate.Status.State = ateapipb.ActorState_ACTOR_STATE_RUNNING return nil }) if err != nil { @@ -181,8 +181,8 @@ func TestUpdateActor_RetriesConcurrentWrite(t *testing.T) { if attempts != 2 { t.Errorf("mutate ran %d times, want 2", attempts) } - if updated.GetStatus() != ateapipb.Actor_STATUS_RUNNING { - t.Errorf("status = %v, want RUNNING", updated.GetStatus()) + if updated.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_RUNNING { + t.Errorf("state = %v, want RUNNING", updated.GetStatus().GetState()) } if got := updated.GetWorkerSelector().GetMatchLabels()["tier"]; got != "paid" { t.Errorf("worker selector tier = %q, want paid: concurrent update was lost", got) @@ -200,7 +200,7 @@ func TestUpdateActor_ExhaustsOptimisticRetries(t *testing.T) { Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "actor-a"}, ActorTemplateNamespace: "default", ActorTemplateName: "template-a", - Status: ateapipb.Actor_STATUS_SUSPENDED, + Status: &ateapipb.ActorStatus{State: ateapipb.ActorState_ACTOR_STATE_SUSPENDED}, }) if err != nil { t.Fatalf("CreateActor failed: %v", err) @@ -211,7 +211,7 @@ func TestUpdateActor_ExhaustsOptimisticRetries(t *testing.T) { _, err = s.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error { attempts++ _, err := s.UpdateActor(ctx, actorRef, func(concurrent *ateapipb.Actor) error { - concurrent.Status = ateapipb.Actor_STATUS_RUNNING + concurrent.Status.State = ateapipb.ActorState_ACTOR_STATE_RUNNING return nil }) return err @@ -242,7 +242,7 @@ func TestUpdateActorTemplate_RetriesConcurrentWrite(t *testing.T) { return fmt.Errorf("concurrent actor template update: %w", err) } } - toUpdate.Phase = &ateapipb.ActorTemplatePhase{Message: "ready"} + toUpdate.Status = &ateapipb.ActorTemplateStatus{Message: "ready"} return nil }) if err != nil { @@ -254,7 +254,7 @@ func TestUpdateActorTemplate_RetriesConcurrentWrite(t *testing.T) { if got := updated.GetWorkerSelector().GetMatchLabels()["tier"]; got != "paid" { t.Errorf("worker selector tier = %q, want paid: concurrent update was lost", got) } - if got := updated.GetPhase().GetMessage(); got != "ready" { + if got := updated.GetStatus().GetMessage(); got != "ready" { t.Errorf("phase message = %q, want ready", got) } } @@ -265,8 +265,8 @@ func TestUpdateActorSnapshotTag_UIDPreventsDeleteRecreateABA(t *testing.T) { createTestAtespace(t, s, "team-a") for _, name := range []string{"snapshot-a", "snapshot-b"} { if _, err := s.CreateActorSnapshot(ctx, &ateapipb.ActorSnapshot{ - Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: name}, - SnapshotUri: "gs://bucket/" + name, + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: name}, + Status: &ateapipb.ActorSnapshotStatus{SnapshotUri: "gs://bucket/" + name}, }); err != nil { t.Fatalf("CreateActorSnapshot(%q) failed: %v", name, err) }