From 68e5c3372212981d05f392e96068f2c3f4eefdd5 Mon Sep 17 00:00:00 2001 From: Samuel Enocsson Date: Mon, 12 Oct 2026 18:42:30 +0200 Subject: [PATCH] feat(master): route jobs to agents and reconcile them on reconnect (#7) --- .../{compose_executor.go => compose.go} | 0 .../updater/{swarm_executor.go => swarm.go} | 0 .../{swarm_executor_test.go => swarm_test.go} | 0 internal/master/hub/hub.go | 104 ++++++- internal/master/hub/hub_test.go | 38 +++ internal/master/hub/jobs.go | 133 +++++++++ internal/master/hub/jobs_test.go | 270 ++++++++++++++++++ internal/master/hub/route.go | 69 +++++ internal/master/hub/route_test.go | 50 ++++ 9 files changed, 653 insertions(+), 11 deletions(-) rename internal/agent/updater/{compose_executor.go => compose.go} (100%) rename internal/agent/updater/{swarm_executor.go => swarm.go} (100%) rename internal/agent/updater/{swarm_executor_test.go => swarm_test.go} (100%) create mode 100644 internal/master/hub/jobs.go create mode 100644 internal/master/hub/jobs_test.go create mode 100644 internal/master/hub/route.go create mode 100644 internal/master/hub/route_test.go diff --git a/internal/agent/updater/compose_executor.go b/internal/agent/updater/compose.go similarity index 100% rename from internal/agent/updater/compose_executor.go rename to internal/agent/updater/compose.go diff --git a/internal/agent/updater/swarm_executor.go b/internal/agent/updater/swarm.go similarity index 100% rename from internal/agent/updater/swarm_executor.go rename to internal/agent/updater/swarm.go diff --git a/internal/agent/updater/swarm_executor_test.go b/internal/agent/updater/swarm_test.go similarity index 100% rename from internal/agent/updater/swarm_executor_test.go rename to internal/agent/updater/swarm_test.go diff --git a/internal/master/hub/hub.go b/internal/master/hub/hub.go index 8270976..da69147 100644 --- a/internal/master/hub/hub.go +++ b/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) - defer h.unregister(agent.ID, c) + ctx, cancel := context.WithCancel(r.Context()) + 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} - err = s.run(r.Context(), c) + s := &session{hub: h, agent: agent, conn: ac} + 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 } diff --git a/internal/master/hub/hub_test.go b/internal/master/hub/hub_test.go index 756a986..4fddadc 100644 --- a/internal/master/hub/hub_test.go +++ b/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() diff --git a/internal/master/hub/jobs.go b/internal/master/hub/jobs.go new file mode 100644 index 0000000..a8c24b3 --- /dev/null +++ b/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 +} diff --git a/internal/master/hub/jobs_test.go b/internal/master/hub/jobs_test.go new file mode 100644 index 0000000..c4e7983 --- /dev/null +++ b/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) +} diff --git a/internal/master/hub/route.go b/internal/master/hub/route.go new file mode 100644 index 0000000..6013ee7 --- /dev/null +++ b/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 +} diff --git a/internal/master/hub/route_test.go b/internal/master/hub/route_test.go new file mode 100644 index 0000000..1ad86d8 --- /dev/null +++ b/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) + }) + } +}