scobleclaw/pool-manager/examples/agent-dispatcher/main.go

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