package broker import ( "context" "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/metrics" "git.warky.dev/wdevs/pgsql-broker/pkg/broker/models" "git.warky.dev/wdevs/pgsql-broker/pkg/broker/queue" ) // DatabaseInstance represents a broker instance for a single database type DatabaseInstance struct { ID int64 Name string DatabaseName string Hostname string PID int Version string config *config.Config dbConfig *config.DatabaseConfig db adapter.DBAdapter logger adapter.Logger queues map[int]*queue.Queue queuesMu sync.RWMutex ctx context.Context shutdown bool shutdownMu sync.RWMutex jobsHandled int64 startTime time.Time metrics *metrics.Metrics // 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 func NewDatabaseInstance(cfg *config.Config, dbCfg *config.DatabaseConfig, db adapter.DBAdapter, logger adapter.Logger, version string, parentCtx context.Context, brokerMetrics ...*metrics.Metrics) (*DatabaseInstance, error) { hostname, err := os.Hostname() if err != nil { hostname = "unknown" } var instanceMetrics *metrics.Metrics if len(brokerMetrics) > 0 { instanceMetrics = brokerMetrics[0] } instance := &DatabaseInstance{ Name: fmt.Sprintf("%s-%s", cfg.Broker.Name, dbCfg.Name), DatabaseName: dbCfg.Name, Hostname: hostname, PID: os.Getpid(), Version: version, config: cfg, dbConfig: dbCfg, db: db, logger: logger.With("component", "database-instance").With("database", dbCfg.Name), queues: make(map[int]*queue.Queue), ctx: parentCtx, startTime: time.Now(), metrics: instanceMetrics, } return instance, nil } // Start begins the database instance func (i *DatabaseInstance) Start() error { i.logger.Info("starting database instance", "name", i.Name, "hostname", i.Hostname, "pid", i.PID) // Connect to database if err := i.db.Connect(i.ctx); err != nil { 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) } i.logger.Info("database instance registered", "id", i.ID) // Start queues if err := i.startQueues(); err != nil { return fmt.Errorf("failed to start queues: %w", err) } // Start listening for notifications if err := i.startListener(); err != nil { return fmt.Errorf("failed to start listener: %w", err) } // 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") if i.metrics != nil { adapter.SupervisedGo(i.logger, "metrics-queue-depth-routine", i.queueDepthRoutine) } return nil } // queueDepthRoutine periodically exports pending jobs grouped by queue. func (i *DatabaseInstance) queueDepthRoutine() { interval := time.Duration(i.config.Broker.QueueDepthPollSec) * time.Second if interval <= 0 { interval = 15 * time.Second } ticker := time.NewTicker(interval) defer ticker.Stop() i.updateQueueDepthMetrics() for { select { case <-ticker.C: i.updateQueueDepthMetrics() case <-i.ctx.Done(): return } } } func (i *DatabaseInstance) updateQueueDepthMetrics() { rows, err := i.db.Query(i.ctx, "SELECT job_queue, COUNT(*) FROM broker.broker_jobs WHERE complete_status = 0 GROUP BY job_queue") if err != nil { i.logger.Warn("failed to collect queue depth metrics", "error", err) return } defer rows.Close() for queueNumber := 1; queueNumber <= i.dbConfig.QueueCount; queueNumber++ { i.metrics.SetJobsQueued(i.DatabaseName, queueNumber, 0) } for rows.Next() { var queueNumber, count int if err := rows.Scan(&queueNumber, &count); err != nil { i.logger.Warn("failed to scan queue depth metric", "error", err) return } i.metrics.SetJobsQueued(i.DatabaseName, queueNumber, float64(count)) } if err := rows.Err(); err != nil { i.logger.Warn("failed to read queue depth metrics", "error", err) } } // 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() if i.shutdown { i.shutdownMu.Unlock() return nil } i.shutdown = true i.shutdownMu.Unlock() i.logger.Info("stopping database instance") // Stop all queues i.queuesMu.Lock() for num, q := range i.queues { i.logger.Info("stopping queue", "number", num) if err := q.Stop(); err != nil { i.logger.Error("failed to stop queue", "number", num, "error", err) } } i.queuesMu.Unlock() // 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 { i.logger.Error("failed to close database", "error", err) } i.logger.Info("database instance stopped") return nil } // 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 i.logger.Debug("registering instance", "name", i.Name, "hostname", i.Hostname, "pid", i.PID, "version", i.Version, "queue_count", i.dbConfig.QueueCount) 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 > 0 { conn.Close() i.logger.Error("broker_register_instance error", "retval", retval, "errmsg", errmsg) return fmt.Errorf("broker_register_instance error: %s", errmsg) } 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 } // startQueues initializes and starts all queues 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, InstanceID: i.ID, WorkerCount: 1, // One worker per queue for now DBAdapter: i.db, Logger: i.logger, BufferSize: i.config.Broker.QueueBufferSize, TimerSeconds: i.config.Broker.QueueTimerSec, FetchSize: i.config.Broker.FetchQueryQueSize, TenantID: i.dbConfig.TenantID, LeaseSeconds: leaseSeconds, Metrics: i.metrics, DatabaseName: i.DatabaseName, } q := queue.New(queueCfg) if err := q.Start(queueCfg); err != nil { return fmt.Errorf("failed to start queue %d: %w", queueNum, err) } i.queues[queueNum] = q i.logger.Info("queue started", "number", queueNum) } i.metrics.SetQueueCount(i.DatabaseName, len(i.queues)) return nil } // startListener starts listening for database notifications func (i *DatabaseInstance) startListener() error { handler := func(n *adapter.Notification) { i.handleNotification(n) } if err := i.db.Listen(i.ctx, "broker.event", handler); err != nil { return fmt.Errorf("failed to start listener: %w", err) } return nil } // 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 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 wake.Queue <= 0 { i.logger.Warn("notification missing queue number", "payload", n.Payload) return } i.queuesMu.RLock() q, exists := i.queues[wake.Queue] i.queuesMu.RUnlock() if !exists { i.logger.Warn("queue not found for notification", "queue", wake.Queue, "job_id", wake.JobID) return } q.Wake() } // pingRoutine periodically updates the instance status in the database func (i *DatabaseInstance) pingRoutine() { ticker := time.NewTicker(30 * time.Second) defer ticker.Stop() for { select { case <-ticker.C: i.shutdownMu.RLock() if i.shutdown { i.shutdownMu.RUnlock() return } i.shutdownMu.RUnlock() if err := i.ping(); err != nil { i.logger.Error("ping failed", "error", err) } case <-i.ctx.Done(): return } } } // 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.broker_ping_instance($1, $2)", i.ID, i.jobsHandled, ).Scan(&retval, &errmsg) if err != nil { return fmt.Errorf("query error: %w", err) } if retval > 0 { return fmt.Errorf("broker_ping_instance error: %s", errmsg) } return nil } // shutdownInstance marks the instance as shutdown in the database func (i *DatabaseInstance) shutdownInstance(ctx context.Context) error { var retval int var errmsg string err := i.db.QueryRow(ctx, "SELECT p_retval, p_errmsg FROM broker.broker_shutdown_instance($1)", i.ID, ).Scan(&retval, &errmsg) if err != nil { return fmt.Errorf("query error: %w", err) } if retval > 0 { return fmt.Errorf("broker_shutdown_instance error: %s", errmsg) } return nil } // GetStats returns instance statistics func (i *DatabaseInstance) GetStats() map[string]interface{} { i.queuesMu.RLock() defer i.queuesMu.RUnlock() stats := map[string]interface{}{ "id": i.ID, "name": i.Name, "database_name": i.DatabaseName, "hostname": i.Hostname, "pid": i.PID, "version": i.Version, "uptime": time.Since(i.startTime).String(), "jobs_handled": i.jobsHandled, "queue_count": len(i.queues), } queueStats := make(map[int]interface{}) for num, q := range i.queues { queueStats[num] = q.GetStats() } stats["queues"] = queueStats return stats }