You've already forked agentic-coding-workflow
feat(jobs): postgres repository with SKIP LOCKED lease + retry/dead-letter/purge
This commit is contained in:
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user