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

2026-05-05 19:40:39 +08:00
parent ff3f27c1dc
commit 2e8135bcdf
2 changed files with 416 additions and 0 deletions
@@ -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)
+}