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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user