feat: ECS agent reboot reconciler — re-adopt running containers on restart (#435)

On startup the agent had no reconcile pass: Run registered, heartbeat, and
immediately polled assignments with an in-memory dedup map. After a restart the
gateway re-delivers the instance's assignments and the agent re-ran them,
orphaning containers still running from before and double-starting tasks.

Add a reconcile() pass that runs once before polling: it lists the runtime's
containers, and for each task still running in this cluster re-registers its IAM
credentials, re-attaches the exit-wait, and refreshes RUNNING on the gateway.
The adopted task set seeds pollAssignments so re-delivered assignments are acked
but not re-run.

- runtime: add a label-agnostic List(ctx) []Container to the Runner seam, with a
  containerd implementation and fake support.
- taskLabels: stamp credID / taskRoleArn / eniMac (when set) so a container is
  self-describing enough to re-register creds and tear down its netns on adoption.
- pollAssignments takes the seed dedup map.

Best-effort: a nil runner or a List error degrades to today's poll+run behaviour.

cmd/ecs-agent/agent.go

@@ -130,7 +130,10 @@ func (a *Agent) Run(ctx context.Context) error {
 	}
 
 	go a.hb.Run(ctx)
-	go a.pollAssignments(ctx)
+
+	// Re-adopt containers still running from before this restart, then poll with
+	// their tasks pre-seeded so re-delivered assignments are acked, not re-run.
+	go a.pollAssignments(ctx, a.reconcile(ctx))
 
 	<-ctx.Done()
 	return a.Stop()

cmd/ecs-agent/reconcile.go

@@ -0,0 +1,72 @@
+package main
+
+import (
+	"context"
+	"log/slog"
+
+	ctrruntime "github.com/mulgadc/spinifex/cmd/ecs-agent/runtime"
+	"github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus"
+)
+
+// reconcile re-adopts containers that survived an agent restart. It lists the
+// runtime's containers and, for each task still running in this cluster,
+// re-registers its IAM credentials, re-attaches the exit-wait, and refreshes its
+// RUNNING state on the gateway. The returned set seeds pollAssignments so the
+// gateway's re-delivered assignments for these tasks are acked but not re-run.
+//
+// It is best-effort: a nil runner or a List error returns an empty set and the
+// agent falls back to its pre-reconcile behaviour (poll + run).
+func (a *Agent) reconcile(ctx context.Context) map[string]bool {
+	adopted := map[string]bool{}
+	if a.runner == nil {
+		return adopted
+	}
+	containers, err := a.runner.List(ctx)
+	if err != nil {
+		slog.Warn("ecs-agent: reconcile list failed", "err", err)
+		return adopted
+	}
+
+	// Group this cluster's running containers back into their tasks. The Assign is
+	// reconstructed from labels — enough to re-register creds and tear down netns.
+	type task struct {
+		as         *bus.Assign
+		containers []ctrruntime.Container
+	}
+	tasks := map[string]*task{}
+	for _, c := range containers {
+		if !c.Running || c.Labels[labelClusterName] != a.id.ClusterName {
+			continue
+		}
+		taskID := c.Labels[labelTaskID]
+		if taskID == "" {
+			continue
+		}
+		t := tasks[taskID]
+		if t == nil {
+			t = &task{as: &bus.Assign{
+				TaskID:        taskID,
+				ClusterName:   a.id.ClusterName,
+				CredID:        c.Labels[labelCredID],
+				TaskRoleARN:   c.Labels[labelTaskRoleARN],
+				ENIMacAddress: c.Labels[labelENIMac],
+			}}
+			tasks[taskID] = t
+		}
+		t.containers = append(t.containers, c)
+	}
+
+	for taskID, t := range tasks {
+		a.registerTaskCreds(t.as)
+		statuses := make([]bus.ContainerStatus, 0, len(t.containers))
+		for _, c := range t.containers {
+			name := c.Labels[labelContainerName]
+			go a.waitContainer(ctx, t.as, name, c.ID)
+			statuses = append(statuses, bus.ContainerStatus{Name: name, Status: bus.TaskStatusRunning, ContainerID: c.ID})
+		}
+		a.reportTaskState(t.as, bus.TaskStatusRunning, "", statuses)
+		adopted[taskID] = true
+		slog.Info("ecs-agent: re-adopted running task", "task", taskID, "containers", len(t.containers))
+	}
+	return adopted
+}

cmd/ecs-agent/reconcile_test.go

@@ -0,0 +1,181 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"slices"
+	"testing"
+	"time"
+
+	ctrruntime "github.com/mulgadc/spinifex/cmd/ecs-agent/runtime"
+	"github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus"
+)
+
+const testTaskRole = "arn:aws:iam::111122223333:role/task"
+
+// adoptedContainer builds a runtime.Container as the reconciler would discover it,
+// stamping only the mulga.ecs.* labels that are non-empty.
+func adoptedContainer(taskID, name, cluster, credID, role, mac string, running bool) ctrruntime.Container {
+	labels := map[string]string{
+		labelTaskID:        taskID,
+		labelContainerName: name,
+		labelClusterName:   cluster,
+	}
+	if credID != "" {
+		labels[labelCredID] = credID
+	}
+	if role != "" {
+		labels[labelTaskRoleARN] = role
+	}
+	if mac != "" {
+		labels[labelENIMac] = mac
+	}
+	return ctrruntime.Container{ID: containerID(taskID, name), Labels: labels, Running: running}
+}
+
+func waitedFor(t *testing.T, rt *ctrruntime.FakePuller, id string) {
+	t.Helper()
+	deadline := time.After(time.Second)
+	for !slices.Contains(rt.Waits(), id) {
+		select {
+		case <-deadline:
+			t.Fatalf("exit-wait never re-attached for %s; waited=%v", id, rt.Waits())
+		case <-time.After(2 * time.Millisecond):
+		}
+	}
+}
+
+// A running, correctly-labeled container is re-adopted: its credentials are
+// re-registered, its RUNNING state refreshed, and its exit-wait re-attached.
+func TestReconcile_AdoptsRunningLabeledContainer(t *testing.T) {
+	cp := &fakeCP{}
+	rt := &ctrruntime.FakePuller{
+		WaitErr: errors.New("blocked"),
+		Containers: []ctrruntime.Container{
+			adoptedContainer("t-001", "web", "default", "t-001", testTaskRole, "52:54:00:de:ad:01", true),
+		},
+	}
+	a := newAgent(config{}, testIdentity(), cp, rt, rt, nil)
+	a.cred = newCredEndpoint(nil, "us-east-1", "https://gw", "", "127.0.0.1", 0, nil)
+
+	adopted := a.reconcile(context.Background())
+
+	if !adopted["t-001"] {
+		t.Fatalf("t-001 not adopted; got %v", adopted)
+	}
+	if got := a.cred.roles["t-001"]; got != testTaskRole {
+		t.Errorf("cred not re-registered: roles[t-001] = %q, want %q", got, testTaskRole)
+	}
+	if n := countStatus(cp.taskStates(), "t-001", bus.TaskStatusRunning); n != 1 {
+		t.Errorf("RUNNING reports for t-001 = %d, want 1", n)
+	}
+	waitedFor(t, rt, containerID("t-001", "web"))
+}
+
+// pollAssignments seeded with an adopted task acks its re-delivered assignment
+// but does not re-run it, while a genuinely new task still dispatches.
+func TestPollAssignments_SeededTaskNotRerun(t *testing.T) {
+	mk := func(id string) bus.Assign {
+		return bus.Assign{
+			AccountID: "123456789012", ClusterName: "default", InstanceID: "i-abc123", TaskID: id,
+			Containers: []bus.AssignContainer{{Name: "web", Image: "registry/web:1", Command: []string{"/bin/true"}}},
+		}
+	}
+	adopted, fresh := mk("t-001"), mk("t-002")
+	cp := &fakeCP{pollReplies: [][]bus.Assign{
+		{adopted, fresh},
+		{adopted, fresh},
+		nil,
+	}}
+	rt := &ctrruntime.FakePuller{WaitErr: errors.New("blocked")}
+	a := newAgent(config{PollInterval: 5 * time.Millisecond}, testIdentity(), cp, rt, rt, nil)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go a.pollAssignments(ctx, map[string]bool{"t-001": true})
+
+	deadline := time.After(time.Second)
+	for countStatus(cp.taskStates(), "t-002", bus.TaskStatusRunning) < 1 {
+		select {
+		case <-deadline:
+			t.Fatal("fresh task t-002 was never dispatched")
+		case <-time.After(2 * time.Millisecond):
+		}
+	}
+	time.Sleep(20 * time.Millisecond)
+	cancel()
+
+	if n := countStatus(cp.taskStates(), "t-001", bus.TaskStatusRunning); n != 0 {
+		t.Errorf("adopted task t-001 was re-run %d times, want 0", n)
+	}
+	if n := countStatus(cp.taskStates(), "t-002", bus.TaskStatusRunning); n != 1 {
+		t.Errorf("fresh task t-002 dispatched %d times, want 1", n)
+	}
+	var t1Acked bool
+	for _, ack := range cp.acks() {
+		for _, id := range ack {
+			if id == "t-001" {
+				t1Acked = true
+			}
+		}
+	}
+	if !t1Acked {
+		t.Errorf("adopted task t-001 was never acked; acks=%v", cp.acks())
+	}
+}
+
+// Containers for another cluster or with a stopped task are not adopted.
+func TestReconcile_FiltersOtherClusterAndStopped(t *testing.T) {
+	rt := &ctrruntime.FakePuller{
+		WaitErr: errors.New("blocked"),
+		Containers: []ctrruntime.Container{
+			adoptedContainer("t-other", "web", "other", "", "", "", true),
+			adoptedContainer("t-stop", "web", "default", "", "", "", false),
+			adoptedContainer("t-ok", "web", "default", "", "", "", true),
+		},
+	}
+	a := newAgent(config{}, testIdentity(), &fakeCP{}, rt, rt, nil)
+
+	adopted := a.reconcile(context.Background())
+
+	if len(adopted) != 1 || !adopted["t-ok"] {
+		t.Fatalf("adopted = %v, want only t-ok", adopted)
+	}
+}
+
+// A nil runner or a List error degrades to an empty set without panicking.
+func TestReconcile_DegradesGracefully(t *testing.T) {
+	noRunner := newAgent(config{}, testIdentity(), &fakeCP{}, nil, nil, nil)
+	if got := noRunner.reconcile(context.Background()); len(got) != 0 {
+		t.Errorf("nil runner: adopted = %v, want empty", got)
+	}
+
+	rt := &ctrruntime.FakePuller{ListErr: errors.New("boom")}
+	listErr := newAgent(config{}, testIdentity(), &fakeCP{}, rt, rt, nil)
+	if got := listErr.reconcile(context.Background()); len(got) != 0 {
+		t.Errorf("list error: adopted = %v, want empty", got)
+	}
+}
+
+// taskLabels emits the optional cred/role/MAC labels only when set.
+func TestTaskLabels_OmitsEmptyOptionalLabels(t *testing.T) {
+	full := taskLabels(&bus.Assign{
+		TaskID: "t1", ClusterName: "c", CredID: "x", TaskRoleARN: testTaskRole, ENIMacAddress: "mac",
+	}, "web")
+	for k, want := range map[string]string{
+		labelCredID: "x", labelTaskRoleARN: testTaskRole, labelENIMac: "mac",
+	} {
+		if full[k] != want {
+			t.Errorf("full labels[%q] = %q, want %q", k, full[k], want)
+		}
+	}
+
+	bare := taskLabels(&bus.Assign{TaskID: "t1", ClusterName: "c"}, "web")
+	if len(bare) != 3 {
+		t.Errorf("bare labels = %v, want only taskID/containerName/clusterName", bare)
+	}
+	for _, k := range []string{labelCredID, labelTaskRoleARN, labelENIMac} {
+		if _, ok := bare[k]; ok {
+			t.Errorf("bare labels unexpectedly carry %q", k)
+		}
+	}
+}

cmd/ecs-agent/runtime/containerd_run.go

@@ -94,6 +94,30 @@ func (p *containerdPuller) Wait(ctx context.Context, containerID string) (RunSta
 	return RunStatus{ExitCode: int(code)}, err
 }
 
+// List enumerates every container in the ecs namespace with its labels and
+// whether its task is running, so the agent can re-adopt live containers after a
+// restart. Per-container label/task errors are treated as "not running" rather
+// than failing the whole list.
+func (p *containerdPuller) List(ctx context.Context) ([]Container, error) {
+	ctx = namespaces.WithNamespace(ctx, ecsNamespace)
+	containers, err := p.client.Containers(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("list containers: %w", err)
+	}
+	out := make([]Container, 0, len(containers))
+	for _, c := range containers {
+		labels, _ := c.Labels(ctx)
+		running := false
+		if task, terr := c.Task(ctx, nil); terr == nil {
+			if status, serr := task.Status(ctx); serr == nil {
+				running = status.Status == containerd.Running
+			}
+		}
+		out = append(out, Container{ID: c.ID(), Labels: labels, Running: running})
+	}
+	return out, nil
+}
+
 // Remove kills and deletes the container's task, then the container + snapshot.
 func (p *containerdPuller) Remove(ctx context.Context, containerID string) error {
 	ctx = namespaces.WithNamespace(ctx, ecsNamespace)

cmd/ecs-agent/runtime/runtime.go

@@ -55,13 +55,24 @@ type RunStatus struct {
 	ExitCode int
 }
 
+// Container is a container the runtime already knows about, discovered by List.
+// Labels carry the mulga.ecs.* task identity so the agent can re-adopt running
+// containers after a restart; Running reports whether its task is live.
+type Container struct {
+	ID      string
+	Labels  map[string]string
+	Running bool
+}
+
 // Runner creates and starts containers from already-pulled images. id is a
 // caller-unique container ID; Wait blocks until the container exits; Remove
-// tears down the container + its task.
+// tears down the container + its task; List enumerates known containers so the
+// reboot reconciler can re-adopt the ones still running.
 type Runner interface {
 	Run(ctx context.Context, id string, spec RunSpec) (containerID string, err error)
 	Wait(ctx context.Context, containerID string) (RunStatus, error)
 	Remove(ctx context.Context, containerID string) error
+	List(ctx context.Context) ([]Container, error)
 }
 
 // Runtime is the full container runtime the ecs-agent drives: pull + run.

cmd/ecs-agent/runtime/runtime_fake.go

@@ -22,8 +22,15 @@ type FakePuller struct {
 	RunErr   error
 	WaitCode int
 	WaitErr  error
+	Waited   []string
 	Removed  []string
 	OnRun    func(string, RunSpec) (string, error)
+
+	// List bookkeeping: Containers is replayed by List; ListErr forces a failure;
+	// Listed records that List was called.
+	Containers []Container
+	ListErr    error
+	Listed     bool
 }
 
 var (
@@ -83,8 +90,11 @@ func (f *FakePuller) Run(_ context.Context, id string, spec RunSpec) (string, er
 	return id, nil
 }
 
-// Wait returns the programmed exit code/error.
-func (f *FakePuller) Wait(_ context.Context, _ string) (RunStatus, error) {
+// Wait records the container ID and returns the programmed exit code/error.
+func (f *FakePuller) Wait(_ context.Context, containerID string) (RunStatus, error) {
+	f.mu.Lock()
+	f.Waited = append(f.Waited, containerID)
+	f.mu.Unlock()
 	return RunStatus{ExitCode: f.WaitCode}, f.WaitErr
 }
 
@@ -95,3 +105,21 @@ func (f *FakePuller) Remove(_ context.Context, containerID string) error {
 	f.Removed = append(f.Removed, containerID)
 	return nil
 }
+
+// Waits returns a copy of the container IDs passed to Wait, for assertions.
+func (f *FakePuller) Waits() []string {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return append([]string(nil), f.Waited...)
+}
+
+// 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 {
+		return nil, f.ListErr
+	}
+	return f.Containers, nil
+}

cmd/ecs-agent/taskrun.go

@@ -15,12 +15,13 @@ import (
 // cancelled. Each assign is dispatched to runTask exactly once; the taskIDs seen
 // on a poll are acked on the next poll so the gateway can drop them — a crash
 // before ack re-delivers (at-least-once), matching ACS. A failed poll is logged
-// and retried on the next tick rather than killing the loop.
-func (a *Agent) pollAssignments(ctx context.Context) {
+// and retried on the next tick rather than killing the loop. dispatched is seeded
+// by the startup reconcile with tasks already adopted from running containers, so
+// their re-delivered assignments are acked but not re-run.
+func (a *Agent) pollAssignments(ctx context.Context, dispatched map[string]bool) {
 	ticker := time.NewTicker(a.cfg.PollInterval)
 	defer ticker.Stop()
 
-	dispatched := map[string]bool{}
 	var ackNext []string
 	for {
 		select {
@@ -199,12 +200,35 @@ func containerID(taskID, name string) string {
 	return fmt.Sprintf("%s-%s", taskID, name)
 }
 
-// taskLabels are the mulga.ecs.* labels stamped on every container so the reboot
-// reconciler can re-associate running containers with their task on restart.
+// mulga.ecs.* label keys stamped on every container. The reboot reconciler reads
+// them back to re-associate running containers with their task on restart; the
+// cred/role/MAC labels make a container self-describing enough to re-register its
+// credentials and tear down its netns without the original assignment.
+const (
+	labelTaskID        = "mulga.ecs.taskID"
+	labelContainerName = "mulga.ecs.containerName"
+	labelClusterName   = "mulga.ecs.clusterName"
+	labelCredID        = "mulga.ecs.credID"
+	labelTaskRoleARN   = "mulga.ecs.taskRoleArn"
+	labelENIMac        = "mulga.ecs.eniMac"
+)
+
+// taskLabels are the mulga.ecs.* labels stamped on a container. The cred/role/MAC
+// labels are omitted when empty so a container carries only what it needs.
 func taskLabels(as *bus.Assign, name string) map[string]string {
-	return map[string]string{
-		"mulga.ecs.taskID":        as.TaskID,
-		"mulga.ecs.containerName": name,
-		"mulga.ecs.clusterName":   as.ClusterName,
+	labels := map[string]string{
+		labelTaskID:        as.TaskID,
+		labelContainerName: name,
+		labelClusterName:   as.ClusterName,
 	}
+	if credID := taskCredID(as); credID != "" {
+		labels[labelCredID] = credID
+	}
+	if as.TaskRoleARN != "" {
+		labels[labelTaskRoleARN] = as.TaskRoleARN
+	}
+	if as.ENIMacAddress != "" {
+		labels[labelENIMac] = as.ENIMacAddress
+	}
+	return labels
 }

cmd/ecs-agent/taskrun_test.go

@@ -136,7 +136,7 @@ func TestPollAssignments_DispatchesOnceAndAcks(t *testing.T) {
 	a := newAgent(config{PollInterval: 5 * time.Millisecond}, testIdentity(), cp, rt, rt, nil)
 
 	ctx, cancel := context.WithCancel(context.Background())
-	go a.pollAssignments(ctx)
+	go a.pollAssignments(ctx, map[string]bool{})
 
 	// runTask reports RUNNING through the (mutex-guarded) control plane. With the
 	// assign returned on two consecutive polls, dedup means exactly one dispatch →