diff --git a/internal/agent/conn/conn.go b/internal/agent/conn/conn.go index dff5318..9eafc2d 100644 --- a/internal/agent/conn/conn.go +++ b/internal/agent/conn/conn.go @@ -1,65 +1,97 @@ // Package conn keeps an agent connected to the master: it dials over pinned -// TLS, announces itself, sends its inventory, and streams changes until the -// connection drops, then reconnects with backoff. +// TLS, announces itself, sends its inventory, runs the jobs the master sends +// and streams changes and job results until the connection drops, then +// reconnects with backoff. package conn import ( "context" + "errors" "fmt" "log/slog" "net/http" + "sync" "time" "github.com/coder/websocket" "github.com/coder/websocket/wsjson" "github.com/shcizo/package-updater/internal/agent/inventory" + "github.com/shcizo/package-updater/internal/agent/updater" "github.com/shcizo/package-updater/internal/pin" "github.com/shcizo/package-updater/internal/protocol" ) -// readLimit caps a single incoming message. The library default (32 KiB) -// is too small for the master's future job messages; matches the master. +// readLimit caps a single incoming message. Matches the master. const readLimit = 1 << 20 // dialTimeout bounds the TLS + WebSocket handshake so a master that accepts // TCP but never answers does not hang the agent forever. const dialTimeout = 30 * time.Second +// writeTimeout bounds one websocket write so a stalled peer ends the session +// instead of blocking the writer forever. +const writeTimeout = 10 * time.Second + +// defaultProgressInterval is how often job output is batched to the master. +const defaultProgressInterval = 500 * time.Millisecond + // Source supplies what the agent reports. *inventory.Collector satisfies it. type Source interface { Hello(ctx context.Context) (protocol.Hello, error) Snapshot(ctx context.Context) ([]protocol.Service, error) } +// Jobs is the agent's job queue as the connection sees it. *updater.Runner +// satisfies it. +type Jobs interface { + // Submit queues a job; a job ID already held is ignored. + Submit(id string, start protocol.JobStart) + // Ack forgets a finished job the master has stored. + Ack(id string) + // Known lists every job ID held: queued, running, or finished unacked. + Known() []string + // TakeProgress returns output produced since the previous call. + TakeProgress() []updater.Progress + // Pending returns results the master has not acknowledged yet. + Pending() []updater.Finished + // Changed receives when there is new progress or a new result. + Changed() <-chan struct{} +} + // Config holds connection settings. Token is sent as a bearer header and must // never be logged. type Config struct { - URL string // wss://host:port/agent - Token string // AGENT_TOKEN - Fingerprint string // MASTER_KEY - Version string - MinBackoff time.Duration // default 1s - MaxBackoff time.Duration // default 30s + URL string // wss://host:port/agent + Token string // AGENT_TOKEN + Fingerprint string // MASTER_KEY + Version string + MinBackoff time.Duration // default 1s + MaxBackoff time.Duration // default 30s + ProgressInterval time.Duration // default 500ms } // Agent is the agent's side of the master connection. type Agent struct { cfg Config src Source + jobs Jobs log *slog.Logger changes chan struct{} } // New returns an Agent. Call Run to connect. -func New(cfg Config, src Source, log *slog.Logger) *Agent { +func New(cfg Config, src Source, jobs Jobs, log *slog.Logger) *Agent { if cfg.MinBackoff == 0 { cfg.MinBackoff = time.Second } if cfg.MaxBackoff == 0 { cfg.MaxBackoff = 30 * time.Second } - return &Agent{cfg: cfg, src: src, log: log, changes: make(chan struct{}, 1)} + if cfg.ProgressInterval == 0 { + cfg.ProgressInterval = defaultProgressInterval + } + return &Agent{cfg: cfg, src: src, jobs: jobs, log: log, changes: make(chan struct{}, 1)} } // Notify tells the agent its inventory may have changed. It never blocks; @@ -94,6 +126,9 @@ func (a *Agent) Run(ctx context.Context) error { // session runs one connection until it fails. connected reports whether the // hello was delivered, so Run resets its backoff only after real progress. +// +// Two goroutines serve the connection: a reader, which handles what the +// master sends, and a writer, the only goroutine that writes to the socket. func (a *Agent) session(parent context.Context) (connected bool, err error) { httpClient := &http.Client{Transport: &http.Transport{TLSClientConfig: pin.ClientConfig(a.cfg.Fingerprint)}} // The dial context only bounds the handshake; coder/websocket does not @@ -107,59 +142,140 @@ func (a *Agent) session(parent context.Context) (connected bool, err error) { if err != nil { return false, fmt.Errorf("dial master: %w", err) } - defer c.CloseNow() c.SetReadLimit(readLimit) - ctx, cancel := context.WithCancel(parent) - defer cancel() - // The master sends nothing yet in this phase, but reading is how a - // dropped connection is noticed, so a reader runs for the whole session. - go func() { - defer cancel() - for { - if _, _, err := c.Read(ctx); err != nil { - return - } + ctx, cancel := context.WithCancelCause(parent) + out := make(chan protocol.Envelope, 64) + var wg sync.WaitGroup + defer func() { + cancel(nil) + if parent.Err() != nil { + _ = c.Close(websocket.StatusNormalClosure, "") } + c.CloseNow() + wg.Wait() }() + wg.Add(2) + go func() { defer wg.Done(); a.writeLoop(ctx, cancel, c, out) }() + go func() { defer wg.Done(); a.readLoop(ctx, cancel, c) }() + + // send queues an envelope for the writer. It gives up when the session ends. + send := func(typ string, payload any) error { + env, err := protocol.Encode(typ, payload) + if err != nil { + return err + } + return enqueue(ctx, out, env) + } + sendJob := func(typ, id string, payload any) error { + env, err := protocol.EncodeJob(typ, id, payload) + if err != nil { + return err + } + return enqueue(ctx, out, env) + } + fail := func(connected bool, err error) (bool, error) { + if cause := context.Cause(ctx); cause != nil && !errors.Is(cause, context.Canceled) { + return connected, cause + } + return connected, err + } + hello, err := a.src.Hello(ctx) if err != nil { - return false, err + return fail(false, err) } hello.AgentVersion = a.cfg.Version - if err := send(ctx, c, protocol.TypeHello, hello); err != nil { - return false, err + hello.ProtocolVersion = protocol.Version + hello.KnownJobs = a.jobs.Known() + if err := send(protocol.TypeHello, hello); err != nil { + return fail(false, err) } last, err := a.src.Snapshot(ctx) if err != nil { - return true, err + return fail(true, err) } - if err := send(ctx, c, protocol.TypeSnapshot, protocol.Snapshot{Services: last}); err != nil { - return true, err + if err := send(protocol.TypeSnapshot, protocol.Snapshot{Services: last}); err != nil { + return fail(true, err) } - a.log.Info("master_connected", "services", len(last)) + a.log.Info("master_connected", "services", len(last), "known_jobs", len(hello.KnownJobs)) + sent := make(map[string]bool) // job results already sent on this connection + flush := func() error { + for _, p := range a.jobs.TakeProgress() { + if err := sendJob(protocol.TypeJobProgress, p.ID, protocol.JobProgress{Output: p.Output}); err != nil { + return err + } + } + for _, f := range a.jobs.Pending() { + if sent[f.ID] { + continue + } + if err := sendJob(protocol.TypeJobResult, f.ID, f.Result); err != nil { + return err + } + sent[f.ID] = true + } + return nil + } + + // flush is throttled to one call per ProgressInterval; a signal that + // arrives early arms the timer for the remainder instead of being lost. + lastFlush := time.Time{} + var timer *time.Timer + var timerC <-chan time.Time + defer func() { + if timer != nil { + timer.Stop() + } + }() + maybeFlush := func() error { + wait := a.cfg.ProgressInterval - time.Since(lastFlush) + if wait > 0 { + if timerC == nil { + timer = time.NewTimer(wait) + timerC = timer.C + } + return nil + } + lastFlush = time.Now() + return flush() + } + if err := maybeFlush(); err != nil { + return fail(true, err) + } + + jobsChanged := a.jobs.Changed() for { select { case <-ctx.Done(): - _ = c.Close(websocket.StatusNormalClosure, "") - return true, ctx.Err() + return fail(true, ctx.Err()) + case <-jobsChanged: + if err := maybeFlush(); err != nil { + return fail(true, err) + } + case <-timerC: + timerC = nil + lastFlush = time.Now() + if err := flush(); err != nil { + return fail(true, err) + } case <-a.changes: next, err := a.src.Snapshot(ctx) if err != nil { - return true, err + return fail(true, err) } upserts, removes := inventory.Diff(last, next) for _, s := range upserts { - if err := send(ctx, c, protocol.TypeUpsert, protocol.Upsert{Service: s}); err != nil { - return true, err + if err := send(protocol.TypeUpsert, protocol.Upsert{Service: s}); err != nil { + return fail(true, err) } } for _, k := range removes { - if err := send(ctx, c, protocol.TypeRemove, protocol.Remove{Key: k}); err != nil { - return true, err + if err := send(protocol.TypeRemove, protocol.Remove{Key: k}); err != nil { + return fail(true, err) } } last = next @@ -167,13 +283,58 @@ func (a *Agent) session(parent context.Context) (connected bool, err error) { } } -func send(ctx context.Context, c *websocket.Conn, typ string, payload any) error { - env, err := protocol.Encode(typ, payload) - if err != nil { - return err +// enqueue hands env to the writer, or gives up when the session ends. +func enqueue(ctx context.Context, out chan<- protocol.Envelope, env protocol.Envelope) error { + select { + case out <- env: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// writeLoop is the only goroutine that writes to c. A failed or timed-out +// write ends the session. +func (a *Agent) writeLoop(ctx context.Context, cancel context.CancelCauseFunc, c *websocket.Conn, out <-chan protocol.Envelope) { + for { + select { + case <-ctx.Done(): + return + case env := <-out: + wctx, wcancel := context.WithTimeout(ctx, writeTimeout) + err := wsjson.Write(wctx, c, env) + wcancel() + if err != nil { + cancel(fmt.Errorf("send %s: %w", env.Type, err)) + return + } + } + } +} + +// readLoop handles what the master sends. A message it cannot parse ends the +// session; a message type it does not know is logged and ignored, so a newer +// master can add message types without breaking this agent. +func (a *Agent) readLoop(ctx context.Context, cancel context.CancelCauseFunc, c *websocket.Conn) { + for { + var env protocol.Envelope + if err := wsjson.Read(ctx, c, &env); err != nil { + cancel(fmt.Errorf("read: %w", err)) + return + } + switch env.Type { + case protocol.TypeJobStart: + start, err := protocol.Decode[protocol.JobStart](env) + if err != nil || env.ID == "" { + cancel(fmt.Errorf("malformed job.start (id %q): %v", env.ID, err)) + return + } + a.jobs.Submit(env.ID, start) + case protocol.TypeJobAck: + // JobAck has no fields, so the payload is not inspected. + a.jobs.Ack(env.ID) + default: + a.log.Warn("unknown_message_type", "type", env.Type) + } } - if err := wsjson.Write(ctx, c, env); err != nil { - return fmt.Errorf("send %s: %w", typ, err) - } - return nil } diff --git a/internal/agent/conn/conn_test.go b/internal/agent/conn/conn_test.go index ac72afb..f4f2208 100644 --- a/internal/agent/conn/conn_test.go +++ b/internal/agent/conn/conn_test.go @@ -19,6 +19,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/shcizo/package-updater/internal/agent/updater" "github.com/shcizo/package-updater/internal/pin" "github.com/shcizo/package-updater/internal/protocol" ) @@ -113,12 +114,111 @@ var ( svcB = protocol.Service{Kind: protocol.KindContainer, ContainerName: "b", Image: "b:1", Status: "running"} ) +// fakeJobs records what the connection tells it and returns scripted state. +type fakeJobs struct { + mu sync.Mutex + known []string + pending []updater.Finished + progress []string // output chunks since the last TakeProgress + submits map[string]protocol.JobStart + acks []string + changed chan struct{} +} + +func newFakeJobs() *fakeJobs { + return &fakeJobs{submits: map[string]protocol.JobStart{}, changed: make(chan struct{}, 1)} +} + +func (f *fakeJobs) Submit(id string, s protocol.JobStart) { + f.mu.Lock() + defer f.mu.Unlock() + f.submits[id] = s +} + +func (f *fakeJobs) Ack(id string) { + f.mu.Lock() + defer f.mu.Unlock() + f.acks = append(f.acks, id) + kept := f.pending[:0:0] + for _, p := range f.pending { + if p.ID != id { + kept = append(kept, p) + } + } + f.pending = kept +} + +func (f *fakeJobs) Known() []string { + f.mu.Lock() + defer f.mu.Unlock() + return f.known +} + +func (f *fakeJobs) TakeProgress() []updater.Progress { + f.mu.Lock() + defer f.mu.Unlock() + if len(f.progress) == 0 { + return nil + } + out := updater.Progress{ID: "1", Output: strings.Join(f.progress, "")} + f.progress = nil + return []updater.Progress{out} +} + +func (f *fakeJobs) Pending() []updater.Finished { + f.mu.Lock() + defer f.mu.Unlock() + return append([]updater.Finished(nil), f.pending...) +} + +func (f *fakeJobs) Changed() <-chan struct{} { return f.changed } + +func (f *fakeJobs) signal() { + select { + case f.changed <- struct{}{}: + default: + } +} + +func (f *fakeJobs) addProgress(s string) { + f.mu.Lock() + f.progress = append(f.progress, s) + f.mu.Unlock() + f.signal() +} + +func (f *fakeJobs) addResult(id string, r protocol.JobResult) { + f.mu.Lock() + f.pending = append(f.pending, updater.Finished{ID: id, Result: r}) + f.mu.Unlock() + f.signal() +} + +func (f *fakeJobs) submitted(id string) (protocol.JobStart, bool) { + f.mu.Lock() + defer f.mu.Unlock() + s, ok := f.submits[id] + return s, ok +} + +func (f *fakeJobs) acked() []string { + f.mu.Lock() + defer f.mu.Unlock() + return append([]string(nil), f.acks...) +} + func startAgent(t *testing.T, fm *fakeMaster, fingerprint string, src Source) *Agent { + t.Helper() + return startAgentWith(t, fm, fingerprint, src, newFakeJobs(), 0) +} + +func startAgentWith(t *testing.T, fm *fakeMaster, fingerprint string, src Source, jobs Jobs, progressInterval time.Duration) *Agent { t.Helper() a := New(Config{ URL: fm.url(), Token: "pua_test", Fingerprint: fingerprint, Version: "v-test", MinBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond, - }, src, slog.New(slog.NewTextHandler(io.Discard, nil))) + ProgressInterval: progressInterval, + }, src, jobs, slog.New(slog.NewTextHandler(io.Discard, nil))) ctx, cancel := context.WithCancel(context.Background()) done := make(chan struct{}) go func() { _ = a.Run(ctx); close(done) }() @@ -137,7 +237,7 @@ func TestAgent_SendsHelloThenSnapshotWithToken(t *testing.T) { require.Equal(t, protocol.TypeHello, hello.Type) h, err := protocol.Decode[protocol.Hello](hello) require.NoError(t, err) - assert.Equal(t, protocol.Hello{Hostname: "node-2", AgentVersion: "v-test"}, h) + assert.Equal(t, protocol.Hello{Hostname: "node-2", AgentVersion: "v-test", ProtocolVersion: protocol.Version}, h) snap := fm.next(t) require.Equal(t, protocol.TypeSnapshot, snap.Type) @@ -198,3 +298,212 @@ func TestAgent_RefusesMasterWithWrongKey(t *testing.T) { time.Sleep(200 * time.Millisecond) assert.Empty(t, fm.auth, "the token must never be sent to a server with the wrong key") } + +func sendEnv(t *testing.T, c *websocket.Conn, env protocol.Envelope) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + require.NoError(t, wsjson.Write(ctx, c, env)) +} + +func TestAgent_HelloCarriesVersionAndKnownJobs(t *testing.T) { + fm := newFakeMaster(t) + jobs := newFakeJobs() + jobs.known = []string{"3", "7"} + startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond) + + h, err := protocol.Decode[protocol.Hello](fm.next(t)) + require.NoError(t, err) + assert.Equal(t, protocol.Version, h.ProtocolVersion) + assert.Equal(t, []string{"3", "7"}, h.KnownJobs) +} + +func TestAgent_JobStartSubmitsAndAckAcks(t *testing.T) { + fm := newFakeMaster(t) + jobs := newFakeJobs() + startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond) + fm.next(t) + fm.next(t) + c := <-fm.conns + + start := protocol.JobStart{Kind: protocol.KindSwarm, SwarmClusterID: "c1", SwarmService: "gitea_gitea", Ref: "gitea/gitea:1.27"} + env, err := protocol.EncodeJob(protocol.TypeJobStart, "42", start) + require.NoError(t, err) + sendEnv(t, c, env) + require.Eventually(t, func() bool { _, ok := jobs.submitted("42"); return ok }, 2*time.Second, 5*time.Millisecond) + got, _ := jobs.submitted("42") + assert.Equal(t, start, got) + + ack, err := protocol.EncodeJob(protocol.TypeJobAck, "42", protocol.JobAck{}) + require.NoError(t, err) + sendEnv(t, c, ack) + require.Eventually(t, func() bool { return len(jobs.acked()) == 1 }, 2*time.Second, 5*time.Millisecond) + assert.Equal(t, []string{"42"}, jobs.acked()) +} + +func TestAgent_SendsPendingResultsAfterSnapshot(t *testing.T) { + fm := newFakeMaster(t) + jobs := newFakeJobs() + jobs.pending = []updater.Finished{ + {ID: "5", Result: protocol.JobResult{Status: protocol.ResultSucceeded, Ref: "a:1@sha256:aa", Output: "ok\n"}}, + {ID: "6", Result: protocol.JobResult{Status: protocol.ResultFailed, Error: "boom"}}, + } + startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond) + + assert.Equal(t, protocol.TypeHello, fm.next(t).Type) + assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type) + r5 := fm.next(t) + assert.Equal(t, protocol.TypeJobResult, r5.Type) + assert.Equal(t, "5", r5.ID) + res, err := protocol.Decode[protocol.JobResult](r5) + require.NoError(t, err) + assert.Equal(t, "a:1@sha256:aa", res.Ref) + r6 := fm.next(t) + assert.Equal(t, "6", r6.ID) +} + +func TestAgent_ResendsUnackedResultAfterReconnect(t *testing.T) { + fm := newFakeMaster(t) + jobs := newFakeJobs() + jobs.addResult("5", protocol.JobResult{Status: protocol.ResultSucceeded}) + startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond) + + fm.next(t) // hello + fm.next(t) // snapshot + assert.Equal(t, "5", fm.next(t).ID) + (<-fm.conns).CloseNow() // dropped without ack + + assert.Equal(t, protocol.TypeHello, fm.next(t).Type) + assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type) + again := fm.next(t) + assert.Equal(t, protocol.TypeJobResult, again.Type) + assert.Equal(t, "5", again.ID) + + c := <-fm.conns + ack, err := protocol.EncodeJob(protocol.TypeJobAck, "5", protocol.JobAck{}) + require.NoError(t, err) + sendEnv(t, c, ack) + require.Eventually(t, func() bool { return len(jobs.acked()) == 1 }, 2*time.Second, 5*time.Millisecond) + c.CloseNow() + + assert.Equal(t, protocol.TypeHello, fm.next(t).Type) + assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type) + select { + case e := <-fm.envs: + t.Fatalf("acked result was sent again: %+v", e) + case <-time.After(100 * time.Millisecond): + } +} + +func TestAgent_ProgressIsThrottled(t *testing.T) { + fm := newFakeMaster(t) + jobs := newFakeJobs() + startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, 100*time.Millisecond) + fm.next(t) + fm.next(t) + + for range 20 { + jobs.addProgress("x") + time.Sleep(2500 * time.Microsecond) + } + var batches int + var total strings.Builder + collect := func(d time.Duration) { + deadline := time.After(d) + for { + select { + case e := <-fm.envs: + if e.Type != protocol.TypeJobProgress { + continue + } + p, err := protocol.Decode[protocol.JobProgress](e) + require.NoError(t, err) + batches++ + total.WriteString(p.Output) + case <-deadline: + return + } + } + } + collect(100 * time.Millisecond) + assert.LessOrEqual(t, batches, 2) + collect(400 * time.Millisecond) + assert.Equal(t, strings.Repeat("x", 20), total.String(), "the last chunk must be delivered") +} + +func TestAgent_IgnoresUnknownType(t *testing.T) { + fm := newFakeMaster(t) + jobs := newFakeJobs() + startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond) + fm.next(t) + fm.next(t) + c := <-fm.conns + + sendEnv(t, c, protocol.Envelope{Type: "future.thing", Payload: []byte(`{"a":1}`)}) + env, err := protocol.EncodeJob(protocol.TypeJobStart, "9", protocol.JobStart{Kind: protocol.KindCompose, Ref: "a:1"}) + require.NoError(t, err) + sendEnv(t, c, env) + require.Eventually(t, func() bool { _, ok := jobs.submitted("9"); return ok }, 2*time.Second, 5*time.Millisecond) + assert.Zero(t, len(fm.conns), "no reconnect expected") +} + +func TestAgent_ClosesOnMalformedJobStart(t *testing.T) { + fm := newFakeMaster(t) + startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, newFakeJobs(), time.Millisecond) + fm.next(t) + fm.next(t) + c := <-fm.conns + + env, err := protocol.Encode(protocol.TypeJobStart, protocol.JobStart{Ref: "a:1"}) // no ID + require.NoError(t, err) + sendEnv(t, c, env) + + select { + case <-fm.conns: + case <-time.After(2 * time.Second): + t.Fatal("agent did not reconnect after a malformed job.start") + } +} + +func TestAgent_InventoryAndProgressConcurrently(t *testing.T) { + fm := newFakeMaster(t) + src := &fakeSource{} + src.set(svcA) + jobs := newFakeJobs() + a := startAgentWith(t, fm, fm.fingerprint, src, jobs, time.Millisecond) + fm.next(t) + fm.next(t) + + done := make(chan struct{}) + go func() { + defer close(done) + for i := range 30 { + jobs.addProgress("y") + if i%2 == 0 { + src.set(svcB) + } else { + src.set(svcA) + } + a.Notify() + time.Sleep(time.Millisecond) + } + }() + <-done + var progress, inventory int + deadline := time.After(500 * time.Millisecond) + for { + select { + case e := <-fm.envs: + switch e.Type { + case protocol.TypeJobProgress: + progress++ + case protocol.TypeUpsert, protocol.TypeRemove: + inventory++ + } + case <-deadline: + assert.Positive(t, progress) + assert.Positive(t, inventory) + return + } + } +}