diff --git a/internal/e2e/agent_master_test.go b/internal/e2e/agent_master_test.go index 36a071b..905b064 100644 --- a/internal/e2e/agent_master_test.go +++ b/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 := ¤tHub{} + 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) }() diff --git a/internal/e2e/jobs_test.go b/internal/e2e/jobs_test.go new file mode 100644 index 0000000..d3f1ae7 --- /dev/null +++ b/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) +}