Postgres `SKIP LOCKED` for High-Throughput Job Queues

CREATE TYPE job_status AS ENUM ('pending', 'processing', 'completed', 'failed');

CREATE TABLE jobs (
    id BIGSERIAL PRIMARY KEY,
    queue_name TEXT NOT NULL DEFAULT 'default',
    payload JSONB NOT NULL,
    status job_status NOT NULL DEFAULT 'pending',
    max_retries INT NOT NULL DEFAULT 3,
    attempts INT NOT NULL DEFAULT 0,
    run_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE INDEX idx_jobs_on_status_run_at ON jobs (status, run_at);

-- ANTI-PATTERN: DO NOT USE IN PRODUCTION FOR CONCURRENT WORKERS
BEGIN;

-- Find a pending job
SELECT id, payload FROM jobs
WHERE status = 'pending' AND run_at <= NOW()
ORDER BY created_at
LIMIT 1
FOR UPDATE;

-- If a job is found, update its status
UPDATE jobs SET status = 'processing', updated_at = NOW() WHERE id = <job_id>;

COMMIT;

-- THE CORRECT PATTERN FOR CONCURRENT DEQUEUING
BEGIN;

-- Atomically find, lock, and update the next available job
UPDATE jobs
SET status = 'processing', attempts = attempts + 1, updated_at = NOW()
WHERE id = (
    SELECT id
    FROM jobs
    WHERE status = 'pending' AND run_at <= NOW()
    ORDER BY created_at
    FOR UPDATE SKIP LOCKED
    LIMIT 1
)
RETURNING id, payload, attempts, max_retries;

COMMIT;

package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"time"
)

type Job struct {
	ID         int64
	Queue      string
	Payload    json.RawMessage
	Status     string
	MaxRetries int
	Attempts   int
	RunAt      time.Time
	CreatedAt  time.Time
}

const dequeueSQL = `
    UPDATE jobs
    SET status = 'processing', attempts = attempts + 1, updated_at = NOW()
    WHERE id = (
        SELECT id
        FROM jobs
        WHERE status = 'pending' AND run_at <= NOW()
        ORDER BY run_at ASC, id ASC
        FOR UPDATE SKIP LOCKED
        LIMIT 1
    )
    RETURNING id, queue_name, payload, attempts, max_retries;`

// Dequeue retrieves and locks the next available job from the database.
func Dequeue(ctx context.Context, db *sql.DB) (*Job, error) {
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback() // Rollback is a no-op if the transaction is committed.

	job := &Job{}
	err = tx.QueryRowContext(ctx, dequeueSQL).Scan(&job.ID, &job.Queue, &job.Payload, &job.Attempts, &job.MaxRetries)
	if err != nil {
		if err == sql.ErrNoRows {
			// This is not an error, just an empty queue.
			return nil, nil
		}
		return nil, err
	}

	if err := tx.Commit(); err != nil {
		return nil, err
	}

	return job, nil
}

// MarkAsCompleted updates a job's status to 'completed'.
func MarkAsCompleted(ctx context.Context, db *sql.DB, jobID int64) error {
	_, err := db.ExecContext(ctx, "UPDATE jobs SET status = 'completed', updated_at = NOW() WHERE id = $1", jobID)
	return err
}

// MarkAsFailed updates a job's status to 'failed' and schedules a retry if applicable.
func MarkAsFailed(ctx context.Context, db *sql.DB, job *Job, jobErr error) error {
	if job.Attempts >= job.MaxRetries {
		// Move to a dead-letter state
		_, err := db.ExecContext(ctx, "UPDATE jobs SET status = 'failed', updated_at = NOW() WHERE id = $1", job.ID)
		return err
	}

	// Calculate next run time with exponential backoff
	backoffDuration := time.Duration(job.Attempts*job.Attempts) * time.Second * 15
	nextRunAt := time.Now().Add(backoffDuration)

	_, err := db.ExecContext(ctx, "UPDATE jobs SET status = 'pending', run_at = $1, updated_at = NOW() WHERE id = $2", nextRunAt, job.ID)
	return err
}

package main

import (
	"context"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

// Processor is a function that processes a job.
// It returns an error if processing fails.
type Processor func(ctx context.Context, payload json.RawMessage) error

type WorkerPool struct {
	db         *sql.DB
	numWorkers int
	pollInterval time.Duration
	processor  Processor
}

func NewWorkerPool(db *sql.DB, numWorkers int, processor Processor) *WorkerPool {
	return &WorkerPool{
		db:         db,
		numWorkers: numWorkers,
		pollInterval: 5 * time.Second, // Poll interval when queue is empty
		processor:  processor,
	}
}

func (wp *WorkerPool) Run() {
	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup

	log.Printf("Starting %d workers...", wp.numWorkers)
	for i := 0; i < wp.numWorkers; i++ {
		wg.Add(1)
		go wp.worker(ctx, &wg, i+1)
	}

	// Wait for termination signal
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	<-sigChan

	log.Println("Shutdown signal received, stopping workers...")
	cancel() // Signal all workers to stop

	wg.Wait() // Wait for all workers to finish their current job
	log.Println("All workers have stopped.")
}

func (wp *WorkerPool) worker(ctx context.Context, wg *sync.WaitGroup, workerID int) {
	defer wg.Done()
	log.Printf("Worker %d started", workerID)

	for {
		select {
		case <-ctx.Done(): // Shutdown signal received
			log.Printf("Worker %d stopping...", workerID)
			return
		default:
			job, err := Dequeue(ctx, wp.db)
			if err != nil {
				log.Printf("Worker %d: Error dequeuing job: %v", workerID, err)
				time.Sleep(wp.pollInterval) // Back off on DB errors
				continue
			}

			if job == nil {
				// Queue is empty, wait before polling again
				time.Sleep(wp.pollInterval)
				continue
			}

			log.Printf("Worker %d: Picked up job %d", workerID, job.ID)
			processingCtx, cancelProcessing := context.WithTimeout(ctx, 30*time.Second) // Job-specific timeout
			
			err = wp.processor(processingCtx, job.Payload)
			cancelProcessing()

			if err != nil {
				log.Printf("Worker %d: Error processing job %d: %v", workerID, job.ID, err)
				if err := MarkAsFailed(ctx, wp.db, job, err); err != nil {
					log.Printf("Worker %d: CRITICAL - Failed to mark job %d as failed: %v", workerID, job.ID, err)
				}
			} else {
				log.Printf("Worker %d: Completed job %d", workerID, job.ID)
				if err := MarkAsCompleted(ctx, wp.db, job.ID); err != nil {
					log.Printf("Worker %d: CRITICAL - Failed to mark job %d as completed: %v", workerID, job.ID, err)
				}
			}
		}
	}
}

-- Drop the old, inefficient index
DROP INDEX IF EXISTS idx_jobs_on_status_run_at;

-- Create a highly-specific, small, and fast partial index
CREATE INDEX idx_jobs_pending_run_at ON jobs (run_at, id)
WHERE status = 'pending';

ALTER TABLE jobs ADD COLUMN locked_until TIMESTAMPTZ;

-- Modified UPDATE query to set a lease time
UPDATE jobs
SET status = 'processing',
    attempts = attempts + 1,
    updated_at = NOW(),
    locked_until = NOW() + INTERVAL '5 minutes'
WHERE id = (
    SELECT id
    FROM jobs
    WHERE status = 'pending' AND run_at <= NOW()
    ORDER BY run_at ASC, id ASC
    FOR UPDATE SKIP LOCKED
    LIMIT 1
)
RETURNING id, payload, attempts, max_retries;

-- Reaper query to find and reset stuck jobs
UPDATE jobs
SET status = 'pending', locked_until = NULL
WHERE status = 'processing' AND locked_until < NOW();

ALTER TABLE jobs ADD COLUMN priority INT NOT NULL DEFAULT 0;

-- New partial index for prioritized queues
CREATE INDEX idx_jobs_pending_priority_run_at ON jobs (priority DESC, run_at ASC, id ASC)
WHERE status = 'pending';

-- New Dequeue query
SELECT id
FROM jobs
WHERE status = 'pending' AND run_at <= NOW()
ORDER BY priority DESC, run_at ASC, id ASC
FOR UPDATE SKIP LOCKED
LIMIT 1;

BEGIN;
-- This will cause contention
UPDATE jobs SET status = 'processing' WHERE id = (SELECT id FROM jobs WHERE status = 'pending' ORDER BY id LIMIT 1 FOR UPDATE) RETURNING id;
COMMIT;

BEGIN;
-- This will avoid contention
UPDATE jobs SET status = 'processing' WHERE id = (SELECT id FROM jobs WHERE status = 'pending' ORDER BY id LIMIT 1 FOR UPDATE SKIP LOCKED) RETURNING id;
COMMIT;

# Naive locking
pgbench -U youruser -d yourdb -f naive_lock.sql -c 16 -T 60

# SKIP LOCKED
pgbench -U youruser -d yourdb -f skip_locked.sql -c 16 -T 60