From 19d82d5c41b0aa395294963a5985cb3eea2714e8 Mon Sep 17 00:00:00 2001 From: Chief of Staff Date: Tue, 2 Jun 2026 00:05:30 -0700 Subject: [PATCH] feat(pool-manager): add pgqueue Go package + dispatcher example - pgqueue package: Go wrapper around the SQL primitive (enqueue, claim, heartbeat, ack, nack, janitor, listener) - examples/agent-dispatcher: reference worker implementation showing end-to-end flow with LISTEN/NOTIFY + heartbeats - All 5 Go tests pass against live pool_manager DB Reference: skills/data-stack-conventions.md Migration: devops/migrations/2026-06-01_pgqueue_primitive.sql --- .../examples/agent-dispatcher/main.go | 145 ++++++++++ pool-manager/pgqueue/pgqueue.go | 256 ++++++++++++++++++ pool-manager/pgqueue/pgqueue_test.go | 204 ++++++++++++++ 3 files changed, 605 insertions(+) create mode 100644 pool-manager/examples/agent-dispatcher/main.go create mode 100644 pool-manager/pgqueue/pgqueue.go create mode 100644 pool-manager/pgqueue/pgqueue_test.go diff --git a/pool-manager/examples/agent-dispatcher/main.go b/pool-manager/examples/agent-dispatcher/main.go new file mode 100644 index 0000000..007aeef --- /dev/null +++ b/pool-manager/examples/agent-dispatcher/main.go @@ -0,0 +1,145 @@ +// Example: a minimal agent_dispatch worker using the pgqueue primitive. +// +// Run as: go run ./examples/agent-dispatcher +// +// This is the "proof" that the new primitive replaces the older poolq package. +// It: +// - Connects to the local pool_manager DB +// - LISTENs on the "pgqueue:agent_dispatch" channel +// - On wakeup, claims the next job +// - Heartbeats every minute for long jobs +// - Acks on success, nacks on failure +// +// This is a reference implementation. Real agents replace the processJob +// function with their actual work. +package main + +import ( + "context" + "encoding/json" + "fmt" + "log" + "os" + "os/signal" + "syscall" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "scobleclaw.dev/pool-manager/pgqueue" +) + +const ( + queueName = "agent_dispatch" + leaseFor = 5 * time.Minute +) + +func main() { + dsn := os.Getenv("PGQUEUE_DSN") + if dsn == "" { + dsn = "postgres://sscoble@/pool_manager?host=/run/postgresql" + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // SIGINT/SIGTERM + sig := make(chan os.Signal, 1) + signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM) + go func() { <-sig; log.Println("shutting down..."); cancel() }() + + workerID := fmt.Sprintf("dispatcher-%d", os.Getpid()) + + pool, err := pgxpool.New(ctx, dsn) + if err != nil { + log.Fatalf("pgxpool: %v", err) + } + defer pool.Close() + + // Optional: enqueue a test job on startup + if os.Getenv("ENQUEUE_DEMO") == "1" { + demo, _ := json.Marshal(map[string]any{"demo": true, "ts": time.Now().Format(time.RFC3339)}) + res, err := pgqueue.Enqueue(ctx, pool, queueName, demo) + if err != nil { + log.Printf("demo enqueue: %v", err) + } else { + log.Printf("demo enqueued: id=%d", res.ID) + } + } + + // LISTEN for wakeups + ln, err := pgqueue.NewListener(ctx, pool, queueName) + if err != nil { + log.Fatalf("listener: %v", err) + } + defer ln.Close() + log.Printf("worker %s: listening on %q (lease %s)", workerID, queueName, leaseFor) + + // Fallback poll every 10s in case we miss a NOTIFY + pollTicker := time.NewTicker(10 * time.Second) + defer pollTicker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-pollTicker.C: + drain(ctx, pool, workerID) + case <-ln.C(): + // Drain until empty (a wakeup might mean multiple jobs) + drain(ctx, pool, workerID) + } + } +} + +// drain claims and processes jobs until the queue is empty (or ctx cancelled). +func drain(ctx context.Context, pool *pgxpool.Pool, workerID string) { + for ctx.Err() == nil { + job, err := pgqueue.Claim(ctx, pool, queueName, workerID, leaseFor) + if err != nil { + log.Printf("claim: %v", err) + return + } + if job == nil { + return + } + if err := processJob(ctx, pool, job, workerID); err != nil { + log.Printf("job %d: %v", job.ID, err) + res, nackErr := pgqueue.Nack(ctx, pool, job.ID, workerID, err.Error(), 30*time.Second) + if nackErr != nil { + log.Printf("nack %d: %v", job.ID, nackErr) + return + } + log.Printf("job %d: nack → %s", job.ID, res) + continue + } + if err := pgqueue.Ack(ctx, pool, job.ID, workerID); err != nil { + log.Printf("ack %d: %v", job.ID, err) + return + } + log.Printf("job %d: ack", job.ID) + } +} + +// processJob is the placeholder. Real agents replace this with actual work. +// For the demo, we just heartbeat a few times to prove the lease works. +func processJob(ctx context.Context, pool *pgxpool.Pool, job *pgqueue.Job, workerID string) error { + log.Printf("job %d: processing %s payload=%s", job.ID, job.Queue, string(job.Payload)) + + // Simulate work with periodic heartbeats + for i := 0; i < 3; i++ { + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(2 * time.Second): + } + ok, err := pgqueue.Heartbeat(ctx, pool, job.ID, workerID, leaseFor) + if err != nil { + return fmt.Errorf("heartbeat: %w", err) + } + if !ok { + return fmt.Errorf("heartbeat lost (job gone or claimed by another)") + } + log.Printf("job %d: heartbeat %d/3", job.ID, i+1) + } + return nil +} diff --git a/pool-manager/pgqueue/pgqueue.go b/pool-manager/pgqueue/pgqueue.go new file mode 100644 index 0000000..09aeccf --- /dev/null +++ b/pool-manager/pgqueue/pgqueue.go @@ -0,0 +1,256 @@ +// Package pgqueue wraps the Postgres-native queue primitive (pgqueue schema). +// +// This is the "v2" queue: SKIP LOCKED + heartbeat-based lease + LISTEN/NOTIFY. +// It coexists with the older poolq package (work_queue table) and supersedes it. +// +// Reference: skills/data-stack-conventions.md +// /home/sscoble/devops/migrations/2026-06-01_pgqueue_primitive.sql +package pgqueue + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +// Job is one queue entry, returned by Claim. +type Job struct { + ID int64 + Queue string + Payload json.RawMessage + Priority int + Attempts int + MaxAttempts int + ClaimedAt time.Time + LeaseUntil time.Time +} + +// EnqueueResult is the return value of Enqueue (idempotency-aware). +type EnqueueResult struct { + ID int64 + Deduplicated bool // true if the id came from an existing idempotency_key +} + +// Enqueue adds a job to a queue. Optionally dedupes by idempotency_key. +// Notifies any worker LISTENing on this queue. +func Enqueue(ctx context.Context, pool *pgxpool.Pool, queue string, payload json.RawMessage, opts ...EnqueueOption) (EnqueueResult, error) { + cfg := enqueueConfig{} + for _, o := range opts { + o(&cfg) + } + + var id int64 + err := pool.QueryRow(ctx, ` + SELECT pgqueue.enqueue($1, $2, $3, $4, $5) + `, queue, payload, cfg.idempotencyKey, cfg.priority, cfg.runAfter).Scan(&id) + if err != nil { + return EnqueueResult{}, fmt.Errorf("pgqueue enqueue: %w", err) + } + // Note: pgqueue.enqueue returns the id of either the new row OR the + // existing matching idempotency_key row. We can't tell from the return + // alone which it was; callers that need that signal should query the + // jobs table directly. For now, the common case is non-dedup and the + // id is just "the job id". + return EnqueueResult{ID: id, Deduplicated: false}, nil +} + +type enqueueConfig struct { + idempotencyKey *string + priority int + runAfter time.Time +} + +type EnqueueOption func(*enqueueConfig) + +func WithIdempotencyKey(key string) EnqueueOption { + return func(c *enqueueConfig) { c.idempotencyKey = &key } +} + +func WithPriority(p int) EnqueueOption { + return func(c *enqueueConfig) { c.priority = p } +} + +func WithRunAfter(t time.Time) EnqueueOption { + return func(c *enqueueConfig) { c.runAfter = t } +} + +// Claim atomically reserves the next ready job in `queue` for `workerID`. +// The lease is valid for `leaseFor`; the worker must Heartbeat or +// janitor_reset_stale_claims will return the job to the pool. +func Claim(ctx context.Context, pool *pgxpool.Pool, queue, workerID string, leaseFor time.Duration) (*Job, error) { + row := pool.QueryRow(ctx, ` + SELECT id, queue, payload, priority, attempts, max_attempts, claimed_at, lease_until + FROM pgqueue.claim_next($1, $2, $3::interval) + `, queue, workerID, fmt.Sprintf("%d seconds", int(leaseFor.Seconds()))) + + var j Job + if err := row.Scan(&j.ID, &j.Queue, &j.Payload, &j.Priority, &j.Attempts, &j.MaxAttempts, &j.ClaimedAt, &j.LeaseUntil); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, nil // no job available — not an error + } + return nil, fmt.Errorf("pgqueue claim: %w", err) + } + return &j, nil +} + +// Heartbeat extends the lease on a job you still hold. +// Returns true if the heartbeat landed; false if the job is gone or +// claimed by someone else. +func Heartbeat(ctx context.Context, pool *pgxpool.Pool, jobID int64, workerID string, extendBy time.Duration) (bool, error) { + var ok bool + err := pool.QueryRow(ctx, ` + SELECT pgqueue.heartbeat($1, $2, $3::interval) + `, jobID, workerID, fmt.Sprintf("%d seconds", int(extendBy.Seconds()))).Scan(&ok) + if err != nil { + return false, fmt.Errorf("pgqueue heartbeat: %w", err) + } + return ok, nil +} + +// Ack marks a job done. +func Ack(ctx context.Context, pool *pgxpool.Pool, jobID int64, workerID string) error { + var ok bool + err := pool.QueryRow(ctx, ` + SELECT pgqueue.ack($1, $2) + `, jobID, workerID).Scan(&ok) + if err != nil { + return fmt.Errorf("pgqueue ack: %w", err) + } + if !ok { + return fmt.Errorf("pgqueue ack: no rows updated (job %d not claimed by %q)", jobID, workerID) + } + return nil +} + +// NackResult tells you what happened. +type NackResult string + +const ( + NackRequeued NackResult = "requeued" // attempts < max; job back in ready + NackDead NackResult = "dead" // attempts >= max; job in dead state + NackLost NackResult = "lost" // job missing or not in claimable state +) + +// Nack either requeues the job (with backoff) or marks it dead if +// attempts >= max_attempts. +func Nack(ctx context.Context, pool *pgxpool.Pool, jobID int64, workerID string, jobErr string, requeueAfter time.Duration) (NackResult, error) { + var r string + err := pool.QueryRow(ctx, ` + SELECT pgqueue.nack($1, $2, $3, $4::interval) + `, jobID, workerID, jobErr, fmt.Sprintf("%d seconds", int(requeueAfter.Seconds()))).Scan(&r) + if err != nil { + return NackLost, fmt.Errorf("pgqueue nack: %w", err) + } + return NackResult(r), nil +} + +// Janitor resets any claimed jobs whose lease has expired. +// Call this from a cron every 1-5 minutes. Returns the number of jobs reset. +func Janitor(ctx context.Context, pool *pgxpool.Pool) (int, error) { + var n int + err := pool.QueryRow(ctx, `SELECT pgqueue.janitor_reset_stale_claims()`).Scan(&n) + if err != nil { + return 0, fmt.Errorf("pgqueue janitor: %w", err) + } + return n, nil +} + +// Listener wraps a pgx connection for LISTEN/NOTIFY on a single queue. +// It sends the queue name on the channel whenever a new job is enqueued. +// +// Usage: +// +// ln, err := pgqueue.NewListener(ctx, pool, "agent_dispatch") +// if err != nil { return err } +// defer ln.Close() +// for { +// select { +// case <-ctx.Done(): return +// case <-ln.C(): +// // try to claim a job +// for { +// job, err := pgqueue.Claim(ctx, pool, "agent_dispatch", workerID, 5*time.Minute) +// if err != nil { log.Printf("claim: %v", err); break } +// if job == nil { break } +// // process... +// } +// } +// } +type Listener struct { + conn *pgx.Conn + ch chan string + closed chan struct{} + queue string +} + +func NewListener(ctx context.Context, pool *pgxpool.Pool, queueName string) (*Listener, error) { + rawConn, err := pool.Acquire(ctx) + if err != nil { + return nil, fmt.Errorf("listener acquire: %w", err) + } + + channel := "pgqueue:" + queueName + if _, err := rawConn.Exec(ctx, "LISTEN \""+channel+"\""); err != nil { + rawConn.Release() + return nil, fmt.Errorf("listen %s: %w", channel, err) + } + + conn := rawConn.Hijack() + + l := &Listener{ + conn: conn, + ch: make(chan string, 16), + closed: make(chan struct{}), + queue: queueName, + } + go l.loop(ctx) + return l, nil +} + +func (l *Listener) C() <-chan string { return l.ch } + +func (l *Listener) Close() error { + select { + case <-l.closed: + return nil + default: + close(l.closed) + close(l.ch) + return l.conn.Close(context.Background()) + } +} + +func (l *Listener) loop(ctx context.Context) { + defer func() { + // recover from a close-while-sending panic + _ = recover() + }() + for { + select { + case <-l.closed: + return + case <-ctx.Done(): + return + default: + } + msg, err := l.conn.WaitForNotification(ctx) + if err != nil { + return + } + if msg == nil { + continue + } + select { + case l.ch <- l.queue: + case <-l.closed: + return + default: + // channel full; drop. Worker will pick up on next poll. + } + } +} diff --git a/pool-manager/pgqueue/pgqueue_test.go b/pool-manager/pgqueue/pgqueue_test.go new file mode 100644 index 0000000..412555f --- /dev/null +++ b/pool-manager/pgqueue/pgqueue_test.go @@ -0,0 +1,204 @@ +package pgqueue + +import ( + "context" + "encoding/json" + "fmt" + "os" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" +) + +// testPool connects to the local pool_manager DB. +// Skips the test if PGBOUNCER or PG is not available. +func testPool(t *testing.T) *pgxpool.Pool { + t.Helper() + dsn := os.Getenv("PGQUEUE_TEST_DSN") + if dsn == "" { + dsn = "postgres://sscoble@/pool_manager?host=/run/postgresql" + } + pool, err := pgxpool.New(context.Background(), dsn) + if err != nil { + t.Skipf("no postgres available: %v", err) + } + return pool +} + +func cleanup(t *testing.T, pool *pgxpool.Pool, queue string) { + t.Helper() + _, err := pool.Exec(context.Background(), + "DELETE FROM pgqueue.jobs WHERE queue = $1", queue) + if err != nil { + t.Logf("cleanup: %v", err) + } +} + +func TestEnqueueClaimAck(t *testing.T) { + pool := testPool(t) + defer pool.Close() + ctx := context.Background() + q := fmt.Sprintf("test_e2e_%d", time.Now().UnixNano()) + defer cleanup(t, pool, q) + + // Enqueue + payload, _ := json.Marshal(map[string]any{"task": "hello", "n": 42}) + res, err := Enqueue(ctx, pool, q, payload) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + if res.ID == 0 { + t.Fatal("got id=0") + } + + // Claim + job, err := Claim(ctx, pool, q, "test-worker", 30*time.Second) + if err != nil { + t.Fatalf("claim: %v", err) + } + if job == nil { + t.Fatal("got nil job") + } + if job.ID != res.ID { + t.Fatalf("id mismatch: enqueue=%d claim=%d", res.ID, job.ID) + } + if job.Attempts != 1 { + t.Fatalf("expected attempts=1, got %d", job.Attempts) + } + + // Ack + if err := Ack(ctx, pool, job.ID, "test-worker"); err != nil { + t.Fatalf("ack: %v", err) + } + + // Verify done + var state string + err = pool.QueryRow(ctx, "SELECT state FROM pgqueue.jobs WHERE id=$1", job.ID).Scan(&state) + if err != nil { + t.Fatalf("verify: %v", err) + } + if state != "done" { + t.Fatalf("expected state=done, got %q", state) + } +} + +func TestNackRequeue(t *testing.T) { + pool := testPool(t) + defer pool.Close() + ctx := context.Background() + q := fmt.Sprintf("test_nack_requeue_%d", time.Now().UnixNano()) + defer cleanup(t, pool, q) + + payload, _ := json.Marshal(map[string]any{"task": "retry"}) + res, _ := Enqueue(ctx, pool, q, payload) + + job, _ := Claim(ctx, pool, q, "w1", 30*time.Second) + if job == nil { + t.Fatal("expected a job") + } + + r, err := Nack(ctx, pool, job.ID, "w1", "transient", 100*time.Millisecond) + if err != nil { + t.Fatalf("nack: %v", err) + } + if r != NackRequeued { + t.Fatalf("expected requeued, got %q", r) + } + + // After requeue, can claim again + job2, _ := Claim(ctx, pool, q, "w1", 30*time.Second) + if job2 == nil { + t.Fatal("expected requeued job to be claimable") + } + if job2.ID != res.ID { + t.Fatal("requeued job has different id") + } +} + +func TestIdempotencyKey(t *testing.T) { + pool := testPool(t) + defer pool.Close() + ctx := context.Background() + q := fmt.Sprintf("test_idem_%d", time.Now().UnixNano()) + defer cleanup(t, pool, q) + + key := fmt.Sprintf("key-%d", time.Now().UnixNano()) + p1, _ := json.Marshal(map[string]any{"v": 1}) + p2, _ := json.Marshal(map[string]any{"v": 2}) + + r1, _ := Enqueue(ctx, pool, q, p1, WithIdempotencyKey(key)) + r2, _ := Enqueue(ctx, pool, q, p2, WithIdempotencyKey(key)) + + if r1.ID != r2.ID { + t.Fatalf("idem: r1=%d r2=%d, expected equal", r1.ID, r2.ID) + } +} + +func TestJanitor(t *testing.T) { + pool := testPool(t) + defer pool.Close() + ctx := context.Background() + q := fmt.Sprintf("test_janitor_%d", time.Now().UnixNano()) + defer cleanup(t, pool, q) + + // Enqueue, claim with very short lease + payload, _ := json.Marshal(map[string]any{"task": "will-stale"}) + Enqueue(ctx, pool, q, payload) + job, _ := Claim(ctx, pool, q, "w1", 100*time.Millisecond) + if job == nil { + t.Fatal("expected job") + } + + // Wait for lease to expire + time.Sleep(200 * time.Millisecond) + + // Run janitor + n, err := Janitor(ctx, pool) + if err != nil { + t.Fatalf("janitor: %v", err) + } + if n < 1 { + t.Fatalf("expected janitor to reset ≥1, got %d", n) + } + + // Verify job is back to ready + var state string + err = pool.QueryRow(ctx, "SELECT state FROM pgqueue.jobs WHERE id=$1", job.ID).Scan(&state) + if err != nil { + t.Fatalf("verify: %v", err) + } + if state != "ready" { + t.Fatalf("expected state=ready, got %q", state) + } +} + +func TestHeartbeat(t *testing.T) { + pool := testPool(t) + defer pool.Close() + ctx := context.Background() + q := fmt.Sprintf("test_hb_%d", time.Now().UnixNano()) + defer cleanup(t, pool, q) + + Enqueue(ctx, pool, q, []byte(`{}`)) + job, _ := Claim(ctx, pool, q, "w1", 100*time.Millisecond) + if job == nil { + t.Fatal("expected job") + } + + // Heartbeat + ok, err := Heartbeat(ctx, pool, job.ID, "w1", 5*time.Second) + if err != nil { + t.Fatalf("heartbeat: %v", err) + } + if !ok { + t.Fatal("heartbeat returned false") + } + + // Verify lease extended + var newLease time.Time + pool.QueryRow(ctx, "SELECT lease_until FROM pgqueue.jobs WHERE id=$1", job.ID).Scan(&newLease) + if newLease.Before(time.Now().Add(4 * time.Second)) { + t.Fatalf("heartbeat did not extend lease: %v", newLease) + } +}