internal/e2e/agent_master_test.go

@@ -21,6 +21,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/shcizo/package-updater/internal/agent/conn"
+	"github.com/shcizo/package-updater/internal/agent/updater"
 	"github.com/shcizo/package-updater/internal/agenttoken"
 	"github.com/shcizo/package-updater/internal/master/events"
 	"github.com/shcizo/package-updater/internal/master/hub"
@@ -35,6 +36,129 @@ type memStore struct {
 	lookups   int
 	snapshots int
 	services  map[string]protocol.Service
+
+	// Jobs: in-memory twins of the store methods in C5, with the same
+	// ownership and "only active jobs change" rules.
+	svcRows       map[int64]model.Service
+	jobs          map[int64]*model.Job
+	nextJob       int64
+	finishApplied int // FinishJob calls that returned true
+}
+
+// seedService adds a service row that jobs can be triggered for.
+func (m *memStore) seedService(svc model.Service) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.svcRows == nil {
+		m.svcRows = map[int64]model.Service{}
+	}
+	m.svcRows[svc.ID] = svc
+}
+
+func (m *memStore) GetService(_ context.Context, id int64) (model.Service, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	svc, ok := m.svcRows[id]
+	if !ok {
+		return model.Service{}, model.ErrNotFound
+	}
+	return svc, nil
+}
+
+func (m *memStore) CreateJob(_ context.Context, nj model.NewJob) (model.Job, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.jobs == nil {
+		m.jobs = map[int64]*model.Job{}
+	}
+	for _, j := range m.jobs {
+		if j.ServiceID == nj.ServiceID && j.Active() {
+			return *j, model.ErrJobActive
+		}
+	}
+	m.nextJob++
+	j := &model.Job{ID: m.nextJob, ServiceID: nj.ServiceID, AgentID: nj.AgentID, FromRef: nj.FromRef, ToRef: nj.ToRef,
+		Trigger: nj.Trigger, RequestedBy: nj.RequestedBy, Status: model.JobQueued, CreatedAt: nj.CreatedAt}
+	m.jobs[j.ID] = j
+	return *j, nil
+}
+
+func (m *memStore) FailJob(_ context.Context, id int64, msg string, at time.Time) (bool, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	j, ok := m.jobs[id]
+	if !ok || !j.Active() {
+		return false, nil
+	}
+	j.Status, j.Error, j.FinishedAt = model.JobFailed, msg, &at
+	return true, nil
+}
+
+func (m *memStore) GetJob(_ context.Context, id int64) (model.Job, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	j, ok := m.jobs[id]
+	if !ok {
+		return model.Job{}, model.ErrNotFound
+	}
+	return *j, nil
+}
+
+func (m *memStore) FailLostJobs(_ context.Context, agentID int64, known []int64, at time.Time) ([]int64, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	var lost []int64
+	for _, j := range m.jobs {
+		if j.AgentID == nil || *j.AgentID != agentID || !j.Active() {
+			continue
+		}
+		isKnown := false
+		for _, k := range known {
+			isKnown = isKnown || k == j.ID
+		}
+		if !isKnown {
+			j.Status, j.Error, j.FinishedAt = model.JobFailed, "lost: the agent did not know this job when it reconnected", &at
+			lost = append(lost, j.ID)
+		}
+	}
+	return lost, nil
+}
+
+func (m *memStore) JobProgress(_ context.Context, agentID, jobID int64, chunk string, at time.Time) (bool, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	j, ok := m.jobs[jobID]
+	if !ok || j.AgentID == nil || *j.AgentID != agentID || !j.Active() {
+		return false, nil
+	}
+	j.Status, j.Output = model.JobRunning, j.Output+chunk
+	if j.StartedAt == nil {
+		j.StartedAt = &at
+	}
+	return true, nil
+}
+
+func (m *memStore) FinishJob(_ context.Context, agentID, jobID int64, f model.JobFinish) (bool, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	j, ok := m.jobs[jobID]
+	if !ok || j.AgentID == nil || *j.AgentID != agentID || !j.Active() {
+		return false, nil
+	}
+	j.Status, j.Error, j.Output, j.OutputTruncated, j.FinishedAt = f.Status, f.Error, f.Output, f.OutputTruncated, &f.FinishedAt
+	if f.Ref != "" {
+		j.ToRef = f.Ref
+	}
+	if j.StartedAt == nil {
+		j.StartedAt = &f.StartedAt
+	}
+	if f.Status == model.JobSucceeded {
+		svc := m.svcRows[j.ServiceID]
+		svc.DesiredRef = f.Ref
+		m.svcRows[j.ServiceID] = svc
+	}
+	m.finishApplied++
+	return true, nil
 }
 
 func (m *memStore) AgentByTokenHash(_ context.Context, hash string) (model.Agent, error) {
@@ -112,8 +236,21 @@ func (s *staticSource) set(v ...protocol.Service) {
 type testMaster struct {
 	url         string
 	fingerprint string
-	hub         *hub.Hub
+	hub         *currentHub   // the hub behind /agent; follows restart()
 	attempts    *atomic.Int64 // TCP connections accepted by the server
+	broker      *events.Broker
+	restart     func() // replaces the hub and drops every connection
+}
+
+// currentHub forwards to whichever hub is live, so a master restart can swap
+// the hub without the mux or the jobs service noticing.
+type currentHub struct{ p atomic.Pointer[hub.Hub] }
+
+func (c *currentHub) ServeHTTP(w http.ResponseWriter, r *http.Request) { c.p.Load().ServeHTTP(w, r) }
+func (c *currentHub) Disconnect(id int64)                              { c.p.Load().Disconnect(id) }
+func (c *currentHub) Route(svc model.Service) (model.Route, error)     { return c.p.Load().Route(svc) }
+func (c *currentHub) StartJob(r model.Route, id int64, s protocol.JobStart) error {
+	return c.p.Load().StartJob(r, id, s)
 }
 
 func startMaster(t *testing.T, store *memStore) testMaster {
@@ -125,7 +262,10 @@ func startMaster(t *testing.T, store *memStore) testMaster {
 	fingerprint, err := pin.Fingerprint(key.Public())
 	require.NoError(t, err)
 
-	h := hub.New(store, events.NewBroker(), slog.New(slog.NewTextHandler(io.Discard, nil)))
+	quiet := slog.New(slog.NewTextHandler(io.Discard, nil))
+	broker := events.NewBroker()
+	h := &currentHub{}
+	h.p.Store(hub.New(store, broker, quiet))
 	mux := http.NewServeMux()
 	mux.Handle("GET /agent", h)
 	attempts := &atomic.Int64{}
@@ -139,14 +279,27 @@ func startMaster(t *testing.T, store *memStore) testMaster {
 	srv.TLS = &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS13}
 	srv.StartTLS()
 	t.Cleanup(srv.Close)
-	return testMaster{url: "wss" + strings.TrimPrefix(srv.URL, "https") + "/agent", fingerprint: fingerprint, hub: h, attempts: attempts}
+	restart := func() {
+		// The hub holds no job state, so a new hub over the same store is a
+		// faithful master restart.
+		old := h.p.Swap(hub.New(store, broker, quiet))
+		// CloseClientConnections skips hijacked (WebSocket) connections, so
+		// cut each agent through the old hub.
+		store.mu.Lock()
+		defer store.mu.Unlock()
+		for _, a := range store.agents {
+			old.Disconnect(a.ID)
+		}
+	}
+	return testMaster{url: "wss" + strings.TrimPrefix(srv.URL, "https") + "/agent", fingerprint: fingerprint,
+		hub: h, attempts: attempts, broker: broker, restart: restart}
 }
 
 func runAgent(t *testing.T, url, fingerprint, token string, src conn.Source) *conn.Agent {
 	t.Helper()
 	a := conn.New(conn.Config{URL: url, Token: token, Fingerprint: fingerprint, Version: "e2e",
 		MinBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond,
-	}, src, slog.New(slog.NewTextHandler(io.Discard, nil)))
+	}, src, updater.NewRunner(nil, time.Minute), slog.New(slog.NewTextHandler(io.Discard, nil)))
 	ctx, cancel := context.WithCancel(context.Background())
 	done := make(chan struct{})
 	go func() { _ = a.Run(ctx); close(done) }()

internal/e2e/jobs_test.go

@@ -0,0 +1,257 @@
+package e2e
+
+import (
+	"context"
+	"crypto/tls"
+	"errors"
+	"io"
+	"log/slog"
+	"net/http"
+	"strings"
+	"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/agent/conn"
+	"github.com/shcizo/package-updater/internal/agent/updater"
+	"github.com/shcizo/package-updater/internal/agenttoken"
+	"github.com/shcizo/package-updater/internal/master/jobs"
+	"github.com/shcizo/package-updater/internal/master/model"
+	"github.com/shcizo/package-updater/internal/protocol"
+)
+
+const (
+	slusk       = "gitea.shcizo.se/shcizo/slusk"
+	agentID     = int64(2)
+	waitFor     = 3 * time.Second
+	pollEvery   = 10 * time.Millisecond
+	lostMessage = "lost: the agent did not know this job when it reconnected"
+)
+
+// gateExecutor writes a line of output, then blocks until release is closed
+// or the job is cancelled.
+type gateExecutor struct {
+	started chan string // receives the job's Ref when it begins
+	release chan struct{}
+}
+
+func newGateExecutor() *gateExecutor {
+	return &gateExecutor{started: make(chan string, 8), release: make(chan struct{})}
+}
+
+func (g *gateExecutor) Execute(ctx context.Context, job protocol.JobStart, out io.Writer) (string, error) {
+	_, _ = io.WriteString(out, "pulling "+job.Ref+"\n")
+	g.started <- job.Ref
+	select {
+	case <-g.release:
+	case <-ctx.Done():
+		return "", ctx.Err()
+	}
+	_, _ = io.WriteString(out, "recreated\n")
+	return job.Ref + "@sha256:e2e", nil
+}
+
+// jobAgent is a connected agent with a real runner and its two lifetimes: the
+// connection and the runner can be stopped separately.
+type jobAgent struct {
+	runner   *updater.Runner
+	stopConn func()
+	stopRun  func()
+}
+
+func startJobAgent(t *testing.T, m testMaster, token string, exec updater.Executor) *jobAgent {
+	t.Helper()
+	src := &staticSource{}
+	src.set(web)
+	runner := updater.NewRunner(exec, time.Minute)
+	runCtx, stopRun := context.WithCancel(context.Background())
+	runDone := make(chan struct{})
+	go func() { runner.Run(runCtx); close(runDone) }()
+
+	a := conn.New(conn.Config{URL: m.url, Token: token, Fingerprint: m.fingerprint, Version: "e2e",
+		MinBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond, ProgressInterval: 10 * time.Millisecond,
+	}, src, runner, slog.New(slog.NewTextHandler(io.Discard, nil)))
+	connCtx, cancelConn := context.WithCancel(context.Background())
+	connDone := make(chan struct{})
+	go func() { _ = a.Run(connCtx); close(connDone) }()
+
+	ja := &jobAgent{
+		runner:   runner,
+		stopConn: func() { cancelConn(); <-connDone },
+		stopRun:  func() { stopRun(); <-runDone },
+	}
+	t.Cleanup(func() { ja.stopConn(); ja.stopRun() })
+	return ja
+}
+
+// jobEnv is a master with one agent (ID 2, token "pua_ok") that owns a Compose
+// service, and a jobs service wired to the master's current hub.
+type jobEnv struct {
+	store *memStore
+	m     testMaster
+	jobs  *jobs.Service
+	exec  *gateExecutor
+	agent *jobAgent
+}
+
+func newJobEnv(t *testing.T) *jobEnv {
+	t.Helper()
+	store := &memStore{agents: map[string]model.Agent{agenttoken.Hash("pua_ok"): {ID: agentID, Name: "node-2"}}}
+	store.seedService(model.Service{ID: 1, Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web",
+		AgentID: ptr(agentID), ActualRef: slusk + ":1.4.2@sha256:old"})
+	m := startMaster(t, store)
+	exec := newGateExecutor()
+	env := &jobEnv{store: store, m: m, exec: exec,
+		jobs:  jobs.New(store, m.hub, m.broker, "gitea.shcizo.se/shcizo/package-updater:latest", slog.New(slog.NewTextHandler(io.Discard, nil))),
+		agent: startJobAgent(t, m, "pua_ok", exec)}
+	env.waitRoutable(t)
+	return env
+}
+
+func ptr[T any](v T) *T { return &v }
+
+// waitRoutable waits until the master can send jobs to the agent.
+func (e *jobEnv) waitRoutable(t *testing.T) {
+	t.Helper()
+	svc, err := e.store.GetService(context.Background(), 1)
+	require.NoError(t, err)
+	require.Eventually(t, func() bool { _, err := e.m.hub.Route(svc); return err == nil }, waitFor, pollEvery)
+}
+
+func (e *jobEnv) job(t *testing.T, id int64) model.Job {
+	t.Helper()
+	j, err := e.store.GetJob(context.Background(), id)
+	require.NoError(t, err)
+	return j
+}
+
+func (e *jobEnv) waitStatus(t *testing.T, id int64, want model.JobStatus) model.Job {
+	t.Helper()
+	require.Eventually(t, func() bool { return e.job(t, id).Status == want }, waitFor, pollEvery,
+		"job %d should become %s, is %s (%s)", id, want, e.job(t, id).Status, e.job(t, id).Error)
+	return e.job(t, id)
+}
+
+func (e *jobEnv) waitStarted(t *testing.T) {
+	t.Helper()
+	select {
+	case <-e.exec.started:
+	case <-time.After(waitFor):
+		t.Fatal("the agent never started the job")
+	}
+}
+
+func (e *jobEnv) trigger(t *testing.T) model.Job {
+	t.Helper()
+	j, err := e.jobs.Trigger(context.Background(), 1, "sub-1")
+	require.NoError(t, err)
+	return j
+}
+
+func TestManualJobFlowsToAgentAndBack(t *testing.T) {
+	e := newJobEnv(t)
+	close(e.exec.release)
+
+	j := e.trigger(t)
+	done := e.waitStatus(t, j.ID, model.JobSucceeded)
+
+	assert.Equal(t, slusk+":1.4.2@sha256:e2e", done.ToRef)
+	assert.Contains(t, done.Output, "pulling "+slusk+":1.4.2")
+	assert.Contains(t, done.Output, "recreated")
+	svc, err := e.store.GetService(context.Background(), 1)
+	require.NoError(t, err)
+	assert.Equal(t, slusk+":1.4.2@sha256:e2e", svc.DesiredRef)
+	require.Eventually(t, func() bool { return len(e.agent.runner.Known()) == 0 }, waitFor, pollEvery, "the result must be acked")
+}
+
+func TestJobFinishedWhileDisconnectedIsRecordedOnce(t *testing.T) {
+	e := newJobEnv(t)
+	j := e.trigger(t)
+	e.waitStarted(t)
+
+	e.m.restart() // drops the agent's connection while the executor blocks
+	close(e.exec.release)
+
+	done := e.waitStatus(t, j.ID, model.JobSucceeded)
+	assert.NotEqual(t, lostMessage, done.Error)
+	assert.GreaterOrEqual(t, e.m.attempts.Load(), int64(2), "the agent must have reconnected")
+	require.Eventually(t, func() bool { return len(e.agent.runner.Known()) == 0 }, waitFor, pollEvery)
+	e.store.mu.Lock()
+	defer e.store.mu.Unlock()
+	assert.Equal(t, 1, e.store.finishApplied)
+}
+
+func TestMasterRestartKeepsJobTheAgentKnows(t *testing.T) {
+	e := newJobEnv(t)
+	j := e.trigger(t)
+	e.waitStarted(t)
+
+	e.m.restart()
+	e.waitRoutable(t) // the agent reconnected and said hello with its known jobs
+	assert.True(t, e.job(t, j.ID).Active(), "the agent still has the job, so it is not lost")
+	close(e.exec.release)
+
+	done := e.waitStatus(t, j.ID, model.JobSucceeded)
+	assert.Empty(t, done.Error)
+}
+
+func TestAgentRestartMarksJobLost(t *testing.T) {
+	e := newJobEnv(t)
+	j := e.trigger(t)
+	e.waitStarted(t)
+
+	// The agent process dies: connection first, so the cancelled job cannot
+	// report a result, then the runner.
+	e.agent.stopConn()
+	e.agent.stopRun()
+	fresh := startJobAgent(t, e.m, "pua_ok", newGateExecutor())
+	_ = fresh
+
+	lost := e.waitStatus(t, j.ID, model.JobFailed)
+	assert.Equal(t, lostMessage, lost.Error)
+}
+
+func TestPhaseOneAgentGetsNoJobs(t *testing.T) {
+	store := &memStore{agents: map[string]model.Agent{agenttoken.Hash("pua_old"): {ID: agentID, Name: "node-2"}}}
+	svcRow := model.Service{ID: 1, Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web",
+		AgentID: ptr(agentID), ActualRef: slusk + ":1.4.2@sha256:old"}
+	store.seedService(svcRow)
+	m := startMaster(t, store)
+	js := jobs.New(store, m.hub, m.broker, "gitea.shcizo.se/shcizo/package-updater:latest", slog.New(slog.NewTextHandler(io.Discard, nil)))
+
+	// A raw client speaking the phase-1 protocol: its hello has no
+	// protocol_version.
+	ctx := context.Background()
+	c, _, err := websocket.Dial(ctx, m.url, &websocket.DialOptions{
+		HTTPHeader: http.Header{"Authorization": {"Bearer pua_old"}},
+		HTTPClient: &http.Client{Transport: &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}},
+	})
+	require.NoError(t, err)
+	defer c.CloseNow()
+	hello, err := protocol.Encode(protocol.TypeHello, map[string]any{"agent_version": "1", "hostname": "node-2", "swarm_manager": false})
+	require.NoError(t, err)
+	require.NoError(t, wsjson.Write(ctx, c, hello))
+	snap, err := protocol.Encode(protocol.TypeSnapshot, protocol.Snapshot{Services: []protocol.Service{web}})
+	require.NoError(t, err)
+	require.NoError(t, wsjson.Write(ctx, c, snap))
+
+	require.Eventually(t, func() bool {
+		_, err := m.hub.Route(svcRow)
+		return errors.Is(err, model.ErrAgentTooOld)
+	}, waitFor, pollEvery)
+
+	j, err := js.Trigger(ctx, 1, "sub-1")
+	require.NoError(t, err)
+	assert.Equal(t, model.JobFailed, j.Status)
+	assert.Equal(t, "agent too old for updates", j.Error)
+
+	rctx, cancel := context.WithTimeout(ctx, 200*time.Millisecond)
+	defer cancel()
+	var env protocol.Envelope
+	err = wsjson.Read(rctx, c, &env)
+	assert.True(t, err != nil && strings.Contains(err.Error(), "deadline"), "no job.start expected, got %q (err %v)", env.Type, err)
+}