fix: gate scheduler task concurrency

internal/runnerfacade/facade.go

@@ -57,6 +57,7 @@ type Facade struct {
 	Registry     *Registry
 	Capabilities Capabilities
 	Upstream     Upstream
+	OnTerminal   func(taskassignment.Assignment)
 }
 
 func (f *Facade) Handler() (string, http.Handler) {
@@ -98,7 +99,11 @@ func (f *Facade) UpdateTask(ctx context.Context, request *connect.Request[runner
 	if request.Msg.GetState().GetId() != assignment.Task.GetId() {
 		return nil, connect.NewError(connect.CodePermissionDenied, errors.New("task update does not match assignment"))
 	}
-	return f.Upstream.UpdateTask(ctx, connect.NewRequest(request.Msg))
+	response, err := f.Upstream.UpdateTask(ctx, connect.NewRequest(request.Msg))
+	if err == nil && request.Msg.GetState().GetResult() != runnerv1.Result_RESULT_UNSPECIFIED && f.OnTerminal != nil {
+		f.OnTerminal(assignment)
+	}
+	return response, err
 }
 
 func (f *Facade) UpdateLog(ctx context.Context, request *connect.Request[runnerv1.UpdateLogRequest]) (*connect.Response[runnerv1.UpdateLogResponse], error) {

internal/runnerfacade/facade_test.go

@@ -163,6 +163,29 @@ func TestFacadeForwardsOnlyMatchingTaskAndLogUpdates(t *testing.T) {
 	}
 }
 
+func TestFacadeSignalsTerminalTaskAfterUpstreamAcceptsIt(t *testing.T) {
+	facade, assignment, token := testFacade(t)
+	ctx := WithSPIFFEID(context.Background(), assignment.Identity.SPIFFEID)
+	if _, err := facade.FetchTask(ctx, authenticatedRequest(&runnerv1.FetchTaskRequest{}, assignment.ID, token)); err != nil {
+		t.Fatal(err)
+	}
+	completed := 0
+	facade.OnTerminal = func(got taskassignment.Assignment) {
+		if got.ID != assignment.ID {
+			t.Fatalf("terminal assignment = %s", got.ID)
+		}
+		completed++
+	}
+	if _, err := facade.UpdateTask(ctx, authenticatedRequest(&runnerv1.UpdateTaskRequest{
+		State: &runnerv1.TaskState{Id: 42, Result: runnerv1.Result_RESULT_SUCCESS},
+	}, assignment.ID, token)); err != nil {
+		t.Fatal(err)
+	}
+	if completed != 1 {
+		t.Fatalf("terminal notifications = %d", completed)
+	}
+}
+
 func TestCapabilitiesAreDeterministicAndAssignmentScoped(t *testing.T) {
 	capabilities, err := NewCapabilities([]byte("0123456789abcdef0123456789abcdef"))
 	if err != nil {

internal/taskscheduler/gate.go

@@ -0,0 +1,31 @@
+package taskscheduler
+
+import "context"
+
+// SingleFlightGate keeps at most one fetched task in flight. Release is
+// idempotent so repeated terminal updates cannot increase capacity.
+type SingleFlightGate struct {
+	token chan struct{}
+}
+
+func NewSingleFlightGate() *SingleFlightGate {
+	gate := &SingleFlightGate{token: make(chan struct{}, 1)}
+	gate.token <- struct{}{}
+	return gate
+}
+
+func (g *SingleFlightGate) Acquire(ctx context.Context) error {
+	select {
+	case <-g.token:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+func (g *SingleFlightGate) Release() {
+	select {
+	case g.token <- struct{}{}:
+	default:
+	}
+}

internal/taskscheduler/gate_test.go

@@ -0,0 +1,29 @@
+package taskscheduler
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+func TestSingleFlightGateBlocksUntilTerminalRelease(t *testing.T) {
+	gate := NewSingleFlightGate()
+	if err := gate.Acquire(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	blocked, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	defer cancel()
+	if err := gate.Acquire(blocked); err == nil {
+		t.Fatal("second task acquired capacity before release")
+	}
+	gate.Release()
+	gate.Release()
+	if err := gate.Acquire(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	blockedAgain, cancelAgain := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	defer cancelAgain()
+	if err := gate.Acquire(blockedAgain); err == nil {
+		t.Fatal("duplicate terminal release increased capacity")
+	}
+}

internal/taskscheduler/poller.go

@@ -26,6 +26,7 @@ type PollerConfig struct {
 type Poller struct {
 	Client    PollClient
 	Scheduler *Scheduler
+	Gate      *SingleFlightGate
 	Config    PollerConfig
 	OnError   func(error)
 }
@@ -50,7 +51,14 @@ func (p Poller) Run(ctx context.Context) error {
 		errorBackoff = 5 * time.Second
 	}
 	var tasksVersion int64
+	haveLease := false
 	for {
+		if p.Gate != nil && !haveLease {
+			if err := p.Gate.Acquire(ctx); err != nil {
+				return nil
+			}
+			haveLease = true
+		}
 		response, err := p.Client.FetchTask(ctx, tasksVersion)
 		if err != nil {
 			if ctx.Err() != nil {
@@ -72,6 +80,10 @@ func (p Poller) Run(ctx context.Context) error {
 		tasksVersion = response.GetTasksVersion()
 		task := response.GetTask()
 		if task == nil {
+			if p.Gate != nil {
+				p.Gate.Release()
+				haveLease = false
+			}
 			if !wait(ctx, emptyBackoff) {
 				return nil
 			}
@@ -79,6 +91,7 @@ func (p Poller) Run(ctx context.Context) error {
 		}
 		for {
 			if err := p.Scheduler.Run(ctx, task); err == nil {
+				haveLease = false
 				break
 			} else {
 				if ctx.Err() != nil {