feat(jobs): postgres repository with SKIP LOCKED lease + retry/dead-letter/purge

internal/jobs/repository.go

@@ -0,0 +1,197 @@
+package jobs
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgtype"
+	"github.com/jackc/pgx/v5/pgxpool"
+
+	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
+	jobssqlc "github.com/yan1h/agent-coding-workflow/internal/jobs/sqlc"
+)
+
+// Repository is the persistence contract for jobs. It is implemented by
+// PgRepository for production and may be faked for unit tests of higher-level
+// components (worker, scheduler).
+type Repository interface {
+	Enqueue(ctx context.Context, typ JobType, payload json.RawMessage, scheduledAt time.Time, maxAttempts int32) (*Job, error)
+	Lease(ctx context.Context, leasedBy string) (*Job, error)
+	Complete(ctx context.Context, id uuid.UUID) error
+	FailWithRetry(ctx context.Context, id uuid.UUID, scheduledAt time.Time, lastErr string) error
+	DeadLetter(ctx context.Context, id uuid.UUID, lastErr string) error
+	Get(ctx context.Context, id uuid.UUID) (*Job, error)
+	ReapStuckRunningJobs(ctx context.Context, leasedBefore time.Time) ([]uuid.UUID, error)
+	PurgeCompleted(ctx context.Context, completedBefore time.Time) (int64, error)
+}
+
+// PgRepository wraps sqlc-generated queries with a pgxpool.Pool. It is the only
+// production implementation of Repository.
+type PgRepository struct {
+	pool *pgxpool.Pool
+	q    *jobssqlc.Queries
+}
+
+// NewPostgresRepository builds a PgRepository on the given pool.
+func NewPostgresRepository(pool *pgxpool.Pool) *PgRepository {
+	return &PgRepository{pool: pool, q: jobssqlc.New(pool)}
+}
+
+var _ Repository = (*PgRepository)(nil)
+
+// ===== type-conversion helpers =====
+
+func toPgUUID(u uuid.UUID) pgtype.UUID { return pgtype.UUID{Bytes: u, Valid: true} }
+
+func fromPgUUID(p pgtype.UUID) uuid.UUID {
+	if !p.Valid {
+		return uuid.Nil
+	}
+	return uuid.UUID(p.Bytes)
+}
+
+// ===== Repository methods =====
+
+// Enqueue persists a new pending job. Caller-supplied payload is rejected if it
+// exceeds MaxPayloadBytes so a runaway producer can't bloat the table.
+func (r *PgRepository) Enqueue(ctx context.Context, typ JobType, payload json.RawMessage, scheduledAt time.Time, maxAttempts int32) (*Job, error) {
+	if len(payload) > MaxPayloadBytes {
+		return nil, errs.New(errs.CodeJobPayloadTooLarge, fmt.Sprintf("payload exceeds %d bytes", MaxPayloadBytes))
+	}
+	row, err := r.q.EnqueueJob(ctx, jobssqlc.EnqueueJobParams{
+		ID:          toPgUUID(uuid.New()),
+		Type:        string(typ),
+		Payload:     payload,
+		ScheduledAt: pgtype.Timestamptz{Time: scheduledAt, Valid: true},
+		MaxAttempts: maxAttempts,
+	})
+	if err != nil {
+		return nil, errs.Wrap(err, errs.CodeInternal, "enqueue job")
+	}
+	return rowToJob(row), nil
+}
+
+// Lease atomically transitions the oldest due pending job to running and
+// returns it. Returns (nil, nil) when no job is ready — callers treat that as
+// "queue empty" and back off.
+func (r *PgRepository) Lease(ctx context.Context, leasedBy string) (*Job, error) {
+	row, err := r.q.LeaseNextJob(ctx, &leasedBy)
+	if err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return nil, nil
+		}
+		return nil, errs.Wrap(err, errs.CodeInternal, "lease job")
+	}
+	return rowToJob(row), nil
+}
+
+// Complete marks a job as completed and clears its lease/error fields.
+func (r *PgRepository) Complete(ctx context.Context, id uuid.UUID) error {
+	if err := r.q.CompleteJob(ctx, toPgUUID(id)); err != nil {
+		return errs.Wrap(err, errs.CodeInternal, "complete job")
+	}
+	return nil
+}
+
+// FailWithRetry returns a job to pending with a new scheduled_at and stores
+// the last error. Attempts is preserved (it was incremented on lease) so the
+// worker can compare against MaxAttempts to decide whether to retry or
+// dead-letter on the next failure.
+func (r *PgRepository) FailWithRetry(ctx context.Context, id uuid.UUID, scheduledAt time.Time, lastErr string) error {
+	if err := r.q.FailJobWithRetry(ctx, jobssqlc.FailJobWithRetryParams{
+		ID:          toPgUUID(id),
+		ScheduledAt: pgtype.Timestamptz{Time: scheduledAt, Valid: true},
+		LastError:   &lastErr,
+	}); err != nil {
+		return errs.Wrap(err, errs.CodeInternal, "fail job with retry")
+	}
+	return nil
+}
+
+// DeadLetter terminally marks a job as dead and stores the last error.
+func (r *PgRepository) DeadLetter(ctx context.Context, id uuid.UUID, lastErr string) error {
+	if err := r.q.DeadLetterJob(ctx, jobssqlc.DeadLetterJobParams{
+		ID:        toPgUUID(id),
+		LastError: &lastErr,
+	}); err != nil {
+		return errs.Wrap(err, errs.CodeInternal, "dead-letter job")
+	}
+	return nil
+}
+
+// Get fetches a job by id. Returns CodeJobNotFound when no row exists.
+func (r *PgRepository) Get(ctx context.Context, id uuid.UUID) (*Job, error) {
+	row, err := r.q.GetJob(ctx, toPgUUID(id))
+	if err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return nil, errs.New(errs.CodeJobNotFound, "job not found")
+		}
+		return nil, errs.Wrap(err, errs.CodeInternal, "get job")
+	}
+	return rowToJob(row), nil
+}
+
+// ReapStuckRunningJobs resets jobs whose lease started before leasedBefore
+// back to pending so another worker can pick them up. Returns the ids of all
+// reaped jobs for logging/metrics.
+func (r *PgRepository) ReapStuckRunningJobs(ctx context.Context, leasedBefore time.Time) ([]uuid.UUID, error) {
+	rows, err := r.q.ReapStuckRunningJobs(ctx, pgtype.Timestamptz{Time: leasedBefore, Valid: true})
+	if err != nil {
+		return nil, errs.Wrap(err, errs.CodeInternal, "reap stuck running jobs")
+	}
+	out := make([]uuid.UUID, 0, len(rows))
+	for _, p := range rows {
+		out = append(out, fromPgUUID(p))
+	}
+	return out, nil
+}
+
+// PurgeCompleted deletes completed and dead jobs whose completed_at is older
+// than completedBefore, returning the row count for metrics.
+func (r *PgRepository) PurgeCompleted(ctx context.Context, completedBefore time.Time) (int64, error) {
+	rows, err := r.q.PurgeCompletedJobs(ctx, pgtype.Timestamptz{Time: completedBefore, Valid: true})
+	if err != nil {
+		return 0, errs.Wrap(err, errs.CodeInternal, "purge completed jobs")
+	}
+	return rows, nil
+}
+
+// rowToJob converts a sqlc-generated row into the domain Job. LeasedBy and
+// LastError are *string in the sqlc layer (emit_pointers_for_null_types), so
+// we copy the value to avoid aliasing the row's pointer; LeasedAt and
+// CompletedAt are pgtype.Timestamptz, so we check Valid.
+func rowToJob(r jobssqlc.Job) *Job {
+	out := &Job{
+		ID:          fromPgUUID(r.ID),
+		Type:        JobType(r.Type),
+		Payload:     r.Payload,
+		Status:      JobStatus(r.Status),
+		ScheduledAt: r.ScheduledAt.Time,
+		Attempts:    r.Attempts,
+		MaxAttempts: r.MaxAttempts,
+		CreatedAt:   r.CreatedAt.Time,
+		UpdatedAt:   r.UpdatedAt.Time,
+	}
+	if r.LeasedAt.Valid {
+		t := r.LeasedAt.Time
+		out.LeasedAt = &t
+	}
+	if r.LeasedBy != nil {
+		s := *r.LeasedBy
+		out.LeasedBy = &s
+	}
+	if r.LastError != nil {
+		s := *r.LastError
+		out.LastError = &s
+	}
+	if r.CompletedAt.Valid {
+		t := r.CompletedAt.Time
+		out.CompletedAt = &t
+	}
+	return out
+}

internal/jobs/repository_test.go

@@ -0,0 +1,219 @@
+//go:build integration
+
+package jobs_test
+
+import (
+	"context"
+	"encoding/json"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	tcpg "github.com/testcontainers/testcontainers-go/modules/postgres"
+
+	"github.com/yan1h/agent-coding-workflow/internal/infra/db"
+	"github.com/yan1h/agent-coding-workflow/internal/jobs"
+)
+
+// setupRepo spins up a throwaway postgres container, runs migrations, and
+// returns a fresh PgRepository plus the underlying pool for tests that need to
+// poke at rows directly. Each test gets its own container so they may run in
+// parallel without bleeding state.
+func setupRepo(t *testing.T) (*jobs.PgRepository, *pgxpool.Pool) {
+	t.Helper()
+	ctx := context.Background()
+	container, err := tcpg.Run(ctx, "postgres:16-alpine",
+		tcpg.WithDatabase("acw"),
+		tcpg.WithUsername("acw"),
+		tcpg.WithPassword("acw"),
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = container.Terminate(ctx) })
+
+	dsn, err := container.ConnectionString(ctx, "sslmode=disable")
+	require.NoError(t, err)
+	require.NoError(t, db.Migrate(ctx, dsn, "file://../../migrations"))
+
+	pool, err := db.NewPool(ctx, dsn)
+	require.NoError(t, err)
+	t.Cleanup(pool.Close)
+
+	return jobs.NewPostgresRepository(pool), pool
+}
+
+func TestEnqueue_PersistsJob(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, _ := setupRepo(t)
+	job, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{"topic":"x"}`), time.Now(), 5)
+	require.NoError(t, err)
+	assert.Equal(t, jobs.StatusPending, job.Status)
+	assert.Equal(t, int32(0), job.Attempts)
+	assert.Equal(t, int32(5), job.MaxAttempts)
+}
+
+func TestLease_PicksOldestPending(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, _ := setupRepo(t)
+	earlier := time.Now().Add(-10 * time.Second)
+	later := time.Now()
+	_, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{"a":1}`), later, 5)
+	require.NoError(t, err)
+	earliest, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{"a":2}`), earlier, 5)
+	require.NoError(t, err)
+
+	got, err := repo.Lease(ctx, "test-host#1")
+	require.NoError(t, err)
+	require.NotNil(t, got)
+	assert.Equal(t, earliest.ID, got.ID)
+	assert.Equal(t, jobs.StatusRunning, got.Status)
+	assert.Equal(t, int32(1), got.Attempts)
+}
+
+func TestLease_ReturnsNilWhenEmpty(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, _ := setupRepo(t)
+	got, err := repo.Lease(ctx, "test-host#1")
+	require.NoError(t, err)
+	assert.Nil(t, got)
+}
+
+// TestLease_SkipLockedConcurrent verifies that with FOR UPDATE SKIP LOCKED in
+// the underlying query, only one of N concurrent leasers wins a single
+// pending row.
+func TestLease_SkipLockedConcurrent(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, _ := setupRepo(t)
+	_, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{"a":1}`), time.Now(), 5)
+	require.NoError(t, err)
+
+	var wg sync.WaitGroup
+	results := make([]*jobs.Job, 5)
+	leaseErrs := make([]error, 5)
+	for i := 0; i < 5; i++ {
+		wg.Add(1)
+		go func(idx int) {
+			defer wg.Done()
+			j, err := repo.Lease(ctx, "test")
+			results[idx] = j
+			leaseErrs[idx] = err
+		}(i)
+	}
+	wg.Wait()
+	got := 0
+	for i := range results {
+		require.NoError(t, leaseErrs[i])
+		if results[i] != nil {
+			got++
+		}
+	}
+	assert.Equal(t, 1, got, "only one worker should win")
+}
+
+func TestComplete_ClearsLease(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, _ := setupRepo(t)
+	job, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{}`), time.Now(), 5)
+	require.NoError(t, err)
+	leased, err := repo.Lease(ctx, "h")
+	require.NoError(t, err)
+	require.NoError(t, repo.Complete(ctx, leased.ID))
+	got, err := repo.Get(ctx, job.ID)
+	require.NoError(t, err)
+	assert.Equal(t, jobs.StatusCompleted, got.Status)
+	assert.Nil(t, got.LeasedAt)
+	assert.Nil(t, got.LeasedBy)
+}
+
+func TestFailWithRetry_ReschedulesAndKeepsAttempts(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, _ := setupRepo(t)
+	_, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{}`), time.Now(), 5)
+	require.NoError(t, err)
+	leased, err := repo.Lease(ctx, "h")
+	require.NoError(t, err)
+	next := time.Now().Add(30 * time.Second)
+	require.NoError(t, repo.FailWithRetry(ctx, leased.ID, next, "boom"))
+	got, err := repo.Get(ctx, leased.ID)
+	require.NoError(t, err)
+	assert.Equal(t, jobs.StatusPending, got.Status)
+	assert.WithinDuration(t, next, got.ScheduledAt, time.Second)
+	require.NotNil(t, got.LastError)
+	assert.Equal(t, "boom", *got.LastError)
+	assert.Equal(t, int32(1), got.Attempts)
+}
+
+func TestDeadLetter_FinalState(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, _ := setupRepo(t)
+	_, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{}`), time.Now(), 5)
+	require.NoError(t, err)
+	leased, err := repo.Lease(ctx, "h")
+	require.NoError(t, err)
+	require.NoError(t, repo.DeadLetter(ctx, leased.ID, "max attempts"))
+	got, err := repo.Get(ctx, leased.ID)
+	require.NoError(t, err)
+	assert.Equal(t, jobs.StatusDead, got.Status)
+	require.NotNil(t, got.CompletedAt)
+}
+
+func TestReapStuck_ResetsLongLeased(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, pool := setupRepo(t)
+	_, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{}`), time.Now(), 5)
+	require.NoError(t, err)
+	leased, err := repo.Lease(ctx, "h")
+	require.NoError(t, err)
+	_, err = pool.Exec(ctx, `UPDATE jobs SET leased_at = now() - INTERVAL '10 minutes' WHERE id=$1`, leased.ID)
+	require.NoError(t, err)
+
+	ids, err := repo.ReapStuckRunningJobs(ctx, time.Now().Add(-5*time.Minute))
+	require.NoError(t, err)
+	require.Len(t, ids, 1)
+	assert.Equal(t, leased.ID, ids[0])
+	got, err := repo.Get(ctx, leased.ID)
+	require.NoError(t, err)
+	assert.Equal(t, jobs.StatusPending, got.Status)
+}
+
+func TestPurge_DeletesOldCompletedAndDead(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, pool := setupRepo(t)
+
+	mkOld := func(status jobs.JobStatus) uuid.UUID {
+		j, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{}`), time.Now(), 5)
+		require.NoError(t, err)
+		_, err = pool.Exec(ctx,
+			`UPDATE jobs SET status=$1, completed_at = now() - INTERVAL '10 days' WHERE id=$2`, string(status), j.ID)
+		require.NoError(t, err)
+		return j.ID
+	}
+	completedOld := mkOld(jobs.StatusCompleted)
+	deadOld := mkOld(jobs.StatusDead)
+	freshComp, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{}`), time.Now(), 5)
+	require.NoError(t, err)
+	require.NoError(t, repo.Complete(ctx, freshComp.ID))
+
+	rows, err := repo.PurgeCompleted(ctx, time.Now().Add(-7*24*time.Hour))
+	require.NoError(t, err)
+	assert.Equal(t, int64(2), rows)
+
+	_, err = repo.Get(ctx, completedOld)
+	assert.Error(t, err)
+	_, err = repo.Get(ctx, deadOld)
+	assert.Error(t, err)
+	_, err = repo.Get(ctx, freshComp.ID)
+	require.NoError(t, err)
+}