From 2e8135bcdff2b8ff592142eb26627cc9409be7cf Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Tue, 5 May 2026 19:40:39 +0800 Subject: [PATCH] feat(jobs): postgres repository with SKIP LOCKED lease + retry/dead-letter/purge --- internal/jobs/repository.go | 197 +++++++++++++++++++++++++++ internal/jobs/repository_test.go | 219 +++++++++++++++++++++++++++++++ 2 files changed, 416 insertions(+) create mode 100644 internal/jobs/repository.go create mode 100644 internal/jobs/repository_test.go diff --git a/internal/jobs/repository.go b/internal/jobs/repository.go new file mode 100644 index 0000000..02b8cb4 --- /dev/null +++ b/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 +} diff --git a/internal/jobs/repository_test.go b/internal/jobs/repository_test.go new file mode 100644 index 0000000..8e1db30 --- /dev/null +++ b/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) +}