internal/master/hub/hub.go

@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/coder/websocket"
@@ -25,6 +26,13 @@ import (
 // small for a full snapshot from a busy machine.
 const readLimit = 1 << 20
 
+// writeTimeout bounds one websocket write. A peer that stops reading is cut
+// off after this long.
+const writeTimeout = 10 * time.Second
+
+// outQueue is how many messages may wait for a connection's writer.
+const outQueue = 64
+
 // Store is the storage the hub writes to. *store.Store satisfies it.
 type Store interface {
 	AgentByTokenHash(ctx context.Context, hash string) (model.Agent, error)
@@ -33,6 +41,9 @@ type Store interface {
 	UpsertService(ctx context.Context, agentID int64, clusterID string, svc protocol.Service, at time.Time) error
 	RemoveService(ctx context.Context, agentID int64, clusterID, agentKey string, at time.Time) error
 	MarkSeen(ctx context.Context, agentID int64, at time.Time) error
+	FailLostJobs(ctx context.Context, agentID int64, known []int64, at time.Time) ([]int64, error)
+	JobProgress(ctx context.Context, agentID, jobID int64, chunk string, at time.Time) (bool, error)
+	FinishJob(ctx context.Context, agentID, jobID int64, f model.JobFinish) (bool, error)
 }
 
 // Publisher receives "something changed" signals. *events.Broker satisfies it.
@@ -47,13 +58,57 @@ type Hub struct {
 	log   *slog.Logger
 	now   func() time.Time
 
+	connSeq atomic.Uint64
+
 	mu    sync.Mutex
-	conns map[int64]*websocket.Conn
+	conns map[int64]*agentConn
+}
+
+// agentConn is one live agent connection. Its writer goroutine is the only
+// code that writes to ws; everything else queues envelopes on out.
+type agentConn struct {
+	ws     *websocket.Conn
+	id     uint64
+	out    chan protocol.Envelope
+	ctx    context.Context
+	cancel context.CancelFunc
+	hello  *protocol.Hello // guarded by Hub.mu; nil until the hello is reconciled
+}
+
+// send queues env for the writer. It returns false if the connection ended
+// before the message could be queued.
+func (c *agentConn) send(env protocol.Envelope) bool {
+	select {
+	case c.out <- env:
+		return true
+	case <-c.ctx.Done():
+		return false
+	}
+}
+
+// writeLoop writes queued envelopes until the connection's context ends. A
+// failed or slow write closes the connection.
+func (c *agentConn) writeLoop() {
+	for {
+		select {
+		case <-c.ctx.Done():
+			return
+		case env := <-c.out:
+			wctx, cancel := context.WithTimeout(c.ctx, writeTimeout)
+			err := wsjson.Write(wctx, c.ws, env)
+			cancel()
+			if err != nil {
+				c.ws.CloseNow()
+				c.cancel()
+				return
+			}
+		}
+	}
 }
 
 // New returns a Hub. Mount it at GET /agent on the TLS listener.
 func New(store Store, pub Publisher, log *slog.Logger) *Hub {
-	return &Hub{store: store, pub: pub, log: log, now: time.Now, conns: make(map[int64]*websocket.Conn)}
+	return &Hub{store: store, pub: pub, log: log, now: time.Now, conns: make(map[int64]*agentConn)}
 }
 
 // ServeHTTP authenticates the agent's token and runs its connection. The token
@@ -78,8 +133,12 @@ func (h *Hub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	c.SetReadLimit(readLimit)
-	h.register(agent, r.RemoteAddr, c)
+	ctx, cancel := context.WithCancel(r.Context())
-	defer h.unregister(agent.ID, c)
+	defer cancel()
+	ac := &agentConn{ws: c, id: h.connSeq.Add(1), out: make(chan protocol.Envelope, outQueue), ctx: ctx, cancel: cancel}
+	go ac.writeLoop()
+	h.register(agent, r.RemoteAddr, ac)
+	defer h.unregister(agent.ID, ac)
 
 	// Revoking commits revoked_at before calling Disconnect. Re-checking after
 	// register closes the window where a revoke lands between the first lookup
@@ -90,8 +149,8 @@ func (h *Hub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	s := &session{hub: h, agent: agent}
+	s := &session{hub: h, agent: agent, conn: ac}
-	err = s.run(r.Context(), c)
+	err = s.run(ctx, c)
 	h.log.Info("agent_disconnected", "agent", agent.Name, "reason", err)
 }
 
@@ -110,25 +169,25 @@ func (h *Hub) Disconnect(agentID int64) {
 	c := h.conns[agentID]
 	h.mu.Unlock()
 	if c != nil {
-		c.CloseNow()
+		c.ws.CloseNow()
 	}
 }
 
 // register makes c the agent's live connection, closing any older one. A
 // replacement is logged: it is normal after a network blip, but repeated
 // replacements mean two agents share a token.
-func (h *Hub) register(agent model.Agent, remoteAddr string, c *websocket.Conn) {
+func (h *Hub) register(agent model.Agent, remoteAddr string, c *agentConn) {
 	h.mu.Lock()
 	old := h.conns[agent.ID]
 	h.conns[agent.ID] = c
 	h.mu.Unlock()
 	if old != nil {
 		h.log.Warn("agent_replaced", "agent_id", agent.ID, "agent", agent.Name, "remote_addr", remoteAddr)
-		old.CloseNow()
+		old.ws.CloseNow()
 	}
 }
 
-func (h *Hub) unregister(id int64, c *websocket.Conn) {
+func (h *Hub) unregister(id int64, c *agentConn) {
 	h.mu.Lock()
 	current := h.conns[id] == c
 	if current {
@@ -148,6 +207,7 @@ func (h *Hub) unregister(id int64, c *websocket.Conn) {
 type session struct {
 	hub   *Hub
 	agent model.Agent
+	conn  *agentConn
 	hello *protocol.Hello
 }
 
@@ -188,7 +248,25 @@ func (s *session) handle(ctx context.Context, env protocol.Envelope) error {
 		if err := h.store.RecordHello(ctx, s.agent.ID, v, h.now()); err != nil {
 			return err
 		}
+		// Reconcile jobs before the connection becomes routable, so a job is
+		// never sent to an agent whose old jobs have not been settled yet.
+		known := parseJobIDs(v.KnownJobs, func(id string) {
+			h.log.Warn("known_job_not_numeric", "agent", s.agent.Name, "job_id", id)
+		})
+		lost, err := h.store.FailLostJobs(ctx, s.agent.ID, known, h.now())
+		if err != nil {
+			return err
+		}
+		if len(lost) > 0 {
+			h.pub.Publish(events.Jobs)
+			for _, id := range lost {
+				h.pub.Publish(events.JobTopic(id))
+			}
+		}
 		s.hello = &v
+		h.mu.Lock()
+		s.conn.hello = &v
+		h.mu.Unlock()
 		h.log.Info("agent_connected", "agent", s.agent.Name, "hostname", v.Hostname, "swarm_manager", v.SwarmManager)
 		h.pub.Publish(events.Agents)
 	case protocol.TypeSnapshot:
@@ -224,8 +302,12 @@ func (s *session) handle(ctx context.Context, env protocol.Envelope) error {
 			return err
 		}
 		h.pub.Publish(events.Services)
+	case protocol.TypeJobProgress:
+		return s.handleJobProgress(ctx, env)
+	case protocol.TypeJobResult:
+		return s.handleJobResult(ctx, env)
 	default:
-		return fmt.Errorf("unknown message type %q", env.Type)
+		h.log.Warn("unknown_message_type", "agent", s.agent.Name, "type", env.Type)
 	}
 	return nil
 }

internal/master/hub/hub_test.go

@@ -3,6 +3,7 @@ package hub
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"io"
 	"log/slog"
 	"net/http"
@@ -35,6 +36,18 @@ type fakeStore struct {
 	// revokeFromLookup, when > 0, makes AgentByTokenHash return the agent as
 	// revoked from that (1-based) lookup onwards.
 	revokeFromLookup int
+
+	// Job calls, recorded; the *Result fields script the answers.
+	lostCalls    [][]int64
+	lostResult   []int64
+	lostErr      error
+	lostGate     chan struct{} // when set, FailLostJobs blocks until it is closed
+	progress     []string      // "agent|job|chunk"
+	progressOK   bool
+	finishes     []model.JobFinish
+	finishAgents []int64
+	finishOK     bool
+	finishErr    error
 }
 
 func newFakeStore(agents ...model.Agent) *fakeStore {
@@ -94,6 +107,31 @@ func (f *fakeStore) MarkSeen(_ context.Context, id int64, _ time.Time) error {
 	return nil
 }
 
+func (f *fakeStore) FailLostJobs(_ context.Context, _ int64, known []int64, _ time.Time) ([]int64, error) {
+	if f.lostGate != nil {
+		<-f.lostGate
+	}
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.lostCalls = append(f.lostCalls, known)
+	return f.lostResult, f.lostErr
+}
+
+func (f *fakeStore) JobProgress(_ context.Context, agentID, jobID int64, chunk string, _ time.Time) (bool, error) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.progress = append(f.progress, fmt.Sprintf("%d|%d|%s", agentID, jobID, chunk))
+	return f.progressOK, nil
+}
+
+func (f *fakeStore) FinishJob(_ context.Context, agentID, _ int64, fin model.JobFinish) (bool, error) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.finishAgents = append(f.finishAgents, agentID)
+	f.finishes = append(f.finishes, fin)
+	return f.finishOK, f.finishErr
+}
+
 func (f *fakeStore) count(get func(*fakeStore) int) func() int {
 	return func() int {
 		f.mu.Lock()

internal/master/hub/jobs.go

@@ -0,0 +1,133 @@
+package hub
+
+import (
+	"context"
+	"strconv"
+	"strings"
+
+	"github.com/shcizo/package-updater/internal/master/events"
+	"github.com/shcizo/package-updater/internal/master/model"
+	"github.com/shcizo/package-updater/internal/protocol"
+)
+
+// StartJob sends job.start over exactly the connection r names. It returns
+// model.ErrAgentOffline if that connection is gone, was replaced by a newer
+// one, or has not finished its hello.
+func (h *Hub) StartJob(r model.Route, jobID int64, start protocol.JobStart) error {
+	h.mu.Lock()
+	c := h.conns[r.AgentID]
+	usable := c != nil && c.id == r.ConnID && c.hello != nil
+	h.mu.Unlock()
+	if !usable {
+		return model.ErrAgentOffline
+	}
+	env, err := protocol.EncodeJob(protocol.TypeJobStart, strconv.FormatInt(jobID, 10), start)
+	if err != nil {
+		return err
+	}
+	// A job whose write later fails stays queued; the agent's next hello
+	// settles it as lost.
+	if !c.send(env) {
+		return model.ErrAgentOffline
+	}
+	return nil
+}
+
+// parseJobIDs converts the agent's job IDs to integers, calling skip for each
+// one that is not a number.
+func parseJobIDs(ids []string, skip func(id string)) []int64 {
+	out := make([]int64, 0, len(ids))
+	for _, s := range ids {
+		id, err := strconv.ParseInt(s, 10, 64)
+		if err != nil {
+			skip(s)
+			continue
+		}
+		out = append(out, id)
+	}
+	return out
+}
+
+// cleanOutput drops NUL bytes, which Postgres text columns reject. Without
+// this a result containing one could never be stored and would be resent
+// forever.
+func cleanOutput(s string) string { return strings.ReplaceAll(s, "\x00", "") }
+
+func (s *session) handleJobProgress(ctx context.Context, env protocol.Envelope) error {
+	h := s.hub
+	v, err := protocol.Decode[protocol.JobProgress](env)
+	if err != nil {
+		return err
+	}
+	id, err := strconv.ParseInt(env.ID, 10, 64)
+	if err != nil {
+		h.log.Warn("job_progress_ignored", "agent", s.agent.Name, "job_id", env.ID, "reason", "not a number")
+		return nil
+	}
+	applied, err := h.store.JobProgress(ctx, s.agent.ID, id, cleanOutput(v.Output), h.now())
+	if err != nil {
+		return err
+	}
+	if applied {
+		h.pub.Publish(events.Jobs)
+		h.pub.Publish(events.JobTopic(id))
+	}
+	return nil
+}
+
+func (s *session) handleJobResult(ctx context.Context, env protocol.Envelope) error {
+	h := s.hub
+	v, err := protocol.Decode[protocol.JobResult](env)
+	if err != nil {
+		return err
+	}
+	ack, err := protocol.EncodeJob(protocol.TypeJobAck, env.ID, protocol.JobAck{})
+	if err != nil {
+		return err
+	}
+	id, err := strconv.ParseInt(env.ID, 10, 64)
+	if err != nil {
+		h.log.Warn("job_result_ignored", "agent", s.agent.Name, "job_id", env.ID, "reason", "not a number")
+		s.conn.send(ack)
+		return nil
+	}
+
+	fin := model.JobFinish{
+		Status:          model.JobStatus(v.Status),
+		Ref:             v.Ref,
+		Error:           v.Error,
+		Output:          cleanOutput(v.Output),
+		OutputTruncated: v.OutputTruncated,
+		StartedAt:       v.StartedAt,
+		FinishedAt:      v.FinishedAt,
+	}
+	if v.Status != protocol.ResultSucceeded && v.Status != protocol.ResultFailed {
+		fin.Status = model.JobFailed
+		fin.Ref = ""
+		fin.Error = "agent sent invalid status " + strconv.Quote(v.Status)
+	}
+	if fin.StartedAt.IsZero() {
+		fin.StartedAt = h.now()
+	}
+	if fin.FinishedAt.IsZero() {
+		fin.FinishedAt = h.now()
+	}
+
+	// A store error closes the connection without an ack, so the agent
+	// resends the result after it reconnects.
+	applied, err := h.store.FinishJob(ctx, s.agent.ID, id, fin)
+	if err != nil {
+		return err
+	}
+	if applied {
+		h.pub.Publish(events.Jobs)
+		h.pub.Publish(events.JobTopic(id))
+		if fin.Status == model.JobSucceeded {
+			h.pub.Publish(events.Services)
+		}
+	} else {
+		h.log.Warn("job_result_ignored", "agent", s.agent.Name, "job_id", env.ID, "reason", "no such active job for this agent")
+	}
+	s.conn.send(ack)
+	return nil
+}

internal/master/hub/jobs_test.go

@@ -0,0 +1,270 @@
+package hub
+
+import (
+	"context"
+	"strconv"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/coder/websocket"
+	"github.com/coder/websocket/wsjson"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/shcizo/package-updater/internal/master/events"
+	"github.com/shcizo/package-updater/internal/master/model"
+	"github.com/shcizo/package-updater/internal/protocol"
+)
+
+var t0 = time.Date(2026, 10, 12, 12, 0, 0, 0, time.UTC)
+
+func sendJob(t *testing.T, c *websocket.Conn, typ, id string, payload any) {
+	t.Helper()
+	env, err := protocol.EncodeJob(typ, id, payload)
+	require.NoError(t, err)
+	require.NoError(t, wsjson.Write(context.Background(), c, env))
+}
+
+func readEnv(t *testing.T, c *websocket.Conn) protocol.Envelope {
+	t.Helper()
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	var env protocol.Envelope
+	require.NoError(t, wsjson.Read(ctx, c, &env))
+	return env
+}
+
+func waitRoute(t *testing.T, h *Hub) model.Route {
+	t.Helper()
+	svc := model.Service{Kind: protocol.KindCompose, AgentID: ptr(node2.ID)}
+	var r model.Route
+	require.Eventually(t, func() bool {
+		var err error
+		r, err = h.Route(svc)
+		return err == nil
+	}, 2*time.Second, 5*time.Millisecond)
+	return r
+}
+
+func TestHub_ReconcilesBeforeRoutable(t *testing.T) {
+	store := newFakeStore(node2)
+	store.lostGate = make(chan struct{})
+	store.lostResult = []int64{5}
+	h, pub, srv := newHub(t, store)
+	c, _, err := dial(t, srv, "pua_node-2")
+	require.NoError(t, err)
+	defer c.CloseNow()
+
+	send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "node-2", ProtocolVersion: 2, KnownJobs: []string{"7", "9", "x"}})
+	svc := model.Service{Kind: protocol.KindCompose, AgentID: ptr(node2.ID)}
+	time.Sleep(50 * time.Millisecond) // hello is now blocked inside FailLostJobs
+	_, err = h.Route(svc)
+	assert.ErrorIs(t, err, model.ErrAgentOffline)
+
+	close(store.lostGate)
+	waitRoute(t, h)
+
+	store.mu.Lock()
+	defer store.mu.Unlock()
+	assert.Equal(t, [][]int64{{7, 9}}, store.lostCalls)
+	assert.True(t, pub.has(events.Jobs))
+	assert.True(t, pub.has(events.JobTopic(5)))
+}
+
+func TestHub_StartJobSendsEnvelope(t *testing.T) {
+	store := newFakeStore(node2)
+	h, _, srv := newHub(t, store)
+	c, _, err := dial(t, srv, "pua_node-2")
+	require.NoError(t, err)
+	defer c.CloseNow()
+	send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
+	r := waitRoute(t, h)
+
+	start := protocol.JobStart{Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web", Ref: "r:1"}
+	require.NoError(t, h.StartJob(r, 42, start))
+
+	env := readEnv(t, c)
+	assert.Equal(t, protocol.TypeJobStart, env.Type)
+	assert.Equal(t, "42", env.ID)
+	got, err := protocol.Decode[protocol.JobStart](env)
+	require.NoError(t, err)
+	assert.Equal(t, start, got)
+}
+
+func TestHub_StartJobOnReplacedConnectionFails(t *testing.T) {
+	store := newFakeStore(node2)
+	h, _, srv := newHub(t, store)
+	first, _, err := dial(t, srv, "pua_node-2")
+	require.NoError(t, err)
+	defer first.CloseNow()
+	send(t, first, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
+	old := waitRoute(t, h)
+
+	second, _, err := dial(t, srv, "pua_node-2")
+	require.NoError(t, err)
+	defer second.CloseNow()
+
+	// Until the new connection says hello, nothing is routable and the old
+	// route is refused.
+	assert.ErrorIs(t, h.StartJob(old, 1, protocol.JobStart{}), model.ErrAgentOffline)
+	send(t, second, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
+	fresh := waitRoute(t, h)
+	require.NotEqual(t, old.ConnID, fresh.ConnID)
+	assert.ErrorIs(t, h.StartJob(old, 1, protocol.JobStart{}), model.ErrAgentOffline)
+
+	require.NoError(t, h.StartJob(fresh, 2, protocol.JobStart{Ref: "r:2"}))
+	env := readEnv(t, second)
+	assert.Equal(t, "2", env.ID, "the new connection got only the job sent to it")
+}
+
+func TestHub_ProgressAndResult(t *testing.T) {
+	store := newFakeStore(node2)
+	store.progressOK, store.finishOK = true, true
+	_, pub, srv := newHub(t, store)
+	c, _, err := dial(t, srv, "pua_node-2")
+	require.NoError(t, err)
+	defer c.CloseNow()
+	send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
+
+	sendJob(t, c, protocol.TypeJobProgress, "42", protocol.JobProgress{Output: "out"})
+	require.Eventually(t, func() bool { return pub.has(events.JobTopic(42)) }, time.Second, 5*time.Millisecond)
+	store.mu.Lock()
+	assert.Equal(t, []string{"2|42|out"}, store.progress)
+	store.mu.Unlock()
+
+	pub.reset()
+	res := protocol.JobResult{Status: protocol.ResultSucceeded, Ref: "r:1@sha256:d", Output: "done", StartedAt: t0, FinishedAt: t0}
+	sendJob(t, c, protocol.TypeJobResult, "42", res)
+	ack := readEnv(t, c)
+	assert.Equal(t, protocol.TypeJobAck, ack.Type)
+	assert.Equal(t, "42", ack.ID)
+	assert.True(t, pub.has(events.Services))
+	assert.True(t, pub.has(events.JobTopic(42)))
+	store.mu.Lock()
+	require.Len(t, store.finishes, 1)
+	assert.Equal(t, model.JobSucceeded, store.finishes[0].Status)
+	assert.Equal(t, "r:1@sha256:d", store.finishes[0].Ref)
+	store.mu.Unlock()
+
+	// Not applied (already recorded, or not this agent's): still acked.
+	store.mu.Lock()
+	store.finishOK = false
+	store.mu.Unlock()
+	pub.reset()
+	sendJob(t, c, protocol.TypeJobResult, "43", res)
+	assert.Equal(t, "43", readEnv(t, c).ID)
+	assert.False(t, pub.has(events.JobTopic(43)))
+
+	// Non-numeric ID: acked, no store call.
+	store.mu.Lock()
+	before := len(store.finishes)
+	store.mu.Unlock()
+	sendJob(t, c, protocol.TypeJobResult, "abc", res)
+	assert.Equal(t, "abc", readEnv(t, c).ID)
+	store.mu.Lock()
+	assert.Len(t, store.finishes, before)
+	store.mu.Unlock()
+
+	// Unknown status is stored as a failure.
+	store.mu.Lock()
+	store.finishOK = true
+	store.mu.Unlock()
+	sendJob(t, c, protocol.TypeJobResult, "44", protocol.JobResult{Status: "weird", Ref: "r:9"})
+	assert.Equal(t, "44", readEnv(t, c).ID)
+	store.mu.Lock()
+	defer store.mu.Unlock()
+	last := store.finishes[len(store.finishes)-1]
+	assert.Equal(t, model.JobFailed, last.Status)
+	assert.Equal(t, `agent sent invalid status "weird"`, last.Error)
+	assert.Empty(t, last.Ref)
+}
+
+func TestHub_ResultStoreErrorClosesConnection(t *testing.T) {
+	store := newFakeStore(node2)
+	store.finishErr = assert.AnError
+	_, _, srv := newHub(t, store)
+	c, _, err := dial(t, srv, "pua_node-2")
+	require.NoError(t, err)
+	defer c.CloseNow()
+	send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
+	sendJob(t, c, protocol.TypeJobResult, "42", protocol.JobResult{Status: protocol.ResultFailed})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	_, _, err = c.Read(ctx)
+	assert.Equal(t, websocket.StatusPolicyViolation, websocket.CloseStatus(err), "no ack, connection closed")
+}
+
+func TestHub_IgnoresUnknownType(t *testing.T) {
+	store := newFakeStore(node2)
+	_, _, srv := newHub(t, store)
+	c, _, err := dial(t, srv, "pua_node-2")
+	require.NoError(t, err)
+	defer c.CloseNow()
+	send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
+	require.NoError(t, wsjson.Write(context.Background(), c, protocol.Envelope{Type: "future.thing"}))
+	send(t, c, protocol.TypeUpsert, protocol.Upsert{Service: protocol.Service{Kind: protocol.KindContainer, ContainerName: "y"}})
+
+	require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.upserts) })() == 1 }, time.Second, 5*time.Millisecond)
+}
+
+func TestHub_LargeResultArrives(t *testing.T) {
+	store := newFakeStore(node2)
+	store.finishOK = true
+	_, _, srv := newHub(t, store)
+	c, _, err := dial(t, srv, "pua_node-2")
+	require.NoError(t, err)
+	defer c.CloseNow()
+	send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
+
+	out := strings.Repeat("\x1b", protocol.MaxJobOutput)
+	sendJob(t, c, protocol.TypeJobResult, "42", protocol.JobResult{Status: protocol.ResultFailed, Output: out})
+	assert.Equal(t, "42", readEnv(t, c).ID)
+	store.mu.Lock()
+	defer store.mu.Unlock()
+	require.Len(t, store.finishes, 1)
+	assert.Equal(t, out, store.finishes[0].Output)
+}
+
+func TestHub_ConcurrentWritesAreSerialised(t *testing.T) {
+	const n = 50
+	store := newFakeStore(node2)
+	store.finishOK = true
+	h, _, srv := newHub(t, store)
+	c, _, err := dial(t, srv, "pua_node-2")
+	require.NoError(t, err)
+	defer c.CloseNow()
+	send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
+	r := waitRoute(t, h)
+
+	var wg sync.WaitGroup
+	for i := 0; i < n; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			assert.NoError(t, h.StartJob(r, int64(1000+i), protocol.JobStart{Ref: "r:1"}))
+		}()
+	}
+	go func() {
+		for i := 0; i < n; i++ {
+			sendJob(t, c, protocol.TypeJobResult, strconv.Itoa(i), protocol.JobResult{Status: protocol.ResultFailed})
+		}
+	}()
+
+	starts, acks := 0, 0
+	for starts+acks < 2*n {
+		switch env := readEnv(t, c); env.Type {
+		case protocol.TypeJobStart:
+			starts++
+		case protocol.TypeJobAck:
+			acks++
+		default:
+			t.Fatalf("unexpected %q", env.Type)
+		}
+	}
+	wg.Wait()
+	assert.Equal(t, n, starts)
+	assert.Equal(t, n, acks)
+}

internal/master/hub/route.go

@@ -0,0 +1,69 @@
+package hub
+
+import (
+	"sort"
+
+	"github.com/shcizo/package-updater/internal/master/model"
+	"github.com/shcizo/package-updater/internal/protocol"
+)
+
+// candidate is a connection that has finished its hello and may get jobs.
+type candidate struct {
+	agentID int64
+	connID  uint64
+	hello   protocol.Hello
+}
+
+// Route picks the connection a job for svc goes to. It returns
+// model.ErrAgentOffline or model.ErrAgentTooOld when there is none.
+func (h *Hub) Route(svc model.Service) (model.Route, error) {
+	h.mu.Lock()
+	conns := make([]candidate, 0, len(h.conns))
+	for id, c := range h.conns {
+		if c.hello != nil {
+			conns = append(conns, candidate{agentID: id, connID: c.id, hello: *c.hello})
+		}
+	}
+	h.mu.Unlock()
+	return route(svc, conns)
+}
+
+// route picks the connection a job for svc goes to. Compose services and
+// containers can only run on the agent that reported them. Swarm services
+// run on the reporting agent if it qualifies, else on any other connected
+// manager of the same cluster (lowest agent ID first). Too-old wins over
+// offline only when a connection would otherwise have qualified.
+func route(svc model.Service, conns []candidate) (model.Route, error) {
+	sort.Slice(conns, func(a, b int) bool { return conns[a].agentID < conns[b].agentID })
+	tooOld := false
+	try := func(c candidate) bool {
+		if svc.Kind == protocol.KindSwarm {
+			if !c.hello.SwarmManager || c.hello.SwarmClusterID == "" || c.hello.SwarmClusterID != svc.SwarmClusterID {
+				return false
+			}
+		}
+		if c.hello.ProtocolVersion < protocol.MinJobVersion {
+			tooOld = true
+			return false
+		}
+		return true
+	}
+	if svc.AgentID != nil {
+		for _, c := range conns {
+			if c.agentID == *svc.AgentID && try(c) {
+				return model.Route{AgentID: c.agentID, ConnID: c.connID}, nil
+			}
+		}
+	}
+	if svc.Kind == protocol.KindSwarm {
+		for _, c := range conns {
+			if try(c) {
+				return model.Route{AgentID: c.agentID, ConnID: c.connID}, nil
+			}
+		}
+	}
+	if tooOld {
+		return model.Route{}, model.ErrAgentTooOld
+	}
+	return model.Route{}, model.ErrAgentOffline
+}

internal/master/hub/route_test.go

@@ -0,0 +1,50 @@
+package hub
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/shcizo/package-updater/internal/master/model"
+	"github.com/shcizo/package-updater/internal/protocol"
+)
+
+func TestRoute(t *testing.T) {
+	composeSvc := model.Service{Kind: protocol.KindCompose, AgentID: ptr(int64(2))}
+	swarmSvc := model.Service{Kind: protocol.KindSwarm, AgentID: ptr(int64(1)), SwarmClusterID: "c1"}
+	v2 := protocol.Hello{ProtocolVersion: protocol.Version}
+	mgrV2 := protocol.Hello{ProtocolVersion: protocol.Version, SwarmManager: true, SwarmClusterID: "c1"}
+
+	tests := []struct {
+		name  string
+		svc   model.Service
+		conns []candidate
+		want  model.Route
+		err   error
+	}{
+		{"compose on own agent", composeSvc, []candidate{{2, 20, v2}}, model.Route{AgentID: 2, ConnID: 20}, nil},
+		{"compose, own agent too old", composeSvc, []candidate{{2, 20, protocol.Hello{}}}, model.Route{}, model.ErrAgentTooOld},
+		{"compose, own agent missing", composeSvc, nil, model.Route{}, model.ErrAgentOffline},
+		{"compose, only another agent online", composeSvc, []candidate{{5, 50, v2}}, model.Route{}, model.ErrAgentOffline},
+		{"swarm on reporting agent", swarmSvc, []candidate{{1, 10, mgrV2}, {3, 30, mgrV2}}, model.Route{AgentID: 1, ConnID: 10}, nil},
+		{"swarm, reporting agent offline, other manager", swarmSvc, []candidate{{3, 30, mgrV2}}, model.Route{AgentID: 3, ConnID: 30}, nil},
+		{"swarm, lowest agent id first", swarmSvc, []candidate{{4, 40, mgrV2}, {3, 30, mgrV2}}, model.Route{AgentID: 3, ConnID: 30}, nil},
+		{"swarm, other manager in another cluster", swarmSvc,
+			[]candidate{{3, 30, protocol.Hello{ProtocolVersion: 2, SwarmManager: true, SwarmClusterID: "c2"}}}, model.Route{}, model.ErrAgentOffline},
+		{"swarm, only a v0 manager of the cluster", swarmSvc,
+			[]candidate{{1, 10, protocol.Hello{SwarmManager: true, SwarmClusterID: "c1"}}}, model.Route{}, model.ErrAgentTooOld},
+		{"swarm, v0 manager but a v2 manager exists", swarmSvc,
+			[]candidate{{1, 10, protocol.Hello{SwarmManager: true, SwarmClusterID: "c1"}}, {3, 30, mgrV2}}, model.Route{AgentID: 3, ConnID: 30}, nil},
+		{"swarm, worker of the cluster", swarmSvc,
+			[]candidate{{1, 10, protocol.Hello{ProtocolVersion: 2, SwarmClusterID: "c1"}}}, model.Route{}, model.ErrAgentOffline},
+		{"swarm, manager with empty cluster id", model.Service{Kind: protocol.KindSwarm},
+			[]candidate{{1, 10, protocol.Hello{ProtocolVersion: 2, SwarmManager: true}}}, model.Route{}, model.ErrAgentOffline},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := route(tt.svc, tt.conns)
+			assert.ErrorIs(t, err, tt.err)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}