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 {