feat: master/agent phase 4 — registry polling and version history (#9) (#23)

internal/e2e/agent_master_test.go

@@ -11,6 +11,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"regexp"
 	"sort"
 	"strings"
 	"sync"
@@ -31,6 +32,9 @@ import (
 	"github.com/shcizo/package-updater/internal/protocol"
 )
 
+// digestPattern is the digest format the real store records.
+var digestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)
+
 type memStore struct {
 	mu        sync.Mutex
 	agents    map[string]model.Agent
@@ -44,6 +48,63 @@ type memStore struct {
 	jobs          map[int64]*model.Job
 	nextJob       int64
 	finishApplied int // FinishJob calls that returned true
+	versions      map[int64]model.ImageVersion
+}
+
+// seedVersion adds an image_versions row that GetImageVersion returns.
+func (m *memStore) seedVersion(v model.ImageVersion) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.versions == nil {
+		m.versions = map[int64]model.ImageVersion{}
+	}
+	m.versions[v.ID] = v
+}
+
+func (m *memStore) GetImageVersion(_ context.Context, id int64) (model.ImageVersion, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	v, ok := m.versions[id]
+	if !ok {
+		return model.ImageVersion{}, model.ErrNotFound
+	}
+	return v, nil
+}
+
+func (m *memStore) ListServices(_ context.Context) ([]model.Service, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	var out []model.Service
+	for _, svc := range m.svcRows {
+		if svc.RemovedAt == nil {
+			out = append(out, svc)
+		}
+	}
+	sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
+	return out, nil
+}
+
+func (m *memStore) SetPollError(_ context.Context, ids []int64, msg string) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	for _, id := range ids {
+		svc := m.svcRows[id]
+		svc.LastPollError = msg
+		m.svcRows[id] = svc
+	}
+	return nil
+}
+
+func (m *memStore) SetLastTagDigest(_ context.Context, id int64, digest string) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	svc, ok := m.svcRows[id]
+	if !ok {
+		return model.ErrNotFound
+	}
+	svc.LastTagDigest = digest
+	m.svcRows[id] = svc
+	return nil
 }
 
 // seedService adds a service row that jobs can be triggered for.
@@ -88,7 +149,7 @@ func (m *memStore) SetAvailable(_ context.Context, id int64, a model.Available)
 	if !ok {
 		return model.ErrNotFound
 	}
-	svc.AvailableRef, svc.AvailableSource, svc.AvailableAt = a.Ref, a.Source, &a.At
+	svc.AvailableRef, svc.AvailableDigest, svc.AvailableSource, svc.AvailableAt = a.Ref, a.Digest, a.Source, &a.At
 	m.svcRows[id] = svc
 	return nil
 }
@@ -184,10 +245,14 @@ func (m *memStore) FinishJob(_ context.Context, agentID, jobID int64, f model.Jo
 	if f.Status == model.JobSucceeded {
 		svc := m.svcRows[j.ServiceID]
 		svc.DesiredRef = f.Ref
+		if _, d, _ := strings.Cut(f.Ref, "@"); !strings.Contains(requestedRef, "@") && digestPattern.MatchString(d) {
+			svc.LastTagDigest = d
+		}
 		// Same rule as the SQL: the job was created for the suggested ref after it was reported.
-		requested, _, _ := strings.Cut(requestedRef, "@")
-		if svc.AvailableAt != nil && svc.AvailableRef == requested && !svc.AvailableAt.After(j.CreatedAt) {
-			svc.AvailableRef, svc.AvailableSource, svc.AvailableAt = "", "", nil
+		requested, digest, _ := strings.Cut(requestedRef, "@")
+		if svc.AvailableAt != nil && svc.AvailableRef == requested && !svc.AvailableAt.After(j.CreatedAt) &&
+			(digest == "" || svc.AvailableDigest == digest) {
+			svc.AvailableRef, svc.AvailableDigest, svc.AvailableSource, svc.AvailableAt = "", "", "", nil
 		}
 		m.svcRows[j.ServiceID] = svc
 	}
@@ -282,7 +347,9 @@ 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) Route(svc model.Service, minProtocol int) (model.Route, error) {
+	return c.p.Load().Route(svc, minProtocol)
+}
 func (c *currentHub) StartJob(r model.Route, id int64, s protocol.JobStart) error {
 	return c.p.Load().StartJob(r, id, s)
 }

internal/e2e/jobs_test.go

@@ -32,6 +32,9 @@ const (
 	lostMessage = "lost: the agent did not know this job when it reconnected"
 )
 
+// e2eDigest is the digest the fake executor pins a tag to.
+const e2eDigest = "sha256:" + "e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0"
+
 // gateExecutor writes a line of output, then blocks until release is closed
 // or the job is cancelled.
 type gateExecutor struct {
@@ -52,7 +55,10 @@ func (g *gateExecutor) Execute(ctx context.Context, job protocol.JobStart, out i
 		return "", ctx.Err()
 	}
 	_, _ = io.WriteString(out, "recreated\n")
-	return job.Ref + "@sha256:e2e", nil
+	if strings.Contains(job.Ref, "@") {
+		return job.Ref, nil
+	}
+	return job.Ref + "@" + e2eDigest, nil
 }
 
 // jobAgent is a connected agent with a real runner and its two lifetimes: the
@@ -119,7 +125,7 @@ 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)
+	require.Eventually(t, func() bool { _, err := e.m.hub.Route(svc, protocol.MinJobVersion); return err == nil }, waitFor, pollEvery)
 }
 
 func (e *jobEnv) job(t *testing.T, id int64) model.Job {
@@ -159,12 +165,12 @@ func TestManualJobFlowsToAgentAndBack(t *testing.T) {
 	j := e.trigger(t)
 	done := e.waitStatus(t, j.ID, model.JobSucceeded)
 
-	assert.Equal(t, slusk+":1.4.2@sha256:e2e", done.ToRef)
+	assert.Equal(t, slusk+":1.4.2@"+e2eDigest, 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)
+	assert.Equal(t, slusk+":1.4.2@"+e2eDigest, svc.DesiredRef)
 	require.Eventually(t, func() bool { return len(e.agent.runner.Known()) == 0 }, waitFor, pollEvery, "the result must be acked")
 }
 
@@ -240,7 +246,7 @@ func TestPhaseOneAgentGetsNoJobs(t *testing.T) {
 	require.NoError(t, wsjson.Write(ctx, c, snap))
 
 	require.Eventually(t, func() bool {
-		_, err := m.hub.Route(svcRow)
+		_, err := m.hub.Route(svcRow, protocol.MinJobVersion)
 		return errors.Is(err, model.ErrAgentTooOld)
 	}, waitFor, pollEvery)
 

internal/e2e/versions_test.go

@@ -0,0 +1,133 @@
+package e2e
+
+import (
+	"context"
+	"io"
+	"log/slog"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/shcizo/package-updater/internal/master/jobs"
+	"github.com/shcizo/package-updater/internal/master/model"
+	"github.com/shcizo/package-updater/internal/master/poller"
+	"github.com/shcizo/package-updater/internal/protocol"
+)
+
+// jobCount is the number of jobs created so far.
+func (m *memStore) jobCount() int {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return len(m.jobs)
+}
+
+// fixedRegistry answers every lookup from a fixed table.
+type fixedRegistry struct {
+	digests map[string]string
+	calls   int
+}
+
+func (f *fixedRegistry) Digest(_ context.Context, ref string) (string, error) {
+	f.calls++
+	return f.digests[ref], nil
+}
+
+func TestPickPinnedVersion(t *testing.T) {
+	env := newJobEnv(t)
+	env.store.seedVersion(model.ImageVersion{ID: 5, Repo: slusk, Tag: "1.3.0", Digest: "sha256:old130", Source: model.SourceRunning})
+
+	job, err := env.jobs.TriggerVersion(context.Background(), 1, 5, "alice")
+
+	require.NoError(t, err)
+	assert.Equal(t, slusk+":1.3.0@sha256:old130", job.ToRef)
+	select {
+	case ref := <-env.exec.started:
+		assert.Equal(t, slusk+":1.3.0@sha256:old130", ref, "the agent receives exactly the picked version")
+	case <-time.After(waitFor):
+		t.Fatal("the agent never received the job")
+	}
+	close(env.exec.release)
+	done := env.waitStatus(t, job.ID, model.JobSucceeded)
+	assert.Equal(t, model.TriggerManual, done.Trigger)
+	assert.Equal(t, "alice", done.RequestedBy)
+}
+
+func TestPickVersionOfAnotherRepoIsRefused(t *testing.T) {
+	env := newJobEnv(t)
+	env.store.seedVersion(model.ImageVersion{ID: 6, Repo: "evil/other", Tag: "1", Digest: "sha256:x", Source: model.SourcePoll})
+
+	_, err := env.jobs.TriggerVersion(context.Background(), 1, 6, "alice")
+
+	assert.ErrorIs(t, err, model.ErrNotFound)
+	assert.Zero(t, env.store.jobCount(), "no job was created")
+}
+
+func TestPollStartsAutoJob(t *testing.T) {
+	env := newJobEnv(t)
+	env.store.seedService(model.Service{ID: 1, Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web",
+		AgentID: ptr(agentID), ImageRepo: slusk, ActualRef: slusk + ":1.4.2@sha256:old", ActualDigest: "sha256:old",
+		UpdateMode: model.ModeAuto})
+	reg := &fixedRegistry{digests: map[string]string{slusk + ":1.4.2": "sha256:new"}}
+	p := poller.New(env.store, env.jobs, reg, env.m.broker, poller.Config{Workers: 1}, slog.New(slog.NewTextHandler(io.Discard, nil)))
+
+	p.Cycle(context.Background())
+
+	select {
+	case ref := <-env.exec.started:
+		assert.Equal(t, slusk+":1.4.2", ref, "the agent gets repo:tag and resolves the digest itself")
+	case <-time.After(waitFor):
+		t.Fatal("the agent never received the job")
+	}
+	require.Equal(t, 1, env.store.jobCount())
+	job := env.job(t, 1)
+	assert.Equal(t, model.TriggerPoll, job.Trigger)
+	assert.Equal(t, slusk+":1.4.2", job.ToRef)
+	svc, err := env.store.GetService(context.Background(), 1)
+	require.NoError(t, err)
+	assert.Equal(t, slusk+":1.4.2", svc.AvailableRef)
+	assert.Equal(t, "sha256:new", svc.AvailableDigest)
+	assert.Equal(t, model.SourcePoll, svc.AvailableSource)
+	assert.Equal(t, "sha256:new", svc.LastTagDigest)
+
+	p.Cycle(context.Background())
+
+	assert.Equal(t, 2, reg.calls, "the registry is asked again")
+	assert.Equal(t, 1, env.store.jobCount(), "the same digest does not start another job")
+	close(env.exec.release)
+	env.waitStatus(t, job.ID, model.JobSucceeded)
+}
+
+// An update started by the webhook, then a rollback by hand: polling must not
+// redo the update while the registry still shows the digest that was deployed.
+func TestPollKeepsARollbackAfterAWebhookUpdate(t *testing.T) {
+	env := newJobEnv(t)
+	ctx := context.Background()
+	close(env.exec.release)
+	env.store.seedService(model.Service{ID: 1, Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web",
+		AgentID: ptr(agentID), ImageRepo: slusk, ActualRef: slusk + ":main@sha256:old", ActualDigest: "sha256:old",
+		UpdateMode: model.ModeAuto})
+	env.store.seedVersion(model.ImageVersion{ID: 5, Repo: slusk, Tag: "main", Digest: "sha256:old", Source: model.SourceRunning})
+	svc, err := env.store.GetService(ctx, 1)
+	require.NoError(t, err)
+
+	hook, err := env.jobs.Start(ctx, svc, jobs.StartRequest{Trigger: model.TriggerWebhook, Ref: slusk + ":main"})
+	require.NoError(t, err)
+	env.waitStatus(t, hook.ID, model.JobSucceeded)
+	svc, err = env.store.GetService(ctx, 1)
+	require.NoError(t, err)
+	assert.Equal(t, e2eDigest, svc.LastTagDigest, "the deployed digest of the tag is remembered")
+
+	pick, err := env.jobs.TriggerVersion(ctx, 1, 5, "alice")
+	require.NoError(t, err)
+	env.waitStatus(t, pick.ID, model.JobSucceeded)
+	jobsBefore := env.store.jobCount()
+
+	reg := &fixedRegistry{digests: map[string]string{slusk + ":main": e2eDigest}}
+	p := poller.New(env.store, env.jobs, reg, env.m.broker, poller.Config{Workers: 1}, slog.New(slog.NewTextHandler(io.Discard, nil)))
+	p.Cycle(ctx)
+
+	assert.Equal(t, 1, reg.calls)
+	assert.Equal(t, jobsBefore, env.store.jobCount(), "the rollback stays")
+}

internal/e2e/webhook_test.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/shcizo/package-updater/internal/master/model"
 	"github.com/shcizo/package-updater/internal/master/webhook"
+	"github.com/shcizo/package-updater/internal/protocol"
 )
 
 // webhookServer mounts the webhook handler on a test server, wired to the
@@ -98,7 +99,7 @@ func TestWebhookWaitsForAutoJob(t *testing.T) {
 	job := e.job(t, r.body.Results[0].JobID)
 	assert.Equal(t, model.TriggerWebhook, job.Trigger)
 	assert.Equal(t, model.JobSucceeded, job.Status)
-	assert.Equal(t, slusk+":2.0.0@sha256:e2e", job.ToRef)
+	assert.Equal(t, slusk+":2.0.0@"+e2eDigest, job.ToRef)
 	svc, err := e.store.GetService(context.Background(), 1)
 	require.NoError(t, err)
 	assert.Empty(t, svc.AvailableRef, "the job deployed the suggestion, so it is cleared")
@@ -111,7 +112,7 @@ func TestWebhookFailedAutoJobIs500(t *testing.T) {
 	e.agent.stopConn() // no agent connection: the job fails at once with "agent offline"
 	require.Eventually(t, func() bool {
 		svc, _ := e.store.GetService(context.Background(), 1)
-		_, err := e.m.hub.Route(svc)
+		_, err := e.m.hub.Route(svc, protocol.MinJobVersion)
 		return err != nil
 	}, waitFor, pollEvery)