From 4e6a2eb46a0e4006f6c8e039b9fdc6368689d2ed Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Mon, 29 Jun 2026 12:13:25 +1000 Subject: [PATCH] =?UTF-8?q?feat:=20ECS=20agent=20reboot=20reconciler=20?= =?UTF-8?q?=E2=80=94=20re-adopt=20running=20containers=20on=20restart=20(#?= =?UTF-8?q?435)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 5 +- cmd/ecs-agent/reconcile.go | 72 ++++++++++ cmd/ecs-agent/reconcile_test.go | 181 ++++++++++++++++++++++++ cmd/ecs-agent/runtime/containerd_run.go | 24 ++++ cmd/ecs-agent/runtime/runtime.go | 13 +- cmd/ecs-agent/runtime/runtime_fake.go | 32 ++++- cmd/ecs-agent/taskrun.go | 42 ++++-- cmd/ecs-agent/taskrun_test.go | 2 +- 8 files changed, 357 insertions(+), 14 deletions(-) create mode 100644 cmd/ecs-agent/reconcile.go create mode 100644 cmd/ecs-agent/reconcile_test.go diff --git a/cmd/ecs-agent/agent.go b/cmd/ecs-agent/agent.go index 80d55caeb..982df46db 100644 --- a/cmd/ecs-agent/agent.go +++ b/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() diff --git a/cmd/ecs-agent/reconcile.go b/cmd/ecs-agent/reconcile.go new file mode 100644 index 000000000..ff153b915 --- /dev/null +++ b/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 +} diff --git a/cmd/ecs-agent/reconcile_test.go b/cmd/ecs-agent/reconcile_test.go new file mode 100644 index 000000000..23e4b14a4 --- /dev/null +++ b/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) + } + } +} diff --git a/cmd/ecs-agent/runtime/containerd_run.go b/cmd/ecs-agent/runtime/containerd_run.go index fde122577..ccec62579 100644 --- a/cmd/ecs-agent/runtime/containerd_run.go +++ b/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) diff --git a/cmd/ecs-agent/runtime/runtime.go b/cmd/ecs-agent/runtime/runtime.go index f27a4a7df..84c5562bd 100644 --- a/cmd/ecs-agent/runtime/runtime.go +++ b/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. diff --git a/cmd/ecs-agent/runtime/runtime_fake.go b/cmd/ecs-agent/runtime/runtime_fake.go index cf09f4ab6..66fa6354a 100644 --- a/cmd/ecs-agent/runtime/runtime_fake.go +++ b/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 +} diff --git a/cmd/ecs-agent/taskrun.go b/cmd/ecs-agent/taskrun.go index 174d4b98e..c9e752d42 100644 --- a/cmd/ecs-agent/taskrun.go +++ b/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 } diff --git a/cmd/ecs-agent/taskrun_test.go b/cmd/ecs-agent/taskrun_test.go index e446a7b90..e88eec539 100644 --- a/cmd/ecs-agent/taskrun_test.go +++ b/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 →