feat(broker): migrations-based install, roles, RLS, and job dependency groups

Replace the ad-hoc tables/procedures install layout with versioned, ordered SQL migrations tracked in broker_schema_migrations. Add optional least-privilege role provisioning (--with-roles), multi-tenant row-level security, lease-based job claiming with stale-lease recovery, and job dependencies -- both by job id and by fan-in job group. Add Docker/Compose support for running the broker and its test suite.
2026-09-17 22:09:41 +02:00
parent 602997bcdb
commit 4c8e1066d4
53 changed files with 2911 additions and 1016 deletions
@@ -2,15 +2,17 @@ package broker

 import (
 	"context"
-	"database/sql" // Import sql package
+	"database/sql"
 	"encoding/json"
 	"fmt"
 	"os"
+	"strings"
 	"sync"
 	"time"

 	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/adapter"
 	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/config"
+	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/install"
 	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/models"
 	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/queue"
 )
@@ -34,6 +36,12 @@ type DatabaseInstance struct {
 	shutdownMu   sync.RWMutex
 	jobsHandled  int64
 	startTime    time.Time
+
+	// sessionConn holds the pg_try_advisory_lock acquired by
+	// registerInstance. The lock is scoped to this one physical connection,
+	// so it must be kept open (never returned to the pool) for the life of
+	// the process and explicitly unlocked on Stop().
+	sessionConn *sql.Conn
 }

 // NewDatabaseInstance creates a new database instance
@@ -70,6 +78,11 @@ func (i *DatabaseInstance) Start() error {
 		return fmt.Errorf("failed to connect to database: %w", err)
 	}

+	// Ensure the schema is up to date before touching any broker objects.
+	if err := i.ensureSchema(); err != nil {
+		return err
+	}
+
 	// Register instance in database
 	if err := i.registerInstance(); err != nil {
 		return fmt.Errorf("failed to register instance: %w", err)
@@ -87,13 +100,48 @@ func (i *DatabaseInstance) Start() error {
 		return fmt.Errorf("failed to start listener: %w", err)
 	}

-	// Start ping routine
-	go i.pingRoutine()
+	// Start ping routine (auto-restarted on panic; must run for the life of
+	// the process)
+	adapter.SupervisedGo(i.logger, "ping-routine", i.pingRoutine)
+
+	// Start stale/expired-lease job recovery routine (auto-restarted on
+	// panic; must run for the life of the process)
+	adapter.SupervisedGo(i.logger, "stale-job-recovery-routine", i.staleJobRecoveryRoutine)

 	i.logger.Info("database instance started successfully")
 	return nil
 }

+// ensureSchema checks the embedded migration set against the database and,
+// depending on dbConfig.AutoMigrate, either applies pending migrations or
+// fails startup fast rather than running against a stale/missing schema.
+func (i *DatabaseInstance) ensureSchema() error {
+	installer := install.New(i.db, i.logger)
+
+	pending, err := installer.PendingMigrations(i.ctx)
+	if err != nil {
+		return fmt.Errorf("failed to check schema migrations: %w", err)
+	}
+
+	if len(pending) == 0 {
+		return nil
+	}
+
+	if !i.dbConfig.AutoMigrate {
+		return fmt.Errorf(
+			"schema is missing or behind: %d migration(s) not applied (%s); either run `pgsql-broker install` or set databases[].auto_migrate: true",
+			len(pending), strings.Join(pending, ", "),
+		)
+	}
+
+	i.logger.Info("auto-migrating database schema", "pending", pending)
+	if err := installer.ApplyMigrations(i.ctx); err != nil {
+		return fmt.Errorf("auto-migration failed: %w", err)
+	}
+
+	return nil
+}
+
 // Stop gracefully stops the database instance
 func (i *DatabaseInstance) Stop() error {
 	i.shutdownMu.Lock()
@@ -116,10 +164,26 @@ func (i *DatabaseInstance) Stop() error {
 	}
 	i.queuesMu.Unlock()

-	// Update instance status in database
-	if err := i.shutdownInstance(); err != nil {
+	// Update instance status in database. i.ctx may already be canceled by
+	// the parent broker's Stop(), so use a fresh short-lived context here.
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	if err := i.shutdownInstance(shutdownCtx); err != nil {
 		i.logger.Error("failed to shutdown instance in database", "error", err)
 	}
+	cancel()
+
+	// Release the advisory lock and close the pinned session connection.
+	if i.sessionConn != nil {
+		if _, err := i.sessionConn.ExecContext(context.Background(),
+			"SELECT pg_advisory_unlock(hashtextextended($1, 0))", "broker:"+i.Name,
+		); err != nil {
+			i.logger.Error("failed to release advisory lock", "error", err)
+		}
+		if err := i.sessionConn.Close(); err != nil {
+			i.logger.Error("failed to close session connection", "error", err)
+		}
+		i.sessionConn = nil
+	}

 	// Close database connection
 	if err := i.db.Close(); err != nil {
@@ -130,72 +194,48 @@ func (i *DatabaseInstance) Stop() error {
 	return nil
 }

-// registerInstance registers the instance in the database
+// registerInstance registers the instance in the database. The advisory
+// lock taken by broker_register_instance is session-scoped, so this runs on
+// a connection pinned out of the pool (i.sessionConn) that is kept open for
+// the life of the process rather than returned after this call.
 func (i *DatabaseInstance) registerInstance() error {
+	conn, err := i.db.Conn(i.ctx)
+	if err != nil {
+		return fmt.Errorf("failed to acquire session connection: %w", err)
+	}
+
 	var retval int
 	var errmsg string
-	var nullableInstanceID sql.NullInt64 // Change to nullable type
+	var nullableInstanceID sql.NullInt64

 	i.logger.Debug("registering instance", "name", i.Name, "hostname", i.Hostname, "pid", i.PID, "version", i.Version, "queue_count", i.dbConfig.QueueCount)
-	err := i.db.QueryRow(i.ctx,
-		"SELECT p_retval, p_errmsg, p_instance_id FROM broker_register_instance($1, $2, $3, $4, $5)",
+	err = conn.QueryRowContext(i.ctx,
+		"SELECT p_retval, p_errmsg, p_instance_id FROM broker.broker_register_instance($1, $2, $3, $4, $5)",
 		i.Name, i.Hostname, i.PID, i.Version, i.dbConfig.QueueCount,
 	).Scan(&retval, &errmsg, &nullableInstanceID)

 	if err != nil {
+		conn.Close()
 		i.logger.Error("query error during instance registration", "error", err)
 		return fmt.Errorf("query error: %w", err)
 	}

-	if retval == 3 {
-		i.logger.Warn("another broker instance is already active, attempting to retrieve ID", "error", errmsg)
-		// Try to retrieve the ID of the active instance
-		var activeID int64
-		err := i.db.QueryRow(i.ctx,
-			"SELECT id_broker_queueinstance FROM broker_queueinstance WHERE name = $1 AND hostname = $2 AND status = 'active' ORDER BY started_at DESC LIMIT 1",
-			i.Name, i.Hostname,
-		).Scan(&activeID)
-		if err != nil {
-			i.logger.Error("failed to retrieve ID of active instance", "error", err)
-			return fmt.Errorf("failed to retrieve ID of active instance: %w", err)
-		}
-		i.ID = activeID
-		i.logger.Info("retrieved active instance ID", "id", i.ID)
-		return nil
-	} else if retval > 0 {
+	if retval > 0 {
+		conn.Close()
 		i.logger.Error("broker_register_instance error", "retval", retval, "errmsg", errmsg)
 		return fmt.Errorf("broker_register_instance error: %s", errmsg)
 	}

-	// If successfully registered, nullableInstanceID.Valid will be true
-	if nullableInstanceID.Valid {
-		i.ID = nullableInstanceID.Int64
-		i.logger.Info("registered new instance", "id", i.ID)
-
-	// Debug logging: Retrieve all entries from broker_queueinstance
-	rows, err := i.db.Query(i.ctx, "SELECT id_broker_queueinstance, name, hostname, status FROM broker_queueinstance")
-	if err != nil {
-		i.logger.Error("debug query failed", "error", err)
-	} else {
-		defer rows.Close()
-		for rows.Next() {
-			var id int64
-			var name, hostname, status string
-			if err := rows.Scan(&id, &name, &hostname, &status); err != nil {
-				i.logger.Error("debug scan failed", "error", err)
-				break
-			}
-			i.logger.Debug("broker_queueinstance entry", "id", id, "name", name, "hostname", hostname, "status", status)
-		}
-	}
-	} else {
-		// This case should ideally not happen if retval is 0 (success)
-		// but if it does, it means p_instance_id was NULL despite success.
-		// This would be an unexpected scenario.
+	if !nullableInstanceID.Valid {
+		conn.Close()
 		i.logger.Error("broker_register_instance returned success but no instance ID", "retval", retval, "errmsg", errmsg)
 		return fmt.Errorf("broker_register_instance returned success but no instance ID")
 	}

+	i.ID = nullableInstanceID.Int64
+	i.sessionConn = conn
+	i.logger.Info("registered new instance", "id", i.ID)
+
 	return nil
 }

@@ -204,6 +244,11 @@ func (i *DatabaseInstance) startQueues() error {
 	i.queuesMu.Lock()
 	defer i.queuesMu.Unlock()

+	leaseSeconds := i.config.Broker.LeaseSeconds
+	if leaseSeconds <= 0 {
+		leaseSeconds = 60
+	}
+
 	for queueNum := 1; queueNum <= i.dbConfig.QueueCount; queueNum++ {
 		queueCfg := queue.Config{
 			Number:       queueNum,
@@ -214,6 +259,8 @@ func (i *DatabaseInstance) startQueues() error {
 			BufferSize:   i.config.Broker.QueueBufferSize,
 			TimerSeconds: i.config.Broker.QueueTimerSec,
 			FetchSize:    i.config.Broker.FetchQueryQueSize,
+			TenantID:     i.dbConfig.TenantID,
+			LeaseSeconds: leaseSeconds,
 		}

 		q := queue.New(queueCfg)
@@ -241,42 +288,38 @@ func (i *DatabaseInstance) startListener() error {
 	return nil
 }

-// handleNotification processes incoming job notifications
+// handleNotification processes incoming wake-up notifications. The payload
+// only carries the queue number (plus a job id kept for logging) --
+// NOTIFY is wake-only, never a hand-off of a job to execute directly, so the
+// woken worker always re-claims via broker_get.
 func (i *DatabaseInstance) handleNotification(n *adapter.Notification) {
+	defer adapter.RecoverAndLog(i.logger, "handle-notification")
+
 	if i.config.Broker.EnableDebug {
 		i.logger.Debug("received notification", "channel", n.Channel, "payload", n.Payload)
 	}

-	var job models.Job
-	if err := json.Unmarshal([]byte(n.Payload), &job); err != nil {
+	var wake models.WakeNotification
+	if err := json.Unmarshal([]byte(n.Payload), &wake); err != nil {
 		i.logger.Error("failed to unmarshal notification", "error", err, "payload", n.Payload)
 		return
 	}

-	if job.ID <= 0 {
-		i.logger.Warn("notification missing job ID", "payload", n.Payload)
+	if wake.Queue <= 0 {
+		i.logger.Warn("notification missing queue number", "payload", n.Payload)
 		return
 	}

-	if job.JobQueue <= 0 {
-		i.logger.Warn("notification missing queue number", "job_id", job.ID)
-		return
-	}
-
-	// Get the queue
 	i.queuesMu.RLock()
-	q, exists := i.queues[job.JobQueue]
+	q, exists := i.queues[wake.Queue]
 	i.queuesMu.RUnlock()

 	if !exists {
-		i.logger.Warn("queue not found for job", "job_id", job.ID, "queue", job.JobQueue)
+		i.logger.Warn("queue not found for notification", "queue", wake.Queue, "job_id", wake.JobID)
 		return
 	}

-	// Add job to queue
-	if err := q.AddJob(job); err != nil {
-		i.logger.Error("failed to add job to queue", "job_id", job.ID, "queue", job.JobQueue, "error", err)
-	}
+	q.Wake()
 }

 // pingRoutine periodically updates the instance status in the database
@@ -304,13 +347,67 @@ func (i *DatabaseInstance) pingRoutine() {
 	}
 }

+// staleJobRecoveryRoutine periodically requeues (or dead-letters) jobs whose
+// lease has expired while still running.
+func (i *DatabaseInstance) staleJobRecoveryRoutine() {
+	interval := time.Duration(i.config.Broker.StaleJobRecoverySec) * time.Second
+	if interval <= 0 {
+		interval = 30 * time.Second
+	}
+
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			i.shutdownMu.RLock()
+			if i.shutdown {
+				i.shutdownMu.RUnlock()
+				return
+			}
+			i.shutdownMu.RUnlock()
+
+			if err := i.recoverStaleJobs(); err != nil {
+				i.logger.Error("stale job recovery failed", "error", err)
+			}
+
+		case <-i.ctx.Done():
+			return
+		}
+	}
+}
+
+// recoverStaleJobs invokes broker_recover_stale_jobs.
+func (i *DatabaseInstance) recoverStaleJobs() error {
+	var retval int
+	var errmsg string
+	var recoveredCount int
+
+	err := i.db.QueryRow(i.ctx, "SELECT p_retval, p_errmsg, p_recovered_count FROM broker.broker_recover_stale_jobs()").
+		Scan(&retval, &errmsg, &recoveredCount)
+	if err != nil {
+		return fmt.Errorf("query error: %w", err)
+	}
+
+	if retval > 0 {
+		return fmt.Errorf("broker_recover_stale_jobs error: %s", errmsg)
+	}
+
+	if recoveredCount > 0 {
+		i.logger.Info("recovered stale jobs", "count", recoveredCount)
+	}
+
+	return nil
+}
+
 // ping updates the instance ping timestamp
 func (i *DatabaseInstance) ping() error {
 	var retval int
 	var errmsg string

 	err := i.db.QueryRow(i.ctx,
-		"SELECT p_retval, p_errmsg FROM broker_ping_instance($1, $2)",
+		"SELECT p_retval, p_errmsg FROM broker.broker_ping_instance($1, $2)",
 		i.ID, i.jobsHandled,
 	).Scan(&retval, &errmsg)

@@ -326,12 +423,12 @@ func (i *DatabaseInstance) ping() error {
 }

 // shutdownInstance marks the instance as shutdown in the database
-func (i *DatabaseInstance) shutdownInstance() error {
+func (i *DatabaseInstance) shutdownInstance(ctx context.Context) error {
 	var retval int
 	var errmsg string

-	err := i.db.QueryRow(i.ctx,
-		"SELECT p_retval, p_errmsg FROM broker_shutdown_instance($1)",
+	err := i.db.QueryRow(ctx,
+		"SELECT p_retval, p_errmsg FROM broker.broker_shutdown_instance($1)",
 		i.ID,
 	).Scan(&retval, &errmsg)

@@ -370,4 +467,4 @@ func (i *DatabaseInstance) GetStats() map[string]interface{} {
 	stats["queues"] = queueStats

 	return stats
-}
+}