diff --git a/cmd/ateapi/internal/controlapi/crash.go b/cmd/ateapi/internal/controlapi/crash.go index e12593c4b0..350f69762d 100644 --- a/cmd/ateapi/internal/controlapi/crash.go +++ b/cmd/ateapi/internal/controlapi/crash.go @@ -31,7 +31,7 @@ import ( // maybeCrashActor inspects err returned by an atelet RPC and crashes the actor // if err carries the actorCrashed=true metadata directive. -func maybeCrashActor(ctx context.Context, st store.Interface, actorRef resources.ActorRef, err error, wrapMsg, opName string) error { +func maybeCrashActor(ctx context.Context, st crashActorStore, actorRef resources.ActorRef, err error, wrapMsg, opName string) error { if err == nil { return nil } @@ -53,7 +53,7 @@ func maybeCrashActor(ctx context.Context, st store.Interface, actorRef resources // crashActor moves the actor to CRASHED state and frees the worker it was // assigned to, if any, so the worker can host other actors. -func crashActor(ctx context.Context, st store.Interface, actorRef resources.ActorRef, opName, reason string) error { +func crashActor(ctx context.Context, st crashActorStore, actorRef resources.ActorRef, opName, reason string) error { actor, err := st.GetActor(ctx, actorRef) if err != nil { return fmt.Errorf("while loading actor to crash: %w", err) @@ -97,10 +97,19 @@ func crashActor(ctx context.Context, st store.Interface, actorRef resources.Acto return errors.Join(errCollected...) } +// crashActorStore encapsulates the subset of store operations needed to crash +// an actor. +type crashActorStore interface { + GetActor(ctx context.Context, actorRef resources.ActorRef) (*ateapipb.Actor, error) + UpdateActor(ctx context.Context, actorRef resources.ActorRef, mutate func(toUpdate *ateapipb.Actor) error) (*ateapipb.Actor, error) + GetWorker(ctx context.Context, namespace, pool, pod string) (*ateapipb.Worker, error) + UpdateWorker(ctx context.Context, worker *ateapipb.Worker, expectedVersion int64) error +} + // releaseWorker clears the worker's assignment if it still points at the given // actor. A missing worker or an already-cleared assignment is not an error. // It returns the worker's sandboxClass if found. -func releaseWorker(ctx context.Context, st store.Interface, actor *ateapipb.Actor) (string, error) { +func releaseWorker(ctx context.Context, st crashActorStore, actor *ateapipb.Actor) (string, error) { assignment := actor.GetWorkerAssignment() if assignment == nil { slog.WarnContext(ctx, "Actor's worker assignment is already cleared") diff --git a/cmd/ateapi/internal/controlapi/service.go b/cmd/ateapi/internal/controlapi/service.go index fd5fd52bfd..da9e140960 100644 --- a/cmd/ateapi/internal/controlapi/service.go +++ b/cmd/ateapi/internal/controlapi/service.go @@ -20,6 +20,7 @@ import ( "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" "github.com/agent-substrate/substrate/cmd/ateapi/internal/workercache" + "github.com/agent-substrate/substrate/internal/resources" "github.com/agent-substrate/substrate/internal/volume" "github.com/agent-substrate/substrate/internal/volume/csi" listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1" @@ -30,7 +31,7 @@ import ( // Service implements ateapipb.Control type Service struct { ateapipb.UnimplementedControlServer - persistence store.Interface + persistence serviceStore workerCache *workercache.Cache dialer *AteletDialer actorTemplateLister listersv1alpha1.ActorTemplateLister @@ -79,6 +80,28 @@ func NewService( return s } +// serviceStore enumerates the exact storage methods needed by +// the control API and nothing more. +type serviceStore interface { + CreateActor(ctx context.Context, actor *ateapipb.Actor) (*ateapipb.Actor, error) + GetActor(ctx context.Context, actorRef resources.ActorRef) (*ateapipb.Actor, error) + UpdateActor(ctx context.Context, actorRef resources.ActorRef, mutate func(toUpdate *ateapipb.Actor) error) (*ateapipb.Actor, error) + ListActors(ctx context.Context, atespace string, opts store.ListOptions) (store.ListResponse[*ateapipb.Actor], error) + GetActorSnapshot(ctx context.Context, atespace, name string) (*ateapipb.ActorSnapshot, error) + ListActorSnapshots(ctx context.Context, atespace string, opts store.ListOptions) (store.ListResponse[*ateapipb.ActorSnapshot], error) + CreateActorSnapshotTag(ctx context.Context, atespace, name string, tag *ateapipb.ActorSnapshotTag) (*ateapipb.ActorSnapshotTag, error) + GetActorSnapshotTag(ctx context.Context, atespace, name string) (*ateapipb.ActorSnapshotTag, error) + UpdateActorSnapshotTag(ctx context.Context, atespace, name string, mutate func(toUpdate *ateapipb.ActorSnapshotTag) error) (*ateapipb.ActorSnapshotTag, error) + DeleteActorSnapshotTag(ctx context.Context, atespace, name string) (*ateapipb.ActorSnapshotTag, error) + AtespaceExists(ctx context.Context, name string) (bool, error) + CreateAtespace(ctx context.Context, atespace *ateapipb.Atespace) (*ateapipb.Atespace, error) + GetAtespace(ctx context.Context, name string) (*ateapipb.Atespace, error) + ListAtespaces(ctx context.Context, opts store.ListOptions) (store.ListResponse[*ateapipb.Atespace], error) + DeleteAtespace(ctx context.Context, name string) (*ateapipb.Atespace, error) + ListWorkers(ctx context.Context, opts store.ListOptions) (store.ListResponse[*ateapipb.Worker], error) + AcquireLock(ctx context.Context, key string) (*store.Lock, error) +} + // GetPlugin retrieves a CSI volume plugin by driver name, dynamically discovering it if not present. func (s *Service) GetPlugin(ctx context.Context, driverName string) (volume.VolumePluginControlPlane, error) { s.mu.RLock() diff --git a/cmd/ateapi/internal/controlapi/syncer.go b/cmd/ateapi/internal/controlapi/syncer.go index 93f3d0507b..ac0d45bb46 100644 --- a/cmd/ateapi/internal/controlapi/syncer.go +++ b/cmd/ateapi/internal/controlapi/syncer.go @@ -54,14 +54,26 @@ type workerKey struct { // key against the current informer cache state, requeuing with rate-limited // backoff on transient failures such as store.ErrVersionConflict. type WorkerPoolSyncer struct { - persistence store.Interface + persistence workerPoolSyncerStore workerInformer cache.SharedIndexInformer workerPoolLister listersv1alpha1.WorkerPoolLister queue workqueue.TypedRateLimitingInterface[workerKey] } +// workerPoolSyncerStore enumerates the exact storage methods needed by +// WorkerPoolSyncer and nothing more. +type workerPoolSyncerStore interface { + GetActor(ctx context.Context, actorRef resources.ActorRef) (*ateapipb.Actor, error) + UpdateActor(ctx context.Context, actorRef resources.ActorRef, mutate func(toUpdate *ateapipb.Actor) error) (*ateapipb.Actor, error) + GetWorker(ctx context.Context, namespace, pool, pod string) (*ateapipb.Worker, error) + CreateWorker(ctx context.Context, worker *ateapipb.Worker) error + UpdateWorker(ctx context.Context, worker *ateapipb.Worker, expectedVersion int64) error + DeleteWorker(ctx context.Context, namespace, pool, pod string) error + ListWorkers(ctx context.Context, opts store.ListOptions) (store.ListResponse[*ateapipb.Worker], error) +} + // NewWorkerPoolSyncer creates a new WorkerPoolSyncer. -func NewWorkerPoolSyncer(persistence store.Interface, workerInformer cache.SharedIndexInformer, workerPoolLister listersv1alpha1.WorkerPoolLister) *WorkerPoolSyncer { +func NewWorkerPoolSyncer(persistence workerPoolSyncerStore, workerInformer cache.SharedIndexInformer, workerPoolLister listersv1alpha1.WorkerPoolLister) *WorkerPoolSyncer { return &WorkerPoolSyncer{ persistence: persistence, workerInformer: workerInformer, diff --git a/cmd/ateapi/internal/controlapi/volumes.go b/cmd/ateapi/internal/controlapi/volumes.go index 10146fda54..54a690c7e1 100644 --- a/cmd/ateapi/internal/controlapi/volumes.go +++ b/cmd/ateapi/internal/controlapi/volumes.go @@ -188,7 +188,7 @@ func actorVolumeID(actorUID string, volumeName string) string { } // detachActorVolumes detaches all mounted external volumes for an actor from its worker node. -func detachActorVolumes(ctx context.Context, st store.Interface, registry VolumePluginRegistry, actor *ateapipb.Actor, template *atev1alpha1.ActorTemplate, action string) error { +func detachActorVolumes(ctx context.Context, st detachActorVolumesStore, registry VolumePluginRegistry, actor *ateapipb.Actor, template *atev1alpha1.ActorTemplate, action string) error { assignment := actor.GetWorkerAssignment() if assignment == nil { slog.WarnContext(ctx, fmt.Sprintf("Actor has no assigned worker pod during %s, skipping detach volumes", action), slog.String("actor_id", actor.GetMetadata().GetName())) @@ -224,3 +224,9 @@ func detachActorVolumes(ctx context.Context, st store.Interface, registry Volume } return nil } + +// detachActorVolumesStore enumerates the subset of store methods needed to +// detach actor volumes. +type detachActorVolumesStore interface { + GetWorker(ctx context.Context, namespace, pool, pod string) (*ateapipb.Worker, error) +} diff --git a/cmd/ateapi/internal/controlapi/workflow.go b/cmd/ateapi/internal/controlapi/workflow.go index 49092f4a59..6e4fd87da2 100644 --- a/cmd/ateapi/internal/controlapi/workflow.go +++ b/cmd/ateapi/internal/controlapi/workflow.go @@ -24,6 +24,7 @@ import ( "github.com/agent-substrate/substrate/cmd/ateapi/internal/workercache" "github.com/agent-substrate/substrate/internal/resources" listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" @@ -67,7 +68,7 @@ func markSkipped(ctx context.Context, reason string) { // ActorWorkflow handles the workflows for actor's resume / suspend operations. type ActorWorkflow struct { - store store.Interface + store actorWorkflowStore workerCache *workercache.Cache scheduler scheduling.Scheduler dialer *AteletDialer @@ -82,7 +83,7 @@ type ActorWorkflow struct { // NewActorWorkflow creates a new ActorWorkflow. instruments may be nil. func NewActorWorkflow( - store store.Interface, + store actorWorkflowStore, workerCache *workercache.Cache, dialer *AteletDialer, actorTemplateLister listersv1alpha1.ActorTemplateLister, @@ -108,6 +109,19 @@ func NewActorWorkflow( } } +// actorWorkflowStore enumerates the exact storage methods needed by +// ActorWorkflow and nothing more. +type actorWorkflowStore interface { + GetActor(ctx context.Context, actorRef resources.ActorRef) (*ateapipb.Actor, error) + UpdateActor(ctx context.Context, actorRef resources.ActorRef, mutate func(toUpdate *ateapipb.Actor) error) (*ateapipb.Actor, error) + DeleteActor(ctx context.Context, actorRef resources.ActorRef) (*ateapipb.Actor, error) + GetWorker(ctx context.Context, namespace, pool, pod string) (*ateapipb.Worker, error) + UpdateWorker(ctx context.Context, worker *ateapipb.Worker, expectedVersion int64) error + GetActorSnapshot(ctx context.Context, atespace, name string) (*ateapipb.ActorSnapshot, error) + CreateActorSnapshot(ctx context.Context, snapshot *ateapipb.ActorSnapshot) (*ateapipb.ActorSnapshot, error) + AcquireLock(ctx context.Context, key string) (*store.Lock, error) +} + func (w *ActorWorkflow) acquireActorLock(ctx context.Context, actorRef resources.ActorRef) (context.Context, *store.Lock, error) { lockKey := "lock:actor:" + actorRef.Atespace + ":" + actorRef.Name diff --git a/cmd/ateapi/internal/debugapi/service.go b/cmd/ateapi/internal/debugapi/service.go index 0f05fb3c24..60ec36c72d 100644 --- a/cmd/ateapi/internal/debugapi/service.go +++ b/cmd/ateapi/internal/debugapi/service.go @@ -15,21 +15,28 @@ package debugapi import ( - "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" + "context" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" ) // Service implements ateapipb.DebugServer. type Service struct { ateapipb.UnimplementedDebugServer - persistence store.Interface + persistence debuggingStore } var _ ateapipb.DebugServer = (*Service)(nil) // NewService creates a new debug service. -func NewService(persistence store.Interface) *Service { +func NewService(persistence debuggingStore) *Service { return &Service{ persistence: persistence, } } + +// debuggingStore enumerates the exact storage methods needed by +// the debug API and nothing more. +type debuggingStore interface { + DebugClearAll(ctx context.Context) error +} diff --git a/cmd/ateapi/internal/workercache/workercache.go b/cmd/ateapi/internal/workercache/workercache.go index 25846e4f87..56b1583b47 100644 --- a/cmd/ateapi/internal/workercache/workercache.go +++ b/cmd/ateapi/internal/workercache/workercache.go @@ -39,7 +39,7 @@ const relistPageSize = 1000 // TODO: add metrics — at minimum a gauge for worker count, a counter for // resync events, and a counter for failed PUBLISH operations (in ateredis). type Cache struct { - store store.Interface + store workerListWatcher relistInterval time.Duration mu sync.RWMutex @@ -51,7 +51,7 @@ type Cache struct { // New creates a Cache backed by a given store. relistInterval controls how // often the cache performs a full ListWorkers to recover from state drifts // caused by missing WorkerWatch events. -func New(store store.Interface, relistInterval time.Duration) *Cache { +func New(store workerListWatcher, relistInterval time.Duration) *Cache { return &Cache{ store: store, relistInterval: relistInterval, @@ -59,6 +59,13 @@ func New(store store.Interface, relistInterval time.Duration) *Cache { } } +// workerListWatcher enumerates the exact storage methods needed by this +// package and nothing more. +type workerListWatcher interface { + WatchWorkers(ctx context.Context) (*store.WorkerWatch, error) + ListWorkers(ctx context.Context, opts store.ListOptions) (store.ListResponse[*ateapipb.Worker], error) +} + // Start syncs the cache synchronously, then spawns a background goroutine // that streams updates, relists periodically, and resyncs on connection loss. // Returns as soon as the initial sync succeeds.