From 484ecd0dfa9b4896abe8bf8323b0b02ade419041 Mon Sep 17 00:00:00 2001 From: Emil Date: Thu, 23 Jul 2026 21:44:44 +0300 Subject: [PATCH] Bind result artifacts to lease attempts --- coordinator/internal/domain/artifact.go | 1 + .../storage/postgres/artifact_repo.go | 6 +-- .../storage/postgres/integration_test.go | 33 ++++++++++++ coordinator/internal/usecase/artifact.go | 15 ++++++ coordinator/internal/usecase/task.go | 6 +-- coordinator/internal/usecase/usecase_test.go | 52 +++++++++++++++++++ .../migrations/0008_artifact_attempt.down.sql | 7 +++ .../migrations/0008_artifact_attempt.up.sql | 34 ++++++++++++ 8 files changed, 148 insertions(+), 6 deletions(-) create mode 100644 coordinator/migrations/0008_artifact_attempt.down.sql create mode 100644 coordinator/migrations/0008_artifact_attempt.up.sql diff --git a/coordinator/internal/domain/artifact.go b/coordinator/internal/domain/artifact.go index 43bc219..dbe287f 100644 --- a/coordinator/internal/domain/artifact.go +++ b/coordinator/internal/domain/artifact.go @@ -23,6 +23,7 @@ type Artifact struct { ID uuid.UUID JobID uuid.UUID TaskID *uuid.UUID // nil for a job-level input + Attempt *int // required for a partial result; nil for non-worker artifacts Kind ArtifactKind Filename string StorageKey string diff --git a/coordinator/internal/storage/postgres/artifact_repo.go b/coordinator/internal/storage/postgres/artifact_repo.go index 771b5ee..f42abac 100644 --- a/coordinator/internal/storage/postgres/artifact_repo.go +++ b/coordinator/internal/storage/postgres/artifact_repo.go @@ -26,14 +26,14 @@ func NewArtifactRepo(pool *pgxpool.Pool) *ArtifactRepo { var _ usecase.ArtifactRepository = (*ArtifactRepo)(nil) var artifactColumns = []string{ - "id", "job_id", "task_id", "kind", "filename", "storage_key", + "id", "job_id", "task_id", "attempt", "kind", "filename", "storage_key", "content_type", "size_bytes", "sha256", "created_at", } func (r *ArtifactRepo) Insert(ctx context.Context, a *domain.Artifact) error { sql, args, err := psql.Insert("artifacts"). Columns(artifactColumns...). - Values(a.ID, a.JobID, a.TaskID, string(a.Kind), a.Filename, a.StorageKey, + Values(a.ID, a.JobID, a.TaskID, a.Attempt, string(a.Kind), a.Filename, a.StorageKey, a.ContentType, a.SizeBytes, a.SHA256, a.CreatedAt). ToSql() if err != nil { @@ -59,7 +59,7 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, kind string ) err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan( - &a.ID, &a.JobID, &a.TaskID, &kind, &a.Filename, &a.StorageKey, + &a.ID, &a.JobID, &a.TaskID, &a.Attempt, &kind, &a.Filename, &a.StorageKey, &a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt) if errors.Is(err, pgx.ErrNoRows) { return nil, domain.ErrArtifactNotFound diff --git a/coordinator/internal/storage/postgres/integration_test.go b/coordinator/internal/storage/postgres/integration_test.go index 05a7e90..fd461bd 100644 --- a/coordinator/internal/storage/postgres/integration_test.go +++ b/coordinator/internal/storage/postgres/integration_test.go @@ -259,6 +259,8 @@ func TestListCompletedIsOrderedByChunkIndex(t *testing.T) { return err } art.SetContent(fmt.Sprintf("rsha-%d", task.ChunkIndex), 1) + attempt := 1 + art.Attempt = &attempt if err := artifacts.Insert(ctx, art); err != nil { return err } @@ -269,6 +271,7 @@ func TestListCompletedIsOrderedByChunkIndex(t *testing.T) { } owner := "worker-1" fresh.Status = domain.TaskLeased + fresh.Attempt = attempt fresh.LeaseOwner = &owner expires := now.Add(time.Minute) fresh.LeaseExpiresAt = &expires @@ -350,6 +353,10 @@ func seedArtifact(t *testing.T, pool *pgxpool.Pool, jobID uuid.UUID, taskID *uui t.Fatalf("build artifact: %v", err) } art.SetContent(fmt.Sprintf("sha-%s", art.ID), 3) + if kind == domain.ArtifactPartialResult { + attempt := 1 + art.Attempt = &attempt + } if err := NewArtifactRepo(pool).Insert(context.Background(), art); err != nil { t.Fatalf("insert artifact: %v", err) } @@ -446,6 +453,32 @@ func TestArtifactRepoRoundTrip(t *testing.T) { } } +func TestPartialResultArtifactRoundTripsAttempt(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, tasks := seedJob(t, pool, 1) + taskID := tasks[0].ID + art, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult, + "result.csv", "text/csv", time.Now().UTC()) + if err != nil { + t.Fatal(err) + } + attempt := 2 + art.Attempt = &attempt + art.SetContent("sha", 3) + repo := NewArtifactRepo(pool) + if err := repo.Insert(ctx, art); err != nil { + t.Fatalf("insert: %v", err) + } + got, err := repo.Get(ctx, art.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if got.Attempt == nil || *got.Attempt != attempt { + t.Fatalf("attempt = %v, want %d", got.Attempt, attempt) + } +} + // A shard task stores its input as an artifact and no URI: this exercises the // nullable input_uri column, the input_artifact_id round-trip, and the // ck_tasks_has_input check that requires one or the other. diff --git a/coordinator/internal/usecase/artifact.go b/coordinator/internal/usecase/artifact.go index dcbd64d..5cb8b8b 100644 --- a/coordinator/internal/usecase/artifact.go +++ b/coordinator/internal/usecase/artifact.go @@ -39,6 +39,8 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) ( if err != nil { return nil, err } + attempt := in.Attempt + art.Attempt = &attempt // Stream to storage first: size and checksum are measured here, by us, not // taken from the worker. A large shard never sits in memory. @@ -48,6 +50,19 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) ( } art.SetContent(sum, size) + // The stream may take longer than the lease. Re-check after it finishes so + // an expired worker cannot leave a durable result record behind. Completion + // performs the same ownership check under its transaction. + current, err := uc.tasks.Get(ctx, in.TaskID) + if err != nil { + _ = uc.blobs.Delete(ctx, art.StorageKey) + return nil, err + } + if !current.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) { + _ = uc.blobs.Delete(ctx, art.StorageKey) + return nil, domain.ErrLeaseConflict + } + // Persist the record. If that fails the blob would be an orphan, so remove it. if err := uc.artifacts.Insert(ctx, art); err != nil { _ = uc.blobs.Delete(ctx, art.StorageKey) diff --git a/coordinator/internal/usecase/task.go b/coordinator/internal/usecase/task.go index 7d01b25..131cffe 100644 --- a/coordinator/internal/usecase/task.go +++ b/coordinator/internal/usecase/task.go @@ -144,7 +144,7 @@ func (uc *CompleteTask) Execute(ctx context.Context, in CompleteTaskInput) (*dom } // Rule 10: never trust a worker-supplied artifact reference. The result // must be an artifact the coordinator itself stored for *this* task. - if err := uc.verifyResultArtifact(ctx, in.TaskID, in.ResultArtifactID); err != nil { + if err := uc.verifyResultArtifact(ctx, in.TaskID, in.Attempt, in.ResultArtifactID); err != nil { return err } now := uc.clock.Now() @@ -176,12 +176,12 @@ func (uc *CompleteTask) Execute(ctx context.Context, in CompleteTaskInput) (*dom // verifyResultArtifact enforces that the referenced artifact was stored by the // coordinator for this exact task. It stops a worker from completing task B with // an artifact it uploaded for task A, and from naming an id that isn't a result. -func (uc *CompleteTask) verifyResultArtifact(ctx context.Context, taskID, artifactID uuid.UUID) error { +func (uc *CompleteTask) verifyResultArtifact(ctx context.Context, taskID uuid.UUID, attempt int, artifactID uuid.UUID) error { art, err := uc.artifacts.Get(ctx, artifactID) if err != nil { return err } - if art.TaskID == nil || *art.TaskID != taskID || art.Kind != domain.ArtifactPartialResult { + if art.TaskID == nil || *art.TaskID != taskID || art.Attempt == nil || *art.Attempt != attempt || art.Kind != domain.ArtifactPartialResult { return domain.ErrResultConflict } return nil diff --git a/coordinator/internal/usecase/usecase_test.go b/coordinator/internal/usecase/usecase_test.go index 254a8ef..4786b53 100644 --- a/coordinator/internal/usecase/usecase_test.go +++ b/coordinator/internal/usecase/usecase_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "io" "strings" "testing" "time" @@ -19,6 +20,17 @@ var ctx = context.Background() const lease = 2 * time.Minute +type expiringBlobStore struct { + *memstore.BlobStore + clock *memstore.Clock +} + +func (s expiringBlobStore) Put(ctx context.Context, key string, body io.Reader) (string, int64, error) { + sum, size, err := s.BlobStore.Put(ctx, key, body) + s.clock.Advance(lease + time.Second) + return sum, size, err +} + // harness wires every use case to in-memory stores so orchestration can be // tested without a database. type harness struct { @@ -239,6 +251,46 @@ func TestCompleteRejectsForeignArtifact(t *testing.T) { } } +func TestCompleteRejectsArtifactFromExpiredAttempt(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) + taskID, attemptOne := h.leaseOne(t, "w1", "w") + staleArtifact := h.uploadResult(t, taskID, "w1", attemptOne) + + h.clk.Advance(lease + time.Second) + if _, err := h.expire.Execute(ctx); err != nil { + t.Fatalf("expire lease: %v", err) + } + _, attemptTwo := h.leaseOne(t, "w2", "w") + if attemptTwo != attemptOne+1 { + t.Fatalf("attempt = %d, want %d", attemptTwo, attemptOne+1) + } + + _, err := h.complete.Execute(ctx, usecase.CompleteTaskInput{ + TaskID: taskID, WorkerID: "w2", Attempt: attemptTwo, ResultArtifactID: staleArtifact, + }) + if !errors.Is(err, domain.ErrResultConflict) { + t.Errorf("stale-attempt artifact: err = %v, want ErrResultConflict", err) + } +} + +func TestUploadRejectsLeaseThatExpiresDuringStreaming(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) + taskID, attempt := h.leaseOne(t, "w1", "w") + h.uploadArt = usecase.NewUploadArtifact( + h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, h.clk, + ) + + _, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{ + TaskID: taskID, WorkerID: "w1", Attempt: attempt, + Filename: "result.csv", ContentType: "text/csv", Body: strings.NewReader("result"), + }) + if !errors.Is(err, domain.ErrLeaseConflict) { + t.Errorf("upload after lease expiry: err = %v, want ErrLeaseConflict", err) + } +} + func TestCompleteIsIdempotentOnReplay(t *testing.T) { h := newHarness() h.seedJob(t, "w", 1) diff --git a/coordinator/migrations/0008_artifact_attempt.down.sql b/coordinator/migrations/0008_artifact_attempt.down.sql new file mode 100644 index 0000000..f0731bb --- /dev/null +++ b/coordinator/migrations/0008_artifact_attempt.down.sql @@ -0,0 +1,7 @@ +BEGIN; + +ALTER TABLE artifacts DROP CONSTRAINT IF EXISTS ck_artifact_attempt_positive; +ALTER TABLE artifacts DROP CONSTRAINT IF EXISTS ck_partial_result_attempt; +ALTER TABLE artifacts DROP COLUMN IF EXISTS attempt; + +COMMIT; diff --git a/coordinator/migrations/0008_artifact_attempt.up.sql b/coordinator/migrations/0008_artifact_attempt.up.sql new file mode 100644 index 0000000..a0edae7 --- /dev/null +++ b/coordinator/migrations/0008_artifact_attempt.up.sql @@ -0,0 +1,34 @@ +BEGIN; + +-- A partial result belongs to the lease attempt that uploaded it. Without this +-- binding a worker holding a later retry could complete a task with stale bytes +-- uploaded by an expired attempt of that same task. +ALTER TABLE artifacts ADD COLUMN attempt integer; + +-- A completed task never gets a later lease, so its current attempt is also +-- the attempt that produced the stored result. +UPDATE artifacts AS a +SET attempt = t.attempt +FROM tasks AS t +WHERE a.task_id = t.id + AND a.kind = 'partial_result'::artifact_kind + AND t.status = 'completed'::task_status + AND a.attempt IS NULL; + +-- For unfinished tasks the old schema cannot tell which attempt uploaded a +-- partial result. Keeping it would let a later retry claim stale bytes, so the +-- worker must upload again. Blob garbage is harmless and follows the existing +-- coordinator-owned storage cleanup policy. +DELETE FROM artifacts AS a +USING tasks AS t +WHERE a.task_id = t.id + AND a.kind = 'partial_result'::artifact_kind + AND t.status <> 'completed'::task_status + AND a.attempt IS NULL; + +ALTER TABLE artifacts ADD CONSTRAINT ck_partial_result_attempt + CHECK (kind <> 'partial_result'::artifact_kind OR attempt IS NOT NULL); +ALTER TABLE artifacts ADD CONSTRAINT ck_artifact_attempt_positive + CHECK (attempt IS NULL OR attempt > 0); + +COMMIT;