feat(master): route jobs to agents and reconcile them on reconnect (#7)
This commit is contained in:
+93
-11
@@ -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
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user