feat(agent): run jobs from the master and resend results until acked (#7)
This commit is contained in:
+200
-39
@@ -1,37 +1,64 @@
|
|||||||
// Package conn keeps an agent connected to the master: it dials over pinned
|
// Package conn keeps an agent connected to the master: it dials over pinned
|
||||||
// TLS, announces itself, sends its inventory, and streams changes until the
|
// TLS, announces itself, sends its inventory, runs the jobs the master sends
|
||||||
// connection drops, then reconnects with backoff.
|
// and streams changes and job results until the connection drops, then
|
||||||
|
// reconnects with backoff.
|
||||||
package conn
|
package conn
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/coder/websocket"
|
"github.com/coder/websocket"
|
||||||
"github.com/coder/websocket/wsjson"
|
"github.com/coder/websocket/wsjson"
|
||||||
|
|
||||||
"github.com/shcizo/package-updater/internal/agent/inventory"
|
"github.com/shcizo/package-updater/internal/agent/inventory"
|
||||||
|
"github.com/shcizo/package-updater/internal/agent/updater"
|
||||||
"github.com/shcizo/package-updater/internal/pin"
|
"github.com/shcizo/package-updater/internal/pin"
|
||||||
"github.com/shcizo/package-updater/internal/protocol"
|
"github.com/shcizo/package-updater/internal/protocol"
|
||||||
)
|
)
|
||||||
|
|
||||||
// readLimit caps a single incoming message. The library default (32 KiB)
|
// readLimit caps a single incoming message. Matches the master.
|
||||||
// is too small for the master's future job messages; matches the master.
|
|
||||||
const readLimit = 1 << 20
|
const readLimit = 1 << 20
|
||||||
|
|
||||||
// dialTimeout bounds the TLS + WebSocket handshake so a master that accepts
|
// dialTimeout bounds the TLS + WebSocket handshake so a master that accepts
|
||||||
// TCP but never answers does not hang the agent forever.
|
// TCP but never answers does not hang the agent forever.
|
||||||
const dialTimeout = 30 * time.Second
|
const dialTimeout = 30 * time.Second
|
||||||
|
|
||||||
|
// writeTimeout bounds one websocket write so a stalled peer ends the session
|
||||||
|
// instead of blocking the writer forever.
|
||||||
|
const writeTimeout = 10 * time.Second
|
||||||
|
|
||||||
|
// defaultProgressInterval is how often job output is batched to the master.
|
||||||
|
const defaultProgressInterval = 500 * time.Millisecond
|
||||||
|
|
||||||
// Source supplies what the agent reports. *inventory.Collector satisfies it.
|
// Source supplies what the agent reports. *inventory.Collector satisfies it.
|
||||||
type Source interface {
|
type Source interface {
|
||||||
Hello(ctx context.Context) (protocol.Hello, error)
|
Hello(ctx context.Context) (protocol.Hello, error)
|
||||||
Snapshot(ctx context.Context) ([]protocol.Service, error)
|
Snapshot(ctx context.Context) ([]protocol.Service, error)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Jobs is the agent's job queue as the connection sees it. *updater.Runner
|
||||||
|
// satisfies it.
|
||||||
|
type Jobs interface {
|
||||||
|
// Submit queues a job; a job ID already held is ignored.
|
||||||
|
Submit(id string, start protocol.JobStart)
|
||||||
|
// Ack forgets a finished job the master has stored.
|
||||||
|
Ack(id string)
|
||||||
|
// Known lists every job ID held: queued, running, or finished unacked.
|
||||||
|
Known() []string
|
||||||
|
// TakeProgress returns output produced since the previous call.
|
||||||
|
TakeProgress() []updater.Progress
|
||||||
|
// Pending returns results the master has not acknowledged yet.
|
||||||
|
Pending() []updater.Finished
|
||||||
|
// Changed receives when there is new progress or a new result.
|
||||||
|
Changed() <-chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
// Config holds connection settings. Token is sent as a bearer header and must
|
// Config holds connection settings. Token is sent as a bearer header and must
|
||||||
// never be logged.
|
// never be logged.
|
||||||
type Config struct {
|
type Config struct {
|
||||||
@@ -41,25 +68,30 @@ type Config struct {
|
|||||||
Version string
|
Version string
|
||||||
MinBackoff time.Duration // default 1s
|
MinBackoff time.Duration // default 1s
|
||||||
MaxBackoff time.Duration // default 30s
|
MaxBackoff time.Duration // default 30s
|
||||||
|
ProgressInterval time.Duration // default 500ms
|
||||||
}
|
}
|
||||||
|
|
||||||
// Agent is the agent's side of the master connection.
|
// Agent is the agent's side of the master connection.
|
||||||
type Agent struct {
|
type Agent struct {
|
||||||
cfg Config
|
cfg Config
|
||||||
src Source
|
src Source
|
||||||
|
jobs Jobs
|
||||||
log *slog.Logger
|
log *slog.Logger
|
||||||
changes chan struct{}
|
changes chan struct{}
|
||||||
}
|
}
|
||||||
|
|
||||||
// New returns an Agent. Call Run to connect.
|
// New returns an Agent. Call Run to connect.
|
||||||
func New(cfg Config, src Source, log *slog.Logger) *Agent {
|
func New(cfg Config, src Source, jobs Jobs, log *slog.Logger) *Agent {
|
||||||
if cfg.MinBackoff == 0 {
|
if cfg.MinBackoff == 0 {
|
||||||
cfg.MinBackoff = time.Second
|
cfg.MinBackoff = time.Second
|
||||||
}
|
}
|
||||||
if cfg.MaxBackoff == 0 {
|
if cfg.MaxBackoff == 0 {
|
||||||
cfg.MaxBackoff = 30 * time.Second
|
cfg.MaxBackoff = 30 * time.Second
|
||||||
}
|
}
|
||||||
return &Agent{cfg: cfg, src: src, log: log, changes: make(chan struct{}, 1)}
|
if cfg.ProgressInterval == 0 {
|
||||||
|
cfg.ProgressInterval = defaultProgressInterval
|
||||||
|
}
|
||||||
|
return &Agent{cfg: cfg, src: src, jobs: jobs, log: log, changes: make(chan struct{}, 1)}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Notify tells the agent its inventory may have changed. It never blocks;
|
// Notify tells the agent its inventory may have changed. It never blocks;
|
||||||
@@ -94,6 +126,9 @@ func (a *Agent) Run(ctx context.Context) error {
|
|||||||
|
|
||||||
// session runs one connection until it fails. connected reports whether the
|
// session runs one connection until it fails. connected reports whether the
|
||||||
// hello was delivered, so Run resets its backoff only after real progress.
|
// hello was delivered, so Run resets its backoff only after real progress.
|
||||||
|
//
|
||||||
|
// Two goroutines serve the connection: a reader, which handles what the
|
||||||
|
// master sends, and a writer, the only goroutine that writes to the socket.
|
||||||
func (a *Agent) session(parent context.Context) (connected bool, err error) {
|
func (a *Agent) session(parent context.Context) (connected bool, err error) {
|
||||||
httpClient := &http.Client{Transport: &http.Transport{TLSClientConfig: pin.ClientConfig(a.cfg.Fingerprint)}}
|
httpClient := &http.Client{Transport: &http.Transport{TLSClientConfig: pin.ClientConfig(a.cfg.Fingerprint)}}
|
||||||
// The dial context only bounds the handshake; coder/websocket does not
|
// The dial context only bounds the handshake; coder/websocket does not
|
||||||
@@ -107,59 +142,140 @@ func (a *Agent) session(parent context.Context) (connected bool, err error) {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return false, fmt.Errorf("dial master: %w", err)
|
return false, fmt.Errorf("dial master: %w", err)
|
||||||
}
|
}
|
||||||
defer c.CloseNow()
|
|
||||||
c.SetReadLimit(readLimit)
|
c.SetReadLimit(readLimit)
|
||||||
|
|
||||||
ctx, cancel := context.WithCancel(parent)
|
ctx, cancel := context.WithCancelCause(parent)
|
||||||
defer cancel()
|
out := make(chan protocol.Envelope, 64)
|
||||||
// The master sends nothing yet in this phase, but reading is how a
|
var wg sync.WaitGroup
|
||||||
// dropped connection is noticed, so a reader runs for the whole session.
|
defer func() {
|
||||||
go func() {
|
cancel(nil)
|
||||||
defer cancel()
|
if parent.Err() != nil {
|
||||||
for {
|
_ = c.Close(websocket.StatusNormalClosure, "")
|
||||||
if _, _, err := c.Read(ctx); err != nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
c.CloseNow()
|
||||||
|
wg.Wait()
|
||||||
}()
|
}()
|
||||||
|
|
||||||
|
wg.Add(2)
|
||||||
|
go func() { defer wg.Done(); a.writeLoop(ctx, cancel, c, out) }()
|
||||||
|
go func() { defer wg.Done(); a.readLoop(ctx, cancel, c) }()
|
||||||
|
|
||||||
|
// send queues an envelope for the writer. It gives up when the session ends.
|
||||||
|
send := func(typ string, payload any) error {
|
||||||
|
env, err := protocol.Encode(typ, payload)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return enqueue(ctx, out, env)
|
||||||
|
}
|
||||||
|
sendJob := func(typ, id string, payload any) error {
|
||||||
|
env, err := protocol.EncodeJob(typ, id, payload)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return enqueue(ctx, out, env)
|
||||||
|
}
|
||||||
|
fail := func(connected bool, err error) (bool, error) {
|
||||||
|
if cause := context.Cause(ctx); cause != nil && !errors.Is(cause, context.Canceled) {
|
||||||
|
return connected, cause
|
||||||
|
}
|
||||||
|
return connected, err
|
||||||
|
}
|
||||||
|
|
||||||
hello, err := a.src.Hello(ctx)
|
hello, err := a.src.Hello(ctx)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return false, err
|
return fail(false, err)
|
||||||
}
|
}
|
||||||
hello.AgentVersion = a.cfg.Version
|
hello.AgentVersion = a.cfg.Version
|
||||||
if err := send(ctx, c, protocol.TypeHello, hello); err != nil {
|
hello.ProtocolVersion = protocol.Version
|
||||||
return false, err
|
hello.KnownJobs = a.jobs.Known()
|
||||||
|
if err := send(protocol.TypeHello, hello); err != nil {
|
||||||
|
return fail(false, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
last, err := a.src.Snapshot(ctx)
|
last, err := a.src.Snapshot(ctx)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return true, err
|
return fail(true, err)
|
||||||
}
|
}
|
||||||
if err := send(ctx, c, protocol.TypeSnapshot, protocol.Snapshot{Services: last}); err != nil {
|
if err := send(protocol.TypeSnapshot, protocol.Snapshot{Services: last}); err != nil {
|
||||||
return true, err
|
return fail(true, err)
|
||||||
}
|
}
|
||||||
a.log.Info("master_connected", "services", len(last))
|
a.log.Info("master_connected", "services", len(last), "known_jobs", len(hello.KnownJobs))
|
||||||
|
|
||||||
|
sent := make(map[string]bool) // job results already sent on this connection
|
||||||
|
flush := func() error {
|
||||||
|
for _, p := range a.jobs.TakeProgress() {
|
||||||
|
if err := sendJob(protocol.TypeJobProgress, p.ID, protocol.JobProgress{Output: p.Output}); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, f := range a.jobs.Pending() {
|
||||||
|
if sent[f.ID] {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if err := sendJob(protocol.TypeJobResult, f.ID, f.Result); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
sent[f.ID] = true
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// flush is throttled to one call per ProgressInterval; a signal that
|
||||||
|
// arrives early arms the timer for the remainder instead of being lost.
|
||||||
|
lastFlush := time.Time{}
|
||||||
|
var timer *time.Timer
|
||||||
|
var timerC <-chan time.Time
|
||||||
|
defer func() {
|
||||||
|
if timer != nil {
|
||||||
|
timer.Stop()
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
maybeFlush := func() error {
|
||||||
|
wait := a.cfg.ProgressInterval - time.Since(lastFlush)
|
||||||
|
if wait > 0 {
|
||||||
|
if timerC == nil {
|
||||||
|
timer = time.NewTimer(wait)
|
||||||
|
timerC = timer.C
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
lastFlush = time.Now()
|
||||||
|
return flush()
|
||||||
|
}
|
||||||
|
if err := maybeFlush(); err != nil {
|
||||||
|
return fail(true, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
jobsChanged := a.jobs.Changed()
|
||||||
for {
|
for {
|
||||||
select {
|
select {
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
_ = c.Close(websocket.StatusNormalClosure, "")
|
return fail(true, ctx.Err())
|
||||||
return true, ctx.Err()
|
case <-jobsChanged:
|
||||||
|
if err := maybeFlush(); err != nil {
|
||||||
|
return fail(true, err)
|
||||||
|
}
|
||||||
|
case <-timerC:
|
||||||
|
timerC = nil
|
||||||
|
lastFlush = time.Now()
|
||||||
|
if err := flush(); err != nil {
|
||||||
|
return fail(true, err)
|
||||||
|
}
|
||||||
case <-a.changes:
|
case <-a.changes:
|
||||||
next, err := a.src.Snapshot(ctx)
|
next, err := a.src.Snapshot(ctx)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return true, err
|
return fail(true, err)
|
||||||
}
|
}
|
||||||
upserts, removes := inventory.Diff(last, next)
|
upserts, removes := inventory.Diff(last, next)
|
||||||
for _, s := range upserts {
|
for _, s := range upserts {
|
||||||
if err := send(ctx, c, protocol.TypeUpsert, protocol.Upsert{Service: s}); err != nil {
|
if err := send(protocol.TypeUpsert, protocol.Upsert{Service: s}); err != nil {
|
||||||
return true, err
|
return fail(true, err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
for _, k := range removes {
|
for _, k := range removes {
|
||||||
if err := send(ctx, c, protocol.TypeRemove, protocol.Remove{Key: k}); err != nil {
|
if err := send(protocol.TypeRemove, protocol.Remove{Key: k}); err != nil {
|
||||||
return true, err
|
return fail(true, err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
last = next
|
last = next
|
||||||
@@ -167,13 +283,58 @@ func (a *Agent) session(parent context.Context) (connected bool, err error) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func send(ctx context.Context, c *websocket.Conn, typ string, payload any) error {
|
// enqueue hands env to the writer, or gives up when the session ends.
|
||||||
env, err := protocol.Encode(typ, payload)
|
func enqueue(ctx context.Context, out chan<- protocol.Envelope, env protocol.Envelope) error {
|
||||||
if err != nil {
|
select {
|
||||||
return err
|
case out <- env:
|
||||||
}
|
|
||||||
if err := wsjson.Write(ctx, c, env); err != nil {
|
|
||||||
return fmt.Errorf("send %s: %w", typ, err)
|
|
||||||
}
|
|
||||||
return nil
|
return nil
|
||||||
|
case <-ctx.Done():
|
||||||
|
return ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// writeLoop is the only goroutine that writes to c. A failed or timed-out
|
||||||
|
// write ends the session.
|
||||||
|
func (a *Agent) writeLoop(ctx context.Context, cancel context.CancelCauseFunc, c *websocket.Conn, out <-chan protocol.Envelope) {
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return
|
||||||
|
case env := <-out:
|
||||||
|
wctx, wcancel := context.WithTimeout(ctx, writeTimeout)
|
||||||
|
err := wsjson.Write(wctx, c, env)
|
||||||
|
wcancel()
|
||||||
|
if err != nil {
|
||||||
|
cancel(fmt.Errorf("send %s: %w", env.Type, err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// readLoop handles what the master sends. A message it cannot parse ends the
|
||||||
|
// session; a message type it does not know is logged and ignored, so a newer
|
||||||
|
// master can add message types without breaking this agent.
|
||||||
|
func (a *Agent) readLoop(ctx context.Context, cancel context.CancelCauseFunc, c *websocket.Conn) {
|
||||||
|
for {
|
||||||
|
var env protocol.Envelope
|
||||||
|
if err := wsjson.Read(ctx, c, &env); err != nil {
|
||||||
|
cancel(fmt.Errorf("read: %w", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
switch env.Type {
|
||||||
|
case protocol.TypeJobStart:
|
||||||
|
start, err := protocol.Decode[protocol.JobStart](env)
|
||||||
|
if err != nil || env.ID == "" {
|
||||||
|
cancel(fmt.Errorf("malformed job.start (id %q): %v", env.ID, err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
a.jobs.Submit(env.ID, start)
|
||||||
|
case protocol.TypeJobAck:
|
||||||
|
// JobAck has no fields, so the payload is not inspected.
|
||||||
|
a.jobs.Ack(env.ID)
|
||||||
|
default:
|
||||||
|
a.log.Warn("unknown_message_type", "type", env.Type)
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -19,6 +19,7 @@ import (
|
|||||||
"github.com/stretchr/testify/assert"
|
"github.com/stretchr/testify/assert"
|
||||||
"github.com/stretchr/testify/require"
|
"github.com/stretchr/testify/require"
|
||||||
|
|
||||||
|
"github.com/shcizo/package-updater/internal/agent/updater"
|
||||||
"github.com/shcizo/package-updater/internal/pin"
|
"github.com/shcizo/package-updater/internal/pin"
|
||||||
"github.com/shcizo/package-updater/internal/protocol"
|
"github.com/shcizo/package-updater/internal/protocol"
|
||||||
)
|
)
|
||||||
@@ -113,12 +114,111 @@ var (
|
|||||||
svcB = protocol.Service{Kind: protocol.KindContainer, ContainerName: "b", Image: "b:1", Status: "running"}
|
svcB = protocol.Service{Kind: protocol.KindContainer, ContainerName: "b", Image: "b:1", Status: "running"}
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// fakeJobs records what the connection tells it and returns scripted state.
|
||||||
|
type fakeJobs struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
known []string
|
||||||
|
pending []updater.Finished
|
||||||
|
progress []string // output chunks since the last TakeProgress
|
||||||
|
submits map[string]protocol.JobStart
|
||||||
|
acks []string
|
||||||
|
changed chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func newFakeJobs() *fakeJobs {
|
||||||
|
return &fakeJobs{submits: map[string]protocol.JobStart{}, changed: make(chan struct{}, 1)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeJobs) Submit(id string, s protocol.JobStart) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
f.submits[id] = s
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeJobs) Ack(id string) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
f.acks = append(f.acks, id)
|
||||||
|
kept := f.pending[:0:0]
|
||||||
|
for _, p := range f.pending {
|
||||||
|
if p.ID != id {
|
||||||
|
kept = append(kept, p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
f.pending = kept
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeJobs) Known() []string {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
return f.known
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeJobs) TakeProgress() []updater.Progress {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
if len(f.progress) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
out := updater.Progress{ID: "1", Output: strings.Join(f.progress, "")}
|
||||||
|
f.progress = nil
|
||||||
|
return []updater.Progress{out}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeJobs) Pending() []updater.Finished {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
return append([]updater.Finished(nil), f.pending...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeJobs) Changed() <-chan struct{} { return f.changed }
|
||||||
|
|
||||||
|
func (f *fakeJobs) signal() {
|
||||||
|
select {
|
||||||
|
case f.changed <- struct{}{}:
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeJobs) addProgress(s string) {
|
||||||
|
f.mu.Lock()
|
||||||
|
f.progress = append(f.progress, s)
|
||||||
|
f.mu.Unlock()
|
||||||
|
f.signal()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeJobs) addResult(id string, r protocol.JobResult) {
|
||||||
|
f.mu.Lock()
|
||||||
|
f.pending = append(f.pending, updater.Finished{ID: id, Result: r})
|
||||||
|
f.mu.Unlock()
|
||||||
|
f.signal()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeJobs) submitted(id string) (protocol.JobStart, bool) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
s, ok := f.submits[id]
|
||||||
|
return s, ok
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeJobs) acked() []string {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
return append([]string(nil), f.acks...)
|
||||||
|
}
|
||||||
|
|
||||||
func startAgent(t *testing.T, fm *fakeMaster, fingerprint string, src Source) *Agent {
|
func startAgent(t *testing.T, fm *fakeMaster, fingerprint string, src Source) *Agent {
|
||||||
|
t.Helper()
|
||||||
|
return startAgentWith(t, fm, fingerprint, src, newFakeJobs(), 0)
|
||||||
|
}
|
||||||
|
|
||||||
|
func startAgentWith(t *testing.T, fm *fakeMaster, fingerprint string, src Source, jobs Jobs, progressInterval time.Duration) *Agent {
|
||||||
t.Helper()
|
t.Helper()
|
||||||
a := New(Config{
|
a := New(Config{
|
||||||
URL: fm.url(), Token: "pua_test", Fingerprint: fingerprint, Version: "v-test",
|
URL: fm.url(), Token: "pua_test", Fingerprint: fingerprint, Version: "v-test",
|
||||||
MinBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond,
|
MinBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond,
|
||||||
}, src, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
ProgressInterval: progressInterval,
|
||||||
|
}, src, jobs, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||||
ctx, cancel := context.WithCancel(context.Background())
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
done := make(chan struct{})
|
done := make(chan struct{})
|
||||||
go func() { _ = a.Run(ctx); close(done) }()
|
go func() { _ = a.Run(ctx); close(done) }()
|
||||||
@@ -137,7 +237,7 @@ func TestAgent_SendsHelloThenSnapshotWithToken(t *testing.T) {
|
|||||||
require.Equal(t, protocol.TypeHello, hello.Type)
|
require.Equal(t, protocol.TypeHello, hello.Type)
|
||||||
h, err := protocol.Decode[protocol.Hello](hello)
|
h, err := protocol.Decode[protocol.Hello](hello)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
assert.Equal(t, protocol.Hello{Hostname: "node-2", AgentVersion: "v-test"}, h)
|
assert.Equal(t, protocol.Hello{Hostname: "node-2", AgentVersion: "v-test", ProtocolVersion: protocol.Version}, h)
|
||||||
|
|
||||||
snap := fm.next(t)
|
snap := fm.next(t)
|
||||||
require.Equal(t, protocol.TypeSnapshot, snap.Type)
|
require.Equal(t, protocol.TypeSnapshot, snap.Type)
|
||||||
@@ -198,3 +298,212 @@ func TestAgent_RefusesMasterWithWrongKey(t *testing.T) {
|
|||||||
time.Sleep(200 * time.Millisecond)
|
time.Sleep(200 * time.Millisecond)
|
||||||
assert.Empty(t, fm.auth, "the token must never be sent to a server with the wrong key")
|
assert.Empty(t, fm.auth, "the token must never be sent to a server with the wrong key")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func sendEnv(t *testing.T, c *websocket.Conn, env protocol.Envelope) {
|
||||||
|
t.Helper()
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
require.NoError(t, wsjson.Write(ctx, c, env))
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgent_HelloCarriesVersionAndKnownJobs(t *testing.T) {
|
||||||
|
fm := newFakeMaster(t)
|
||||||
|
jobs := newFakeJobs()
|
||||||
|
jobs.known = []string{"3", "7"}
|
||||||
|
startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
|
||||||
|
|
||||||
|
h, err := protocol.Decode[protocol.Hello](fm.next(t))
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, protocol.Version, h.ProtocolVersion)
|
||||||
|
assert.Equal(t, []string{"3", "7"}, h.KnownJobs)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgent_JobStartSubmitsAndAckAcks(t *testing.T) {
|
||||||
|
fm := newFakeMaster(t)
|
||||||
|
jobs := newFakeJobs()
|
||||||
|
startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
|
||||||
|
fm.next(t)
|
||||||
|
fm.next(t)
|
||||||
|
c := <-fm.conns
|
||||||
|
|
||||||
|
start := protocol.JobStart{Kind: protocol.KindSwarm, SwarmClusterID: "c1", SwarmService: "gitea_gitea", Ref: "gitea/gitea:1.27"}
|
||||||
|
env, err := protocol.EncodeJob(protocol.TypeJobStart, "42", start)
|
||||||
|
require.NoError(t, err)
|
||||||
|
sendEnv(t, c, env)
|
||||||
|
require.Eventually(t, func() bool { _, ok := jobs.submitted("42"); return ok }, 2*time.Second, 5*time.Millisecond)
|
||||||
|
got, _ := jobs.submitted("42")
|
||||||
|
assert.Equal(t, start, got)
|
||||||
|
|
||||||
|
ack, err := protocol.EncodeJob(protocol.TypeJobAck, "42", protocol.JobAck{})
|
||||||
|
require.NoError(t, err)
|
||||||
|
sendEnv(t, c, ack)
|
||||||
|
require.Eventually(t, func() bool { return len(jobs.acked()) == 1 }, 2*time.Second, 5*time.Millisecond)
|
||||||
|
assert.Equal(t, []string{"42"}, jobs.acked())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgent_SendsPendingResultsAfterSnapshot(t *testing.T) {
|
||||||
|
fm := newFakeMaster(t)
|
||||||
|
jobs := newFakeJobs()
|
||||||
|
jobs.pending = []updater.Finished{
|
||||||
|
{ID: "5", Result: protocol.JobResult{Status: protocol.ResultSucceeded, Ref: "a:1@sha256:aa", Output: "ok\n"}},
|
||||||
|
{ID: "6", Result: protocol.JobResult{Status: protocol.ResultFailed, Error: "boom"}},
|
||||||
|
}
|
||||||
|
startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
|
||||||
|
|
||||||
|
assert.Equal(t, protocol.TypeHello, fm.next(t).Type)
|
||||||
|
assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type)
|
||||||
|
r5 := fm.next(t)
|
||||||
|
assert.Equal(t, protocol.TypeJobResult, r5.Type)
|
||||||
|
assert.Equal(t, "5", r5.ID)
|
||||||
|
res, err := protocol.Decode[protocol.JobResult](r5)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, "a:1@sha256:aa", res.Ref)
|
||||||
|
r6 := fm.next(t)
|
||||||
|
assert.Equal(t, "6", r6.ID)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgent_ResendsUnackedResultAfterReconnect(t *testing.T) {
|
||||||
|
fm := newFakeMaster(t)
|
||||||
|
jobs := newFakeJobs()
|
||||||
|
jobs.addResult("5", protocol.JobResult{Status: protocol.ResultSucceeded})
|
||||||
|
startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
|
||||||
|
|
||||||
|
fm.next(t) // hello
|
||||||
|
fm.next(t) // snapshot
|
||||||
|
assert.Equal(t, "5", fm.next(t).ID)
|
||||||
|
(<-fm.conns).CloseNow() // dropped without ack
|
||||||
|
|
||||||
|
assert.Equal(t, protocol.TypeHello, fm.next(t).Type)
|
||||||
|
assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type)
|
||||||
|
again := fm.next(t)
|
||||||
|
assert.Equal(t, protocol.TypeJobResult, again.Type)
|
||||||
|
assert.Equal(t, "5", again.ID)
|
||||||
|
|
||||||
|
c := <-fm.conns
|
||||||
|
ack, err := protocol.EncodeJob(protocol.TypeJobAck, "5", protocol.JobAck{})
|
||||||
|
require.NoError(t, err)
|
||||||
|
sendEnv(t, c, ack)
|
||||||
|
require.Eventually(t, func() bool { return len(jobs.acked()) == 1 }, 2*time.Second, 5*time.Millisecond)
|
||||||
|
c.CloseNow()
|
||||||
|
|
||||||
|
assert.Equal(t, protocol.TypeHello, fm.next(t).Type)
|
||||||
|
assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type)
|
||||||
|
select {
|
||||||
|
case e := <-fm.envs:
|
||||||
|
t.Fatalf("acked result was sent again: %+v", e)
|
||||||
|
case <-time.After(100 * time.Millisecond):
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgent_ProgressIsThrottled(t *testing.T) {
|
||||||
|
fm := newFakeMaster(t)
|
||||||
|
jobs := newFakeJobs()
|
||||||
|
startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, 100*time.Millisecond)
|
||||||
|
fm.next(t)
|
||||||
|
fm.next(t)
|
||||||
|
|
||||||
|
for range 20 {
|
||||||
|
jobs.addProgress("x")
|
||||||
|
time.Sleep(2500 * time.Microsecond)
|
||||||
|
}
|
||||||
|
var batches int
|
||||||
|
var total strings.Builder
|
||||||
|
collect := func(d time.Duration) {
|
||||||
|
deadline := time.After(d)
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case e := <-fm.envs:
|
||||||
|
if e.Type != protocol.TypeJobProgress {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
p, err := protocol.Decode[protocol.JobProgress](e)
|
||||||
|
require.NoError(t, err)
|
||||||
|
batches++
|
||||||
|
total.WriteString(p.Output)
|
||||||
|
case <-deadline:
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
collect(100 * time.Millisecond)
|
||||||
|
assert.LessOrEqual(t, batches, 2)
|
||||||
|
collect(400 * time.Millisecond)
|
||||||
|
assert.Equal(t, strings.Repeat("x", 20), total.String(), "the last chunk must be delivered")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgent_IgnoresUnknownType(t *testing.T) {
|
||||||
|
fm := newFakeMaster(t)
|
||||||
|
jobs := newFakeJobs()
|
||||||
|
startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
|
||||||
|
fm.next(t)
|
||||||
|
fm.next(t)
|
||||||
|
c := <-fm.conns
|
||||||
|
|
||||||
|
sendEnv(t, c, protocol.Envelope{Type: "future.thing", Payload: []byte(`{"a":1}`)})
|
||||||
|
env, err := protocol.EncodeJob(protocol.TypeJobStart, "9", protocol.JobStart{Kind: protocol.KindCompose, Ref: "a:1"})
|
||||||
|
require.NoError(t, err)
|
||||||
|
sendEnv(t, c, env)
|
||||||
|
require.Eventually(t, func() bool { _, ok := jobs.submitted("9"); return ok }, 2*time.Second, 5*time.Millisecond)
|
||||||
|
assert.Zero(t, len(fm.conns), "no reconnect expected")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgent_ClosesOnMalformedJobStart(t *testing.T) {
|
||||||
|
fm := newFakeMaster(t)
|
||||||
|
startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, newFakeJobs(), time.Millisecond)
|
||||||
|
fm.next(t)
|
||||||
|
fm.next(t)
|
||||||
|
c := <-fm.conns
|
||||||
|
|
||||||
|
env, err := protocol.Encode(protocol.TypeJobStart, protocol.JobStart{Ref: "a:1"}) // no ID
|
||||||
|
require.NoError(t, err)
|
||||||
|
sendEnv(t, c, env)
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-fm.conns:
|
||||||
|
case <-time.After(2 * time.Second):
|
||||||
|
t.Fatal("agent did not reconnect after a malformed job.start")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgent_InventoryAndProgressConcurrently(t *testing.T) {
|
||||||
|
fm := newFakeMaster(t)
|
||||||
|
src := &fakeSource{}
|
||||||
|
src.set(svcA)
|
||||||
|
jobs := newFakeJobs()
|
||||||
|
a := startAgentWith(t, fm, fm.fingerprint, src, jobs, time.Millisecond)
|
||||||
|
fm.next(t)
|
||||||
|
fm.next(t)
|
||||||
|
|
||||||
|
done := make(chan struct{})
|
||||||
|
go func() {
|
||||||
|
defer close(done)
|
||||||
|
for i := range 30 {
|
||||||
|
jobs.addProgress("y")
|
||||||
|
if i%2 == 0 {
|
||||||
|
src.set(svcB)
|
||||||
|
} else {
|
||||||
|
src.set(svcA)
|
||||||
|
}
|
||||||
|
a.Notify()
|
||||||
|
time.Sleep(time.Millisecond)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
<-done
|
||||||
|
var progress, inventory int
|
||||||
|
deadline := time.After(500 * time.Millisecond)
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case e := <-fm.envs:
|
||||||
|
switch e.Type {
|
||||||
|
case protocol.TypeJobProgress:
|
||||||
|
progress++
|
||||||
|
case protocol.TypeUpsert, protocol.TypeRemove:
|
||||||
|
inventory++
|
||||||
|
}
|
||||||
|
case <-deadline:
|
||||||
|
assert.Positive(t, progress)
|
||||||
|
assert.Positive(t, inventory)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user