[ACTP] keep executor alive during key sync

Author: embeaken

pkg/privateactionrunner/executor/server.go

@@ -51,6 +51,7 @@ type Server struct {
 
 	ready  atomic.Bool
 	active atomic.Int32
+	busy   atomic.Int32
 
 	lastActivity atomic.Int64
 	clock        clock.Clock
@@ -72,8 +73,17 @@ func (s *Server) touch() {
 	s.lastActivity.Store(s.clock.Now().UnixNano())
 }
 
+func (s *Server) beginActivity() func() {
+	s.touch()
+	s.busy.Add(1)
+	return func() {
+		s.touch()
+		s.busy.Add(-1)
+	}
+}
+
 func (s *Server) idleFor() time.Duration {
-	if s.active.Load() > 0 {
+	if s.busy.Load() > 0 {
 		return 0
 	}
 	return s.clock.Since(time.Unix(0, s.lastActivity.Load()))
@@ -109,7 +119,8 @@ func (s *Server) SyncKeys(ctx context.Context, req *pb.SyncKeysRequest) (*pb.Syn
 			Key:     append([]byte(nil), key.GetKey()...),
 		})
 	}
-	s.touch()
+	finishActivity := s.beginActivity()
+	defer finishActivity()
 	if err := s.keysManager.Seed(seed); err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid signing-key seed: %v", err)
 	}
@@ -143,11 +154,11 @@ func (s *Server) RunAction(req *pb.RunActionRequest, stream pb.Executor_RunActio
 		))
 	}
 
-	s.touch()
+	finishActivity := s.beginActivity()
 	s.active.Add(1)
 	defer func() {
-		s.touch()
 		s.active.Add(-1)
+		finishActivity()
 	}()
 
 	// Raw bytes must stay unmodified for signature verification.

pkg/privateactionrunner/executor/server_test.go

@@ -71,13 +71,27 @@ func (f *fakeExecutor) RunPrepared(ctx context.Context, _ *runners.PreparedWorkf
 }
 
 type fakeKeysManager struct {
-	seed     []taskverifier.SigningKey
-	snapshot []taskverifier.SigningKey
+	seed        []taskverifier.SigningKey
+	snapshot    []taskverifier.SigningKey
+	readyGate   <-chan struct{}
+	waitStarted chan<- struct{}
 }
 
-func (f *fakeKeysManager) Start(context.Context)              {}
-func (f *fakeKeysManager) GetKey(string) types.DecodedKey     { return nil }
-func (f *fakeKeysManager) WaitForReady(context.Context) error { return nil }
+func (f *fakeKeysManager) Start(context.Context)          {}
+func (f *fakeKeysManager) GetKey(string) types.DecodedKey { return nil }
+func (f *fakeKeysManager) WaitForReady(ctx context.Context) error {
+	if f.waitStarted != nil {
+		f.waitStarted <- struct{}{}
+	}
+	if f.readyGate != nil {
+		select {
+		case <-f.readyGate:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+	return nil
+}
 func (f *fakeKeysManager) Seed(keys []taskverifier.SigningKey) error {
 	f.seed = keys
 	return nil
@@ -262,6 +276,34 @@ func TestServeSyncKeysSeedsAndReturnsCurrentSnapshot(t *testing.T) {
 	require.True(t, health.GetReady())
 }
 
+func TestSyncKeysSuppressesIdleExitWhileWaitingForKeys(t *testing.T) {
+	readyGate := make(chan struct{})
+	waitStarted := make(chan struct{}, 1)
+	keysManager := &fakeKeysManager{readyGate: readyGate, waitStarted: waitStarted}
+	srv := NewServer(&fakeExecutor{}, "test-version", keysManager)
+	mockClock := clock.NewMock()
+	srv.clock = mockClock
+	srv.touch()
+
+	done := make(chan error, 1)
+	go func() {
+		_, err := srv.SyncKeys(context.Background(), &pb.SyncKeysRequest{})
+		done <- err
+	}()
+	<-waitStarted
+
+	mockClock.Add(2 * time.Minute)
+	assert.Zero(t, srv.idleFor(), "key synchronization is executor activity")
+	health, err := srv.Health(context.Background(), &pb.HealthRequest{})
+	require.NoError(t, err)
+	assert.Zero(t, health.ActiveActions, "key synchronization is not an action")
+
+	close(readyGate)
+	require.NoError(t, <-done)
+	mockClock.Add(time.Minute)
+	assert.Equal(t, time.Minute, srv.idleFor())
+}
+
 func TestServeRunActionStreamsOutputAndForwardsRawTask(t *testing.T) {
 	fake := &fakeExecutor{
 		prepared: &runners.PreparedWorkflowTask{Task: &types.Task{}},
@@ -595,11 +637,13 @@ func TestIdleTracking(t *testing.T) {
 	mockClock.Add(time.Second)
 	assert.Equal(t, time.Second, srv.idleFor(), "health checks should not count as activity")
 
-	srv.active.Add(1)
+	finishActivity := srv.beginActivity()
 	mockClock.Add(2 * timeout)
-	assert.Zero(t, srv.idleFor(), "an active action should suppress the idle state")
-	srv.active.Add(-1)
-	srv.touch()
+	assert.Zero(t, srv.idleFor(), "activity should suppress the idle state")
+	health, err := srv.Health(context.Background(), &pb.HealthRequest{})
+	require.NoError(t, err)
+	assert.Zero(t, health.ActiveActions, "non-action activity must not affect action accounting")
+	finishActivity()
 
 	mockClock.Add(timeout)
 	assert.Equal(t, timeout, srv.idleFor())