Bind result artifacts to lease attempts

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

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

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.

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)

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

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)

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;

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;