diff --git a/cmd/atelet/actor_operation_locks.go b/cmd/atelet/actor_operation_locks.go new file mode 100644 index 000000000..81f1d723b --- /dev/null +++ b/cmd/atelet/actor_operation_locks.go @@ -0,0 +1,142 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "errors" + "fmt" + "sync" + + "github.com/agent-substrate/substrate/internal/actoroperation" + "github.com/agent-substrate/substrate/internal/ateompath" + "github.com/google/uuid" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// actorOperationLocks prevents overlapping lifecycle operations from mutating +// the same actor's node-local files. Its zero value is ready for use. +type actorOperationLocks struct { + mu sync.Mutex + locked map[string]*actorOperationLease +} + +// actorOperationLease gives each acquisition its own identity. A stale or +// duplicate release therefore cannot unlock a newer operation for the same +// actor. +// +// The byte makes this type non-zero-sized. Go permits pointers to distinct +// zero-sized variables to compare equal, which would defeat the identity +// check below. +type actorOperationLease struct { + _ byte +} + +func (l *actorOperationLocks) tryLock(actorUID string) (func(), bool) { + l.mu.Lock() + defer l.mu.Unlock() + + if _, ok := l.locked[actorUID]; ok { + return nil, false + } + if l.locked == nil { + l.locked = make(map[string]*actorOperationLease) + } + lease := &actorOperationLease{} + l.locked[actorUID] = lease + + return func() { + l.mu.Lock() + defer l.mu.Unlock() + + if l.locked[actorUID] == lease { + delete(l.locked, actorUID) + } + }, true +} + +type actorOperation struct { + actorUID string + id string + fileLock *actoroperation.Lock + endProcess func() +} + +func (s *AteomHerder) beginActorOperation(actorUID string) (*actorOperation, error) { + endProcess, ok := s.actorOperations.tryLock(actorUID) + if !ok { + return nil, status.Error(codes.Aborted, "another operation is in progress for this actor") + } + + fileLock, err := actoroperation.TryAcquire(ateompath.ActorPath(actorUID)) + if err != nil { + endProcess() + if errors.Is(err, actoroperation.ErrLocked) { + return nil, status.Error(codes.Aborted, "another operation is in progress for this actor") + } + return nil, fmt.Errorf("while acquiring actor file lock: %w", err) + } + + operation := &actorOperation{ + actorUID: actorUID, + id: uuid.NewString(), + fileLock: fileLock, + endProcess: endProcess, + } + if err := fileLock.SetOperationID(operation.id); err != nil { + _ = fileLock.Release() + endProcess() + return nil, fmt.Errorf("while recording actor operation: %w", err) + } + return operation, nil +} + +func (o *actorOperation) releaseFiles() error { + if o.fileLock == nil { + return nil + } + err := o.fileLock.Release() + o.fileLock = nil + return err +} + +func (o *actorOperation) reacquireFiles(ctx context.Context) error { + if o.fileLock != nil { + return errors.New("actor file lock is already held") + } + + fileLock, err := actoroperation.Acquire(ctx, ateompath.ActorPath(o.actorUID)) + if err != nil { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return status.FromContextError(err).Err() + } + return fmt.Errorf("while reacquiring actor file lock: %w", err) + } + if err := fileLock.CheckOperationID(o.id); err != nil { + _ = fileLock.Release() + if errors.Is(err, actoroperation.ErrOperationChanged) { + return status.Error(codes.Aborted, "actor operation was superseded") + } + return fmt.Errorf("while checking actor operation: %w", err) + } + o.fileLock = fileLock + return nil +} + +func (o *actorOperation) end() { + _ = o.releaseFiles() + o.endProcess() +} diff --git a/cmd/atelet/actor_operation_locks_test.go b/cmd/atelet/actor_operation_locks_test.go new file mode 100644 index 000000000..9cf4c0685 --- /dev/null +++ b/cmd/atelet/actor_operation_locks_test.go @@ -0,0 +1,203 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "runtime" + "strings" + "sync" + "testing" + + "github.com/agent-substrate/substrate/internal/proto/ateletpb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func TestActorOperationLocks(t *testing.T) { + var locks actorOperationLocks + + releaseActor1, ok := locks.tryLock("actor-1") + if !ok { + t.Fatal("first lock for actor-1 failed") + } + if _, ok := locks.tryLock("actor-1"); ok { + t.Error("overlapping lock for actor-1 succeeded") + } + releaseActor2, ok := locks.tryLock("actor-2") + if !ok { + t.Error("lock for a different actor failed") + } + + releaseActor1() + releaseActor1Again, ok := locks.tryLock("actor-1") + if !ok { + t.Error("actor-1 could not be locked again after unlock") + } + releaseActor1Again() + releaseActor2() + if len(locks.locked) != 0 { + t.Errorf("lock map contains %d entries after unlocks, want 0", len(locks.locked)) + } +} + +func TestActorOperationLocksStaleReleaseDoesNotUnlockNewLease(t *testing.T) { + var locks actorOperationLocks + + staleRelease, ok := locks.tryLock("actor-1") + if !ok { + t.Fatal("first lock for actor-1 failed") + } + staleRelease() + + currentRelease, ok := locks.tryLock("actor-1") + if !ok { + t.Fatal("second lock for actor-1 failed") + } + staleRelease() + if overlappingRelease, acquired := locks.tryLock("actor-1"); acquired { + overlappingRelease() + t.Error("stale release unlocked the current lease") + } + + currentRelease() +} + +func TestActorOperationLockHeldUntilExplicitRelease(t *testing.T) { + var locks actorOperationLocks + + ctx, cancel := context.WithCancel(context.Background()) + release, ok := locks.tryLock("actor-1") + if !ok { + t.Fatal("first lock for actor-1 failed") + } + + cancel() + <-ctx.Done() + if overlappingRelease, acquired := locks.tryLock("actor-1"); acquired { + overlappingRelease() + t.Error("context cancellation released the actor lock") + } + + release() + reacquiredRelease, ok := locks.tryLock("actor-1") + if !ok { + t.Error("actor-1 could not be locked after explicit release") + return + } + reacquiredRelease() +} + +func TestActorOperationLocksConcurrent(t *testing.T) { + var locks actorOperationLocks + const callers = 64 + + start := make(chan struct{}) + release := make(chan struct{}) + results := make(chan bool, callers) + var wg sync.WaitGroup + for range callers { + wg.Add(1) + go func() { + defer wg.Done() + <-start + endOperation, acquired := locks.tryLock("actor-1") + results <- acquired + if acquired { + <-release + endOperation() + } + }() + } + + close(start) + acquired := 0 + for range callers { + if <-results { + acquired++ + } + } + close(release) + wg.Wait() + + if acquired != 1 { + t.Errorf("%d concurrent callers acquired the same actor lock, want 1", acquired) + } + if len(locks.locked) != 0 { + t.Errorf("lock map contains %d entries after concurrent unlock, want 0", len(locks.locked)) + } +} + +func TestRPCBoundariesRejectOverlappingActorOperation(t *testing.T) { + s := &AteomHerder{} + ctx := context.Background() + + tests := []struct { + name string + request func() error + }{ + { + name: "Run", + request: func() error { + req := validRunRequest() + req.SandboxAssets = &ateletpb.SandboxAssets{ + SandboxClass: "gvisor", + Assets: map[string]*ateletpb.ArchAssets{ + runtime.GOARCH: { + Files: map[string]*ateletpb.AssetFile{ + "runsc": { + Url: "gs://bucket/runsc", + Sha256: strings.Repeat("a", 64), + }, + }, + }, + }, + } + _, err := s.Run(ctx, req) + return err + }, + }, + { + name: "Checkpoint", + request: func() error { + _, err := s.Checkpoint(ctx, validCheckpointRequest()) + return err + }, + }, + { + name: "Restore", + request: func() error { + _, err := s.Restore(ctx, validRestoreRequest()) + return err + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + const actorUID = "123e4567-e89b-12d3-a456-426614174000" + endOperation, ok := s.actorOperations.tryLock(actorUID) + if !ok { + t.Fatalf("failed to set up lock for %s", actorUID) + } + t.Cleanup(endOperation) + + err := tt.request() + if got := status.Code(err); got != codes.Aborted { + t.Errorf("status.Code(err) = %v, want %v (err: %v)", got, codes.Aborted, err) + } + }) + } +} diff --git a/cmd/atelet/main.go b/cmd/atelet/main.go index a826361d5..77ce9181a 100644 --- a/cmd/atelet/main.go +++ b/cmd/atelet/main.go @@ -191,10 +191,11 @@ func main() { type AteomHerder struct { ateletpb.UnimplementedAteomHerderServer - ateomDialer *AteomDialer - imageCache *imagecache.Store - anonGCSClient ategcs.ObjectStorage - gcsClient ategcs.ObjectStorage + ateomDialer *AteomDialer + imageCache *imagecache.Store + anonGCSClient ategcs.ObjectStorage + gcsClient ategcs.ObjectStorage + actorOperations actorOperationLocks } var _ ateletpb.AteomHerderServer = (*AteomHerder)(nil) @@ -229,6 +230,12 @@ func (s *AteomHerder) Run(ctx context.Context, req *ateletpb.RunRequest) (*atele if err != nil { return nil, status.Error(codes.InvalidArgument, err.Error()) } + operation, err := s.beginActorOperation(actorUID) + if err != nil { + return nil, err + } + defer operation.end() + assetPaths, err := s.ensureSandboxAssets(ctx, sandboxRec) if err != nil { return nil, err @@ -255,6 +262,9 @@ func (s *AteomHerder) Run(ctx context.Context, req *ateletpb.RunRequest) (*atele if err != nil { return nil, err } + if err := operation.releaseFiles(); err != nil { + return nil, fmt.Errorf("while releasing actor file lock before ateom RunWorkload: %w", err) + } // Tell ateom to start the workload. gVisor uses RunscPath; the micro-VM // runtime uses the full RuntimeAssetPaths set. @@ -267,6 +277,7 @@ func (s *AteomHerder) Run(ctx context.Context, req *ateletpb.RunRequest) (*atele RuntimeAssetPaths: assetPaths, Spec: buildAteomWorkloadSpec(req.GetSpec()), ActorUid: actorUID, + OperationId: operation.id, }); err != nil { return nil, fmt.Errorf("while calling ateom.RunWorkload: %w", err) } @@ -317,6 +328,12 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe actorUID, atespace, actorName := req.GetActorUid(), req.GetAtespace(), req.GetActorName() + operation, err := s.beginActorOperation(actorUID) + if err != nil { + return nil, err + } + defer operation.end() + // Checkpoint requests no longer carry the sandbox config; recover the // version this actor was started with from the on-node record and re-fetch // it (a cache hit) so ateom can drive runsc, and so we can pin it into the @@ -336,6 +353,9 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe if err != nil { return nil, err } + if err := operation.releaseFiles(); err != nil { + return nil, fmt.Errorf("while releasing actor file lock before ateom CheckpointWorkload: %w", err) + } // Tell ateom to take the checkpoint and delete containers. ateom reports the // exact files it wrote so we ship precisely that set (gVisor's image files, @@ -350,12 +370,16 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe Spec: buildAteomWorkloadSpec(req.GetSpec()), Scope: toAteomSnapshotScope(req.GetScope()), ActorUid: actorUID, + OperationId: operation.id, }) if err != nil { // TODO: Ateom should classify checkpoint failures, and set "should-crash" // in the metadata if the error is not retriable. return nil, fmt.Errorf("while calling ateom.CheckpointWorkload: %w", err) } + if err := operation.reacquireFiles(ctx); err != nil { + return nil, err + } sandboxRec.SnapshotFiles = resp.GetSnapshotFiles() if len(sandboxRec.SnapshotFiles) == 0 { return nil, ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonInvalidCheckpointResult, ateerrors.ActorCrashedMetadata(), errors.New("ateom reported no snapshot files for checkpoint")) @@ -462,6 +486,12 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest) actorUID, atespace, actorName := req.GetActorUid(), req.GetAtespace(), req.GetActorName() + operation, err := s.beginActorOperation(actorUID) + if err != nil { + return nil, err + } + defer operation.end() + // Not crashing the actor, because terminal errors here indicate problems with atelet, // node or the disk itself. if err := resetActorDirs(actorUID); err != nil { @@ -551,6 +581,9 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest) if err != nil { return nil, err } + if err := operation.releaseFiles(); err != nil { + return nil, fmt.Errorf("while releasing actor file lock before ateom RestoreWorkload: %w", err) + } // Tell ateom to do runsc create + runsc restore for pause container and // all application containers. @@ -565,11 +598,15 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest) Spec: buildAteomWorkloadSpec(req.GetSpec()), Scope: toAteomSnapshotScope(req.GetScope()), ActorUid: req.GetActorUid(), + OperationId: operation.id, }); err != nil { // TODO: classify the errors returned by Ateom and crash the actor if needed. return nil, fmt.Errorf("while calling ateom.RestoreWorkload: %w", err) } dAteom = time.Since(tAteom) + if err := operation.reacquireFiles(ctx); err != nil { + return nil, err + } // Record the (manifest-pinned) sandbox binaries on-node so a subsequent // Checkpoint of this restored actor can re-pin the same version. diff --git a/cmd/ateom-gvisor/actor_operation.go b/cmd/ateom-gvisor/actor_operation.go new file mode 100644 index 000000000..897b13bef --- /dev/null +++ b/cmd/ateom-gvisor/actor_operation.go @@ -0,0 +1,73 @@ +//go:build linux + +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "errors" + "fmt" + "log/slog" + + "github.com/agent-substrate/substrate/internal/actoroperation" + "github.com/agent-substrate/substrate/internal/ateompath" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func acquireActorOperation(ctx context.Context, actorUID, operationID string) (*actoroperation.Lock, error) { + if err := ctx.Err(); err != nil { + return nil, status.FromContextError(err).Err() + } + + fileLock, err := actoroperation.TryAcquire(ateompath.ActorPath(actorUID)) + if err != nil { + if errors.Is(err, actoroperation.ErrLocked) { + return nil, status.Error(codes.Aborted, "another operation is in progress for this actor") + } + return nil, fmt.Errorf("while acquiring actor file lock: %w", err) + } + + // Empty operation IDs come from an older atelet during a rolling upgrade. + // Accept one only if no generation-aware atelet has prepared this actor; + // otherwise a queued legacy request could mutate a newer operation's files. + if operationID == "" { + hasOperationID, err := fileLock.HasOperationID() + if err != nil { + _ = fileLock.Release() + return nil, err + } + if hasOperationID { + _ = fileLock.Release() + return nil, status.Error(codes.Aborted, "actor operation ID is missing") + } + return fileLock, nil + } + if err := fileLock.CheckOperationID(operationID); err != nil { + _ = fileLock.Release() + if errors.Is(err, actoroperation.ErrOperationChanged) { + return nil, status.Error(codes.Aborted, "actor operation was superseded") + } + return nil, fmt.Errorf("while checking actor operation: %w", err) + } + return fileLock, nil +} + +func releaseActorOperation(ctx context.Context, fileLock *actoroperation.Lock) { + if err := fileLock.Release(); err != nil { + slog.ErrorContext(ctx, "Failed to release actor file lock", "err", err) + } +} diff --git a/cmd/ateom-gvisor/main.go b/cmd/ateom-gvisor/main.go index 4d245dc6f..5d79f2f13 100644 --- a/cmd/ateom-gvisor/main.go +++ b/cmd/ateom-gvisor/main.go @@ -182,6 +182,12 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload s.lock.Lock() defer s.lock.Unlock() + operation, err := acquireActorOperation(ctx, req.GetActorUid(), req.GetOperationId()) + if err != nil { + return nil, err + } + defer releaseActorOperation(ctx, operation) + s.actorLogger.EmitLifecycleLog("Actor starting", req.GetAtespace(), req.GetActorName(), req.GetActorUid(), req.GetActorTemplateNamespace(), req.GetActorTemplateName()) // Contract with atelet: @@ -260,6 +266,12 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec s.lock.Lock() defer s.lock.Unlock() + operation, err := acquireActorOperation(ctx, req.GetActorUid(), req.GetOperationId()) + if err != nil { + return nil, err + } + defer releaseActorOperation(ctx, operation) + s.actorLogger.EmitLifecycleLog("Actor checkpointing", req.GetAtespace(), req.GetActorName(), req.GetActorUid(), req.GetActorTemplateNamespace(), req.GetActorTemplateName()) // Contract with atelet: @@ -383,6 +395,12 @@ func (s *AteomService) RestoreWorkload(ctx context.Context, req *ateompb.Restore s.lock.Lock() defer s.lock.Unlock() + operation, err := acquireActorOperation(ctx, req.GetActorUid(), req.GetOperationId()) + if err != nil { + return nil, err + } + defer releaseActorOperation(ctx, operation) + s.actorLogger.EmitLifecycleLog("Actor restoring", req.GetAtespace(), req.GetActorName(), req.GetActorUid(), req.GetActorTemplateNamespace(), req.GetActorTemplateName()) // Contract with atelet: diff --git a/cmd/ateom-microvm/actor_operation.go b/cmd/ateom-microvm/actor_operation.go new file mode 100644 index 000000000..897b13bef --- /dev/null +++ b/cmd/ateom-microvm/actor_operation.go @@ -0,0 +1,73 @@ +//go:build linux + +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "errors" + "fmt" + "log/slog" + + "github.com/agent-substrate/substrate/internal/actoroperation" + "github.com/agent-substrate/substrate/internal/ateompath" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func acquireActorOperation(ctx context.Context, actorUID, operationID string) (*actoroperation.Lock, error) { + if err := ctx.Err(); err != nil { + return nil, status.FromContextError(err).Err() + } + + fileLock, err := actoroperation.TryAcquire(ateompath.ActorPath(actorUID)) + if err != nil { + if errors.Is(err, actoroperation.ErrLocked) { + return nil, status.Error(codes.Aborted, "another operation is in progress for this actor") + } + return nil, fmt.Errorf("while acquiring actor file lock: %w", err) + } + + // Empty operation IDs come from an older atelet during a rolling upgrade. + // Accept one only if no generation-aware atelet has prepared this actor; + // otherwise a queued legacy request could mutate a newer operation's files. + if operationID == "" { + hasOperationID, err := fileLock.HasOperationID() + if err != nil { + _ = fileLock.Release() + return nil, err + } + if hasOperationID { + _ = fileLock.Release() + return nil, status.Error(codes.Aborted, "actor operation ID is missing") + } + return fileLock, nil + } + if err := fileLock.CheckOperationID(operationID); err != nil { + _ = fileLock.Release() + if errors.Is(err, actoroperation.ErrOperationChanged) { + return nil, status.Error(codes.Aborted, "actor operation was superseded") + } + return nil, fmt.Errorf("while checking actor operation: %w", err) + } + return fileLock, nil +} + +func releaseActorOperation(ctx context.Context, fileLock *actoroperation.Lock) { + if err := fileLock.Release(); err != nil { + slog.ErrorContext(ctx, "Failed to release actor file lock", "err", err) + } +} diff --git a/cmd/ateom-microvm/checkpoint.go b/cmd/ateom-microvm/checkpoint.go index f95e5a7ee..e8c004222 100644 --- a/cmd/ateom-microvm/checkpoint.go +++ b/cmd/ateom-microvm/checkpoint.go @@ -46,6 +46,12 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec s.lock.Lock() defer s.lock.Unlock() + operation, err := acquireActorOperation(ctx, req.GetActorUid(), req.GetOperationId()) + if err != nil { + return nil, err + } + defer releaseActorOperation(ctx, operation) + atespace := req.GetAtespace() name := req.GetActorName() actorUID := req.GetActorUid() diff --git a/cmd/ateom-microvm/restore.go b/cmd/ateom-microvm/restore.go index 3a98860e1..fc19f2339 100644 --- a/cmd/ateom-microvm/restore.go +++ b/cmd/ateom-microvm/restore.go @@ -54,6 +54,12 @@ func (s *AteomService) RestoreWorkload(ctx context.Context, req *ateompb.Restore s.lock.Lock() defer s.lock.Unlock() + operation, err := acquireActorOperation(ctx, req.GetActorUid(), req.GetOperationId()) + if err != nil { + return nil, err + } + defer releaseActorOperation(ctx, operation) + atespace := req.GetAtespace() name := req.GetActorName() actorUID := req.GetActorUid() diff --git a/cmd/ateom-microvm/run.go b/cmd/ateom-microvm/run.go index 7c323e787..598673586 100644 --- a/cmd/ateom-microvm/run.go +++ b/cmd/ateom-microvm/run.go @@ -189,6 +189,12 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload s.lock.Lock() defer s.lock.Unlock() + operation, err := acquireActorOperation(ctx, req.GetActorUid(), req.GetOperationId()) + if err != nil { + return nil, err + } + defer releaseActorOperation(ctx, operation) + atespace := req.GetAtespace() name := req.GetActorName() actorUID := req.GetActorUid() diff --git a/internal/actoroperation/actoroperation.go b/internal/actoroperation/actoroperation.go new file mode 100644 index 000000000..114d459b4 --- /dev/null +++ b/internal/actoroperation/actoroperation.go @@ -0,0 +1,163 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package actoroperation coordinates access to an actor's node-local files +// across atelet and ateom processes. +package actoroperation + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "sync" + "time" + + "golang.org/x/sys/unix" +) + +const ( + lockFileName = ".operation.lock" + operationIDName = ".operation-id" + lockRetryInterval = 10 * time.Millisecond +) + +var ( + // ErrLocked means another process currently owns the actor's file lock. + ErrLocked = errors.New("actor files are locked") + + // ErrOperationChanged means the files were prepared by a different + // operation than the caller expected. + ErrOperationChanged = errors.New("actor operation changed") +) + +// Lock is an exclusive advisory lock for one actor directory. +type Lock struct { + file *os.File + actorPath string + + releaseOnce sync.Once + releaseErr error +} + +// TryAcquire acquires actorPath's lock without waiting. +func TryAcquire(actorPath string) (*Lock, error) { + if err := os.MkdirAll(actorPath, 0o700); err != nil { + return nil, fmt.Errorf("while creating actor directory: %w", err) + } + + file, err := os.OpenFile(filepath.Join(actorPath, lockFileName), os.O_CREATE|os.O_RDWR, 0o600) + if err != nil { + return nil, fmt.Errorf("while opening actor operation lock: %w", err) + } + if err := unix.Flock(int(file.Fd()), unix.LOCK_EX|unix.LOCK_NB); err != nil { + _ = file.Close() + if errors.Is(err, unix.EWOULDBLOCK) || errors.Is(err, unix.EAGAIN) { + return nil, ErrLocked + } + return nil, fmt.Errorf("while locking actor operation file: %w", err) + } + + return &Lock{file: file, actorPath: actorPath}, nil +} + +// Acquire waits until actorPath's lock is available or ctx is done. +func Acquire(ctx context.Context, actorPath string) (*Lock, error) { + ticker := time.NewTicker(lockRetryInterval) + defer ticker.Stop() + + for { + if err := ctx.Err(); err != nil { + return nil, err + } + lock, err := TryAcquire(actorPath) + if !errors.Is(err, ErrLocked) { + return lock, err + } + + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-ticker.C: + } + } +} + +// SetOperationID records which operation prepared the actor's files. The +// caller must hold l until this method returns. +func (l *Lock) SetOperationID(operationID string) error { + if operationID == "" { + return errors.New("operation ID is empty") + } + + file, err := os.CreateTemp(l.actorPath, ".operation-id-") + if err != nil { + return fmt.Errorf("while creating temporary operation ID file: %w", err) + } + tmpPath := file.Name() + defer os.Remove(tmpPath) + + if _, err := fmt.Fprintln(file, operationID); err != nil { + _ = file.Close() + return fmt.Errorf("while writing operation ID: %w", err) + } + if err := file.Close(); err != nil { + return fmt.Errorf("while closing operation ID: %w", err) + } + if err := os.Rename(tmpPath, filepath.Join(l.actorPath, operationIDName)); err != nil { + return fmt.Errorf("while installing operation ID: %w", err) + } + return nil +} + +// CheckOperationID verifies that no newer operation prepared the actor's +// files. The caller must hold l until this method returns. +func (l *Lock) CheckOperationID(operationID string) error { + got, err := os.ReadFile(filepath.Join(l.actorPath, operationIDName)) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("%w: operation ID is missing", ErrOperationChanged) + } + return fmt.Errorf("while reading operation ID: %w", err) + } + if string(got) != operationID+"\n" { + return fmt.Errorf("%w: got %q, want %q", ErrOperationChanged, string(got), operationID) + } + return nil +} + +// HasOperationID reports whether a generation-aware atelet has prepared this +// actor before. The caller must hold l until this method returns. +func (l *Lock) HasOperationID() (bool, error) { + _, err := os.Stat(filepath.Join(l.actorPath, operationIDName)) + switch { + case err == nil: + return true, nil + case errors.Is(err, os.ErrNotExist): + return false, nil + default: + return false, fmt.Errorf("while checking operation ID file: %w", err) + } +} + +// Release unlocks and closes the lock. It is safe to call more than once. +func (l *Lock) Release() error { + l.releaseOnce.Do(func() { + unlockErr := unix.Flock(int(l.file.Fd()), unix.LOCK_UN) + closeErr := l.file.Close() + l.releaseErr = errors.Join(unlockErr, closeErr) + }) + return l.releaseErr +} diff --git a/internal/actoroperation/actoroperation_test.go b/internal/actoroperation/actoroperation_test.go new file mode 100644 index 000000000..cd9e26fad --- /dev/null +++ b/internal/actoroperation/actoroperation_test.go @@ -0,0 +1,247 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package actoroperation + +import ( + "context" + "errors" + "os" + "os/exec" + "path/filepath" + "testing" + "time" +) + +func TestTryAcquireExcludesAndReleases(t *testing.T) { + actorPath := filepath.Join(t.TempDir(), "actor") + + lock, err := TryAcquire(actorPath) + if err != nil { + t.Fatalf("TryAcquire() failed: %v", err) + } + if _, err := TryAcquire(actorPath); !errors.Is(err, ErrLocked) { + t.Fatalf("overlapping TryAcquire() error = %v, want ErrLocked", err) + } + if err := lock.Release(); err != nil { + t.Fatalf("Release() failed: %v", err) + } + if err := lock.Release(); err != nil { + t.Fatalf("second Release() failed: %v", err) + } + + reacquired, err := TryAcquire(actorPath) + if err != nil { + t.Fatalf("TryAcquire() after release failed: %v", err) + } + if err := reacquired.Release(); err != nil { + t.Fatalf("reacquired Release() failed: %v", err) + } +} + +func TestAcquireHonorsContext(t *testing.T) { + actorPath := filepath.Join(t.TempDir(), "actor") + lock, err := TryAcquire(actorPath) + if err != nil { + t.Fatalf("TryAcquire() failed: %v", err) + } + t.Cleanup(func() { + if err := lock.Release(); err != nil { + t.Errorf("Release() failed: %v", err) + } + }) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + if _, err := Acquire(ctx, actorPath); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("Acquire() error = %v, want DeadlineExceeded", err) + } +} + +func TestAcquireDoesNotTakeFreeLockAfterCancellation(t *testing.T) { + actorPath := filepath.Join(t.TempDir(), "actor") + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + if _, err := Acquire(ctx, actorPath); !errors.Is(err, context.Canceled) { + t.Fatalf("Acquire() error = %v, want Canceled", err) + } + + lock, err := TryAcquire(actorPath) + if err != nil { + t.Fatalf("TryAcquire() after canceled Acquire failed: %v", err) + } + if err := lock.Release(); err != nil { + t.Fatalf("Release() failed: %v", err) + } +} + +func TestOperationIDRejectsStaleOperation(t *testing.T) { + actorPath := filepath.Join(t.TempDir(), "actor") + lock, err := TryAcquire(actorPath) + if err != nil { + t.Fatalf("TryAcquire() failed: %v", err) + } + defer lock.Release() + + if err := lock.SetOperationID("operation-1"); err != nil { + t.Fatalf("SetOperationID(operation-1) failed: %v", err) + } + if exists, err := lock.HasOperationID(); err != nil || !exists { + t.Fatalf("HasOperationID() = %v, %v, want true, nil", exists, err) + } + if err := lock.CheckOperationID("operation-1"); err != nil { + t.Fatalf("CheckOperationID(operation-1) failed: %v", err) + } + if err := lock.SetOperationID("operation-2"); err != nil { + t.Fatalf("SetOperationID(operation-2) failed: %v", err) + } + if err := lock.CheckOperationID("operation-1"); !errors.Is(err, ErrOperationChanged) { + t.Fatalf("stale CheckOperationID() error = %v, want ErrOperationChanged", err) + } +} + +func TestMissingOperationIDIsStale(t *testing.T) { + actorPath := filepath.Join(t.TempDir(), "actor") + lock, err := TryAcquire(actorPath) + if err != nil { + t.Fatalf("TryAcquire() failed: %v", err) + } + defer lock.Release() + + if err := lock.CheckOperationID("operation-1"); !errors.Is(err, ErrOperationChanged) { + t.Fatalf("CheckOperationID() error = %v, want ErrOperationChanged", err) + } +} + +func TestHasOperationIDIsFalseBeforeFirstGeneration(t *testing.T) { + actorPath := filepath.Join(t.TempDir(), "actor") + lock, err := TryAcquire(actorPath) + if err != nil { + t.Fatalf("TryAcquire() failed: %v", err) + } + defer lock.Release() + + if exists, err := lock.HasOperationID(); err != nil || exists { + t.Fatalf("HasOperationID() = %v, %v, want false, nil", exists, err) + } +} + +func TestTryAcquireExcludesAnotherProcess(t *testing.T) { + if os.Getenv("ACTOR_OPERATION_LOCK_HELPER") == "1" { + actorPath := os.Getenv("ACTOR_OPERATION_LOCK_PATH") + readyPath := os.Getenv("ACTOR_OPERATION_READY_PATH") + releasePath := os.Getenv("ACTOR_OPERATION_RELEASE_PATH") + + lock, err := TryAcquire(actorPath) + if err != nil { + os.Exit(2) + } + if err := os.WriteFile(readyPath, nil, 0o600); err != nil { + os.Exit(3) + } + if os.Getenv("ACTOR_OPERATION_EXIT_AFTER_READY") == "1" { + os.Exit(0) + } + for { + if _, err := os.Stat(releasePath); err == nil { + break + } else if !errors.Is(err, os.ErrNotExist) { + os.Exit(4) + } + time.Sleep(time.Millisecond) + } + if err := lock.Release(); err != nil { + os.Exit(5) + } + os.Exit(0) + } + + tmpDir := t.TempDir() + actorPath := filepath.Join(tmpDir, "actor") + readyPath := filepath.Join(tmpDir, "ready") + releasePath := filepath.Join(tmpDir, "release") + cmd := exec.Command(os.Args[0], "-test.run=TestTryAcquireExcludesAnotherProcess") + cmd.Env = append(os.Environ(), + "ACTOR_OPERATION_LOCK_HELPER=1", + "ACTOR_OPERATION_LOCK_PATH="+actorPath, + "ACTOR_OPERATION_READY_PATH="+readyPath, + "ACTOR_OPERATION_RELEASE_PATH="+releasePath, + ) + if err := cmd.Start(); err != nil { + t.Fatalf("starting helper: %v", err) + } + t.Cleanup(func() { + _ = os.WriteFile(releasePath, nil, 0o600) + _ = cmd.Wait() + }) + + deadline := time.Now().Add(5 * time.Second) + for { + if _, err := os.Stat(readyPath); err == nil { + break + } else if !errors.Is(err, os.ErrNotExist) { + t.Fatalf("checking helper readiness: %v", err) + } + if time.Now().After(deadline) { + t.Fatal("timed out waiting for helper to acquire lock") + } + time.Sleep(time.Millisecond) + } + + if _, err := TryAcquire(actorPath); !errors.Is(err, ErrLocked) { + t.Fatalf("TryAcquire() while helper holds lock error = %v, want ErrLocked", err) + } + if err := os.WriteFile(releasePath, nil, 0o600); err != nil { + t.Fatalf("releasing helper: %v", err) + } + if err := cmd.Wait(); err != nil { + t.Fatalf("helper failed: %v", err) + } + + lock, err := TryAcquire(actorPath) + if err != nil { + t.Fatalf("TryAcquire() after helper exited failed: %v", err) + } + if err := lock.Release(); err != nil { + t.Fatalf("Release() failed: %v", err) + } +} + +func TestLockReleasedWhenProcessExits(t *testing.T) { + tmpDir := t.TempDir() + actorPath := filepath.Join(tmpDir, "actor") + readyPath := filepath.Join(tmpDir, "ready") + cmd := exec.Command(os.Args[0], "-test.run=TestTryAcquireExcludesAnotherProcess") + cmd.Env = append(os.Environ(), + "ACTOR_OPERATION_LOCK_HELPER=1", + "ACTOR_OPERATION_EXIT_AFTER_READY=1", + "ACTOR_OPERATION_LOCK_PATH="+actorPath, + "ACTOR_OPERATION_READY_PATH="+readyPath, + ) + if err := cmd.Run(); err != nil { + t.Fatalf("helper failed: %v", err) + } + if _, err := os.Stat(readyPath); err != nil { + t.Fatalf("helper did not report lock acquisition: %v", err) + } + + lock, err := TryAcquire(actorPath) + if err != nil { + t.Fatalf("TryAcquire() after helper exit failed: %v", err) + } + if err := lock.Release(); err != nil { + t.Fatalf("Release() failed: %v", err) + } +} diff --git a/internal/proto/ateompb/ateom.pb.go b/internal/proto/ateompb/ateom.pb.go index 12c4fa1a7..6592b72a9 100644 --- a/internal/proto/ateompb/ateom.pb.go +++ b/internal/proto/ateompb/ateom.pb.go @@ -104,8 +104,12 @@ type RunWorkloadRequest struct { // to the local on-disk path atelet fetched it to (content-addressed, like // runsc_path). Empty for the gVisor runtime, which uses runsc_path. RuntimeAssetPaths map[string]string `protobuf:"bytes,8,rep,name=runtime_asset_paths,json=runtimeAssetPaths,proto3" json:"runtime_asset_paths,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // operation_id identifies the atelet operation that prepared the actor's + // node-local files. Ateom rejects a stale request if a newer operation has + // replaced those files before this RPC starts. + OperationId string `protobuf:"bytes,9,opt,name=operation_id,json=operationId,proto3" json:"operation_id,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *RunWorkloadRequest) Reset() { @@ -194,6 +198,13 @@ func (x *RunWorkloadRequest) GetRuntimeAssetPaths() map[string]string { return nil } +func (x *RunWorkloadRequest) GetOperationId() string { + if x != nil { + return x.OperationId + } + return "" +} + // WorkloadSpec parallels Pod, but with far fewer configurable fields. type WorkloadSpec struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -458,7 +469,10 @@ type CheckpointWorkloadRequest struct { // atelet fetched it to (see RunWorkloadRequest). Empty for gVisor. RuntimeAssetPaths map[string]string `protobuf:"bytes,9,rep,name=runtime_asset_paths,json=runtimeAssetPaths,proto3" json:"runtime_asset_paths,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` // What content to include in the checkpoint. - Scope SnapshotScope `protobuf:"varint,10,opt,name=scope,proto3,enum=ateom.SnapshotScope" json:"scope,omitempty"` + Scope SnapshotScope `protobuf:"varint,10,opt,name=scope,proto3,enum=ateom.SnapshotScope" json:"scope,omitempty"` + // operation_id identifies the atelet operation that owns the actor's + // node-local files. + OperationId string `protobuf:"bytes,11,opt,name=operation_id,json=operationId,proto3" json:"operation_id,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -563,6 +577,13 @@ func (x *CheckpointWorkloadRequest) GetScope() SnapshotScope { return SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED } +func (x *CheckpointWorkloadRequest) GetOperationId() string { + if x != nil { + return x.OperationId + } + return "" +} + type CheckpointWorkloadResponse struct { state protoimpl.MessageState `protogen:"open.v1"` // snapshot_files lists the files ateom wrote into the checkpoint directory @@ -625,7 +646,10 @@ type RestoreWorkloadRequest struct { // atelet fetched it to (see RunWorkloadRequest). Empty for gVisor. RuntimeAssetPaths map[string]string `protobuf:"bytes,9,rep,name=runtime_asset_paths,json=runtimeAssetPaths,proto3" json:"runtime_asset_paths,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` // What content to restore from the snapshot. - Scope SnapshotScope `protobuf:"varint,10,opt,name=scope,proto3,enum=ateom.SnapshotScope" json:"scope,omitempty"` + Scope SnapshotScope `protobuf:"varint,10,opt,name=scope,proto3,enum=ateom.SnapshotScope" json:"scope,omitempty"` + // operation_id identifies the atelet operation that prepared the actor's + // node-local files. + OperationId string `protobuf:"bytes,11,opt,name=operation_id,json=operationId,proto3" json:"operation_id,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -730,6 +754,13 @@ func (x *RestoreWorkloadRequest) GetScope() SnapshotScope { return SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED } +func (x *RestoreWorkloadRequest) GetOperationId() string { + if x != nil { + return x.OperationId + } + return "" +} + type RestoreWorkloadResponse struct { state protoimpl.MessageState `protogen:"open.v1"` unknownFields protoimpl.UnknownFields @@ -770,7 +801,7 @@ var File_ateom_proto protoreflect.FileDescriptor const file_ateom_proto_rawDesc = "" + "\n" + - "\vateom.proto\x12\x05ateom\"\xc6\x03\n" + + "\vateom.proto\x12\x05ateom\"\xe9\x03\n" + "\x12RunWorkloadRequest\x12\x1a\n" + "\batespace\x18\x01 \x01(\tR\batespace\x12\x1d\n" + "\n" + @@ -781,7 +812,8 @@ const file_ateom_proto_rawDesc = "" + "\n" + "runsc_path\x18\x06 \x01(\tR\trunscPath\x12'\n" + "\x04spec\x18\a \x01(\v2\x13.ateom.WorkloadSpecR\x04spec\x12`\n" + - "\x13runtime_asset_paths\x18\b \x03(\v20.ateom.RunWorkloadRequest.RuntimeAssetPathsEntryR\x11runtimeAssetPaths\x1aD\n" + + "\x13runtime_asset_paths\x18\b \x03(\v20.ateom.RunWorkloadRequest.RuntimeAssetPathsEntryR\x11runtimeAssetPaths\x12!\n" + + "\foperation_id\x18\t \x01(\tR\voperationId\x1aD\n" + "\x16RuntimeAssetPathsEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"@\n" + @@ -798,7 +830,7 @@ const file_ateom_proto_rawDesc = "" + "\rHTTPGetAction\x12\x12\n" + "\x04path\x18\x01 \x01(\tR\x04path\x12\x12\n" + "\x04port\x18\x02 \x01(\x05R\x04port\"\x15\n" + - "\x13RunWorkloadResponse\"\xb0\x04\n" + + "\x13RunWorkloadResponse\"\xd3\x04\n" + "\x19CheckpointWorkloadRequest\x12\x1a\n" + "\batespace\x18\x01 \x01(\tR\batespace\x12\x1d\n" + "\n" + @@ -812,12 +844,13 @@ const file_ateom_proto_rawDesc = "" + "\x13snapshot_uri_prefix\x18\b \x01(\tR\x11snapshotUriPrefix\x12g\n" + "\x13runtime_asset_paths\x18\t \x03(\v27.ateom.CheckpointWorkloadRequest.RuntimeAssetPathsEntryR\x11runtimeAssetPaths\x12*\n" + "\x05scope\x18\n" + - " \x01(\x0e2\x14.ateom.SnapshotScopeR\x05scope\x1aD\n" + + " \x01(\x0e2\x14.ateom.SnapshotScopeR\x05scope\x12!\n" + + "\foperation_id\x18\v \x01(\tR\voperationId\x1aD\n" + "\x16RuntimeAssetPathsEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"C\n" + "\x1aCheckpointWorkloadResponse\x12%\n" + - "\x0esnapshot_files\x18\x01 \x03(\tR\rsnapshotFiles\"\xaa\x04\n" + + "\x0esnapshot_files\x18\x01 \x03(\tR\rsnapshotFiles\"\xcd\x04\n" + "\x16RestoreWorkloadRequest\x12\x1a\n" + "\batespace\x18\x01 \x01(\tR\batespace\x12\x1d\n" + "\n" + @@ -831,7 +864,8 @@ const file_ateom_proto_rawDesc = "" + "\x13snapshot_uri_prefix\x18\b \x01(\tR\x11snapshotUriPrefix\x12d\n" + "\x13runtime_asset_paths\x18\t \x03(\v24.ateom.RestoreWorkloadRequest.RuntimeAssetPathsEntryR\x11runtimeAssetPaths\x12*\n" + "\x05scope\x18\n" + - " \x01(\x0e2\x14.ateom.SnapshotScopeR\x05scope\x1aD\n" + + " \x01(\x0e2\x14.ateom.SnapshotScopeR\x05scope\x12!\n" + + "\foperation_id\x18\v \x01(\tR\voperationId\x1aD\n" + "\x16RuntimeAssetPathsEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\x19\n" + diff --git a/internal/proto/ateompb/ateom.proto b/internal/proto/ateompb/ateom.proto index a282d81cb..b67d65d11 100644 --- a/internal/proto/ateompb/ateom.proto +++ b/internal/proto/ateompb/ateom.proto @@ -64,6 +64,11 @@ message RunWorkloadRequest { // to the local on-disk path atelet fetched it to (content-addressed, like // runsc_path). Empty for the gVisor runtime, which uses runsc_path. map runtime_asset_paths = 8; + + // operation_id identifies the atelet operation that prepared the actor's + // node-local files. Ateom rejects a stale request if a newer operation has + // replaced those files before this RPC starts. + string operation_id = 9; } // WorkloadSpec parallels Pod, but with far fewer configurable fields. @@ -134,6 +139,10 @@ message CheckpointWorkloadRequest { // What content to include in the checkpoint. SnapshotScope scope = 10; + + // operation_id identifies the atelet operation that owns the actor's + // node-local files. + string operation_id = 11; } message CheckpointWorkloadResponse { @@ -164,6 +173,10 @@ message RestoreWorkloadRequest { // What content to restore from the snapshot. SnapshotScope scope = 10; + + // operation_id identifies the atelet operation that prepared the actor's + // node-local files. + string operation_id = 11; } message RestoreWorkloadResponse {