From f1c3163be4315ae0261c01ac927280e79498358f Mon Sep 17 00:00:00 2001 From: Efremenko Arhip Date: Wed, 22 Jul 2026 14:48:58 +0300 Subject: [PATCH] feat(coordinator): implement PostgreSQL repositories MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the repository stubs with real pgx queries, so the queue now works end to end: a job is split into tasks, leased to workers one at a time, heartbeated, completed, and reflected in job progress. Task claiming is a single statement — SELECT ... FOR UPDATE SKIP LOCKED feeding an UPDATE — so concurrent coordinators lease different rows instead of blocking on the same one. Writes use optimistic concurrency: the entity increments its version in memory, and the UPDATE guards on the previous value. Retries moved to the transaction level. Once Postgres aborts a transaction with a serialization failure, replaying one statement inside it cannot help; the unit of retry is Begin -> fn -> Commit, which is safe because each attempt re-reads its rows through GetForUpdate. Adds integration tests behind the `integration` build tag, run against a real PostgreSQL through TEST_DATABASE_URL: concurrent claiming hands each task to exactly one worker, job creation rolls back whole, stale writes are refused, completed results keep chunk order, and expired leases return to the queue. Two bugs the tests caught: - a nil parameters map reached a NOT NULL jsonb column as SQL NULL, since pgx sends NULL rather than omitting the column and letting DEFAULT '{}' apply; - replaying an already-recorded result returned 409. The idempotent path leaves the entity untouched, so the version guard matched nothing and a successful no-op looked like a conflict. CompleteTask now skips the write when the entity did not change. --- coordinator/Makefile | 9 +- coordinator/README.md | 19 +- .../storage/postgres/integration_test.go | 360 ++++++++++++++++++ .../internal/storage/postgres/job_repo.go | 49 ++- .../internal/storage/postgres/task_repo.go | 257 ++++++++++--- coordinator/internal/storage/postgres/tx.go | 34 +- coordinator/internal/usecase/task.go | 11 +- 7 files changed, 668 insertions(+), 71 deletions(-) create mode 100644 coordinator/internal/storage/postgres/integration_test.go diff --git a/coordinator/Makefile b/coordinator/Makefile index d85c19e..38e4b54 100644 --- a/coordinator/Makefile +++ b/coordinator/Makefile @@ -10,6 +10,11 @@ run: test: go test ./... +# Needs a running PostgreSQL; the spec forbids mocks for these guarantees. +# make test-integration TEST_DATABASE_URL='postgres://...' +test-integration: + TEST_DATABASE_URL="$(TEST_DATABASE_URL)" go test -tags=integration ./... -v + vet: go vet ./... @@ -18,8 +23,8 @@ vet: LINT_VERSION := v2.12.2 lint: @command -v golangci-lint >/dev/null 2>&1 \ - && golangci-lint run ./... \ - || go run github.com/golangci/golangci-lint/v2/cmd/golangci-lint@$(LINT_VERSION) run ./... + && golangci-lint run --build-tags=integration ./... \ + || go run github.com/golangci/golangci-lint/v2/cmd/golangci-lint@$(LINT_VERSION) run --build-tags=integration ./... tidy: go mod tidy diff --git a/coordinator/README.md b/coordinator/README.md index f83a9be..ef54ec1 100644 --- a/coordinator/README.md +++ b/coordinator/README.md @@ -113,20 +113,21 @@ See `.env.example`; only `DATABASE_URL` is required. ## Status -Scaffold with a **complete, tested domain**. Layers, wiring, routing, auth, -access logging, error mapping, transactions, migrations, and graceful shutdown -are in place. Repository methods are stubs returning `ErrNotImplemented` -(→ HTTP 501); the SQL for claiming and lease expiry is written and ready to wire. +The queue works end to end: a job can be submitted, split into tasks, leased to +workers one at a time, heartbeated, completed, and reflected in job progress. Roadmap: 1. schema + migrations ✅ -2. `ClaimNext`, `InsertBatch` — atomic claim via `FOR UPDATE SKIP LOCKED` -3. `GetForUpdate`, `Update`, `CountByStatus` — completes the result/failure paths -4. file upload / chunk download -5. `ExpireLeases` — SQL is written, needs wiring +2. `ClaimNext`, `InsertBatch` — atomic claim via `FOR UPDATE SKIP LOCKED` ✅ +3. `GetForUpdate`, `Update`, `CountByStatus` — result/failure paths ✅ +4. file upload / chunk download — **next** +5. `ExpireLeases` ✅ (reaper + a sweep before every claim) 6. stitcher: merge per-chunk top-k into the final CSV -7. integration tests against real Postgres via `TEST_DATABASE_URL` +7. more integration coverage as features land + +Still stubbed: `StitchJob.Execute`, and there is no `POST /upload` or +`GET /download_chunk` yet — so chunk files must be referenced by URI for now. ## Tests diff --git a/coordinator/internal/storage/postgres/integration_test.go b/coordinator/internal/storage/postgres/integration_test.go new file mode 100644 index 0000000..d6ec8a7 --- /dev/null +++ b/coordinator/internal/storage/postgres/integration_test.go @@ -0,0 +1,360 @@ +//go:build integration + +// Integration tests run against a real PostgreSQL instance supplied through +// TEST_DATABASE_URL. The spec forbids mocks or SQLite here: the guarantees +// being verified — FOR UPDATE SKIP LOCKED, optimistic concurrency, transaction +// rollback — are properties of Postgres, not of our Go code. +// +// docker compose up -d +// TEST_DATABASE_URL='postgres://scimesh:scimesh@localhost:5432/scimesh?sslmode=disable' \ +// go test -tags=integration ./internal/storage/postgres/ -v +package postgres + +import ( + "context" + "errors" + "fmt" + "os" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +func testPool(t *testing.T) *pgxpool.Pool { + t.Helper() + url := os.Getenv("TEST_DATABASE_URL") + if url == "" { + t.Skip("TEST_DATABASE_URL is not set") + } + pool, err := pgxpool.New(context.Background(), url) + if err != nil { + t.Fatalf("connect: %v", err) + } + t.Cleanup(pool.Close) + return pool +} + +// seedJob creates a job with n pending tasks and removes them afterwards, so +// tests stay independent of each other and of leftovers from earlier runs. +func seedJob(t *testing.T, pool *pgxpool.Pool, n int) (*domain.Job, []*domain.Task) { + t.Helper() + ctx := context.Background() + + chunks := make([]domain.ChunkSpec, 0, n) + for i := 0; i < n; i++ { + chunks = append(chunks, domain.ChunkSpec{ + ChunkIndex: i, + InputURI: fmt.Sprintf("s3://chunk-%d", i), + InputSHA256: fmt.Sprintf("sha-%d", i), + }) + } + job, tasks, err := domain.NewJobWithTasks("similarity_search", "s3://ds", nil, chunks, time.Now().UTC()) + if err != nil { + t.Fatalf("build job: %v", err) + } + + jobs, taskRepo, tx := NewJobRepo(pool), NewTaskRepo(pool), NewTxManager(pool) + err = tx.WithinTx(ctx, func(ctx context.Context) error { + if err := jobs.Insert(ctx, job); err != nil { + return err + } + return taskRepo.InsertBatch(ctx, tasks) + }) + if err != nil { + t.Fatalf("seed: %v", err) + } + + t.Cleanup(func() { + // ON DELETE CASCADE removes the tasks with it. + _, _ = pool.Exec(context.Background(), `DELETE FROM jobs WHERE id = $1`, job.ID) + }) + return job, tasks +} + +func TestCreateJobPersistsEveryTask(t *testing.T) { + pool := testPool(t) + job, _ := seedJob(t, pool, 3) + + counts, err := NewTaskRepo(pool).CountByStatus(context.Background(), job.ID) + if err != nil { + t.Fatalf("count: %v", err) + } + if counts[domain.TaskPending] != 3 { + t.Errorf("pending = %d, want 3", counts[domain.TaskPending]) + } +} + +// A job must land whole or not at all: a half-created job leaves chunks no +// worker could ever complete. +func TestCreateJobRollsBackOnFailure(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + + chunks := []domain.ChunkSpec{{ChunkIndex: 0, InputURI: "s3://c0", InputSHA256: "sha0"}} + job, tasks, err := domain.NewJobWithTasks("similarity_search", "s3://ds", nil, chunks, time.Now().UTC()) + if err != nil { + t.Fatal(err) + } + + jobs, taskRepo, tx := NewJobRepo(pool), NewTaskRepo(pool), NewTxManager(pool) + boom := errors.New("boom") + err = tx.WithinTx(ctx, func(ctx context.Context) error { + if err := jobs.Insert(ctx, job); err != nil { + return err + } + if err := taskRepo.InsertBatch(ctx, tasks); err != nil { + return err + } + return boom // fail after both writes + }) + if !errors.Is(err, boom) { + t.Fatalf("err = %v, want boom", err) + } + + if _, err := jobs.Get(ctx, job.ID); !errors.Is(err, domain.ErrJobNotFound) { + t.Errorf("job survived the rollback: %v", err) + } +} + +// The acceptance criterion: N workers claiming at once must each get a +// different task, and no task may be handed out twice. +func TestConcurrentClaimGivesEachTaskToExactlyOneWorker(t *testing.T) { + pool := testPool(t) + const tasks = 8 + job, _ := seedJob(t, pool, tasks) + + repo := NewTaskRepo(pool) + now := time.Now().UTC() + + var ( + mu sync.Mutex + claimed = make(map[uuid.UUID]string) + wg sync.WaitGroup + ) + // More workers than tasks, so the surplus must come back empty rather than + // steal an already-leased row. + for i := 0; i < tasks*2; i++ { + wg.Add(1) + go func(n int) { + defer wg.Done() + task, err := repo.ClaimNext(context.Background(), usecase.ClaimFilter{ + Owner: fmt.Sprintf("worker-%d", n), + Now: now, + LeaseUntil: now.Add(time.Minute), + }) + if err != nil { + t.Errorf("claim: %v", err) + return + } + if task == nil || task.JobID != job.ID { + return // empty queue, or a task from another test's job + } + mu.Lock() + defer mu.Unlock() + if prev, dup := claimed[task.ID]; dup { + t.Errorf("task %s handed to both %s and worker-%d", task.ID, prev, n) + } + claimed[task.ID] = fmt.Sprintf("worker-%d", n) + }(i) + } + wg.Wait() + + if len(claimed) != tasks { + t.Errorf("claimed %d tasks, want %d", len(claimed), tasks) + } +} + +func TestClaimNextReturnsNilOnEmptyQueue(t *testing.T) { + pool := testPool(t) + now := time.Now().UTC() + + // Drain everything first, then ask once more. + repo := NewTaskRepo(pool) + for { + task, err := repo.ClaimNext(context.Background(), usecase.ClaimFilter{ + Owner: "drainer", Now: now, LeaseUntil: now.Add(time.Minute), + }) + if err != nil { + t.Fatalf("drain: %v", err) + } + if task == nil { + break + } + } + + task, err := repo.ClaimNext(context.Background(), usecase.ClaimFilter{ + Owner: "worker-1", Now: now, LeaseUntil: now.Add(time.Minute), + }) + if err != nil { + t.Fatalf("claim: %v", err) + } + if task != nil { + t.Errorf("expected nil on an empty queue, got %s", task.ID) + } +} + +func TestUpdateRejectsStaleVersion(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, _ := seedJob(t, pool, 1) + + repo, tx := NewTaskRepo(pool), NewTxManager(pool) + now := time.Now().UTC() + + task, err := repo.ClaimNext(ctx, usecase.ClaimFilter{ + Owner: "worker-1", Now: now, LeaseUntil: now.Add(time.Minute), + }) + if err != nil || task == nil || task.JobID != job.ID { + t.Skipf("could not claim this job's task (got %v, %v)", task, err) + } + + // A stale copy: same row, but the version it remembers is behind. + stale := *task + stale.Version = task.Version // pretend the caller mutated it once + + err = tx.WithinTx(ctx, func(ctx context.Context) error { + fresh, err := repo.GetForUpdate(ctx, task.ID) + if err != nil { + return err + } + if err := fresh.RenewLease("worker-1", fresh.Attempt, now.Add(2*time.Minute)); err != nil { + return err + } + return repo.Update(ctx, fresh) + }) + if err != nil { + t.Fatalf("legitimate update failed: %v", err) + } + + // Now the stale copy's version is behind by one; its write must be refused. + stale.Version++ // as a domain method would have done + if err := repo.Update(ctx, &stale); !errors.Is(err, domain.ErrLeaseConflict) { + t.Errorf("stale update err = %v, want ErrLeaseConflict", err) + } +} + +func TestListCompletedIsOrderedByChunkIndex(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, tasks := seedJob(t, pool, 4) + + repo, tx := NewTaskRepo(pool), NewTxManager(pool) + now := time.Now().UTC() + + // Complete them out of order to prove the ordering comes from SQL. + for _, i := range []int{2, 0, 3, 1} { + task := tasks[i] + err := tx.WithinTx(ctx, func(ctx context.Context) error { + fresh, err := repo.GetForUpdate(ctx, task.ID) + if err != nil { + return err + } + owner := "worker-1" + fresh.Status = domain.TaskLeased + fresh.LeaseOwner = &owner + expires := now.Add(time.Minute) + fresh.LeaseExpiresAt = &expires + if err := fresh.CompleteWith( + fmt.Sprintf("s3://result-%d", fresh.ChunkIndex), + fmt.Sprintf("rsha-%d", fresh.ChunkIndex), + nil, owner, fresh.Attempt, now, + ); err != nil { + return err + } + return repo.Update(ctx, fresh) + }) + if err != nil { + t.Fatalf("complete chunk %d: %v", i, err) + } + } + + done, err := repo.ListCompleted(ctx, job.ID) + if err != nil { + t.Fatalf("list: %v", err) + } + if len(done) != 4 { + t.Fatalf("got %d completed, want 4", len(done)) + } + for i, task := range done { + if task.ChunkIndex != i { + t.Errorf("position %d holds chunk_index %d — order is not deterministic", i, task.ChunkIndex) + } + } +} + +// A worker whose network dropped resends the same manifest. That must succeed: +// the entity is unchanged, so nothing is written, and the optimistic-concurrency +// guard must not turn the replay into a conflict. +func TestCompleteTaskReplayIsIdempotent(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, _ := seedJob(t, pool, 1) + + tasks, jobs, tx := NewTaskRepo(pool), NewJobRepo(pool), NewTxManager(pool) + clk := fixedClock{now: time.Now().UTC()} + uc := usecase.NewCompleteTask(tasks, jobs, tx, clk) + + claimed, err := tasks.ClaimNext(ctx, usecase.ClaimFilter{ + Owner: "worker-1", Now: clk.now, LeaseUntil: clk.now.Add(time.Minute), + }) + if err != nil || claimed == nil || claimed.JobID != job.ID { + t.Skipf("could not claim this job's task (got %v, %v)", claimed, err) + } + + in := usecase.CompleteTaskInput{ + TaskID: claimed.ID, WorkerID: "worker-1", Attempt: claimed.Attempt, + ResultURI: "s3://r0", ResultSHA256: "rsha", + } + if _, err := uc.Execute(ctx, in); err != nil { + t.Fatalf("first submission: %v", err) + } + if _, err := uc.Execute(ctx, in); err != nil { + t.Errorf("replay must be idempotent, got %v", err) + } + + // A different manifest for the same task is a genuine conflict. + other := in + other.ResultURI = "s3://different" + if _, err := uc.Execute(ctx, other); !errors.Is(err, domain.ErrResultConflict) { + t.Errorf("err = %v, want ErrResultConflict", err) + } +} + +type fixedClock struct{ now time.Time } + +func (c fixedClock) Now() time.Time { return c.now } + +func TestExpireLeasesRequeuesElapsedTasks(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, _ := seedJob(t, pool, 1) + + repo := NewTaskRepo(pool) + past := time.Now().UTC().Add(-time.Hour) + + // Lease it with an expiry already in the past. + task, err := repo.ClaimNext(ctx, usecase.ClaimFilter{ + Owner: "dead-worker", Now: past, LeaseUntil: past.Add(time.Minute), + }) + if err != nil || task == nil || task.JobID != job.ID { + t.Skipf("could not claim this job's task (got %v, %v)", task, err) + } + + if _, err := repo.ExpireLeases(ctx, time.Now().UTC()); err != nil { + t.Fatalf("expire: %v", err) + } + + counts, err := repo.CountByStatus(ctx, job.ID) + if err != nil { + t.Fatalf("count: %v", err) + } + if counts[domain.TaskPending] != 1 { + t.Errorf("pending = %d, want 1 — a dead worker must not strand its task", counts[domain.TaskPending]) + } +} diff --git a/coordinator/internal/storage/postgres/job_repo.go b/coordinator/internal/storage/postgres/job_repo.go index 3751801..1945a07 100644 --- a/coordinator/internal/storage/postgres/job_repo.go +++ b/coordinator/internal/storage/postgres/job_repo.go @@ -2,9 +2,11 @@ package postgres import ( "context" + "errors" "time" "github.com/google/uuid" + "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" "github.com/emil28092005/SciMesh/coordinator/internal/domain" @@ -22,19 +24,50 @@ func NewJobRepo(pool *pgxpool.Pool) *JobRepo { var _ usecase.JobRepository = (*JobRepo)(nil) -// TODO(phase 2-3): replace stubs with real pgx queries. +const jobColumns = `id, workload, input_uri, parameters, status, created_at, completed_at` +const insertJobSQL = ` +INSERT INTO jobs (id, workload, input_uri, parameters, status, created_at) +VALUES ($1, $2, $3, $4, $5, $6)` + +// Insert runs inside the caller's transaction, alongside the job's tasks — that +// is what makes "all tasks or none" hold. func (r *JobRepo) Insert(ctx context.Context, j *domain.Job) error { - // Phase 2: INSERT INTO jobs ...; runs inside the caller's transaction. - return usecase.ErrNotImplemented + _, err := conn(ctx, r.pool).Exec(ctx, insertJobSQL, + j.ID, j.Workload, j.InputURI, jsonbOrEmpty(j.Parameters), string(j.Status), j.CreatedAt) + return err } +const getJobSQL = `SELECT ` + jobColumns + ` FROM jobs WHERE id = $1` + func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) { - // Phase 3: SELECT ... WHERE id = $1; no rows -> domain.ErrJobNotFound. - return nil, usecase.ErrNotImplemented + var ( + j domain.Job + status string + ) + err := conn(ctx, r.pool).QueryRow(ctx, getJobSQL, id).Scan( + &j.ID, &j.Workload, &j.InputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt) + if errors.Is(err, pgx.ErrNoRows) { + return nil, domain.ErrJobNotFound + } + if err != nil { + return nil, err + } + j.Status = domain.JobStatus(status) + return &j, nil } -func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID, status domain.JobStatus, completedAt *time.Time) error { - // Phase 3: UPDATE jobs SET status = $2, completed_at = $3 WHERE id = $1. - return usecase.ErrNotImplemented +const updateJobStatusSQL = `UPDATE jobs SET status = $2, completed_at = $3 WHERE id = $1` + +func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID, + status domain.JobStatus, completedAt *time.Time) error { + + tag, err := conn(ctx, r.pool).Exec(ctx, updateJobStatusSQL, id, string(status), completedAt) + if err != nil { + return err + } + if tag.RowsAffected() == 0 { + return domain.ErrJobNotFound + } + return nil } diff --git a/coordinator/internal/storage/postgres/task_repo.go b/coordinator/internal/storage/postgres/task_repo.go index 631a88b..4c0abe1 100644 --- a/coordinator/internal/storage/postgres/task_repo.go +++ b/coordinator/internal/storage/postgres/task_repo.go @@ -2,9 +2,11 @@ package postgres import ( "context" + "errors" "time" "github.com/google/uuid" + "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" "github.com/emil28092005/SciMesh/coordinator/internal/domain" @@ -22,6 +24,37 @@ func NewTaskRepo(pool *pgxpool.Pool) *TaskRepo { var _ usecase.TaskRepository = (*TaskRepo)(nil) +// taskColumns is the single source of truth for the shape scanTask expects. +// Every query that returns a task selects exactly this list, in this order — +// three hand-written column lists would drift apart within a week. +const taskColumns = `id, job_id, chunk_index, workload, input_uri, input_sha256, + parameters, status, attempt, max_attempts, lease_owner, lease_expires_at, + result_uri, result_sha256, metrics, error_code, error_message, + created_at, started_at, completed_at, version` + +// scanTask maps one row onto an entity. +// +// status is read into a plain string rather than domain.TaskStatus: pgx does +// not know the task_status enum, and going through string keeps the driver out +// of the domain's type system. +func scanTask(row pgx.Row) (*domain.Task, error) { + var ( + t domain.Task + status string + ) + err := row.Scan( + &t.ID, &t.JobID, &t.ChunkIndex, &t.Workload, &t.InputURI, &t.InputSHA256, + &t.Parameters, &status, &t.Attempt, &t.MaxAttempts, &t.LeaseOwner, &t.LeaseExpiresAt, + &t.ResultURI, &t.ResultSHA256, &t.Metrics, &t.ErrorCode, &t.ErrorMessage, + &t.CreatedAt, &t.StartedAt, &t.CompletedAt, &t.Version, + ) + if err != nil { + return nil, err + } + t.Status = domain.TaskStatus(status) + return &t, nil +} + // claimNextSQL leases one task in a single statement. // // FOR UPDATE SKIP LOCKED is what makes concurrent coordinators safe: each @@ -29,10 +62,11 @@ var _ usecase.TaskRepository = (*TaskRepo)(nil) // so no task is ever handed to two workers and no claim blocks behind another. // Splitting this into SELECT + UPDATE would reintroduce exactly that race. // -//nolint:unused // wired up in phase 2; kept beside the repository it belongs to +// The CTE column is aliased to cid so the RETURNING list below can use bare +// column names without colliding with the candidate relation. const claimNextSQL = ` WITH candidate AS ( - SELECT id + SELECT id AS cid FROM tasks WHERE status = 'pending' AND attempt < max_attempts @@ -49,17 +83,174 @@ SET status = 'leased', started_at = COALESCE(started_at, $4), version = version + 1 FROM candidate -WHERE tasks.id = candidate.id -RETURNING tasks.id, tasks.job_id, tasks.chunk_index, tasks.workload, - tasks.input_uri, tasks.input_sha256, tasks.parameters, - tasks.status, tasks.attempt, tasks.max_attempts, - tasks.lease_owner, tasks.lease_expires_at, tasks.version; -` +WHERE tasks.id = candidate.cid +RETURNING ` + taskColumns + +// ClaimNext atomically leases the next eligible task. +func (r *TaskRepo) ClaimNext(ctx context.Context, f usecase.ClaimFilter) (*domain.Task, error) { + workloads := f.Workloads + if workloads == nil { + workloads = []string{} // NULL would make the cardinality() guard fail + } + + var task *domain.Task + err := withRetry(ctx, func(ctx context.Context) error { + row := conn(ctx, r.pool).QueryRow(ctx, claimNextSQL, workloads, f.Owner, f.LeaseUntil, f.Now) + t, err := scanTask(row) + if errors.Is(err, pgx.ErrNoRows) { + task = nil + return nil // an empty queue is a normal state, not a failure + } + if err != nil { + return err + } + task = t + return nil + }) + if err != nil { + return nil, err + } + return task, nil +} + +const getForUpdateSQL = `SELECT ` + taskColumns + ` FROM tasks WHERE id = $1 FOR UPDATE` + +// GetForUpdate reads a task and holds its row lock until the caller's +// transaction ends, so read-modify-write use cases cannot interleave. +func (r *TaskRepo) GetForUpdate(ctx context.Context, id uuid.UUID) (*domain.Task, error) { + t, err := scanTask(conn(ctx, r.pool).QueryRow(ctx, getForUpdateSQL, id)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, domain.ErrTaskNotFound + } + if err != nil { + return nil, err + } + return t, nil +} + +// updateTaskSQL writes the mutated entity back under optimistic concurrency. +// +// The entity has already incremented its Version in memory, so the new value +// goes into SET while the guard in WHERE compares against the previous one. +const updateTaskSQL = ` +UPDATE tasks +SET status = $2, + attempt = $3, + lease_owner = $4, + lease_expires_at = $5, + result_uri = $6, + result_sha256 = $7, + metrics = $8, + error_code = $9, + error_message = $10, + started_at = $11, + completed_at = $12, + version = $13 +WHERE id = $1 AND version = $13 - 1` + +func (r *TaskRepo) Update(ctx context.Context, t *domain.Task) error { + tag, err := conn(ctx, r.pool).Exec(ctx, updateTaskSQL, + t.ID, string(t.Status), t.Attempt, t.LeaseOwner, t.LeaseExpiresAt, + t.ResultURI, t.ResultSHA256, t.Metrics, t.ErrorCode, t.ErrorMessage, + t.StartedAt, t.CompletedAt, t.Version, + ) + if err != nil { + return err + } + if tag.RowsAffected() == 0 { + // Either the row vanished or someone else advanced its version while we + // held a stale copy. Both mean this write must not land. + return domain.ErrLeaseConflict + } + return nil +} + +const insertTaskSQL = ` +INSERT INTO tasks (id, job_id, chunk_index, workload, input_uri, input_sha256, + parameters, status, attempt, max_attempts, created_at, version) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)` + +// InsertBatch writes every task in one round trip. It runs inside the caller's +// transaction, which is what makes "all tasks or none" hold. +func (r *TaskRepo) InsertBatch(ctx context.Context, tasks []*domain.Task) error { + if len(tasks) == 0 { + return nil + } + + batch := &pgx.Batch{} + for _, t := range tasks { + batch.Queue(insertTaskSQL, + t.ID, t.JobID, t.ChunkIndex, t.Workload, t.InputURI, t.InputSHA256, + jsonbOrEmpty(t.Parameters), string(t.Status), t.Attempt, t.MaxAttempts, t.CreatedAt, t.Version, + ) + } + + results := conn(ctx, r.pool).SendBatch(ctx, batch) + for range tasks { + if _, err := results.Exec(); err != nil { + _ = results.Close() + return err + } + } + return results.Close() +} + +const listCompletedSQL = ` +SELECT ` + taskColumns + ` +FROM tasks +WHERE job_id = $1 AND status = 'completed' +ORDER BY chunk_index` + +// ListCompleted returns results in chunk order, which the stitcher relies on: +// a non-deterministic order would make the merged output depend on which worker +// happened to finish first. +func (r *TaskRepo) ListCompleted(ctx context.Context, jobID uuid.UUID) ([]*domain.Task, error) { + rows, err := conn(ctx, r.pool).Query(ctx, listCompletedSQL, jobID) + if err != nil { + return nil, err + } + defer rows.Close() + + var tasks []*domain.Task + for rows.Next() { + t, err := scanTask(rows) + if err != nil { + return nil, err + } + tasks = append(tasks, t) + } + return tasks, rows.Err() +} + +const countByStatusSQL = `SELECT status, count(*) FROM tasks WHERE job_id = $1 GROUP BY status` + +func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[domain.TaskStatus]int, error) { + rows, err := conn(ctx, r.pool).Query(ctx, countByStatusSQL, jobID) + if err != nil { + return nil, err + } + defer rows.Close() + + counts := make(map[domain.TaskStatus]int) + for rows.Next() { + var ( + status string + n int + ) + if err := rows.Scan(&status, &n); err != nil { + return nil, err + } + counts[domain.TaskStatus(status)] = n + } + return counts, rows.Err() +} // expireLeasesSQL applies the lease-expiry rule set-based, mirroring // domain.Task.ExpireLease: requeue while attempts remain, otherwise fail. // -//nolint:unused // wired up in phase 5; mirrors domain.Task.ExpireLease +// It is one statement rather than a load-decide-save loop because several +// coordinators run it concurrently; an atomic UPDATE makes the duplicate work +// harmless — the loser simply updates zero rows. const expireLeasesSQL = ` UPDATE tasks SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_status @@ -72,43 +263,17 @@ SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_sta ELSE error_message END, completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END, version = version + 1 -WHERE status = 'leased' AND lease_expires_at < $1; -` - -// TODO(phase 2-6): replace stubs with real pgx queries; the SQL above is ready -// to wire up. Each method maps 1:1 to a roadmap phase. - -func (r *TaskRepo) ClaimNext(ctx context.Context, f usecase.ClaimFilter) (*domain.Task, error) { - // Phase 2: run claimNextSQL; pgx.ErrNoRows -> (nil, nil). - return nil, usecase.ErrNotImplemented -} - -func (r *TaskRepo) GetForUpdate(ctx context.Context, id uuid.UUID) (*domain.Task, error) { - // Phase 3: SELECT ... WHERE id = $1 FOR UPDATE; no rows -> domain.ErrTaskNotFound. - return nil, usecase.ErrNotImplemented -} - -func (r *TaskRepo) Update(ctx context.Context, t *domain.Task) error { - // Phase 3: UPDATE ... WHERE id = $1 AND version = $2 (optimistic concurrency). - return usecase.ErrNotImplemented -} - -func (r *TaskRepo) InsertBatch(ctx context.Context, tasks []*domain.Task) error { - // Phase 2: pgx.Batch or COPY; runs inside the caller's transaction. - return usecase.ErrNotImplemented -} - -func (r *TaskRepo) ListCompleted(ctx context.Context, jobID uuid.UUID) ([]*domain.Task, error) { - // Phase 6: WHERE job_id = $1 AND status = 'completed' ORDER BY chunk_index. - return nil, usecase.ErrNotImplemented -} - -func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[domain.TaskStatus]int, error) { - // Phase 3: SELECT status, count(*) ... GROUP BY status. - return nil, usecase.ErrNotImplemented -} +WHERE status = 'leased' AND lease_expires_at < $1` func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) { - // Phase 5: run expireLeasesSQL, return the affected row count. - return 0, usecase.ErrNotImplemented + var affected int64 + err := withRetry(ctx, func(ctx context.Context) error { + tag, err := conn(ctx, r.pool).Exec(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired) + if err != nil { + return err + } + affected = tag.RowsAffected() + return nil + }) + return affected, err } diff --git a/coordinator/internal/storage/postgres/tx.go b/coordinator/internal/storage/postgres/tx.go index 72dfe9b..4b70e17 100644 --- a/coordinator/internal/storage/postgres/tx.go +++ b/coordinator/internal/storage/postgres/tx.go @@ -12,12 +12,11 @@ import ( // querier is satisfied by both *pgxpool.Pool and pgx.Tx, letting every // repository method run identically inside or outside a transaction. -// -//nolint:unused // used by repository methods once phase 2 replaces the stubs type querier interface { Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) + SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults } // txKey is an unexported struct type, so no other package can collide with it @@ -39,11 +38,28 @@ func NewTxManager(pool *pgxpool.Pool) *TxManager { // The transaction travels in the context rather than in fn's signature, which // is what lets the usecase layer express "do these repository calls atomically" // without its port ever mentioning pgx. +// Retrying happens here, around the whole transaction, and deliberately not +// inside the repositories. Once Postgres aborts a transaction with a +// serialization failure or deadlock, every further statement in it fails too — +// replaying a single query would accomplish nothing. The unit of retry is +// Begin → fn → Commit. +// +// This is safe because fn re-reads its rows (via GetForUpdate) on each attempt, +// so a retry starts from the current state rather than stale entities. func (m *TxManager) WithinTx(ctx context.Context, fn func(ctx context.Context) error) error { if _, ok := ctx.Value(txKey{}).(pgx.Tx); ok { - return fn(ctx) // already inside a transaction — join it, don't nest + // Already inside a transaction — join it. Retrying here would be wrong + // twice over: the outer transaction owns the retry, and re-running fn + // alone cannot undo what the outer one already wrote. + return fn(ctx) } + return withRetry(ctx, func(ctx context.Context) error { + return m.runTx(ctx, fn) + }) +} + +func (m *TxManager) runTx(ctx context.Context, fn func(ctx context.Context) error) error { tx, err := m.pool.Begin(ctx) if err != nil { return err @@ -58,9 +74,17 @@ func (m *TxManager) WithinTx(ctx context.Context, fn func(ctx context.Context) e return tx.Commit(ctx) } +// jsonbOrEmpty keeps a nil map from reaching a NOT NULL jsonb column. pgx +// encodes a nil map as SQL NULL rather than omitting the column, so the +// DEFAULT '{}' never gets a chance to apply. +func jsonbOrEmpty(m map[string]any) map[string]any { + if m == nil { + return map[string]any{} + } + return m +} + // conn returns the transaction bound to ctx, or the pool when there is none. -// -//nolint:unused // every repository method will route through this in phase 2 func conn(ctx context.Context, pool *pgxpool.Pool) querier { if tx, ok := ctx.Value(txKey{}).(pgx.Tx); ok { return tx diff --git a/coordinator/internal/usecase/task.go b/coordinator/internal/usecase/task.go index c9dc801..da177bf 100644 --- a/coordinator/internal/usecase/task.go +++ b/coordinator/internal/usecase/task.go @@ -130,14 +130,23 @@ func (uc *CompleteTask) Execute(ctx context.Context, in CompleteTaskInput) (*dom return err } now := uc.clock.Now() + before := task.Version if err := task.CompleteWith(in.ResultURI, in.ResultSHA256, in.Metrics, in.WorkerID, in.Attempt, now); err != nil { return err } + out = task + + // A replay of an already-recorded result leaves the entity untouched. + // Writing anyway would fail the optimistic-concurrency guard (the stored + // version already equals ours) and turn an idempotent call into a 409. + if task.Version == before { + return nil + } + if err := uc.tasks.Update(ctx, task); err != nil { return err } - out = task return syncJobStatus(ctx, uc.jobs, uc.tasks, task.JobID, now) }) if err != nil {