Emil_Shanaty/SciMesh
1
0
Fork 0
Code Issues Pull Requests Actions 121 Packages Projects Releases Wiki Activity

Bind result artifacts to lease attempts

Browse Source
This commit is contained in:
Emil
2026-07-23 21:44:44 +03:00
parent 983c5843ec
commit 484ecd0dfa
8 changed files with 148 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coordinator/internal/domain/artifact.go
+1
View File
@@ -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
+3 -3
View File
@@ -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
+33
View File
@@ -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
+15
View File
@@ -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
+3 -3
View File
@@ -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
+52
View File
@@ -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
+7
View File
@@ -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
+34
View File
@@ -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;