Bind result artifacts to lease attempts

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)