Harden PostgreSQL concurrency, integrity, and error handling

cmd/ateapi/internal/controlapi/actor.go

@@ -105,6 +105,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)
 	}
 
@@ -243,7 +246,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,

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
 }

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

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) {

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,

cmd/ateapi/internal/controlapi/functional_test.go

@@ -2970,16 +2970,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")

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
+}

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,

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
 	}
 

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)