fix: ack an ECS stop directive only after the reap reports STOPPED (#1085)

The agent acked a stop in the same loop iteration that launched the reap, so
the directive was deleted before the reap had done anything. When the reap
could not run — no container runtime, or a failing List — nothing re-posted it,
and the task held its ENI, its capacity reservation and its target-group
registrations forever. The in-memory `stopping` map made that permanent: it was
set at dispatch and never cleared, so the agent would not re-dispatch even if
the directive did come back.

A stop ledger now separates a reap in flight from a reap that reported STOPPED.
Only the second is acked, so a failed reap leaves the directive in the inbox to
be redelivered and retried, while the assign loop keeps reading both halves and
so still refuses to start a task that is stopping or stopped.

cmd/ecs-agent/registration_test.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"errors"
+	"slices"
 	"sync"
 	"testing"
 	"time"
@@ -24,6 +25,7 @@ type fakeCP struct {
 	pollStopAcks [][]string
 	pollReplies  [][]bus.Assign
 	stopReplies  [][]bus.StopDirective
+	stopPending  []bus.StopDirective
 	pollCalls    int
 
 	registerErr bool
@@ -72,6 +74,12 @@ func (f *fakeCP) PollAssignments(_, _ string, ackAssigns, ackStops []string) ([]
 	if f.pollCalls < len(f.stopReplies) {
 		stops = f.stopReplies[f.pollCalls]
 	}
+	// stopPending is the durable inbox the gateway keeps: an acked directive is
+	// deleted, and everything still in it is redelivered on every poll.
+	f.stopPending = slices.DeleteFunc(f.stopPending, func(sd bus.StopDirective) bool {
+		return slices.Contains(ackStops, sd.TaskID)
+	})
+	stops = slices.Concat(stops, f.stopPending)
 	f.pollCalls++
 	return out, stops, nil
 }

cmd/ecs-agent/runtime/runtime_fake.go

@@ -34,11 +34,14 @@ type FakePuller struct {
 	StopTimeouts []time.Duration
 	StopErr      error
 
-	// List bookkeeping: Containers is replayed by List; ListErr forces a failure;
-	// Listed records that List was called.
+	// List bookkeeping: Containers is replayed by List; ListErr forces a failure,
+	// limited to the first ListErrFor calls when that is non-zero, so a test can
+	// model a runtime that recovers. Listed and ListCalls record the calls.
 	Containers []Container
 	ListErr    error
+	ListErrFor int
 	Listed     bool
+	ListCalls  int
 }
 
 var (
@@ -130,13 +133,36 @@ func (f *FakePuller) Waits() []string {
 	return append([]string(nil), f.Waited...)
 }
 
+// RunCalls returns a copy of the specs passed to Run, for assertions made while
+// the agent's goroutines may still be live.
+func (f *FakePuller) RunCalls() []RunSpec {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return append([]RunSpec(nil), f.Runs...)
+}
+
+// Reaped returns a copy of the container IDs passed to Stop.
+func (f *FakePuller) Reaped() []string {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return append([]string(nil), f.Stopped...)
+}
+
 // List replays the programmed container set (or error).
 func (f *FakePuller) List(_ context.Context) ([]Container, error) {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 	f.Listed = true
-	if f.ListErr != nil {
+	f.ListCalls++
+	if f.ListErr != nil && (f.ListErrFor == 0 || f.ListCalls <= f.ListErrFor) {
 		return nil, f.ListErr
 	}
 	return f.Containers, nil
 }
+
+// Lists reports how many times List was called.
+func (f *FakePuller) Lists() int {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.ListCalls
+}

cmd/ecs-agent/taskrun.go

@@ -6,6 +6,7 @@ import (
 	"log/slog"
 	"maps"
 	"strconv"
+	"sync"
 	"time"
 
 	ctrruntime "github.com/mulgadc/spinifex/cmd/ecs-agent/runtime"
@@ -28,7 +29,7 @@ func (a *Agent) pollAssignments(ctx context.Context, dispatched map[string]bool)
 	ticker := time.NewTicker(a.cfg.PollInterval)
 	defer ticker.Stop()
 
-	stopping := map[string]bool{}
+	reaps := newStopLedger()
 	var ackAssigns, ackStops []string
 	for {
 		select {
@@ -45,17 +46,18 @@ func (a *Agent) pollAssignments(ctx context.Context, dispatched map[string]bool)
 			ackStops = ackStops[:0]
 			for i := range stops {
 				sd := stops[i]
-				if !stopping[sd.TaskID] {
-					stopping[sd.TaskID] = true
+				switch {
+				case reaps.done(sd.TaskID):
+					ackStops = append(ackStops, sd.TaskID)
+				case reaps.begin(sd.TaskID):
 					slog.Info("ecs-agent: task stop requested", "task", sd.TaskID, "reason", sd.Reason)
-					go a.stopTask(ctx, sd)
+					go func() { reaps.end(sd.TaskID, a.stopTask(ctx, sd)) }()
 				}
-				ackStops = append(ackStops, sd.TaskID)
 			}
 			ackAssigns = ackAssigns[:0]
 			for i := range assigns {
 				as := assigns[i]
-				if !dispatched[as.TaskID] && !stopping[as.TaskID] {
+				if !dispatched[as.TaskID] && !reaps.seen(as.TaskID) {
 					dispatched[as.TaskID] = true
 					slog.Info("ecs-agent: task assigned", "task", as.TaskID, "containers", len(as.Containers))
 					go a.runTask(ctx, &as)
@@ -66,20 +68,74 @@ func (a *Agent) pollAssignments(ctx context.Context, dispatched map[string]bool)
 	}
 }
 
+// stopLedger tracks reaps in flight and reaps that reported STOPPED. The two
+// are separate because only the second may be acked: acking the first deletes a
+// directive that was never performed, and nothing re-posts it.
+type stopLedger struct {
+	mu       sync.Mutex
+	inflight map[string]bool
+	stopped  map[string]bool
+}
+
+func newStopLedger() *stopLedger {
+	return &stopLedger{inflight: map[string]bool{}, stopped: map[string]bool{}}
+}
+
+// begin claims a task for a reap, reporting false when one is already running.
+func (l *stopLedger) begin(taskID string) bool {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	if l.inflight[taskID] {
+		return false
+	}
+	l.inflight[taskID] = true
+	return true
+}
+
+// end releases the claim, recording the task as stopped only if the reap got
+// that far. A failed reap leaves nothing behind, so the next poll redelivers
+// the directive and dispatches it again.
+func (l *stopLedger) end(taskID string, reported bool) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	delete(l.inflight, taskID)
+	if reported {
+		l.stopped[taskID] = true
+	}
+}
+
+// done reports whether the task's reap reported STOPPED, which is the only
+// condition under which its directive may be acked.
+func (l *stopLedger) done(taskID string) bool {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return l.stopped[taskID]
+}
+
+// seen reports whether a stop for the task is in flight or complete, so an
+// assignment arriving alongside its own stop is never started.
+func (l *stopLedger) seen(taskID string) bool {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return l.inflight[taskID] || l.stopped[taskID]
+}
+
 // stopTask reaps a task's containers on a control-plane stop directive: it
 // lists the runtime's containers, gracefully stops (SIGTERM, wait up to the
 // container's stopTimeout, then force kill + delete) the ones labeled with
 // this task, then reports the task STOPPED with the directive's reason.
 // Idempotent — a task with no live containers still reports STOPPED so the
-// scheduler releases its capacity.
-func (a *Agent) stopTask(ctx context.Context, sd bus.StopDirective) {
+// scheduler releases its capacity. Reports whether it got as far as saying
+// STOPPED: a caller may only ack the directive when it did.
+func (a *Agent) stopTask(ctx context.Context, sd bus.StopDirective) bool {
 	if a.runner == nil {
-		return
+		slog.Warn("ecs-agent: stop deferred, no container runtime", "task", sd.TaskID)
+		return false
 	}
 	containers, err := a.runner.List(ctx)
 	if err != nil {
 		slog.Warn("ecs-agent: stop list failed", "task", sd.TaskID, "err", err)
-		return
+		return false
 	}
 	statuses := make([]bus.ContainerStatus, 0)
 	for _, c := range containers {
@@ -99,6 +155,7 @@ func (a *Agent) stopTask(ctx context.Context, sd bus.StopDirective) {
 	}
 	a.reportTaskState(&bus.Assign{TaskID: sd.TaskID}, bus.TaskStatusStopped, reason, statuses)
 	slog.Info("ecs-agent: task reaped", "task", sd.TaskID, "containers", len(statuses))
+	return true
 }
 
 // runTask pulls each container image, starts the containers, and reports task

cmd/ecs-agent/taskrun_test.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"errors"
+	"slices"
 	"testing"
 	"time"
 
@@ -532,7 +533,7 @@ func TestPollAssignments_StopReapsAndSuppressesAssign(t *testing.T) {
 	as := testAssign()
 	cp := &fakeCP{
 		pollReplies: [][]bus.Assign{{*as}, {*as}, nil},
-		stopReplies: [][]bus.StopDirective{{{TaskID: as.TaskID, Reason: "bye"}}},
+		stopPending: []bus.StopDirective{{TaskID: as.TaskID, Reason: "bye"}},
 	}
 	rt := &ctrruntime.FakePuller{Containers: []ctrruntime.Container{
 		{ID: "t-001-web", Running: true, Labels: map[string]string{"mulga.ecs.taskID": "t-001"}},
@@ -576,6 +577,122 @@ func TestPollAssignments_StopReapsAndSuppressesAssign(t *testing.T) {
 	}
 }
 
+// countStopAcks reports how many polls carried taskID in their stop-ack list.
+func countStopAcks(acks [][]string, taskID string) int {
+	n := 0
+	for _, ack := range acks {
+		if slices.Contains(ack, taskID) {
+			n++
+		}
+	}
+	return n
+}
+
+// A reap that cannot run is never acked, so the directive stays in the gateway's
+// inbox and is redelivered and retried until the runtime comes back.
+func TestPollAssignments_UnperformableStopIsNotAckedAndRetried(t *testing.T) {
+	cp := &fakeCP{stopPending: []bus.StopDirective{{TaskID: "t-001", Reason: "bye"}}}
+	rt := &ctrruntime.FakePuller{ListErr: errors.New("containerd is down")}
+	a := newAgent(config{PollInterval: 5 * time.Millisecond}, testIdentity(), cp, rt, rt, nil)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go a.pollAssignments(ctx, map[string]bool{})
+
+	deadline := time.After(time.Second)
+	for rt.Lists() < 3 {
+		select {
+		case <-deadline:
+			t.Fatalf("failed reap was not retried; List calls=%d", rt.Lists())
+		case <-time.After(2 * time.Millisecond):
+		}
+	}
+	cancel()
+
+	if n := countStopAcks(cp.stopAcks(), "t-001"); n != 0 {
+		t.Errorf("a reap that never ran was acked %d times; stopAcks=%v", n, cp.stopAcks())
+	}
+	if stopped := countStatus(cp.taskStates(), "t-001", bus.TaskStatusStopped); stopped != 0 {
+		t.Errorf("a reap that never ran reported STOPPED %d times", stopped)
+	}
+}
+
+// Once the runtime recovers, the redelivered directive is reaped, reported
+// STOPPED and acked exactly once, which deletes it from the inbox.
+func TestPollAssignments_RecoveredStopIsAckedExactlyOnce(t *testing.T) {
+	cp := &fakeCP{stopPending: []bus.StopDirective{{TaskID: "t-001", Reason: "bye"}}}
+	rt := &ctrruntime.FakePuller{
+		ListErr:    errors.New("containerd is down"),
+		ListErrFor: 1,
+		Containers: []ctrruntime.Container{
+			{ID: "t-001-web", Running: true, Labels: map[string]string{"mulga.ecs.taskID": "t-001"}},
+		},
+	}
+	a := newAgent(config{PollInterval: 5 * time.Millisecond}, testIdentity(), cp, rt, rt, nil)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go a.pollAssignments(ctx, map[string]bool{})
+
+	deadline := time.After(2 * time.Second)
+	for countStopAcks(cp.stopAcks(), "t-001") == 0 {
+		select {
+		case <-deadline:
+			t.Fatalf("recovered reap was never acked; states=%+v stopAcks=%v", cp.taskStates(), cp.stopAcks())
+		case <-time.After(2 * time.Millisecond):
+		}
+	}
+	time.Sleep(30 * time.Millisecond)
+	cancel()
+
+	if n := countStopAcks(cp.stopAcks(), "t-001"); n != 1 {
+		t.Errorf("stop acked %d times, want exactly 1; stopAcks=%v", n, cp.stopAcks())
+	}
+	if stopped := countStatus(cp.taskStates(), "t-001", bus.TaskStatusStopped); stopped != 1 {
+		t.Errorf("task reported STOPPED %d times, want exactly 1", stopped)
+	}
+	if reaped := rt.Reaped(); len(reaped) != 1 || reaped[0] != "t-001-web" {
+		t.Errorf("want t-001-web reaped once, got %+v", reaped)
+	}
+}
+
+// A task whose reap completed is never started again by an assign the gateway
+// is still redelivering.
+func TestPollAssignments_StoppedTaskIsNotRerunByRedeliveredAssign(t *testing.T) {
+	as := testAssign()
+	cp := &fakeCP{
+		pollReplies: [][]bus.Assign{{*as}, {*as}, {*as}, {*as}, {*as}, {*as}},
+		stopPending: []bus.StopDirective{{TaskID: as.TaskID, Reason: "bye"}},
+	}
+	rt := &ctrruntime.FakePuller{Containers: []ctrruntime.Container{
+		{ID: "t-001-web", Running: true, Labels: map[string]string{"mulga.ecs.taskID": "t-001"}},
+	}}
+	a := newAgent(config{PollInterval: 5 * time.Millisecond}, testIdentity(), cp, rt, rt, nil)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go a.pollAssignments(ctx, map[string]bool{})
+
+	deadline := time.After(time.Second)
+	for countStopAcks(cp.stopAcks(), as.TaskID) == 0 {
+		select {
+		case <-deadline:
+			t.Fatalf("stop never completed; states=%+v", cp.taskStates())
+		case <-time.After(2 * time.Millisecond):
+		}
+	}
+	// Let the remaining scripted assigns arrive after the reap has finished.
+	time.Sleep(40 * time.Millisecond)
+	cancel()
+
+	if runs := rt.RunCalls(); len(runs) != 0 {
+		t.Errorf("stopped task was started by a redelivered assign: %+v", runs)
+	}
+	if running := countStatus(cp.taskStates(), as.TaskID, bus.TaskStatusRunning); running != 0 {
+		t.Errorf("stopped task reported RUNNING %d times", running)
+	}
+}
+
 // stopTask with no runtime is a no-op: nothing to reap, no state reported.
 func TestStopTask_NilRunnerNoOp(t *testing.T) {
 	cp := &fakeCP{}