pgsql-broker/pkg/broker/database_instance.go

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
}