scobleclaw/pool-manager/pgqueue/pgqueue.go

// 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.
		}
	}
}