Decouple actor lock TTL from workflow deadline via heartbeat (#14)

ActorWorkflow.ResumeActor and SuspendActor used to derive their workflow
ctx from the Redis lock TTL via acquireActorLock(ctx, id, 30s, 2s) — the
workflow deadline and the lock TTL were a single 28s knob. That meant
image pulls / restores that legitimately need more than 28s death-looped
forever, while raising the knob also raised how long peers wait to retry
an actor after a crashed ateapi replica.

Split the two concerns:

- Lock TTL stays short (30s constant, internal). Bounds peer failover.
- Workflow deadline is a separate operator-configurable knob via the
  new --actor-workflow-deadline pflag (default 5m). Bounds a single
  Resume/Suspend.
- A heartbeat goroutine refreshes the lock every lockTTL/3 (~10s) for
  the full workflow duration. On RefreshLock=false or any Redis error
  (peer stole the lock, Redis blip), the workflow ctx is cancelled with
  errLostActorLock as the cause so in-flight steps unwind cleanly and
  the mutual-exclusion invariant is preserved.
- The release function stops the heartbeat (waits for goroutine exit)
  before best-effort ReleaseLock.

Adds store.Interface.RefreshLock with a Redis CAS Lua script mirroring
the existing ReleaseLock script.

Author: jjamroga

cmd/ateapi/internal/controlapi/functional_test.go

@@ -367,7 +367,7 @@ func setupTest(t *testing.T, ns string) *testContext {
 	volPlugins := map[string]volume.VolumePluginControlPlane{
 		mockDriverName: mockPlugin,
 	}
-	service := NewService(persistence, wc, actorTemplateLister, workerPoolLister, sandboxConfigLister, csiDriverConfigLister, scLister, dialer, instruments, "", volPlugins)
+	service := NewService(persistence, wc, actorTemplateLister, workerPoolLister, sandboxConfigLister, csiDriverConfigLister, scLister, dialer, instruments, "", 30*time.Second, volPlugins)
 
 	// 5. Start REAL gRPC Server for ATE API
 	grpcServer := grpc.NewServer(grpc.UnaryInterceptor(ateinterceptors.ServerUnaryInterceptor))

cmd/ateapi/internal/controlapi/service.go

@@ -17,6 +17,7 @@ package controlapi
 import (
 	"context"
 	"sync"
+	"time"
 
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/workercache"
@@ -50,7 +51,8 @@ type VolumePluginRegistry interface {
 	GetPlugin(ctx context.Context, name string) (volume.VolumePluginControlPlane, error)
 }
 
-// NewService creates a service. instruments may be nil; the record helpers no-op.
+// NewService creates a service. actorWorkflowDeadline bounds how long a single
+// Resume/Suspend workflow can run end-to-end. instruments may be nil.
 func NewService(
 	persistence store.Interface,
 	workerCache *workercache.Cache,
@@ -62,6 +64,7 @@ func NewService(
 	dialer *AteletDialer,
 	instruments *Instruments,
 	egressGatewayAddress string,
+	actorWorkflowDeadline time.Duration,
 	volumePlugins map[string]volume.VolumePluginControlPlane,
 ) *Service {
 	s := &Service{
@@ -75,7 +78,7 @@ func NewService(
 		instruments:           instruments,
 		volumePlugins:         volumePlugins,
 	}
-	s.actorWorkflow = NewActorWorkflow(persistence, workerCache, dialer, actorTemplateLister, workerPoolLister, sandboxConfigLister, storageClassLister, instruments, egressGatewayAddress, s)
+	s.actorWorkflow = NewActorWorkflow(persistence, workerCache, dialer, actorTemplateLister, workerPoolLister, sandboxConfigLister, storageClassLister, instruments, egressGatewayAddress, s, actorWorkflowDeadline)
 	return s
 }
 

cmd/ateapi/internal/controlapi/workflow.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/scheduling"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
@@ -78,9 +79,12 @@ type ActorWorkflow struct {
 	instruments          *Instruments
 	egressGatewayAddress string
 	pluginRegistry       VolumePluginRegistry
+	// workflowDeadline is the maximum duration of a single actor workflow.
+	workflowDeadline time.Duration
 }
 
-// NewActorWorkflow creates a new ActorWorkflow. instruments may be nil.
+// NewActorWorkflow creates a new ActorWorkflow. workflowDeadline bounds how
+// long a single Resume/Suspend can run end-to-end; instruments may be nil.
 func NewActorWorkflow(
 	store store.Interface,
 	workerCache *workercache.Cache,
@@ -92,6 +96,7 @@ func NewActorWorkflow(
 	instruments *Instruments,
 	egressGatewayAddress string,
 	pluginRegistry VolumePluginRegistry,
+	workflowDeadline time.Duration,
 ) *ActorWorkflow {
 	return &ActorWorkflow{
 		store:                store,
@@ -105,19 +110,23 @@ func NewActorWorkflow(
 		instruments:          instruments,
 		egressGatewayAddress: egressGatewayAddress,
 		pluginRegistry:       pluginRegistry,
+		workflowDeadline:     workflowDeadline,
 	}
 }
 
 func (w *ActorWorkflow) acquireActorLock(ctx context.Context, actorRef resources.ActorRef) (context.Context, *store.Lock, error) {
 	lockKey := "lock:actor:" + actorRef.Atespace + ":" + actorRef.Name
+	workflowCtx, cancel := context.WithTimeout(ctx, w.workflowDeadline)
 
-	lock, err := w.store.AcquireLock(ctx, lockKey)
+	lock, err := w.store.AcquireLock(workflowCtx, lockKey)
 	if err != nil {
+		cancel()
 		if errors.Is(err, store.ErrLockConflict) {
 			return nil, nil, status.Error(grpcCodes.Aborted, "another operation is in progress for this actor")
 		}
 		return nil, nil, fmt.Errorf("while acquiring lock: %w", err)
 	}
 
+	context.AfterFunc(lock.Context(), cancel)
 	return lock.Context(), lock, nil
 }

cmd/ateapi/internal/controlapi/workflow_lock_test.go

@@ -0,0 +1,49 @@
+// 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 controlapi
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store/ateredis"
+	"github.com/agent-substrate/substrate/internal/resources"
+	"github.com/alicebob/miniredis/v2"
+	"github.com/redis/go-redis/v9"
+)
+
+func TestAcquireActorLockWorkflowDeadline(t *testing.T) {
+	mr := miniredis.RunT(t)
+	rdb := redis.NewClusterClient(&redis.ClusterOptions{Addrs: []string{mr.Addr()}})
+	t.Cleanup(func() { _ = rdb.Close() })
+	w := &ActorWorkflow{store: ateredis.NewPersistence(rdb), workflowDeadline: 20 * time.Millisecond}
+
+	ctx, lock, err := w.acquireActorLock(context.Background(), resources.ActorRef{Atespace: "space", Name: "actor"})
+	if err != nil {
+		t.Fatalf("acquireActorLock: %v", err)
+	}
+	t.Cleanup(lock.Close)
+
+	select {
+	case <-ctx.Done():
+		if !errors.Is(ctx.Err(), context.DeadlineExceeded) {
+			t.Fatalf("context error = %v, want DeadlineExceeded", ctx.Err())
+		}
+	case <-time.After(time.Second):
+		t.Fatal("workflow context did not reach its deadline")
+	}
+}

cmd/ateapi/internal/controlapi/workflow_suspend_test.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 	"testing"
+	"time"
 
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store/ateredis"
@@ -679,7 +680,7 @@ func TestSuspendActor_PausedWithoutLocalSnapshotCrashes(t *testing.T) {
 	}); err != nil {
 		t.Fatalf("add template to indexer: %v", err)
 	}
-	w := NewActorWorkflow(st, nil, nil, listersv1alpha1.NewActorTemplateLister(indexer), nil, nil, nil, nil, "", nil)
+	w := NewActorWorkflow(st, nil, nil, listersv1alpha1.NewActorTemplateLister(indexer), nil, nil, nil, nil, "", nil, time.Minute)
 
 	seedWorkflowActor(t, ctx, st, resources.ActorRef{Atespace: "team-a", Name: "id1"}, "ns", "tmpl1", ateapipb.Actor_STATUS_PAUSED)
 

cmd/ateapi/internal/controlapi/workflow_testutil_test.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"slices"
 	"testing"
+	"time"
 
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
 	"github.com/agent-substrate/substrate/internal/resources"
@@ -42,7 +43,7 @@ func newTestActorWorkflow(t *testing.T, st store.Interface, tmplNamespace, tmplN
 	}); err != nil {
 		t.Fatalf("add template to indexer: %v", err)
 	}
-	return NewActorWorkflow(st, nil, nil, listersv1alpha1.NewActorTemplateLister(indexer), nil, nil, nil, nil, "", nil)
+	return NewActorWorkflow(st, nil, nil, listersv1alpha1.NewActorTemplateLister(indexer), nil, nil, nil, nil, "", nil, time.Minute)
 }
 
 // seedWorkflowActor stores an actor with the given status, bound to the given

cmd/ateapi/main.go

@@ -32,7 +32,6 @@ import (
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/actoridentity"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/controlapi"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/debugapi"
-	"github.com/agent-substrate/substrate/cmd/ateapi/internal/k8sjwt"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store/ateredis"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/workercache"
 	"github.com/agent-substrate/substrate/internal/ateapiauth"
@@ -88,6 +87,8 @@ var (
 	drainDelay   = pflag.Duration("drain-delay", 13*time.Second, "How long to keep accepting new work after SIGTERM, before starting the gRPC drain.")
 	drainTimeout = pflag.Duration("drain-timeout", 15*time.Second, "Deadline for the graceful gRPC drain on shutdown. In-flight RPCs still running past it are forcefully cancelled.")
 
+	actorWorkflowDeadline = pflag.Duration("actor-workflow-deadline", 5*time.Minute, "Maximum wall-clock duration of a single Resume/Suspend workflow; raise it for slow image registries.")
+
 	showVersion     = pflag.Bool("version", false, "Print version and exit.")
 	logLevelFlag    = pflag.String("log-level", "info", "Minimum log level: debug, info, warn, or error.")
 	clientJWTCAFile = pflag.String("client-jwt-ca-cert", ateapiauth.DefaultServiceAccountCAFile, "CA cert file used to verify TLS when fetching the OIDC discovery document and JWKS for JWT authentication. Defaults to the in-cluster service account CA.")
@@ -196,7 +197,7 @@ func main() {
 
 	volPlugins := make(map[string]volume.VolumePluginControlPlane)
 	ateletDialer := controlapi.NewAteletDialer(workerPodInformer.GetIndexer(), ateletPodInformer.GetIndexer(), *ateletClientCredBundle, *podIdentityCACerts, *ateletInsecure)
-	sm := controlapi.NewService(redisPersistence, workerCache, actorTemplateLister, workerPoolLister, sandboxConfigLister, csiDriverConfigLister, storageClassLister, ateletDialer, instruments, *egressGatewayAddress, volPlugins)
+	sm := controlapi.NewService(redisPersistence, workerCache, actorTemplateLister, workerPoolLister, sandboxConfigLister, csiDriverConfigLister, storageClassLister, ateletDialer, instruments, *egressGatewayAddress, *actorWorkflowDeadline, volPlugins)
 
 	jwtIssuerDiscoveryClient := buildK8sServiceAccountIssuerDiscoveryClient(ctx, *clientJWTCAFile, *clientJWTIssuer)
 
@@ -328,6 +329,7 @@ func logFlagValues(ctx context.Context) {
 		slog.Bool("atelet-insecure", *ateletInsecure),
 		slog.Duration("drain-delay", *drainDelay),
 		slog.Duration("drain-timeout", *drainTimeout),
+		slog.Duration("actor-workflow-deadline", *actorWorkflowDeadline),
 	)
 }