feat(agent): run jobs from the master and resend results until acked (#7)

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
 }