diff --git a/cmd/ecs-agent/agent.go b/cmd/ecs-agent/agent.go index b95b9e70f..676becb62 100644 --- a/cmd/ecs-agent/agent.go +++ b/cmd/ecs-agent/agent.go @@ -12,8 +12,6 @@ import ( "strings" "time" - "github.com/nats-io/nats.go" - "github.com/mulgadc/spinifex/cmd/ecs-agent/credentials" ctrruntime "github.com/mulgadc/spinifex/cmd/ecs-agent/runtime" "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" @@ -23,14 +21,15 @@ import ( // Overridable via -ldflags "-X main.version=...". var version = "dev" -// Agent wires the ecs-agent's runtime seams: a NATS publisher for the Layer-2 -// bus, a container runtime, and an ECR resolver. It registers the host as a -// container instance at boot, heartbeats while alive, and (Sprint 4e) runs -// assigned tasks through containerd, reporting their state back on the bus. +// Agent wires the ecs-agent's runtime seams: a gateway control-plane client, a +// container runtime, and an ECR resolver. It registers the host as a container +// instance at boot, heartbeats (re-registers) while alive, polls the gateway for +// task assignments, and runs them through containerd, reporting state back over +// the gateway. It never connects to NATS — the bus stays host-internal. type Agent struct { cfg config id identity - pub publisher + cp controlPlane puller ctrruntime.ImagePuller runner ctrruntime.Runner resolver ctrruntime.Resolver @@ -41,7 +40,6 @@ type Agent struct { // netns builds awsvpc task network namespaces (nil-safe; bridge/host skip it). netns *taskNetns - nc *nats.Conn closers []func() error } @@ -49,26 +47,26 @@ type Agent struct { // with fakes; New builds the production seams and delegates here. runner may be // nil when containerd is unavailable; the assign path then reports the task // STOPPED rather than crashing. -func newAgent(cfg config, id identity, pub publisher, puller ctrruntime.ImagePuller, runner ctrruntime.Runner, resolver ctrruntime.Resolver) *Agent { +func newAgent(cfg config, id identity, cp controlPlane, puller ctrruntime.ImagePuller, runner ctrruntime.Runner, resolver ctrruntime.Resolver) *Agent { return &Agent{ cfg: cfg, id: id, - pub: pub, + cp: cp, puller: puller, runner: runner, resolver: resolver, - reg: newRegistrar(pub, id), - hb: newHeartbeater(pub, id, cfg.Heartbeat, nil), + reg: newRegistrar(cp, id), + hb: newHeartbeater(cp, id, cfg.Heartbeat), netns: newTaskNetns(execNetRunner{}), } } // New builds an Agent from config: it resolves the host identity from IMDS, -// connects to NATS, builds the ECR resolver and (best-effort) the containerd -// runtime. A containerd connect failure is logged, not fatal — registration and -// heartbeat still run so the instance is visible while the runtime recovers. -// The ECR gateway client is built lazily on first image pull (not here), so a -// missing or malformed gateway CA does not stop the agent from registering. +// builds the gateway control-plane client, the ECR resolver and (best-effort) +// the containerd runtime. A containerd connect failure is logged, not fatal — +// registration and heartbeat still run so the instance is visible while the +// runtime recovers. The ECR gateway client is built lazily on first image pull +// (not here), so a missing or malformed gateway CA does not stop registration. func New(cfg config) (*Agent, error) { imdsClient := &http.Client{Timeout: 5 * time.Second} @@ -99,41 +97,35 @@ func New(cfg config) (*Agent, error) { runner = rt } - nc, err := nats.Connect(cfg.NATSURL, nats.Name("ecs-agent/"+id.InstanceID)) + cp, err := newGatewayControlPlane(cfg) if err != nil { - return nil, fmt.Errorf("connect nats %s: %w", cfg.NATSURL, err) + return nil, fmt.Errorf("build gateway control-plane: %w", err) } - a := newAgent(cfg, id, nc, puller, runner, resolver) - a.nc = nc + a := newAgent(cfg, id, cp, puller, runner, resolver) if puller != nil { a.closers = append(a.closers, puller.Close) } return a, nil } -// Run registers the instance, starts the heartbeat loop, and blocks until ctx is -// cancelled, then tears down. +// Run registers the instance, starts the heartbeat and assignment-poll loops, +// and blocks until ctx is cancelled, then tears down. func (a *Agent) Run(ctx context.Context) error { if err := a.reg.Register(); err != nil { slog.Error("ecs-agent: initial registration failed", "err", err) } else { - slog.Info("ecs-agent: registered", - "cluster", a.id.ClusterName, "instance", a.id.InstanceID, - "subject", bus.RegisterSubject(a.id.AccountID, a.id.ClusterName, a.id.InstanceID)) + slog.Info("ecs-agent: registered", "cluster", a.id.ClusterName, "instance", a.id.InstanceID) } go a.hb.Run(ctx) - - if err := a.subscribeAssign(ctx); err != nil { - slog.Error("ecs-agent: assign subscription failed; task assignment disabled", "err", err) - } + go a.pollAssignments(ctx) <-ctx.Done() return a.Stop() } -// Stop drains NATS and closes the runtime. Safe to call once. +// Stop closes the runtime. Safe to call once. func (a *Agent) Stop() error { var firstErr error for _, c := range a.closers { @@ -141,10 +133,6 @@ func (a *Agent) Stop() error { firstErr = err } } - if a.nc != nil { - _ = a.nc.Drain() - a.nc.Close() - } return firstErr } diff --git a/cmd/ecs-agent/agent_test.go b/cmd/ecs-agent/agent_test.go index c15d2489e..d076608a9 100644 --- a/cmd/ecs-agent/agent_test.go +++ b/cmd/ecs-agent/agent_test.go @@ -6,22 +6,20 @@ import ( "time" ctrruntime "github.com/mulgadc/spinifex/cmd/ecs-agent/runtime" - "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" ) func TestAgent_RunRegistersThenStopsOnContext(t *testing.T) { - pub := newCapturePub() - cfg := config{Heartbeat: 5 * time.Millisecond} - id := testIdentity() + cp := &fakeCP{} + cfg := config{Heartbeat: 5 * time.Millisecond, PollInterval: 5 * time.Millisecond} puller := &ctrruntime.FakePuller{} - a := newAgent(cfg, id, pub, puller, puller, nil) + a := newAgent(cfg, testIdentity(), cp, puller, puller, nil) a.closers = append(a.closers, puller.Close) ctx, cancel := context.WithCancel(context.Background()) done := make(chan error, 1) go func() { done <- a.Run(ctx) }() - // Give it time to register + emit a heartbeat or two. + // Give it time to register + re-register (heartbeat) a few times. time.Sleep(30 * time.Millisecond) cancel() @@ -34,13 +32,9 @@ func TestAgent_RunRegistersThenStopsOnContext(t *testing.T) { t.Fatal("Run did not return after cancel") } - regSubj := bus.RegisterSubject(id.AccountID, id.ClusterName, id.InstanceID) - if _, ok := pub.msgs[regSubj]; !ok { - t.Errorf("no register message on %s", regSubj) - } - hbSubj := bus.HeartbeatSubject(id.AccountID, id.ClusterName, id.InstanceID) - if _, ok := pub.msgs[hbSubj]; !ok { - t.Errorf("no heartbeat message on %s", hbSubj) + // One boot register + at least one heartbeat re-register. + if cp.registerCount() < 2 { + t.Errorf("registers = %d, want >= 2 (boot + heartbeat)", cp.registerCount()) } if !puller.Closed { t.Error("Stop did not close the puller") diff --git a/cmd/ecs-agent/config.go b/cmd/ecs-agent/config.go index 4054fcdd9..143dfa798 100644 --- a/cmd/ecs-agent/config.go +++ b/cmd/ecs-agent/config.go @@ -12,22 +12,26 @@ const ( defaultGatewayCA = "/etc/spinifex-ecs/gateway-ca.pem" defaultIMDSBase = "http://169.254.169.254/latest" defaultContainerdSocket = "/run/containerd/containerd.sock" - defaultNATSURL = "nats://127.0.0.1:4222" defaultHeartbeat = 30 * time.Second + defaultPollInterval = 5 * time.Second ) // config holds the static settings the agent reads at boot. Cluster identity // (account ID, instance ID) is discovered at runtime from IMDS, not configured; -// only the cluster *name* the instance was launched into is static. +// only the cluster *name* the instance was launched into is static. AccessKey / +// SecretKey are the instance's seeded IAM credentials used to SigV4-sign gateway +// calls — the agent never holds a NATS token. type config struct { GatewayURL string GatewayCA string Region string ClusterName string - NATSURL string + AccessKey string + SecretKey string IMDSBase string ContainerdSocket string Heartbeat time.Duration + PollInterval time.Duration } // loadConfig reads the cloud-init env file then lets real env vars override. @@ -45,7 +49,8 @@ func loadConfig(envFile string) config { GatewayCA: get("ECS_GATEWAY_CA"), Region: get("ECS_REGION"), ClusterName: get("ECS_CLUSTER"), - NATSURL: get("ECS_NATS_URL"), + AccessKey: get("ECS_ACCESS_KEY"), + SecretKey: get("ECS_SECRET_KEY"), IMDSBase: get("ECS_IMDS_BASE"), ContainerdSocket: get("ECS_CONTAINERD_SOCKET"), } @@ -58,9 +63,6 @@ func loadConfig(envFile string) config { if cfg.ContainerdSocket == "" { cfg.ContainerdSocket = defaultContainerdSocket } - if cfg.NATSURL == "" { - cfg.NATSURL = defaultNATSURL - } if cfg.ClusterName == "" { cfg.ClusterName = "default" } @@ -70,6 +72,12 @@ func loadConfig(envFile string) config { cfg.Heartbeat = d } } + cfg.PollInterval = defaultPollInterval + if v := get("ECS_POLL_INTERVAL"); v != "" { + if d, err := time.ParseDuration(v); err == nil && d > 0 { + cfg.PollInterval = d + } + } return cfg } diff --git a/cmd/ecs-agent/config_test.go b/cmd/ecs-agent/config_test.go index 7145ec677..c2454989a 100644 --- a/cmd/ecs-agent/config_test.go +++ b/cmd/ecs-agent/config_test.go @@ -10,8 +10,8 @@ import ( func TestLoadConfig_Defaults(t *testing.T) { // Ensure env vars don't leak in from the host. for _, k := range []string{"ECS_GATEWAY_URL", "ECS_GATEWAY_CA", "ECS_REGION", - "ECS_CLUSTER", "ECS_NATS_URL", "ECS_IMDS_BASE", "ECS_CONTAINERD_SOCKET", - "ECS_HEARTBEAT_INTERVAL"} { + "ECS_CLUSTER", "ECS_ACCESS_KEY", "ECS_SECRET_KEY", "ECS_IMDS_BASE", + "ECS_CONTAINERD_SOCKET", "ECS_HEARTBEAT_INTERVAL", "ECS_POLL_INTERVAL"} { t.Setenv(k, "") } cfg := loadConfig(filepath.Join(t.TempDir(), "absent.env")) @@ -24,15 +24,15 @@ func TestLoadConfig_Defaults(t *testing.T) { if cfg.ContainerdSocket != defaultContainerdSocket { t.Errorf("ContainerdSocket = %q, want default", cfg.ContainerdSocket) } - if cfg.NATSURL != defaultNATSURL { - t.Errorf("NATSURL = %q, want default", cfg.NATSURL) - } if cfg.ClusterName != "default" { t.Errorf("ClusterName = %q, want default", cfg.ClusterName) } if cfg.Heartbeat != defaultHeartbeat { t.Errorf("Heartbeat = %v, want %v", cfg.Heartbeat, defaultHeartbeat) } + if cfg.PollInterval != defaultPollInterval { + t.Errorf("PollInterval = %v, want %v", cfg.PollInterval, defaultPollInterval) + } } func TestLoadConfig_FileThenEnvOverride(t *testing.T) { @@ -43,7 +43,7 @@ func TestLoadConfig_FileThenEnvOverride(t *testing.T) { t.Fatal(err) } t.Setenv("ECS_GATEWAY_URL", "https://gw.env") - t.Setenv("ECS_NATS_URL", "nats://10.0.0.1:4222") + t.Setenv("ECS_ACCESS_KEY", "AKIAENV") cfg := loadConfig(envFile) if cfg.GatewayURL != "https://gw.env" { @@ -55,8 +55,8 @@ func TestLoadConfig_FileThenEnvOverride(t *testing.T) { if cfg.ClusterName != "prod" { t.Errorf("ClusterName = %q, want from file", cfg.ClusterName) } - if cfg.NATSURL != "nats://10.0.0.1:4222" { - t.Errorf("NATSURL = %q, want env", cfg.NATSURL) + if cfg.AccessKey != "AKIAENV" { + t.Errorf("AccessKey = %q, want env", cfg.AccessKey) } } diff --git a/cmd/ecs-agent/controlplane.go b/cmd/ecs-agent/controlplane.go new file mode 100644 index 000000000..b1178325d --- /dev/null +++ b/cmd/ecs-agent/controlplane.go @@ -0,0 +1,114 @@ +package main + +import ( + "encoding/json" + "fmt" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ecs" + + "github.com/mulgadc/spinifex/internal/ecsgw" + handlers_ecs "github.com/mulgadc/spinifex/spinifex/handlers/ecs" + "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" +) + +// controlPlane is the agent's channel to the AWS gateway. It replaces the direct +// NATS publisher: every method is a SigV4-signed gateway request, so the NATS bus +// stays host-internal. Tests inject a fake; production uses gatewayControlPlane. +type controlPlane interface { + // Register (re-)registers the container instance. Called at boot and on every + // heartbeat tick — the gateway's RegisterContainerInstance is idempotent and + // refreshes LastSeen, folding liveness into registration (no bus.Heartbeat). + Register(id identity) error + // SubmitTaskState reports a task's RUNNING/STOPPED transition. + SubmitTaskState(st bus.TaskState) error + // PollAssignments drains this instance's assignment inbox, acking the taskIDs + // accepted on the previous poll. Returns the still-pending assignments. + PollAssignments(cluster, instance string, ack []string) ([]bus.Assign, error) +} + +// gatewayControlPlane implements controlPlane over the SigV4 ECS gateway client. +type gatewayControlPlane struct { + client *ecsgw.Client +} + +var _ controlPlane = (*gatewayControlPlane)(nil) + +// pollTimeout bounds a single PollAssignments request; larger than the register +// timeout to leave room for a future server-side long-poll. +const pollTimeout = 30 * time.Second + +// newGatewayControlPlane builds the SigV4 client from seeded creds + pinned CA. +func newGatewayControlPlane(cfg config) (*gatewayControlPlane, error) { + client, err := ecsgw.New(cfg.GatewayURL, cfg.GatewayCA, cfg.AccessKey, cfg.SecretKey, cfg.Region, pollTimeout) + if err != nil { + return nil, err + } + return &gatewayControlPlane{client: client}, nil +} + +func (g *gatewayControlPlane) Register(id identity) error { + in := &ecs.RegisterContainerInstanceInput{ + Cluster: aws.String(id.ClusterName), + InstanceIdentityDocument: aws.String(id.InstanceID), + TotalResources: []*ecs.Resource{ + {Name: aws.String("CPU"), Type: aws.String("INTEGER"), IntegerValue: aws.Int64(int64(id.Capacity.CPU))}, + {Name: aws.String("MEMORY"), Type: aws.String("INTEGER"), IntegerValue: aws.Int64(int64(id.Capacity.MemoryMiB))}, + }, + VersionInfo: &ecs.VersionInfo{AgentVersion: aws.String(id.AgentVersion)}, + } + body, err := json.Marshal(in) + if err != nil { + return fmt.Errorf("marshal register: %w", err) + } + if _, err := g.client.Call("RegisterContainerInstance", body); err != nil { + return fmt.Errorf("register: %w", err) + } + return nil +} + +func (g *gatewayControlPlane) SubmitTaskState(st bus.TaskState) error { + in := &ecs.SubmitTaskStateChangeInput{ + Cluster: aws.String(st.ClusterName), + Task: aws.String(st.TaskID), + Status: aws.String(st.LastStatus), + Reason: aws.String(st.Reason), + } + for _, c := range st.Containers { + csc := &ecs.ContainerStateChange{ + ContainerName: aws.String(c.Name), + Status: aws.String(c.Status), + RuntimeId: aws.String(c.ContainerID), + } + if c.ExitCode != nil { + csc.ExitCode = aws.Int64(int64(*c.ExitCode)) + } + in.Containers = append(in.Containers, csc) + } + body, err := json.Marshal(in) + if err != nil { + return fmt.Errorf("marshal task-state: %w", err) + } + if _, err := g.client.Call("SubmitTaskStateChange", body); err != nil { + return fmt.Errorf("submit task-state: %w", err) + } + return nil +} + +func (g *gatewayControlPlane) PollAssignments(cluster, instance string, ack []string) ([]bus.Assign, error) { + in := &handlers_ecs.PollAssignmentsInput{Cluster: cluster, ContainerInstance: instance, AckTaskIDs: ack} + body, err := json.Marshal(in) + if err != nil { + return nil, fmt.Errorf("marshal poll: %w", err) + } + resp, err := g.client.Call("PollAssignments", body) + if err != nil { + return nil, fmt.Errorf("poll: %w", err) + } + var out handlers_ecs.PollAssignmentsOutput + if err := json.Unmarshal(resp, &out); err != nil { + return nil, fmt.Errorf("decode poll: %w", err) + } + return out.Assignments, nil +} diff --git a/cmd/ecs-agent/heartbeat.go b/cmd/ecs-agent/heartbeat.go index 4d3b9edcc..dc8e67827 100644 --- a/cmd/ecs-agent/heartbeat.go +++ b/cmd/ecs-agent/heartbeat.go @@ -2,54 +2,30 @@ package main import ( "context" - "encoding/json" - "fmt" "log/slog" "time" - - "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" ) -// heartbeater periodically emits instance-heartbeat messages so the scheduler -// can tell live instances from dead ones. +// heartbeater keeps the container instance alive by periodically re-registering +// it through the gateway. Real ECS folds liveness into ACS; v1 folds it into an +// idempotent RegisterContainerInstance, which refreshes the instance's LastSeen +// so the scheduler can tell live instances from dead ones (no bus.Heartbeat). type heartbeater struct { - pub publisher + cp controlPlane id identity interval time.Duration - // runningTasks reports the current running-task count; nil means zero. - runningTasks func() int } -func newHeartbeater(pub publisher, id identity, interval time.Duration, running func() int) *heartbeater { +func newHeartbeater(cp controlPlane, id identity, interval time.Duration) *heartbeater { if interval <= 0 { interval = defaultHeartbeat } - return &heartbeater{pub: pub, id: id, interval: interval, runningTasks: running} + return &heartbeater{cp: cp, id: id, interval: interval} } -// beat publishes a single Heartbeat message. +// beat re-registers the instance once, refreshing its LastSeen. func (h *heartbeater) beat() error { - n := 0 - if h.runningTasks != nil { - n = h.runningTasks() - } - msg := bus.Heartbeat{ - AccountID: h.id.AccountID, - ClusterName: h.id.ClusterName, - InstanceID: h.id.InstanceID, - Status: bus.StatusActive, - RunningTasks: n, - SentAt: time.Now().UTC(), - } - data, err := json.Marshal(&msg) - if err != nil { - return fmt.Errorf("marshal heartbeat: %w", err) - } - subj := bus.HeartbeatSubject(h.id.AccountID, h.id.ClusterName, h.id.InstanceID) - if err := h.pub.Publish(subj, data); err != nil { - return fmt.Errorf("publish heartbeat %s: %w", subj, err) - } - return nil + return h.cp.Register(h.id) } // Run beats every interval until ctx is cancelled. A failed beat is logged and @@ -63,7 +39,7 @@ func (h *heartbeater) Run(ctx context.Context) { return case <-ticker.C: if err := h.beat(); err != nil { - slog.Warn("ecs-agent: heartbeat failed", "err", err) + slog.Warn("ecs-agent: heartbeat (re-register) failed", "err", err) } } } diff --git a/cmd/ecs-agent/main.go b/cmd/ecs-agent/main.go index 108c2862a..a87f7cce9 100644 --- a/cmd/ecs-agent/main.go +++ b/cmd/ecs-agent/main.go @@ -1,13 +1,13 @@ // ecs-agent runs inside a Spinifex ECS container instance (a guest VM booted from // the ECS-AMI, with containerd baked in). It registers the host with the ECS -// scheduler over the Layer-2 NATS bus, heartbeats while alive, and pulls task -// images from the internal ECR registry through containerd. Task assignment and -// container lifecycle land in Sprint 4e. +// scheduler through the AWS gateway over TLS+SigV4 (never NATS), heartbeats by +// re-registering, polls the gateway for task assignments, runs them through +// containerd, and reports state back over the gateway. // -// Static config (gateway URL, CA, region, cluster, NATS, containerd socket) is -// read from the cloud-init env file /etc/spinifex-ecs/agent.env (KEY=value); -// real env vars override it. Host identity (account, instance, AZ) comes from -// IMDS at boot. +// Static config (gateway URL, CA, region, cluster, seeded IAM creds, containerd +// socket) is read from the cloud-init env file /etc/spinifex-ecs/agent.env +// (KEY=value); real env vars override it. Host identity (account, instance, AZ) +// comes from IMDS at boot. package main import ( diff --git a/cmd/ecs-agent/registration.go b/cmd/ecs-agent/registration.go index 5665e2908..2b077ea72 100644 --- a/cmd/ecs-agent/registration.go +++ b/cmd/ecs-agent/registration.go @@ -1,11 +1,7 @@ package main import ( - "encoding/json" "fmt" - "time" - - "github.com/nats-io/nats.go" "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" ) @@ -23,44 +19,21 @@ type identity struct { AgentVersion string } -// publisher is the subset of *nats.Conn the agent uses to emit bus messages. -// Narrowed to an interface so tests can drive it without a live server. -type publisher interface { - Publish(subject string, data []byte) error -} - -var _ publisher = (*nats.Conn)(nil) - -// registrar emits the boot-time instance-register message. +// registrar registers the host as a container instance through the gateway. type registrar struct { - pub publisher - id identity + cp controlPlane + id identity } -func newRegistrar(pub publisher, id identity) *registrar { - return ®istrar{pub: pub, id: id} +func newRegistrar(cp controlPlane, id identity) *registrar { + return ®istrar{cp: cp, id: id} } -// Register publishes the RegisterInstance message on the cluster's register -// subject. The scheduler records the container instance on receipt (Sprint 4e). +// Register calls the gateway's RegisterContainerInstance. The scheduler records +// (or refreshes) the container instance on receipt. func (r *registrar) Register() error { - msg := bus.RegisterInstance{ - AccountID: r.id.AccountID, - ClusterName: r.id.ClusterName, - InstanceID: r.id.InstanceID, - AZ: r.id.AZ, - Hostname: r.id.Hostname, - Capacity: r.id.Capacity, - AgentVersion: r.id.AgentVersion, - RegisteredAt: time.Now().UTC(), - } - data, err := json.Marshal(&msg) - if err != nil { - return fmt.Errorf("marshal register: %w", err) - } - subj := bus.RegisterSubject(r.id.AccountID, r.id.ClusterName, r.id.InstanceID) - if err := r.pub.Publish(subj, data); err != nil { - return fmt.Errorf("publish register %s: %w", subj, err) + if err := r.cp.Register(r.id); err != nil { + return fmt.Errorf("register instance %s: %w", r.id.InstanceID, err) } return nil } diff --git a/cmd/ecs-agent/registration_test.go b/cmd/ecs-agent/registration_test.go index a46ced13c..3b4e34e4e 100644 --- a/cmd/ecs-agent/registration_test.go +++ b/cmd/ecs-agent/registration_test.go @@ -2,7 +2,6 @@ package main import ( "context" - "encoding/json" "errors" "sync" "testing" @@ -11,33 +10,72 @@ import ( "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" ) -// capturePub records published messages; ErrOn forces Publish to fail. -type capturePub struct { - mu sync.Mutex - msgs map[string][]byte - count int - errOn bool +// fakeCP is a test double for the gateway control-plane. It records registers and +// task-state reports, and replays a scripted queue of poll responses. +type fakeCP struct { + mu sync.Mutex + + registers int + states []bus.TaskState + pollAcks [][]string + pollReplies [][]bus.Assign + pollCalls int + + registerErr bool + submitErr bool } -func newCapturePub() *capturePub { return &capturePub{msgs: map[string][]byte{}} } +var _ controlPlane = (*fakeCP)(nil) -func (c *capturePub) Publish(subject string, data []byte) error { - c.mu.Lock() - defer c.mu.Unlock() - if c.errOn { - return errors.New("publish boom") +func (f *fakeCP) Register(_ identity) error { + f.mu.Lock() + defer f.mu.Unlock() + if f.registerErr { + return errors.New("register boom") } - c.count++ - cp := make([]byte, len(data)) - copy(cp, data) - c.msgs[subject] = cp + f.registers++ return nil } -func (c *capturePub) calls() int { - c.mu.Lock() - defer c.mu.Unlock() - return c.count +func (f *fakeCP) SubmitTaskState(st bus.TaskState) error { + f.mu.Lock() + defer f.mu.Unlock() + if f.submitErr { + return errors.New("submit boom") + } + f.states = append(f.states, st) + return nil +} + +func (f *fakeCP) PollAssignments(_, _ string, ack []string) ([]bus.Assign, error) { + f.mu.Lock() + defer f.mu.Unlock() + cp := append([]string(nil), ack...) + f.pollAcks = append(f.pollAcks, cp) + var out []bus.Assign + if f.pollCalls < len(f.pollReplies) { + out = f.pollReplies[f.pollCalls] + } + f.pollCalls++ + return out, nil +} + +func (f *fakeCP) registerCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.registers +} + +func (f *fakeCP) taskStates() []bus.TaskState { + f.mu.Lock() + defer f.mu.Unlock() + return append([]bus.TaskState(nil), f.states...) +} + +func (f *fakeCP) acks() [][]string { + f.mu.Lock() + defer f.mu.Unlock() + return append([][]string(nil), f.pollAcks...) } func testIdentity() identity { @@ -53,67 +91,42 @@ func testIdentity() identity { } func TestRegistrar_Register(t *testing.T) { - pub := newCapturePub() - id := testIdentity() - if err := newRegistrar(pub, id).Register(); err != nil { + cp := &fakeCP{} + if err := newRegistrar(cp, testIdentity()).Register(); err != nil { t.Fatalf("Register: %v", err) } - subj := bus.RegisterSubject(id.AccountID, id.ClusterName, id.InstanceID) - raw, ok := pub.msgs[subj] - if !ok { - t.Fatalf("no message on %s; got %v", subj, pub.msgs) - } - var got bus.RegisterInstance - if err := json.Unmarshal(raw, &got); err != nil { - t.Fatalf("unmarshal: %v", err) - } - if got.InstanceID != id.InstanceID || got.Capacity.CPU != 2048 { - t.Errorf("payload mismatch: %+v", got) - } - if got.RegisteredAt.IsZero() { - t.Error("RegisteredAt not set") + if cp.registerCount() != 1 { + t.Errorf("registers = %d, want 1", cp.registerCount()) } } -func TestRegistrar_PublishError(t *testing.T) { - pub := newCapturePub() - pub.errOn = true - if err := newRegistrar(pub, testIdentity()).Register(); err == nil { - t.Fatal("expected publish error") +func TestRegistrar_RegisterError(t *testing.T) { + cp := &fakeCP{registerErr: true} + if err := newRegistrar(cp, testIdentity()).Register(); err == nil { + t.Fatal("expected register error") } } func TestHeartbeater_Beat(t *testing.T) { - pub := newCapturePub() - id := testIdentity() - h := newHeartbeater(pub, id, time.Second, func() int { return 3 }) - if err := h.beat(); err != nil { + cp := &fakeCP{} + if err := newHeartbeater(cp, testIdentity(), time.Second).beat(); err != nil { t.Fatalf("beat: %v", err) } - subj := bus.HeartbeatSubject(id.AccountID, id.ClusterName, id.InstanceID) - var got bus.Heartbeat - if err := json.Unmarshal(pub.msgs[subj], &got); err != nil { - t.Fatalf("unmarshal: %v", err) - } - if got.Status != bus.StatusActive || got.RunningTasks != 3 { - t.Errorf("payload mismatch: %+v", got) + if cp.registerCount() != 1 { + t.Errorf("beat should re-register once, got %d", cp.registerCount()) } } -func TestHeartbeater_NilRunningTasksIsZero(t *testing.T) { - pub := newCapturePub() - h := newHeartbeater(pub, testIdentity(), 0, nil) +func TestHeartbeater_DefaultInterval(t *testing.T) { + h := newHeartbeater(&fakeCP{}, testIdentity(), 0) if h.interval != defaultHeartbeat { t.Errorf("interval = %v, want default", h.interval) } - if err := h.beat(); err != nil { - t.Fatal(err) - } } func TestHeartbeater_RunStopsOnContext(t *testing.T) { - pub := newCapturePub() - h := newHeartbeater(pub, testIdentity(), 5*time.Millisecond, nil) + cp := &fakeCP{} + h := newHeartbeater(cp, testIdentity(), 5*time.Millisecond) ctx, cancel := context.WithCancel(context.Background()) done := make(chan struct{}) @@ -126,7 +139,7 @@ func TestHeartbeater_RunStopsOnContext(t *testing.T) { case <-time.After(time.Second): t.Fatal("Run did not stop on context cancel") } - if pub.calls() == 0 { - t.Error("expected at least one heartbeat before cancel") + if cp.registerCount() == 0 { + t.Error("expected at least one re-register before cancel") } } diff --git a/cmd/ecs-agent/taskrun.go b/cmd/ecs-agent/taskrun.go index d8d4e3ef2..10e845d67 100644 --- a/cmd/ecs-agent/taskrun.go +++ b/cmd/ecs-agent/taskrun.go @@ -2,40 +2,47 @@ package main import ( "context" - "encoding/json" "fmt" "log/slog" "time" - "github.com/nats-io/nats.go" - ctrruntime "github.com/mulgadc/spinifex/cmd/ecs-agent/runtime" "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" ) -// subscribeAssign wires the per-instance assign subject to the task runner. It is -// a no-op when the agent has no live NATS connection (unit tests drive runTask -// directly). -func (a *Agent) subscribeAssign(ctx context.Context) error { - if a.nc == nil { - return nil - } - subj := bus.AssignSubject(a.id.AccountID, a.id.ClusterName, a.id.InstanceID) - sub, err := a.nc.Subscribe(subj, func(msg *nats.Msg) { - var as bus.Assign - if err := json.Unmarshal(msg.Data, &as); err != nil { - slog.Error("ecs-agent: bad assign payload", "err", err) +// pollAssignments drains the instance's assignment inbox on a ticker until ctx is +// 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) { + ticker := time.NewTicker(a.cfg.PollInterval) + defer ticker.Stop() + + dispatched := map[string]bool{} + var ackNext []string + for { + select { + case <-ctx.Done(): return + case <-ticker.C: + assigns, err := a.cp.PollAssignments(a.id.ClusterName, a.id.InstanceID, ackNext) + if err != nil { + slog.Warn("ecs-agent: poll assignments failed", "err", err) + continue + } + ackNext = ackNext[:0] + for i := range assigns { + as := assigns[i] + if !dispatched[as.TaskID] { + dispatched[as.TaskID] = true + slog.Info("ecs-agent: task assigned", "task", as.TaskID, "containers", len(as.Containers)) + go a.runTask(ctx, &as) + } + ackNext = append(ackNext, as.TaskID) + } } - slog.Info("ecs-agent: task assigned", "task", as.TaskID, "containers", len(as.Containers)) - go a.runTask(ctx, &as) - }) - if err != nil { - return fmt.Errorf("subscribe %s: %w", subj, err) } - a.closers = append(a.closers, sub.Unsubscribe) - slog.Info("ecs-agent: assign subscription active", "subject", subj) - return nil } // runTask pulls each container image, starts the containers, and reports task @@ -102,9 +109,8 @@ func (a *Agent) waitContainer(ctx context.Context, as *bus.Assign, name, contain slog.Info("ecs-agent: task stopped", "task", as.TaskID, "container", name, "exitCode", exit) } -// setupTaskNetns builds the awsvpc task netns when the assign carries an ENI MAC. -// Returns an empty path (host netns) for bridge/host tasks or when no netns -// manager is configured. +// setupTaskNetns builds the awsvpc task netns from the hot-plugged ENI. Bridge/ +// host tasks (no ENI MAC) and a nil controller are no-ops returning an empty path. func (a *Agent) setupTaskNetns(as *bus.Assign) (string, error) { if as.ENIMacAddress == "" || a.netns == nil { return "", nil @@ -112,7 +118,7 @@ func (a *Agent) setupTaskNetns(as *bus.Assign) (string, error) { return a.netns.Setup(as.TaskID, as.ENIMacAddress) } -// teardownTaskNetns removes the awsvpc task netns; a no-op for bridge/host tasks. +// teardownTaskNetns removes the awsvpc task netns; no-op for bridge/host tasks. func (a *Agent) teardownTaskNetns(as *bus.Assign) { if as.ENIMacAddress == "" || a.netns == nil { return @@ -122,9 +128,10 @@ func (a *Agent) teardownTaskNetns(as *bus.Assign) { } } -// reportTaskState publishes a TaskState message on the task's state subject. +// reportTaskState reports a task transition through the gateway's +// SubmitTaskStateChange action. func (a *Agent) reportTaskState(as *bus.Assign, status, reason string, containers []bus.ContainerStatus) { - if a.pub == nil { + if a.cp == nil { return } msg := bus.TaskState{ @@ -137,14 +144,8 @@ func (a *Agent) reportTaskState(as *bus.Assign, status, reason string, container Reason: reason, ReportedAt: time.Now().UTC(), } - data, err := json.Marshal(&msg) - if err != nil { - slog.Error("ecs-agent: marshal task-state failed", "task", as.TaskID, "err", err) - return - } - subj := bus.TaskStateSubject(a.id.AccountID, a.id.ClusterName, as.TaskID) - if err := a.pub.Publish(subj, data); err != nil { - slog.Error("ecs-agent: publish task-state failed", "task", as.TaskID, "subject", subj, "err", err) + if err := a.cp.SubmitTaskState(msg); err != nil { + slog.Error("ecs-agent: report task-state failed", "task", as.TaskID, "err", err) } } diff --git a/cmd/ecs-agent/taskrun_test.go b/cmd/ecs-agent/taskrun_test.go index fda0c0f88..e446a7b90 100644 --- a/cmd/ecs-agent/taskrun_test.go +++ b/cmd/ecs-agent/taskrun_test.go @@ -2,9 +2,7 @@ package main import ( "context" - "encoding/json" "errors" - "sync" "testing" "time" @@ -12,35 +10,6 @@ import ( "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" ) -// historyPub records every published message in order (capturePub keeps only the -// last per subject, which loses the RUNNING→STOPPED transition on one subject). -type historyPub struct { - mu sync.Mutex - msgs [][]byte -} - -func (h *historyPub) Publish(_ string, data []byte) error { - h.mu.Lock() - defer h.mu.Unlock() - cp := make([]byte, len(data)) - copy(cp, data) - h.msgs = append(h.msgs, cp) - return nil -} - -func (h *historyPub) states() []bus.TaskState { - h.mu.Lock() - defer h.mu.Unlock() - out := make([]bus.TaskState, 0, len(h.msgs)) - for _, raw := range h.msgs { - var ts bus.TaskState - if err := json.Unmarshal(raw, &ts); err == nil { - out = append(out, ts) - } - } - return out -} - func testAssign() *bus.Assign { return &bus.Assign{ AccountID: "123456789012", @@ -54,17 +23,17 @@ func testAssign() *bus.Assign { } } -func newRunAgent(pub publisher, rt *ctrruntime.FakePuller) *Agent { - return newAgent(config{}, testIdentity(), pub, rt, rt, nil) +func newRunAgent(cp controlPlane, rt *ctrruntime.FakePuller) *Agent { + return newAgent(config{}, testIdentity(), cp, rt, rt, nil) } // runTask with no runtime reports the task STOPPED instead of crashing. func TestRunTask_NoRuntimeReportsStopped(t *testing.T) { - pub := &historyPub{} - a := newAgent(config{}, testIdentity(), pub, nil, nil, nil) + cp := &fakeCP{} + a := newAgent(config{}, testIdentity(), cp, nil, nil, nil) a.runTask(context.Background(), testAssign()) - st := pub.states() + st := cp.taskStates() if len(st) != 1 || st[0].LastStatus != bus.TaskStatusStopped { t.Fatalf("want one STOPPED state, got %+v", st) } @@ -76,9 +45,9 @@ func TestRunTask_NoRuntimeReportsStopped(t *testing.T) { // Happy path: pull + run reports RUNNING with the container ID; Wait blocking // (forced error) means no STOPPED overwrite, so we can assert the RUNNING frame. func TestRunTask_PullRunReportsRunning(t *testing.T) { - pub := &historyPub{} + cp := &fakeCP{} rt := &ctrruntime.FakePuller{WaitErr: errors.New("blocked")} - a := newRunAgent(pub, rt) + a := newRunAgent(cp, rt) as := testAssign() a.runTask(context.Background(), as) @@ -92,7 +61,7 @@ func TestRunTask_PullRunReportsRunning(t *testing.T) { t.Errorf("missing task label: %+v", rt.Runs[0].Labels) } - st := pub.states() + st := cp.taskStates() if len(st) == 0 || st[len(st)-1].LastStatus != bus.TaskStatusRunning { t.Fatalf("want RUNNING final state, got %+v", st) } @@ -103,15 +72,15 @@ func TestRunTask_PullRunReportsRunning(t *testing.T) { // A pull failure reports STOPPED and never starts the container. func TestRunTask_PullFailureReportsStopped(t *testing.T) { - pub := &historyPub{} + cp := &fakeCP{} rt := &ctrruntime.FakePuller{Err: errors.New("pull denied")} - a := newRunAgent(pub, rt) + a := newRunAgent(cp, rt) a.runTask(context.Background(), testAssign()) if len(rt.Runs) != 0 { t.Errorf("container should not run after pull failure, got %+v", rt.Runs) } - st := pub.states() + st := cp.taskStates() if len(st) != 1 || st[0].LastStatus != bus.TaskStatusStopped { t.Fatalf("want STOPPED, got %+v", st) } @@ -119,12 +88,12 @@ func TestRunTask_PullFailureReportsStopped(t *testing.T) { // A run failure reports STOPPED. func TestRunTask_RunFailureReportsStopped(t *testing.T) { - pub := &historyPub{} + cp := &fakeCP{} rt := &ctrruntime.FakePuller{RunErr: errors.New("start boom")} - a := newRunAgent(pub, rt) + a := newRunAgent(cp, rt) a.runTask(context.Background(), testAssign()) - st := pub.states() + st := cp.taskStates() if len(st) != 1 || st[0].LastStatus != bus.TaskStatusStopped { t.Fatalf("want STOPPED, got %+v", st) } @@ -132,14 +101,14 @@ func TestRunTask_RunFailureReportsStopped(t *testing.T) { // On container exit, waitContainer reports STOPPED with the exit code. func TestRunTask_ExitReportsStoppedWithCode(t *testing.T) { - pub := &historyPub{} + cp := &fakeCP{} rt := &ctrruntime.FakePuller{WaitCode: 7} - a := newRunAgent(pub, rt) + a := newRunAgent(cp, rt) a.runTask(context.Background(), testAssign()) deadline := time.After(time.Second) for { - st := pub.states() + st := cp.taskStates() if len(st) > 0 && st[len(st)-1].LastStatus == bus.TaskStatusStopped { last := st[len(st)-1] if len(last.Containers) != 1 || last.Containers[0].ExitCode == nil || *last.Containers[0].ExitCode != 7 { @@ -149,18 +118,76 @@ func TestRunTask_ExitReportsStoppedWithCode(t *testing.T) { } select { case <-deadline: - t.Fatalf("no STOPPED state after exit; got %+v", pub.states()) + t.Fatalf("no STOPPED state after exit; got %+v", cp.taskStates()) case <-time.After(2 * time.Millisecond): } } } +// pollAssignments dispatches each assign once and acks it on the next poll. +func TestPollAssignments_DispatchesOnceAndAcks(t *testing.T) { + as := testAssign() + cp := &fakeCP{pollReplies: [][]bus.Assign{ + {*as}, // poll 1: one pending assign + {*as}, // poll 2: still present (not yet acked when poll 2 was issued) + nil, // poll 3: gateway dropped it after the ack + }} + 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) + + // runTask reports RUNNING through the (mutex-guarded) control plane. With the + // assign returned on two consecutive polls, dedup means exactly one dispatch → + // exactly one RUNNING report. + deadline := time.After(time.Second) + for { + if running := countStatus(cp.taskStates(), as.TaskID, bus.TaskStatusRunning); running >= 1 { + break + } + select { + case <-deadline: + t.Fatal("assign was never dispatched (no RUNNING report)") + case <-time.After(2 * time.Millisecond): + } + } + time.Sleep(20 * time.Millisecond) + cancel() + + if running := countStatus(cp.taskStates(), as.TaskID, bus.TaskStatusRunning); running != 1 { + t.Errorf("task dispatched %d times, want exactly 1 (dedup on taskID)", running) + } + // The taskID must appear in a later poll's ack list. + acked := false + for _, ack := range cp.acks() { + for _, id := range ack { + if id == as.TaskID { + acked = true + } + } + } + if !acked { + t.Errorf("task %s was never acked; acks=%v", as.TaskID, cp.acks()) + } +} + +func countStatus(states []bus.TaskState, taskID, status string) int { + n := 0 + for _, st := range states { + if st.TaskID == taskID && st.LastStatus == status { + n++ + } + } + return n +} + // awsvpc: an assign carrying an ENI MAC builds the task netns and passes its // path to the container RunSpec. func TestRunTask_AwsvpcBuildsNetns(t *testing.T) { - pub := &historyPub{} + cp := &fakeCP{} rt := &ctrruntime.FakePuller{WaitErr: errors.New("blocked")} - a := newRunAgent(pub, rt) + a := newRunAgent(cp, rt) f := &fakeNetRunner{linkOut: linkWithENI} a.netns = newTestNetns(f) @@ -179,9 +206,9 @@ func TestRunTask_AwsvpcBuildsNetns(t *testing.T) { // bridge/host: no ENI MAC means no netns is built and the container runs in the // host (VM) netns (empty NetnsPath). func TestRunTask_NoMacSkipsNetns(t *testing.T) { - pub := &historyPub{} + cp := &fakeCP{} rt := &ctrruntime.FakePuller{WaitErr: errors.New("blocked")} - a := newRunAgent(pub, rt) + a := newRunAgent(cp, rt) f := &fakeNetRunner{linkOut: linkWithENI} a.netns = newTestNetns(f) diff --git a/internal/ecsgw/client.go b/internal/ecsgw/client.go new file mode 100644 index 000000000..51e9fee12 --- /dev/null +++ b/internal/ecsgw/client.go @@ -0,0 +1,112 @@ +// Package ecsgw is the on-VM SigV4 client the ecs-agent uses to reach the AWS +// gateway over HTTPS instead of connecting to NATS directly. It signs requests +// for service "ecs" with the instance's seeded IAM credentials, keeping the NATS +// bus host-internal (the gateway relays agent calls onto ecs.bus.* host-side). +package ecsgw + +import ( + "bytes" + "crypto/sha256" + "crypto/tls" + "crypto/x509" + "encoding/hex" + "fmt" + "io" + "net/http" + "os" + "strings" + "time" + + "github.com/mulgadc/predastore/auth" + "github.com/mulgadc/spinifex/internal/tlsconfig" +) + +// targetPrefix is the X-Amz-Target service prefix the gateway's ECS dispatch +// strips to the bare action. Any prefix works (ecsActionFromTarget takes the +// suffix); this one mirrors the AWS JSON 1.1 convention. +const targetPrefix = "AmazonEC2ContainerServiceV20141113." + +// Client posts SigV4-signed AWS JSON 1.1 requests to the gateway for service +// "ecs". One client is reused for register/heartbeat/state/poll. +type Client struct { + baseURL string + accessKey string + secretKey string + region string + httpClient *http.Client +} + +// New builds a client. caPath optionally pins the gateway TLS CA; empty relies on +// the system trust store. region defaults to us-east-1 when empty (SigV4 requires +// a non-empty region). timeout bounds a single call — callers needing a long-poll +// pass a larger value than the default register/heartbeat timeout. +func New(baseURL, caPath, accessKey, secretKey, region string, timeout time.Duration) (*Client, error) { + if baseURL == "" { + return nil, fmt.Errorf("ecsgw: baseURL is required") + } + if accessKey == "" || secretKey == "" { + return nil, fmt.Errorf("ecsgw: access key and secret key are required") + } + if region == "" { + region = "us-east-1" + } + if timeout <= 0 { + timeout = 10 * time.Second + } + + tlsCfg := &tls.Config{ + MinVersion: tls.VersionTLS13, + CurvePreferences: tlsconfig.Curves, + } + if caPath != "" { + pem, err := os.ReadFile(caPath) + if err != nil { + return nil, fmt.Errorf("ecsgw: read gateway CA %q: %w", caPath, err) + } + pool := x509.NewCertPool() + if !pool.AppendCertsFromPEM(pem) { + return nil, fmt.Errorf("ecsgw: gateway CA %q has no usable certificates", caPath) + } + tlsCfg.RootCAs = pool + } + + return &Client{ + baseURL: strings.TrimRight(baseURL, "/"), + accessKey: accessKey, + secretKey: secretKey, + region: region, + httpClient: &http.Client{ + Timeout: timeout, + Transport: &http.Transport{TLSClientConfig: tlsCfg, MaxIdleConns: 2, IdleConnTimeout: 30 * time.Second}, + }, + }, nil +} + +// Call SigV4-signs and POSTs body under X-Amz-Target "" to the +// gateway root, returning the 2xx response body. Errors carry the gateway status +// and body. No retry; callers wrap as needed. +func (c *Client) Call(action string, body []byte) ([]byte, error) { + req, err := http.NewRequest(http.MethodPost, c.baseURL+"/", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("create request: %w", err) + } + req.Header.Set("Content-Type", "application/x-amz-json-1.1") + req.Header.Set("X-Amz-Target", targetPrefix+action) + + sum := sha256.Sum256(body) + if err := auth.SignReq(req, c.accessKey, c.secretKey, hex.EncodeToString(sum[:]), "ecs", c.region); err != nil { + return nil, fmt.Errorf("sign request: %w", err) + } + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("send request: %w", err) + } + defer resp.Body.Close() + + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("gateway returned %d: %s", resp.StatusCode, strings.TrimSpace(string(respBody))) + } + return respBody, nil +} diff --git a/internal/ecsgw/client_test.go b/internal/ecsgw/client_test.go new file mode 100644 index 000000000..793b15340 --- /dev/null +++ b/internal/ecsgw/client_test.go @@ -0,0 +1,76 @@ +package ecsgw + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +func newTestClient(t *testing.T, srv *httptest.Server) *Client { + t.Helper() + c, err := New(srv.URL, "", "AKIDTEST", "secret", "ap-southeast-2", time.Second) + if err != nil { + t.Fatalf("New: %v", err) + } + c.httpClient = srv.Client() + return c +} + +// Call signs for service "ecs", sets the X-Amz-Target action, and returns the +// 2xx body. +func TestClient_CallSignsAndTargets(t *testing.T) { + var gotTarget, gotAuth, gotBody string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotTarget = r.Header.Get("X-Amz-Target") + gotAuth = r.Header.Get("Authorization") + b, _ := io.ReadAll(r.Body) + gotBody = string(b) + _, _ = w.Write([]byte(`{"ok":true}`)) + })) + defer srv.Close() + + c := newTestClient(t, srv) + out, err := c.Call("RegisterContainerInstance", []byte(`{"instanceId":"i-1"}`)) + if err != nil { + t.Fatalf("Call: %v", err) + } + if string(out) != `{"ok":true}` { + t.Errorf("body = %q", out) + } + if gotTarget != targetPrefix+"RegisterContainerInstance" { + t.Errorf("target = %q", gotTarget) + } + if gotBody != `{"instanceId":"i-1"}` { + t.Errorf("relayed body = %q", gotBody) + } + if !strings.Contains(gotAuth, "AWS4-HMAC-SHA256") || !strings.Contains(gotAuth, "/ecs/aws4_request") { + t.Errorf("auth header not SigV4 for ecs: %q", gotAuth) + } +} + +// A non-2xx gateway response surfaces as an error carrying status + body. +func TestClient_CallErrorStatus(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"__type":"AccessDenied"}`)) + })) + defer srv.Close() + + c := newTestClient(t, srv) + if _, err := c.Call("Heartbeat", []byte(`{}`)); err == nil || + !strings.Contains(err.Error(), "403") || !strings.Contains(err.Error(), "AccessDenied") { + t.Fatalf("want 403 AccessDenied error, got %v", err) + } +} + +func TestNew_Validation(t *testing.T) { + if _, err := New("", "", "a", "b", "r", 0); err == nil { + t.Error("want error for empty baseURL") + } + if _, err := New("https://gw:9999", "", "", "b", "r", 0); err == nil { + t.Error("want error for empty access key") + } +} diff --git a/scripts/ecs-node-bringup.sh b/scripts/ecs-node-bringup.sh new file mode 100755 index 000000000..5c7681d0e --- /dev/null +++ b/scripts/ecs-node-bringup.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# ecs-node-bringup.sh — launch a Spinifex ECS container-instance VM. +# +# ECS container instances are plain EC2 instances booted from the spinifex-ecs-node +# AMI (Alpine + containerd + ecs-agent), AWS-faithful: there is no ECS-specific +# launch API. This script renders the cloud-init user-data that seeds the agent's +# control-plane config + IAM credentials into /etc/spinifex-ecs/agent.env and the +# gateway TLS CA into /etc/spinifex-ecs/gateway-ca.pem, then runs the instance. +# +# The agent then registers, heartbeats, polls assignments, and reports task state +# entirely through the AWS gateway over TLS+SigV4 — it never touches NATS. +# +# Usage: +# AWS_PROFILE=spinifex scripts/ecs-node-bringup.sh +# +# Env overrides: +# GATEWAY_IP gateway/mgmt IP the guest dials (default 10.15.8.1) +# GATEWAY_CA host path to the gateway CA PEM (default /etc/spinifex/ca.pem) +# AMI_ID ECS-node AMI (default: latest spinifex-ecs-node) +# SUBNET_ID / SG_ID / KEY_NAME / INSTANCE_TYPE (default: first available / t3.small) +set -euo pipefail + +CLUSTER="${1:?usage: ecs-node-bringup.sh }" +PROFILE="${AWS_PROFILE:-spinifex}" +REGION="$(aws configure get region --profile "$PROFILE")" +GATEWAY_IP="${GATEWAY_IP:-10.15.8.1}" +GATEWAY_CA="${GATEWAY_CA:-/etc/spinifex/ca.pem}" +INSTANCE_TYPE="${INSTANCE_TYPE:-t3.small}" + +awscli() { aws --profile "$PROFILE" "$@"; } + +AMI_ID="${AMI_ID:-$(awscli ec2 describe-images \ + --filters 'Name=name,Values=spinifex-ecs-node' \ + --query 'sort_by(Images,&CreationDate)[-1].ImageId' --output text)}" +SUBNET_ID="${SUBNET_ID:-$(awscli ec2 describe-subnets \ + --query 'Subnets[0].SubnetId' --output text)}" +SG_ID="${SG_ID:-$(awscli ec2 describe-security-groups \ + --query 'SecurityGroups[0].GroupId' --output text)}" +KEY_NAME="${KEY_NAME:-$(awscli ec2 describe-key-pairs \ + --query 'KeyPairs[0].KeyName' --output text)}" + +ACCESS_KEY="$(aws configure get aws_access_key_id --profile "$PROFILE")" +SECRET_KEY="$(aws configure get aws_secret_access_key --profile "$PROFILE")" + +echo "cluster=$CLUSTER ami=$AMI_ID subnet=$SUBNET_ID sg=$SG_ID key=$KEY_NAME gw=https://$GATEWAY_IP:9999" >&2 + +# Indent a file's content by six spaces for a cloud-init "content: |" block. +indent6() { sed 's/^/ /' "$1"; } + +USERDATA="$(cat < /etc/resolv.conf +write_files: + - path: /etc/spinifex-ecs/agent.env + permissions: '0600' + content: | + ECS_GATEWAY_URL=https://$GATEWAY_IP:9999 + ECS_GATEWAY_CA=/etc/spinifex-ecs/gateway-ca.pem + ECS_REGION=$REGION + ECS_CLUSTER=$CLUSTER + ECS_ACCESS_KEY=$ACCESS_KEY + ECS_SECRET_KEY=$SECRET_KEY + - path: /etc/spinifex-ecs/gateway-ca.pem + permissions: '0644' + content: | +$(indent6 "$GATEWAY_CA") +EOF +)" + +INSTANCE_ID="$(awscli ec2 run-instances \ + --image-id "$AMI_ID" --instance-type "$INSTANCE_TYPE" \ + --key-name "$KEY_NAME" --subnet-id "$SUBNET_ID" --security-group-ids "$SG_ID" \ + --user-data "$USERDATA" \ + --tag-specifications "ResourceType=instance,Tags=[{Key=Name,Value=ecs-node-$CLUSTER}]" \ + --query 'Instances[0].InstanceId' --output text)" + +echo "$INSTANCE_ID" diff --git a/scripts/images/ecs-agent/ecs-agent.initd b/scripts/images/ecs-agent/ecs-agent.initd index 09cee953d..092d781c4 100644 --- a/scripts/images/ecs-agent/ecs-agent.initd +++ b/scripts/images/ecs-agent/ecs-agent.initd @@ -1,10 +1,10 @@ #!/sbin/openrc-run # ecs-agent — Spinifex ECS container-instance agent. Resolves host identity from -# IMDSv2, connects to the cluster NATS bus, registers + heartbeats on the Layer-2 -# subjects, and pulls task images through containerd. +# IMDSv2, registers + heartbeats + polls task assignments through the AWS gateway +# over TLS+SigV4 (no NATS), and pulls task images through containerd. # -# Static config (gateway URL, CA, region, cluster, NATS, containerd socket) comes -# from /etc/spinifex-ecs/agent.env, dropped by cloud-init user-data at launch. +# Static config (gateway URL, CA, region, cluster, seeded IAM creds, containerd +# socket) comes from /etc/spinifex-ecs/agent.env, dropped by cloud-init at launch. description="Spinifex ECS container-instance agent" command="/usr/local/bin/ecs-agent" @@ -23,11 +23,12 @@ depend() { start_pre() { # Cloud-init drops the agent config here; OpenRC injects it into the agent's # environment. Absent file is tolerated — the agent falls back to its - # built-in defaults (cluster=default, local NATS, default IMDS). + # built-in defaults (cluster=default, default IMDS). if [ -f /etc/spinifex-ecs/agent.env ]; then # shellcheck disable=SC1091 . /etc/spinifex-ecs/agent.env export ECS_GATEWAY_URL ECS_GATEWAY_CA ECS_REGION ECS_CLUSTER \ - ECS_NATS_URL ECS_IMDS_BASE ECS_CONTAINERD_SOCKET ECS_HEARTBEAT_INTERVAL + ECS_ACCESS_KEY ECS_SECRET_KEY ECS_IMDS_BASE ECS_CONTAINERD_SOCKET \ + ECS_HEARTBEAT_INTERVAL ECS_POLL_INTERVAL fi } diff --git a/scripts/images/ecs-agent/manifest.conf b/scripts/images/ecs-agent/manifest.conf index 5ddbb2fe9..39950185c 100644 --- a/scripts/images/ecs-agent/manifest.conf +++ b/scripts/images/ecs-agent/manifest.conf @@ -27,7 +27,8 @@ scripts/images/ecs-agent/10-bridge.conflist:/etc/cni/net.d/10-bridge.conflist" # ecs-ca-install (first boot) installs the cloud-init gateway CA into the system # trust before containerd starts; containerd then trusts the ECR registry TLS. -# ecs-agent registers + heartbeats on the NATS bus once net + containerd are up. +# ecs-agent registers + heartbeats + polls the AWS gateway (TLS+SigV4, no NATS) +# once net + containerd are up. ENABLE_SERVICES="containerd ecs-ca-install ecs-agent" SETUP_SCRIPT="scripts/images/ecs-agent/setup.sh" diff --git a/spinifex/daemon/daemon.go b/spinifex/daemon/daemon.go index 60d67c34e..bff7dc316 100644 --- a/spinifex/daemon/daemon.go +++ b/spinifex/daemon/daemon.go @@ -965,6 +965,8 @@ func (d *Daemon) subscribeAll() error { natsSub{"ecs.RunTask", d.handleECSRunTask, "spinifex-workers"}, natsSub{"ecs.DescribeTasks", d.handleECSDescribeTasks, "spinifex-workers"}, natsSub{"ecs.ListTasks", d.handleECSListTasks, "spinifex-workers"}, + natsSub{"ecs.SubmitTaskStateChange", d.handleECSSubmitTaskStateChange, "spinifex-workers"}, + natsSub{"ecs.PollAssignments", d.handleECSPollAssignments, "spinifex-workers"}, ) } diff --git a/spinifex/daemon/daemon_handlers_ecs.go b/spinifex/daemon/daemon_handlers_ecs.go index b15d9383e..d4eb017a0 100644 --- a/spinifex/daemon/daemon_handlers_ecs.go +++ b/spinifex/daemon/daemon_handlers_ecs.go @@ -54,6 +54,14 @@ func (d *Daemon) handleECSDescribeTasks(msg *nats.Msg) { handleNATSRequest(msg, d.ecsService.DescribeTasks) } +func (d *Daemon) handleECSSubmitTaskStateChange(msg *nats.Msg) { + handleNATSRequest(msg, d.ecsService.SubmitTaskStateChange) +} + func (d *Daemon) handleECSListTasks(msg *nats.Msg) { handleNATSRequest(msg, d.ecsService.ListTasks) } + +func (d *Daemon) handleECSPollAssignments(msg *nats.Msg) { + handleNATSRequest(msg, d.ecsService.PollAssignments) +} diff --git a/spinifex/gateway/ecs.go b/spinifex/gateway/ecs.go index 075950eb5..c46aec18a 100644 --- a/spinifex/gateway/ecs.go +++ b/spinifex/gateway/ecs.go @@ -57,6 +57,10 @@ func (gw *GatewayConfig) ECS_Request(w http.ResponseWriter, r *http.Request) err return err } + if gateway_ecs.RawJSONActions[action] { + gateway_ecs.WriteRawJSONResponse(w, output) + return nil + } gateway_ecs.WriteJSONResponse(w, output) return nil } diff --git a/spinifex/gateway/ecs/actions.go b/spinifex/gateway/ecs/actions.go index c7e635f88..d493bd2d7 100644 --- a/spinifex/gateway/ecs/actions.go +++ b/spinifex/gateway/ecs/actions.go @@ -119,3 +119,26 @@ func ListTasks(nc *nats.Conn, accountID string, body []byte) (any, error) { } return handlers_ecs.NewNATSECSService(nc).ListTasks(input, accountID) } + +// SubmitTaskStateChange is the agent's task-state report path over the gateway +// (replaces the Layer-2 bus publish). The account is the SigV4 caller, so an +// instance cannot report state for another account's task. +func SubmitTaskStateChange(nc *nats.Conn, accountID string, body []byte) (any, error) { + input := new(ecs.SubmitTaskStateChangeInput) + if err := unmarshalIfBody(body, input); err != nil { + return nil, err + } + return handlers_ecs.NewNATSECSService(nc).SubmitTaskStateChange(input, accountID) +} + +// PollAssignments drains the calling instance's assignment inbox (replaces the +// Layer-2 assign subscribe). Internal agent↔gateway action, not an AWS SDK +// shape; the response carries bus.Assign with an RFC3339 time, so the gateway +// encodes it with encoding/json (RawJSONActions), not the jsonutil marshaler. +func PollAssignments(nc *nats.Conn, accountID string, body []byte) (any, error) { + input := new(handlers_ecs.PollAssignmentsInput) + if err := unmarshalIfBody(body, input); err != nil { + return nil, err + } + return handlers_ecs.NewNATSECSService(nc).PollAssignments(input, accountID) +} diff --git a/spinifex/gateway/ecs/handler.go b/spinifex/gateway/ecs/handler.go index b52b9e37f..a2eeae57e 100644 --- a/spinifex/gateway/ecs/handler.go +++ b/spinifex/gateway/ecs/handler.go @@ -9,6 +9,7 @@ package gateway_ecs import ( + "encoding/json" "errors" "log/slog" "net/http" @@ -62,6 +63,12 @@ var Actions = map[string]Handler{ "DescribeTasks": DescribeTasks, "ListTasks": ListTasks, + // Agent task-state reporting (agent → gateway → recordTaskState). + "SubmitTaskStateChange": SubmitTaskStateChange, + + // Agent assignment polling (agent drains its inbox over the gateway). + "PollAssignments": PollAssignments, + // Service. ListServicesByNamespace is a no-op stub in v1. "CreateService": NotImplemented, "UpdateService": NotImplemented, @@ -87,6 +94,14 @@ var Actions = map[string]Handler{ "ListTagsForResource": NotImplemented, } +// RawJSONActions encode their response with encoding/json instead of jsonutil. +// These are internal agent↔gateway shapes (not AWS SDK types) whose payloads +// carry RFC3339 time fields the on-VM agent decodes with encoding/json; the +// jsonutil marshaler would emit epoch-second floats and break the decode. +var RawJSONActions = map[string]bool{ + "PollAssignments": true, +} + // WriteJSONResponse serialises obj as a 200 AWS JSON 1.1 response using the // SDK's jsonutil marshaler (epoch-second time encoding), matching ECR/ACM/EKS. func WriteJSONResponse(w http.ResponseWriter, obj any) { @@ -96,6 +111,22 @@ func WriteJSONResponse(w http.ResponseWriter, obj any) { http.Error(w, awserrors.ErrorInternalError, http.StatusInternalServerError) return } + writeJSONBody(w, body) +} + +// WriteRawJSONResponse serialises obj with encoding/json (RFC3339 times) for the +// internal agent↔gateway actions in RawJSONActions. +func WriteRawJSONResponse(w http.ResponseWriter, obj any) { + body, err := json.Marshal(obj) + if err != nil { + slog.Error("ECS: failed to marshal raw response JSON", "err", err) + http.Error(w, awserrors.ErrorInternalError, http.StatusInternalServerError) + return + } + writeJSONBody(w, body) +} + +func writeJSONBody(w http.ResponseWriter, body []byte) { w.Header().Set("Content-Type", JSONContentType) w.WriteHeader(http.StatusOK) if _, err := w.Write(body); err != nil { diff --git a/spinifex/gateway/ecs/handler_test.go b/spinifex/gateway/ecs/handler_test.go index 2fa499ced..e0204552e 100644 --- a/spinifex/gateway/ecs/handler_test.go +++ b/spinifex/gateway/ecs/handler_test.go @@ -25,6 +25,8 @@ var wiredActions = map[string]bool{ "RegisterTaskDefinition": true, "DescribeTaskDefinition": true, "ListTaskDefinitions": true, "RegisterContainerInstance": true, "DescribeContainerInstances": true, "ListContainerInstances": true, "RunTask": true, "DescribeTasks": true, "ListTasks": true, + "SubmitTaskStateChange": true, + "PollAssignments": true, } func TestActions_StubsReturnNotImplemented(t *testing.T) { diff --git a/spinifex/handlers/ecs/conv.go b/spinifex/handlers/ecs/conv.go index 67708b286..05a018d69 100644 --- a/spinifex/handlers/ecs/conv.go +++ b/spinifex/handlers/ecs/conv.go @@ -39,6 +39,15 @@ func containerInstanceShortID(ref string) string { return ref } +// taskShortID extracts the task UUID from a bare id or a task ARN +// (arn:...:task/{cluster}/{taskID}); the final path segment is the task id. +func taskShortID(ref string) string { + if i := strings.LastIndexByte(ref, '/'); i >= 0 { + return ref[i+1:] + } + return ref +} + // containerDefsFromAWS maps SDK container definitions to the persisted subset. func containerDefsFromAWS(in []*ecs.ContainerDefinition) []ContainerDef { out := make([]ContainerDef, 0, len(in)) diff --git a/spinifex/handlers/ecs/eni_test.go b/spinifex/handlers/ecs/eni_test.go index 48f1a921c..13b3d7373 100644 --- a/spinifex/handlers/ecs/eni_test.go +++ b/spinifex/handlers/ecs/eni_test.go @@ -92,7 +92,7 @@ func awsvpcRunInput() *ecs.RunTaskInput { } func TestRunTask_Awsvpc_AllocatesAttachesAssigns(t *testing.T) { - svc, nc := newTestService(t) + svc, _ := newTestService(t) eni := &stubENI{} svc.eni = eni _, err := svc.CreateCluster(&ecs.CreateClusterInput{ClusterName: aws.String("web")}, testAccountID) @@ -100,9 +100,6 @@ func TestRunTask_Awsvpc_AllocatesAttachesAssigns(t *testing.T) { registerAwsvpcTaskDef(t, svc, "app", 128, 256) registerInstance(t, svc, "web", "i-1", 1024, 2048) - sub, err := nc.SubscribeSync(bus.AssignSubject(testAccountID, "web", "i-1")) - require.NoError(t, err) - out, err := svc.RunTask(awsvpcRunInput(), testAccountID) require.NoError(t, err) require.Len(t, out.Tasks, 1) @@ -120,11 +117,11 @@ func TestRunTask_Awsvpc_AllocatesAttachesAssigns(t *testing.T) { assert.Equal(t, "eni-stub", detailValue(att, "networkInterfaceId")) assert.Equal(t, "172.31.0.50", detailValue(att, "privateIPv4Address")) - // ENI plumbed into the assign payload. - msg, err := sub.NextMsg(2 * time.Second) + // ENI plumbed into the assign payload, delivered via the instance KV inbox. + poll, err := svc.PollAssignments(&PollAssignmentsInput{Cluster: "web", ContainerInstance: "i-1"}, testAccountID) require.NoError(t, err) - var as bus.Assign - require.NoError(t, json.Unmarshal(msg.Data, &as)) + require.Len(t, poll.Assignments, 1) + as := poll.Assignments[0] assert.Equal(t, "eni-stub", as.ENIID) assert.Equal(t, "02:aa:bb:cc:dd:ee", as.ENIMacAddress) assert.Equal(t, "172.31.0.50", as.ENIPrivateIP) diff --git a/spinifex/handlers/ecs/instances.go b/spinifex/handlers/ecs/instances.go index 03c74c032..1e7aee0d3 100644 --- a/spinifex/handlers/ecs/instances.go +++ b/spinifex/handlers/ecs/instances.go @@ -227,12 +227,44 @@ func (s *Service) recordTaskState(msg *bus.TaskState) error { // Release capacity + reclaim the task ENI once, on the transition into STOPPED. if msg.LastStatus == TaskStatusStopped && prev != TaskStatusStopped { + s.reclaimAssignInbox(kv, msg.ClusterName, task.ContainerInstanceID, msg.TaskID) s.reclaimTaskENI(msg.AccountID, &task) return s.releaseReservation(kv, msg.ClusterName, task.ContainerInstanceID, msg.TaskID, task.ReservedCPU, task.ReservedMemoryMiB) } return nil } +// SubmitTaskStateChange is the AWS-API task-state path (agent → gateway → here). +// It maps the SDK input onto the same bus.TaskState shape the Layer-2 bus +// delivers and converges on recordTaskState, so a gateway-routed agent reports +// state without touching NATS. The account is authoritative from accountID. +func (s *Service) SubmitTaskStateChange(input *ecs.SubmitTaskStateChangeInput, accountID string) (*ecs.SubmitTaskStateChangeOutput, error) { + msg := bus.TaskState{ + AccountID: accountID, + ClusterName: clusterShortName(aws.StringValue(input.Cluster)), + TaskID: taskShortID(aws.StringValue(input.Task)), + LastStatus: aws.StringValue(input.Status), + Reason: aws.StringValue(input.Reason), + ReportedAt: time.Now().UTC(), + } + for _, c := range input.Containers { + cs := bus.ContainerStatus{ + Name: aws.StringValue(c.ContainerName), + Status: aws.StringValue(c.Status), + ContainerID: aws.StringValue(c.RuntimeId), + } + if c.ExitCode != nil { + code := int(aws.Int64Value(c.ExitCode)) + cs.ExitCode = &code + } + msg.Containers = append(msg.Containers, cs) + } + if err := s.recordTaskState(&msg); err != nil { + return nil, err + } + return &ecs.SubmitTaskStateChangeOutput{Acknowledgment: aws.String("OK")}, nil +} + // releaseReservation returns a stopped task's capacity to its instance under CAS. func (s *Service) releaseReservation(kv nats.KeyValue, cluster, instanceID, taskID string, cpu, mem int) error { for range reservePlacementRetries { diff --git a/spinifex/handlers/ecs/poll.go b/spinifex/handlers/ecs/poll.go new file mode 100644 index 000000000..33e535648 --- /dev/null +++ b/spinifex/handlers/ecs/poll.go @@ -0,0 +1,69 @@ +package handlers_ecs + +import ( + "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" + "github.com/nats-io/nats.go" +) + +// PollAssignmentsInput is the agent's poll request. ContainerInstance names the +// instance whose inbox to drain; AckTaskIDs are assignments the agent accepted on +// a prior poll and that may now be removed (at-least-once: an unacked assign is +// re-delivered after an agent crash). This is an internal agent↔gateway shape, +// not an AWS SDK action. +type PollAssignmentsInput struct { + Cluster string `json:"cluster"` + ContainerInstance string `json:"containerInstance"` + AckTaskIDs []string `json:"ackTaskIds,omitempty"` +} + +// PollAssignmentsOutput carries the instance's currently pending assignments. +type PollAssignmentsOutput struct { + Assignments []bus.Assign `json:"assignments,omitempty"` +} + +// PollAssignments acks any completed assignments (deletes their inbox keys) then +// returns the instance's still-pending assignments. The scheduler is the only +// writer of the inbox; the agent is the only reader/acker. Account is +// authoritative from accountID (SigV4 caller), so an instance cannot drain +// another account's inbox. +func (s *Service) PollAssignments(input *PollAssignmentsInput, accountID string) (*PollAssignmentsOutput, error) { + cluster := clusterShortName(input.Cluster) + instanceID := containerInstanceShortID(input.ContainerInstance) + kv, err := s.bucket(accountID) + if err != nil { + return nil, err + } + + for _, taskID := range input.AckTaskIDs { + if taskID == "" { + continue + } + _ = kv.Delete(AssignmentKey(cluster, instanceID, taskShortID(taskID))) + } + + keys, err := keysWithPrefix(kv, AssignmentsPrefix(cluster, instanceID)) + if err != nil { + return nil, err + } + out := &PollAssignmentsOutput{} + for _, key := range keys { + var as bus.Assign + found, err := getJSON(kv, key, &as) + if err != nil { + return nil, err + } + if found { + out.Assignments = append(out.Assignments, as) + } + } + return out, nil +} + +// reclaimAssignInbox removes a task's inbox entry on the STOPPED transition so a +// stopped task is never re-delivered to the agent. Best-effort. +func (s *Service) reclaimAssignInbox(kv nats.KeyValue, cluster, instanceID, taskID string) { + if instanceID == "" || taskID == "" { + return + } + _ = kv.Delete(AssignmentKey(cluster, instanceID, taskID)) +} diff --git a/spinifex/handlers/ecs/service.go b/spinifex/handlers/ecs/service.go index d3d6a2095..d36f59505 100644 --- a/spinifex/handlers/ecs/service.go +++ b/spinifex/handlers/ecs/service.go @@ -34,6 +34,11 @@ type ECSService interface { RunTask(input *ecs.RunTaskInput, accountID string) (*ecs.RunTaskOutput, error) DescribeTasks(input *ecs.DescribeTasksInput, accountID string) (*ecs.DescribeTasksOutput, error) ListTasks(input *ecs.ListTasksInput, accountID string) (*ecs.ListTasksOutput, error) + + SubmitTaskStateChange(input *ecs.SubmitTaskStateChangeInput, accountID string) (*ecs.SubmitTaskStateChangeOutput, error) + + // PollAssignments drains an instance's task-assignment inbox (agent → gateway). + PollAssignments(input *PollAssignmentsInput, accountID string) (*PollAssignmentsOutput, error) } // Service is the daemon-side ECS control-plane implementation, backed by the diff --git a/spinifex/handlers/ecs/service_nats.go b/spinifex/handlers/ecs/service_nats.go index f791063c0..750729d5c 100644 --- a/spinifex/handlers/ecs/service_nats.go +++ b/spinifex/handlers/ecs/service_nats.go @@ -78,3 +78,11 @@ func (s *NATSECSService) DescribeTasks(input *ecs.DescribeTasksInput, accountID func (s *NATSECSService) ListTasks(input *ecs.ListTasksInput, accountID string) (*ecs.ListTasksOutput, error) { return utils.NATSRequest[ecs.ListTasksOutput](s.natsConn, "ecs.ListTasks", input, defaultTimeout, accountID) } + +func (s *NATSECSService) SubmitTaskStateChange(input *ecs.SubmitTaskStateChangeInput, accountID string) (*ecs.SubmitTaskStateChangeOutput, error) { + return utils.NATSRequest[ecs.SubmitTaskStateChangeOutput](s.natsConn, "ecs.SubmitTaskStateChange", input, defaultTimeout, accountID) +} + +func (s *NATSECSService) PollAssignments(input *PollAssignmentsInput, accountID string) (*PollAssignmentsOutput, error) { + return utils.NATSRequest[PollAssignmentsOutput](s.natsConn, "ecs.PollAssignments", input, defaultTimeout, accountID) +} diff --git a/spinifex/handlers/ecs/service_test.go b/spinifex/handlers/ecs/service_test.go index 1daebe0cb..917b77b5d 100644 --- a/spinifex/handlers/ecs/service_test.go +++ b/spinifex/handlers/ecs/service_test.go @@ -1,7 +1,6 @@ package handlers_ecs import ( - "encoding/json" "testing" "time" @@ -116,15 +115,12 @@ func registerInstance(t *testing.T, svc *Service, cluster, id string, cpu, mem i } func TestService_RunTask_PlacesAndAssigns(t *testing.T) { - svc, nc := newTestService(t) + svc, _ := newTestService(t) _, err := svc.CreateCluster(&ecs.CreateClusterInput{ClusterName: aws.String("web")}, testAccountID) require.NoError(t, err) registerTaskDef(t, svc, "app", 128, 256) registerInstance(t, svc, "web", "i-1", 1024, 2048) - sub, err := nc.SubscribeSync(bus.AssignSubject(testAccountID, "web", "i-1")) - require.NoError(t, err) - out, err := svc.RunTask(&ecs.RunTaskInput{ Cluster: aws.String("web"), TaskDefinition: aws.String("app"), @@ -135,11 +131,11 @@ func TestService_RunTask_PlacesAndAssigns(t *testing.T) { assert.Empty(t, out.Failures) assert.Equal(t, "PENDING", aws.StringValue(out.Tasks[0].LastStatus)) - // Assign published to the instance's agent. - msg, err := sub.NextMsg(2 * time.Second) + // Assign written to the instance's KV inbox, drained by polling the gateway. + poll, err := svc.PollAssignments(&PollAssignmentsInput{Cluster: "web", ContainerInstance: "i-1"}, testAccountID) require.NoError(t, err) - var as bus.Assign - require.NoError(t, json.Unmarshal(msg.Data, &as)) + require.Len(t, poll.Assignments, 1) + as := poll.Assignments[0] assert.Equal(t, "i-1", as.InstanceID) require.Len(t, as.Containers, 1) assert.Equal(t, "registry/app:1", as.Containers[0].Image) @@ -159,6 +155,46 @@ func TestService_RunTask_PlacesAndAssigns(t *testing.T) { assert.Equal(t, as.TaskID, containerInstanceShortID(aws.StringValue(out.Tasks[0].TaskArn))) } +// PollAssignments is at-least-once: re-poll without ack re-delivers the assign; +// acking it (then STOPPED) drains the inbox so it is never re-delivered. +func TestService_PollAssignments_AckAndReclaim(t *testing.T) { + svc, _ := newTestService(t) + _, err := svc.CreateCluster(&ecs.CreateClusterInput{ClusterName: aws.String("web")}, testAccountID) + require.NoError(t, err) + registerTaskDef(t, svc, "app", 128, 256) + registerInstance(t, svc, "web", "i-1", 1024, 2048) + out, err := svc.RunTask(&ecs.RunTaskInput{Cluster: aws.String("web"), TaskDefinition: aws.String("app")}, testAccountID) + require.NoError(t, err) + taskID := containerInstanceShortID(aws.StringValue(out.Tasks[0].TaskArn)) + + // Unacked re-poll re-delivers. + p1, err := svc.PollAssignments(&PollAssignmentsInput{Cluster: "web", ContainerInstance: "i-1"}, testAccountID) + require.NoError(t, err) + require.Len(t, p1.Assignments, 1) + p2, err := svc.PollAssignments(&PollAssignmentsInput{Cluster: "web", ContainerInstance: "i-1"}, testAccountID) + require.NoError(t, err) + require.Len(t, p2.Assignments, 1) + + // Ack drains the inbox. + p3, err := svc.PollAssignments(&PollAssignmentsInput{ + Cluster: "web", ContainerInstance: "i-1", AckTaskIDs: []string{taskID}, + }, testAccountID) + require.NoError(t, err) + assert.Empty(t, p3.Assignments) + + // A fresh RunTask + STOPPED reclaims its inbox entry without an explicit ack. + out2, err := svc.RunTask(&ecs.RunTaskInput{Cluster: aws.String("web"), TaskDefinition: aws.String("app")}, testAccountID) + require.NoError(t, err) + task2 := containerInstanceShortID(aws.StringValue(out2.Tasks[0].TaskArn)) + require.NoError(t, svc.recordTaskState(&bus.TaskState{ + AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1", TaskID: task2, + LastStatus: bus.TaskStatusStopped, + })) + p4, err := svc.PollAssignments(&PollAssignmentsInput{Cluster: "web", ContainerInstance: "i-1"}, testAccountID) + require.NoError(t, err) + assert.Empty(t, p4.Assignments) +} + func TestService_RunTask_ClusterNotFound(t *testing.T) { svc, _ := newTestService(t) registerTaskDef(t, svc, "app", 1, 1) @@ -224,6 +260,44 @@ func TestService_RecordTaskState_ReleasesCapacityOnStop(t *testing.T) { } } +// The AWS-API SubmitTaskStateChange path (gateway-routed agent) converges on the +// same task record + capacity release as the bus path, and resolves a task ARN. +func TestService_SubmitTaskStateChange_StopsTask(t *testing.T) { + svc, _ := newTestService(t) + _, err := svc.CreateCluster(&ecs.CreateClusterInput{ClusterName: aws.String("web")}, testAccountID) + require.NoError(t, err) + registerTaskDef(t, svc, "app", 128, 256) + registerInstance(t, svc, "web", "i-1", 1024, 2048) + out, err := svc.RunTask(&ecs.RunTaskInput{Cluster: aws.String("web"), TaskDefinition: aws.String("app")}, testAccountID) + require.NoError(t, err) + taskARN := aws.StringValue(out.Tasks[0].TaskArn) // full ARN exercises taskShortID + + exit := int64(0) + ack, err := svc.SubmitTaskStateChange(&ecs.SubmitTaskStateChangeInput{ + Cluster: aws.String("web"), Task: aws.String(taskARN), + Status: aws.String(bus.TaskStatusStopped), Reason: aws.String("exited"), + Containers: []*ecs.ContainerStateChange{ + {ContainerName: aws.String("app"), Status: aws.String(bus.TaskStatusStopped), + ExitCode: &exit, RuntimeId: aws.String("ctr-1")}, + }, + }, testAccountID) + require.NoError(t, err) + require.NotNil(t, ack.Acknowledgment) + + dt, err := svc.DescribeTasks(&ecs.DescribeTasksInput{ + Cluster: aws.String("web"), Tasks: []*string{aws.String(taskARN)}, + }, testAccountID) + require.NoError(t, err) + require.Len(t, dt.Tasks, 1) + assert.Equal(t, "STOPPED", aws.StringValue(dt.Tasks[0].LastStatus)) + + di, err := svc.DescribeContainerInstances(&ecs.DescribeContainerInstancesInput{ + Cluster: aws.String("web"), ContainerInstances: []*string{aws.String("i-1")}, + }, testAccountID) + require.NoError(t, err) + assert.Equal(t, int64(0), aws.Int64Value(di.ContainerInstances[0].RunningTasksCount)) +} + func TestService_RecordHeartbeat_UpdatesStatus(t *testing.T) { svc, _ := newTestService(t) registerInstance(t, svc, "web", "i-1", 1024, 2048) diff --git a/spinifex/handlers/ecs/store.go b/spinifex/handlers/ecs/store.go index 24408f616..6b2a6916a 100644 --- a/spinifex/handlers/ecs/store.go +++ b/spinifex/handlers/ecs/store.go @@ -68,6 +68,19 @@ func TaskKey(cluster, taskID string) string { return TasksPrefix(cluster) + taskID } +// AssignmentsPrefix returns the KV key prefix under which a container instance's +// pending task assignments (its "inbox") live. Kept in a sibling path to +// instances/ + tasks/ so listing those records never picks up assignment keys. +// The agent drains this inbox by polling the gateway (no direct NATS). +func AssignmentsPrefix(cluster, instanceID string) string { + return fmt.Sprintf("clusters/%s/assignments/%s/", cluster, instanceID) +} + +// AssignmentKey returns the KV key for one task's assignment in an instance inbox. +func AssignmentKey(cluster, instanceID, taskID string) string { + return AssignmentsPrefix(cluster, instanceID) + taskID +} + // ServicesPrefix returns the KV key prefix under which all of a cluster's service // records live. Used by ListServices to enumerate. func ServicesPrefix(cluster string) string { diff --git a/spinifex/handlers/ecs/tasks.go b/spinifex/handlers/ecs/tasks.go index e2baa78b7..5ff20cb8c 100644 --- a/spinifex/handlers/ecs/tasks.go +++ b/spinifex/handlers/ecs/tasks.go @@ -81,7 +81,7 @@ func (s *Service) RunTask(input *ecs.RunTaskInput, accountID string) (*ecs.RunTa if err := putJSON(kv, TaskKey(cluster, taskID), rec); err != nil { return nil, err } - if err := s.publishAssign(accountID, cluster, inst.InstanceID, rec, taskDef); err != nil { + if err := s.publishAssign(kv, accountID, cluster, inst.InstanceID, rec, taskDef); err != nil { slog.Error("ECS RunTask: failed to publish assign", "task", taskID, "instance", inst.InstanceID, "err", err) } out.Tasks = append(out.Tasks, s.taskToAWS(accountID, rec)) @@ -184,8 +184,11 @@ func (s *Service) newTaskRecord(accountID, cluster, taskID string, td *TaskDefRe return rec } -// publishAssign sends the task assignment to the instance's agent over the bus. -func (s *Service) publishAssign(accountID, cluster, instanceID string, rec *TaskRecord, td *TaskDefRecord) error { +// publishAssign writes the task assignment into the instance's KV inbox. The +// agent drains it by polling the gateway (PollAssignments) rather than +// subscribing to NATS, so the bus stays host-internal. Durable + restart-safe: +// an unacked assign survives an agent crash and is re-delivered on the next poll. +func (s *Service) publishAssign(kv nats.KeyValue, accountID, cluster, instanceID string, rec *TaskRecord, td *TaskDefRecord) error { msg := bus.Assign{ AccountID: accountID, ClusterName: cluster, @@ -203,12 +206,7 @@ func (s *Service) publishAssign(accountID, cluster, instanceID string, rec *Task for _, c := range td.Containers { msg.Containers = append(msg.Containers, c.toAssignContainer()) } - data, err := json.Marshal(&msg) - if err != nil { - return err - } - subj := bus.AssignSubject(accountID, cluster, instanceID) - return s.nc.Publish(subj, data) + return putJSON(kv, AssignmentKey(cluster, instanceID, rec.TaskID), &msg) } // DescribeTasks returns task records for the named tasks in a cluster.