Scaling Distributed Task Scheduling via PostgreSQL Advisory Locks and Bounded Queue Workers — code snippets

snippet-1.sql

-- Robust schema for high-throughput scheduled tasks
CREATE TABLE scheduled_tasks (
    id BIGSERIAL PRIMARY KEY,
    task_type VARCHAR(64) NOT NULL,
    payload JSONB NOT NULL DEFAULT '{}',
    run_at TIMESTAMPTZ NOT NULL,
    status VARCHAR(20) NOT NULL DEFAULT 'queued',
    max_retries INT NOT NULL DEFAULT 3,
    retry_count INT NOT NULL DEFAULT 0,
    error_log TEXT,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

-- Index specifically tailored for the queue polling query
CREATE INDEX idx_tasks_poll ON scheduled_tasks (run_at, status) 
WHERE status = 'queued';

-- Separate index to optimize the cleaning of stuck/orphaned tasks
CREATE INDEX idx_tasks_stuck ON scheduled_tasks (status, updated_at) 
WHERE status = 'processing';

snippet-2.sql

-- Atomic poll and lock query using transaction-level advisory locks
WITH candidate_task AS (
    SELECT id 
    FROM scheduled_tasks
    WHERE status = 'queued' AND run_at <= NOW()
    ORDER BY run_at ASC, id ASC
    LIMIT 10
    FOR UPDATE SKIP LOCKED
)
UPDATE scheduled_tasks
SET status = 'processing', updated_at = NOW()
WHERE id = (
    SELECT id FROM candidate_task
    WHERE pg_try_advisory_xact_lock(1337, (id % 2147483647)::integer)
    LIMIT 1
)
RETURNING id, task_type, payload, retry_count;

snippet-3.go

package scheduler

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

type Task struct {
	ID         int64
	TaskType   string
	Payload    json.RawMessage
	RetryCount int
}

type Config struct {
	PollInterval    time.Duration
	MaxConcurrency  int
	DBTimeout       time.Duration
	NamespaceID     int32
}

type DistributedScheduler struct {
	db        *pgxpool.Pool
	cfg       Config
	semaphore chan struct{}
	wg        sync.WaitGroup
	ctx       context.Context
	cancel    context.CancelFunc
}

func NewScheduler(db *pgxpool.Pool, cfg Config) *DistributedScheduler {
	ctx, cancel := context.WithCancel(context.Background())
	return &DistributedScheduler{
		db:        db,
		cfg:       cfg,
		semaphore: make(chan struct{}, cfg.MaxConcurrency),
		ctx:       ctx,
		cancel:    cancel,
	}
}

snippet-4.go

// Start begins the scheduling polling loop
func (s *DistributedScheduler) Start() {
	ticker := time.NewTicker(s.cfg.PollInterval)
	defer ticker.Stop()

	log.Printf("Scheduler started with concurrency limit: %d", s.cfg.MaxConcurrency)

	for {
		select {
		case <-s.ctx.Done():
			log.Println("Shutting down poll loop...")
			return
		case <-ticker.C:
			// Check if we have available worker slots before polling the DB
			select {
			case s.semaphore <- struct{}{}:
				// Slot acquired, proceed to poll
				go func() {
					defer func() { <-s.semaphore }()
					s.pollAndExecute()
				}()
			default:
				// No slots available, skip this tick to avoid DB overloading
				continue
			}
		}
	}
}

// Stop initiates graceful shutdown, waiting for active workers to drain
func (s *DistributedScheduler) Stop() {
	s.cancel()
	s.wg.Wait()
	log.Println("Scheduler gracefully stopped.")
}

snippet-5.go

// pollAndExecute fetches a task inside a transaction and processes it
func (s *DistributedScheduler) pollAndExecute() {
	s.wg.Add(1)
	defer s.wg.Done()

	ctx, cancel := context.WithTimeout(s.ctx, s.cfg.DBTimeout)
	defer cancel()

	tx, err := s.db.Begin(ctx)
	if err != nil {
		log.Printf("failed to begin transaction: %v", err)
		return
	}
	defer tx.Rollback(ctx)

	// Fetch task and acquire transaction-level advisory lock
	query := `
		WITH candidate_task AS (
			SELECT id 
			FROM scheduled_tasks
			WHERE status = 'queued' AND run_at <= NOW()
			ORDER BY run_at ASC, id ASC
			LIMIT 10
			FOR UPDATE SKIP LOCKED
		)
		UPDATE scheduled_tasks
		SET status = 'processing', updated_at = NOW()
		WHERE id = (
			SELECT id FROM candidate_task
			WHERE pg_try_advisory_xact_lock($1, (id % 2147483647)::integer)
			LIMIT 1
		)
		RETURNING id, task_type, payload, retry_count;`

	var task Task
	err = tx.QueryRow(ctx, query, s.cfg.NamespaceID).Scan(&task.ID, &task.TaskType, &task.Payload, &task.RetryCount)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) || err.Error() == "no rows in result set" {
			// No tasks currently ready for processing
			return
		}
		log.Printf("error querying database for task: %v", err)
		return
	}

	// Commit transaction immediately to release the advisory lock.
	// This prevents holding DB connections open during long-running tasks.
	if err := tx.Commit(ctx); err != nil {
		log.Printf("failed to commit task acquisition: %v", err)
		return
	}

	// Process the task payload asynchronously/locally
	s.executeTask(task)
}

func (s *DistributedScheduler) executeTask(task Task) {
	// Recover from panics in dynamic task handlers
	defer func() {
		if r := recover(); r != nil {
			s.handleFailure(task.ID, fmt.Errorf("panic in execution: %v", r), task.RetryCount)
		}
	}()

	log.Printf("processing task %d of type %s", task.ID, task.TaskType)

	// Simulate work execution (replace with dynamic routing based on TaskType)
	time.Sleep(500 * time.Millisecond) 

	// Update status to completed
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	
	_, err := s.db.Exec(ctx, "UPDATE scheduled_tasks SET status = 'completed', updated_at = NOW() WHERE id = $1", task.ID)
	if err != nil {
		log.Printf("failed to mark task %d completed: %v", task.ID, err)
	}
}

func (s *DistributedScheduler) handleFailure(taskID int64, err error, currentRetries int) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	log.Printf("task %d failed: %v", taskID, err)
	
	query := `
		UPDATE scheduled_tasks 
		SET status = CASE WHEN retry_count < max_retries THEN 'queued' ELSE 'failed' END,
		    retry_count = retry_count + 1,
		    error_log = $2,
		    run_at = NOW() + (INTERVAL '1 minute' * (retry_count + 1)), -- exponential backoff
		    updated_at = NOW()
		WHERE id = $1`
	
	_, dbErr := s.db.Exec(ctx, query, taskID, err.Error())
	if dbErr != nil {
		log.Printf("failed to register failure for task %d: %v", taskID, dbErr)
	}
}

snippet-6.go

package scheduler

import (
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

var (
	taskExecutionDuration = promauto.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "scheduler_task_execution_duration_seconds",
			Help:    "Time taken to execute a task, partitioned by type.",
			Buckets: prometheus.DefBuckets,
		},
		[]string{"task_type", "status"},
	)

	schedulingLatency = promauto.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "scheduler_task_latency_seconds",
			Help:    "Delay between run_at timestamp and actual execution start.",
			Buckets: []float64{0.1, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0},
		},
		[]string{"task_type"},
	)

	activeWorkers = promauto.NewGauge(
		prometheus.GaugeOpts{
			Name: "scheduler_active_workers",
			Help: "Number of concurrent workers currently running tasks.",
		},
	)

	databaseLockAttempts = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "scheduler_lock_attempts_total",
			Help: "Total number of advisory lock attempts and their results.",
		},
		[]string{"result"}, // "acquired", "skipped", "failed"
	)
)

snippet-7.sql

-- Clean up orphaned tasks that are stuck in the 'processing' state
UPDATE scheduled_tasks
SET status = 'queued',
    retry_count = retry_count + 1,
    error_log = 'Reaper detected orphaned task: worker node terminated unexpectedly',
    updated_at = NOW()
WHERE status = 'processing'
  AND updated_at < NOW() - INTERVAL '5 minutes'
  AND pg_try_advisory_xact_lock(1337, (id % 2147483647)::integer);