From c09e774d743aa3389618984ddf014f9fee6aa792 Mon Sep 17 00:00:00 2001 From: Samuel Enocsson Date: Mon, 12 Oct 2026 19:14:32 +0200 Subject: [PATCH] fix(master): refuse jobs for revoked agents, signal revoked jobs, agent start time wins (#7) --- internal/master/store/agents.go | 30 +++++--- .../master/store/agents_integration_test.go | 19 +++-- internal/master/store/jobs.go | 73 +++++++++++++------ .../master/store/jobs_integration_test.go | 38 +++++++++- internal/master/web/server.go | 9 ++- internal/master/web/server_test.go | 26 ++++--- 6 files changed, 138 insertions(+), 57 deletions(-) diff --git a/internal/master/store/agents.go b/internal/master/store/agents.go index 1d363f3..79ff91f 100644 --- a/internal/master/store/agents.go +++ b/internal/master/store/agents.go @@ -63,33 +63,39 @@ func (s *Store) AgentByTokenHash(ctx context.Context, hash string) (model.Agent, return scanAgent(s.pool.QueryRow(ctx, `SELECT `+agentColumns+` FROM agents WHERE token_hash = $1`, hash)) } -// RevokeAgent marks the agent revoked and its Compose and plain-container -// services removed and its queued and running jobs failed, in one transaction. Swarm rows are shared with other -// managers and left alone. Its token stops working immediately; closing a -// live connection is the hub's job. -func (s *Store) RevokeAgent(ctx context.Context, id int64, at time.Time) error { +// RevokeAgent marks the agent revoked, its Compose and plain-container +// services removed and its queued and running jobs failed, in one +// transaction, and returns the IDs of the jobs it failed. Swarm rows are +// shared with other managers and left alone. Its token stops working +// immediately; closing a live connection is the hub's job. +func (s *Store) RevokeAgent(ctx context.Context, id int64, at time.Time) ([]int64, error) { tx, err := s.pool.Begin(ctx) if err != nil { - return fmt.Errorf("revoke agent: begin: %w", err) + return nil, fmt.Errorf("revoke agent: begin: %w", err) } defer func() { _ = tx.Rollback(ctx) }() if _, err := tx.Exec(ctx, `UPDATE agents SET revoked_at = $2 WHERE id = $1 AND revoked_at IS NULL`, id, at); err != nil { - return fmt.Errorf("revoke agent: %w", err) + return nil, fmt.Errorf("revoke agent: %w", err) } if _, err := tx.Exec(ctx, `UPDATE services SET removed_at = $2 WHERE agent_id = $1 AND kind IN ('compose', 'container') AND removed_at IS NULL`, id, at); err != nil { - return fmt.Errorf("revoke agent: remove services: %w", err) + return nil, fmt.Errorf("revoke agent: remove services: %w", err) } - if _, err := tx.Exec(ctx, failActiveJobsSQL, id, "agent revoked", at); err != nil { - return fmt.Errorf("revoke agent: fail jobs: %w", err) + rows, err := tx.Query(ctx, failActiveJobsSQL, id, "agent revoked", at) + if err != nil { + return nil, fmt.Errorf("revoke agent: fail jobs: %w", err) + } + failed, err := pgx.CollectRows(rows, pgx.RowTo[int64]) + if err != nil { + return nil, fmt.Errorf("revoke agent: fail jobs: %w", err) } if err := tx.Commit(ctx); err != nil { - return fmt.Errorf("revoke agent: commit: %w", err) + return nil, fmt.Errorf("revoke agent: commit: %w", err) } - return nil + return failed, nil } // RecordHello stores what the agent reported about itself on connect. diff --git a/internal/master/store/agents_integration_test.go b/internal/master/store/agents_integration_test.go index 7227d93..9cc22b9 100644 --- a/internal/master/store/agents_integration_test.go +++ b/internal/master/store/agents_integration_test.go @@ -46,7 +46,8 @@ func TestAgents_CreateLookupRevoke(t *testing.T) { _, err = s.AgentByTokenHash(ctx, "nope") assert.ErrorIs(t, err, model.ErrNotFound) - require.NoError(t, s.RevokeAgent(ctx, created.ID, time.Now())) + _, err = s.RevokeAgent(ctx, created.ID, time.Now()) + require.NoError(t, err) found, err = s.AgentByTokenHash(ctx, "hash-2") require.NoError(t, err) assert.NotNil(t, found.RevokedAt) @@ -88,7 +89,8 @@ func TestAgents_RevokedNameCanBeReused(t *testing.T) { ctx := context.Background() old, err := s.CreateAgent(ctx, "node-2", "h1") require.NoError(t, err) - require.NoError(t, s.RevokeAgent(ctx, old.ID, time.Now())) + _, err = s.RevokeAgent(ctx, old.ID, time.Now()) + require.NoError(t, err) fresh, err := s.CreateAgent(ctx, "node-2", "h2") require.NoError(t, err) @@ -104,7 +106,8 @@ func TestAgents_RevokeRemovesOwnComposeAndContainerServices(t *testing.T) { id := newAgent(t, s, "swarm-mgr") require.NoError(t, s.ApplySnapshot(ctx, id, manager, []protocol.Service{slusk, adhoc, gitea}, t0)) - require.NoError(t, s.RevokeAgent(ctx, id, t0.Add(time.Hour))) + _, err := s.RevokeAgent(ctx, id, t0.Add(time.Hour)) + require.NoError(t, err) got, err := s.ListServices(ctx) require.NoError(t, err) @@ -115,7 +118,8 @@ func TestRevokedAgent_InventoryWritesRefused(t *testing.T) { s := openTestStore(t) ctx := context.Background() id := newAgent(t, s, "node-2") - require.NoError(t, s.RevokeAgent(ctx, id, t0)) + _, err := s.RevokeAgent(ctx, id, t0) + require.NoError(t, err) assert.ErrorIs(t, s.ApplySnapshot(ctx, id, worker, []protocol.Service{slusk}, t0.Add(time.Minute)), model.ErrRevoked) assert.ErrorIs(t, s.UpsertService(ctx, id, "", adhoc, t0.Add(time.Minute)), model.ErrRevoked) @@ -130,7 +134,8 @@ func TestRevokedAgent_UpsertDoesNotResurrectRemovedRow(t *testing.T) { ctx := context.Background() id := newAgent(t, s, "node-2") require.NoError(t, s.ApplySnapshot(ctx, id, worker, []protocol.Service{slusk}, t0)) - require.NoError(t, s.RevokeAgent(ctx, id, t0.Add(time.Hour))) + _, err := s.RevokeAgent(ctx, id, t0.Add(time.Hour)) + require.NoError(t, err) assert.ErrorIs(t, s.UpsertService(ctx, id, "", slusk, t0.Add(2*time.Hour)), model.ErrRevoked) @@ -149,7 +154,9 @@ func TestRevokeAgent_FailsActiveJobs(t *testing.T) { jb, err := s.CreateJob(ctx, newJob(svcB, b)) require.NoError(t, err) - require.NoError(t, s.RevokeAgent(ctx, a, t0)) + failed, err := s.RevokeAgent(ctx, a, t0) + require.NoError(t, err) + assert.Equal(t, []int64{ja.ID}, failed, "the failed job IDs are returned so the UI can signal them") got, err := s.GetJob(ctx, ja.ID) require.NoError(t, err) diff --git a/internal/master/store/jobs.go b/internal/master/store/jobs.go index 4f708c6..a118c43 100644 --- a/internal/master/store/jobs.go +++ b/internal/master/store/jobs.go @@ -19,7 +19,8 @@ const lostJobMessage = "lost: the agent did not know this job when it reconnecte // message $2 at time $3. const failActiveJobsSQL = ` UPDATE update_jobs SET status = 'failed', error = $2, finished_at = $3 -WHERE agent_id = $1 AND status IN ('queued', 'running')` +WHERE agent_id = $1 AND status IN ('queued', 'running') +RETURNING id` const jobSelectSQL = ` SELECT j.id, j.service_id, s.kind, s.compose_project, s.compose_service, s.swarm_service, s.container_name, @@ -56,27 +57,44 @@ func (s *Store) getJob(ctx context.Context, id int64) (model.Job, error) { // CreateJob inserts a queued job. If the service already has a queued or // running job, it returns that job and model.ErrJobActive (detected via the -// unique index, so it holds under concurrent calls). +// unique index, so it holds under concurrent calls). A job for a revoked +// agent is refused with model.ErrAgentOffline: the agent's revoke fails its +// jobs once, so a job created afterwards would never be settled and would +// block the service. The share lock on the agent row makes a concurrent +// revoke wait for this insert, so its job cleanup sees the new row. func (s *Store) CreateJob(ctx context.Context, j model.NewJob) (model.Job, error) { - var id int64 - err := s.pool.QueryRow(ctx, ` - INSERT INTO update_jobs (service_id, agent_id, from_ref, to_ref, trigger, requested_by, status, created_at) - VALUES ($1, $2, $3, $4, $5, $6, 'queued', $7) - RETURNING id`, - j.ServiceID, j.AgentID, j.FromRef, j.ToRef, string(j.Trigger), j.RequestedBy, j.CreatedAt).Scan(&id) - var pgErr *pgconn.PgError - if errors.As(err, &pgErr) && pgErr.Code == "23505" && pgErr.ConstraintName == "update_jobs_one_active" { - active, err := scanJob(s.pool.QueryRow(ctx, fmt.Sprintf(jobSelectSQL, "j.output")+ - ` WHERE j.service_id = $1 AND j.status IN ('queued', 'running')`, j.ServiceID)) - if err != nil { - return model.Job{}, fmt.Errorf("find active job: %w", err) + // The active job can finish between the unique violation and the lookup + // below; then the insert is simply tried once more. + for attempt := 0; attempt < 2; attempt++ { + var id int64 + err := s.pool.QueryRow(ctx, ` + INSERT INTO update_jobs (service_id, agent_id, from_ref, to_ref, trigger, requested_by, status, created_at) + SELECT $1::bigint, $2::bigint, $3::text, $4::text, $5::text, $6::text, 'queued', $7::timestamptz + WHERE $2::bigint IS NULL + OR EXISTS (SELECT 1 FROM agents WHERE id = $2 AND revoked_at IS NULL FOR SHARE) + RETURNING id`, + j.ServiceID, j.AgentID, j.FromRef, j.ToRef, string(j.Trigger), j.RequestedBy, j.CreatedAt).Scan(&id) + if errors.Is(err, pgx.ErrNoRows) { + return model.Job{}, model.ErrAgentOffline } - return active, model.ErrJobActive + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) && pgErr.Code == "23505" && pgErr.ConstraintName == "update_jobs_one_active" { + active, err := scanJob(s.pool.QueryRow(ctx, fmt.Sprintf(jobSelectSQL, "j.output")+ + ` WHERE j.service_id = $1 AND j.status IN ('queued', 'running')`, j.ServiceID)) + if errors.Is(err, pgx.ErrNoRows) { + continue + } + if err != nil { + return model.Job{}, fmt.Errorf("find active job: %w", err) + } + return active, model.ErrJobActive + } + if err != nil { + return model.Job{}, fmt.Errorf("create job: %w", err) + } + return s.getJob(ctx, id) } - if err != nil { - return model.Job{}, fmt.Errorf("create job: %w", err) - } - return s.getJob(ctx, id) + return model.Job{}, fmt.Errorf("create job: active job vanished twice") } // FailJob sets a queued or running job to failed with msg and finished_at = @@ -167,9 +185,10 @@ func (s *Store) JobProgress(ctx context.Context, agentID, jobID int64, chunk str // FinishJob records the agent's result for a queued or running job sent to // agentID: status, error, output (replacing what progress assembled), -// started_at (kept if already set), finished_at; to_ref = f.Ref when set. On -// success it also sets the service's desired_ref = f.Ref, in the same -// transaction. Returns false if no such active job belongs to agentID. +// started_at (the agent's, or the existing one if the agent sent none), +// finished_at; to_ref = f.Ref when set. On success it also sets the +// service's desired_ref = f.Ref, in the same transaction. Returns false if no +// such active job belongs to agentID. func (s *Store) FinishJob(ctx context.Context, agentID, jobID int64, f model.JobFinish) (bool, error) { tx, err := s.pool.Begin(ctx) if err != nil { @@ -177,15 +196,21 @@ func (s *Store) FinishJob(ctx context.Context, agentID, jobID int64, f model.Job } defer func() { _ = tx.Rollback(ctx) }() + // The agent's clock wins when it sent a start time; otherwise the time the + // first progress arrived stays. + var startedAt *time.Time + if !f.StartedAt.IsZero() { + startedAt = &f.StartedAt + } var serviceID int64 err = tx.QueryRow(ctx, ` UPDATE update_jobs SET status = $3, error = $4, output = $5, output_truncated = $6, to_ref = CASE WHEN $7::text <> '' THEN $7::text ELSE to_ref END, - started_at = COALESCE(started_at, $8), finished_at = $9 + started_at = COALESCE($8, started_at), finished_at = $9 WHERE id = $1 AND agent_id = $2 AND status IN ('queued', 'running') RETURNING service_id`, - jobID, agentID, string(f.Status), f.Error, f.Output, f.OutputTruncated, f.Ref, f.StartedAt, f.FinishedAt).Scan(&serviceID) + jobID, agentID, string(f.Status), f.Error, f.Output, f.OutputTruncated, f.Ref, startedAt, f.FinishedAt).Scan(&serviceID) if errors.Is(err, pgx.ErrNoRows) { return false, nil } diff --git a/internal/master/store/jobs_integration_test.go b/internal/master/store/jobs_integration_test.go index 2e7a54d..1a1c1b0 100644 --- a/internal/master/store/jobs_integration_test.go +++ b/internal/master/store/jobs_integration_test.go @@ -218,7 +218,7 @@ func TestFinishJob_OnlyOnceAndOnlyByOwner(t *testing.T) { assert.Equal(t, model.JobSucceeded, got.Status) assert.Equal(t, "r:1@sha256:d", got.ToRef) assert.Equal(t, "full", got.Output) - assert.True(t, t0.Equal(*got.StartedAt), "started_at set by the first progress is kept") + assert.True(t, fin.StartedAt.Equal(*got.StartedAt), "the agent's started_at wins over the master's progress clock") assert.True(t, fin.FinishedAt.Equal(*got.FinishedAt)) service, err := s.GetService(ctx, svc) require.NoError(t, err) @@ -303,3 +303,39 @@ func TestGetJobAndListJobs(t *testing.T) { require.NoError(t, err) assert.Len(t, list, 1) } + +func TestFinishJob_ZeroStartedAtKeepsProgressClock(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + a, svc := jobFixture(t, s, "a") + j, err := s.CreateJob(ctx, newJob(svc, a)) + require.NoError(t, err) + _, err = s.JobProgress(ctx, a, j.ID, "x", t0) + require.NoError(t, err) + + ok, err := s.FinishJob(ctx, a, j.ID, model.JobFinish{Status: model.JobFailed, FinishedAt: t0.Add(time.Hour)}) + require.NoError(t, err) + require.True(t, ok) + got, err := s.GetJob(ctx, j.ID) + require.NoError(t, err) + assert.True(t, t0.Equal(*got.StartedAt)) +} + +func TestCreateJob_RefusesRevokedAgent(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + a, svc := jobFixture(t, s, "a") + _, err := s.RevokeAgent(ctx, a, t0) + require.NoError(t, err) + + _, err = s.CreateJob(ctx, newJob(svc, a)) + assert.ErrorIs(t, err, model.ErrAgentOffline) + list, err := s.ListJobs(ctx, model.JobFilter{}) + require.NoError(t, err) + assert.Empty(t, list, "no job is left behind to block the service") + + nj := newJob(svc, a) + nj.AgentID = nil + _, err = s.CreateJob(ctx, nj) + assert.NoError(t, err, "a job that was never routed has no agent to check") +} diff --git a/internal/master/web/server.go b/internal/master/web/server.go index ce5c4e0..c1cef5b 100644 --- a/internal/master/web/server.go +++ b/internal/master/web/server.go @@ -36,7 +36,8 @@ type Store interface { ListJobs(ctx context.Context, f model.JobFilter) ([]model.Job, error) ListAgents(ctx context.Context) ([]model.Agent, error) CreateAgent(ctx context.Context, name, tokenHash string) (model.Agent, error) - RevokeAgent(ctx context.Context, id int64, at time.Time) error + // RevokeAgent returns the IDs of the jobs it failed. + RevokeAgent(ctx context.Context, id int64, at time.Time) ([]int64, error) } // Hub reports and controls live agent connections. *hub.Hub satisfies it. @@ -389,13 +390,17 @@ func (s *Server) revokeAgent(w http.ResponseWriter, r *http.Request) { http.Error(w, "invalid agent id", http.StatusBadRequest) return } - if err := s.store.RevokeAgent(r.Context(), id, s.now()); err != nil { + failedJobs, err := s.store.RevokeAgent(r.Context(), id, s.now()) + if err != nil { s.fail(w, err) return } s.hub.Disconnect(id) s.broker.Publish(events.Services) // the agent's services were marked removed s.broker.Publish(events.Jobs) // the agent's active jobs were failed + for _, jobID := range failedJobs { + s.broker.Publish(events.JobTopic(jobID)) // open job pages show the failure + } http.Redirect(w, r, "/agents", http.StatusSeeOther) } diff --git a/internal/master/web/server_test.go b/internal/master/web/server_test.go index ba4686b..ec21f14 100644 --- a/internal/master/web/server_test.go +++ b/internal/master/web/server_test.go @@ -27,12 +27,13 @@ import ( // fakeStore is written from the HTTP handler goroutine and read by the test // goroutine, so every access goes through mu. type fakeStore struct { - mu sync.Mutex - services []model.Service - jobs []model.Job - agents []model.Agent - created map[string]string // name → token hash - revokedID int64 + mu sync.Mutex + services []model.Service + jobs []model.Job + agents []model.Agent + created map[string]string // name → token hash + revokedID int64 + revokeFailed []int64 // job IDs RevokeAgent reports as failed } func (f *fakeStore) ListServices(context.Context) ([]model.Service, error) { return f.services, nil } @@ -76,11 +77,11 @@ func (f *fakeStore) CreateAgent(_ context.Context, name, hash string) (model.Age return model.Agent{ID: 9, Name: name}, nil } -func (f *fakeStore) RevokeAgent(_ context.Context, id int64, _ time.Time) error { +func (f *fakeStore) RevokeAgent(_ context.Context, id int64, _ time.Time) ([]int64, error) { f.mu.Lock() defer f.mu.Unlock() f.revokedID = id - return nil + return f.revokeFailed, nil } func (f *fakeStore) createdHash(name string) string { @@ -276,6 +277,7 @@ func TestCreateAgent_Validation(t *testing.T) { func TestRevokeAgent_RevokesAndDisconnects(t *testing.T) { srv, store, hub, broker := newTestServer(t) + store.revokeFailed = []int64{5, 6} ch, cancel := broker.Subscribe() defer cancel() resp, err := noRedirectClient().Post(srv.URL+"/agents/2/revoke", "application/x-www-form-urlencoded", nil) @@ -287,16 +289,16 @@ func TestRevokeAgent_RevokesAndDisconnects(t *testing.T) { assert.Equal(t, int64(2), store.revoked()) assert.Equal(t, int64(2), hub.lastDisconnected(), "a revoked agent must be cut off immediately") var got []events.Topic - for len(got) < 2 { + for len(got) < 4 { select { case topic := <-ch: got = append(got, topic) default: - t.Fatalf("revoke published %v, want services and jobs", got) + t.Fatalf("revoke published %v, want services, jobs and one topic per failed job", got) } } - assert.ElementsMatch(t, []events.Topic{events.Services, events.Jobs}, got, - "the overview drops the agent's services and job lists show its failed jobs") + assert.ElementsMatch(t, []events.Topic{events.Services, events.Jobs, events.JobTopic(5), events.JobTopic(6)}, got, + "the overview drops the agent's services; job lists and open job pages show the failures") } func TestRevokeAgent_BadID(t *testing.T) {