pgsql-broker/pkg/broker/worker/worker.go

package worker

import (
	"context"
	"database/sql"
	"fmt"
	"runtime/debug"
	"sync"
	"time"

	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/adapter"
	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/metrics"
)

// Worker represents a single job processing worker
type Worker struct {
	ID           int
	QueueNumber  int
	InstanceID   int64
	db           adapter.DBAdapter
	logger       adapter.Logger
	wakeChan     chan struct{}
	shutdown     chan struct{}
	wg           *sync.WaitGroup
	running      bool
	mu           sync.RWMutex
	lastActivity time.Time
	jobsHandled  int64
	timerSeconds int
	fetchSize    int
	tenantID     string
	leaseSeconds int
	metrics      *metrics.Metrics
	databaseName string
}

// Stats holds worker statistics
type Stats struct {
	LastActivity time.Time
	JobsHandled  int64
	Running      bool
}

// Config holds worker configuration
type Config struct {
	ID           int
	QueueNumber  int
	InstanceID   int64
	DBAdapter    adapter.DBAdapter
	Logger       adapter.Logger
	BufferSize   int
	TimerSeconds int
	FetchSize    int
	TenantID     string
	LeaseSeconds int
	Metrics      *metrics.Metrics
	DatabaseName string
}

// New creates a new worker
func New(cfg Config) *Worker {
	leaseSeconds := cfg.LeaseSeconds
	if leaseSeconds <= 0 {
		leaseSeconds = 60
	}

	return &Worker{
		ID:           cfg.ID,
		QueueNumber:  cfg.QueueNumber,
		InstanceID:   cfg.InstanceID,
		db:           cfg.DBAdapter,
		logger:       cfg.Logger.With("worker_id", cfg.ID).With("queue", cfg.QueueNumber),
		wakeChan:     make(chan struct{}, 1),
		shutdown:     make(chan struct{}),
		wg:           &sync.WaitGroup{},
		timerSeconds: cfg.TimerSeconds,
		fetchSize:    cfg.FetchSize,
		tenantID:     cfg.TenantID,
		leaseSeconds: leaseSeconds,
		metrics:      cfg.Metrics,
		databaseName: cfg.DatabaseName,
	}
}

// Start begins the worker processing loop
func (w *Worker) Start(ctx context.Context) error {
	w.mu.Lock()
	if w.running {
		w.mu.Unlock()
		return fmt.Errorf("worker %d already running", w.ID)
	}
	w.running = true
	w.mu.Unlock()

	w.logger.Info("worker starting")

	w.wg.Add(1)
	go w.superviseProcessLoop(ctx)

	return nil
}

// superviseProcessLoop runs processLoop for the life of the worker,
// restarting it (after a short backoff) if it ever panics, so a bug in job
// processing can never permanently kill this worker's goroutine.
func (w *Worker) superviseProcessLoop(ctx context.Context) {
	defer w.wg.Done()

	for {
		select {
		case <-w.shutdown:
			return
		case <-ctx.Done():
			return
		default:
		}

		if w.runProcessLoopOnce(ctx) {
			return
		}

		time.Sleep(time.Second)
	}
}

// runProcessLoopOnce runs processLoop, recovering any panic. It returns true
// if processLoop returned normally (shutdown/context done, no restart
// needed) and false if it panicked (caller should restart it).
func (w *Worker) runProcessLoopOnce(ctx context.Context) (clean bool) {
	defer func() {
		if r := recover(); r != nil {
			w.logger.Error("worker panic recovered, restarting", "panic", r, "stack", string(debug.Stack()))
			clean = false
		}
	}()
	w.processLoop(ctx)
	return true
}

// Stop gracefully stops the worker
func (w *Worker) Stop() error {
	w.mu.Lock()
	if !w.running {
		w.mu.Unlock()
		return nil
	}
	w.mu.Unlock()

	w.logger.Info("worker stopping")
	close(w.shutdown)
	w.wg.Wait()

	w.mu.Lock()
	w.running = false
	w.mu.Unlock()

	w.logger.Info("worker stopped")
	return nil
}

// Wake signals the worker to check for available jobs immediately, instead
// of waiting for the next timer tick. Fetching is always DB-driven (via
// broker_get with FOR UPDATE SKIP LOCKED), so a redundant or coalesced wake
// is harmless.
func (w *Worker) Wake() {
	select {
	case w.wakeChan <- struct{}{}:
	default:
	}
}

// processLoop is the main worker processing loop
func (w *Worker) processLoop(ctx context.Context) {
	timer := time.NewTimer(time.Duration(w.timerSeconds) * time.Second)
	defer timer.Stop()

	for {
		select {
		case <-w.wakeChan:
			w.updateActivity()
			w.processJobs(ctx)

		case <-timer.C:
			// Timer expired - fetch jobs from database
			if w.timerSeconds > 0 {
				w.updateActivity()
				w.processJobs(ctx)
			}
			timer.Reset(time.Duration(w.timerSeconds) * time.Second)

		case <-w.shutdown:
			w.logger.Info("worker shutdown signal received")
			return

		case <-ctx.Done():
			w.logger.Info("worker context cancelled")
			return
		}
	}
}

// processJobs processes jobs from the queue within a transaction
func (w *Worker) processJobs(ctx context.Context) {
	defer w.recoverPanic()

	for i := 0; i < w.fetchSize; i++ {
		tx, err := w.db.Begin(ctx)
		if err != nil {
			w.logger.Error("failed to begin transaction", "error", err)
			return
		}

		if err := w.setTenantTx(ctx, tx); err != nil {
			if rbErr := tx.Rollback(); rbErr != nil {
				w.logger.Error("failed to rollback transaction", "error", rbErr)
			}
			w.logger.Error("failed to set tenant", "error", err)
			return
		}

		jobID, leaseToken, err := w.fetchNextJobTx(ctx, tx)
		if err != nil {
			if rbErr := tx.Rollback(); rbErr != nil {
				w.logger.Error("failed to rollback transaction", "error", rbErr)
			}
			w.logger.Error("failed to fetch job", "error", err)
			return
		}

		if jobID <= 0 {
			// No job found, rollback
			if rbErr := tx.Rollback(); rbErr != nil {
				w.logger.Error("failed to rollback transaction", "error", rbErr)
			}
			return // No more jobs
		}

		jobName, jobGroup, err := w.fetchJobLabelsTx(ctx, tx, jobID)
		if err != nil {
			w.logger.Warn("failed to fetch job labels for metrics", "job_id", jobID, "error", err)
		}

		// Run the job
		start := time.Now()
		jobStatus, err := w.runJobTx(ctx, tx, jobID, leaseToken)
		duration := time.Since(start)
		if err == nil {
			w.recordJobMetric(jobStatus, jobGroup, jobName, duration)
		}
		if err != nil {
			// Rollback on genuine infra failure
			if rbErr := tx.Rollback(); rbErr != nil {
				w.logger.Error("failed to rollback transaction", "error", rbErr)
			}
			w.logger.Error("failed to run job", "job_id", jobID, "error", err)
		} else {
			if err := tx.Commit(); err != nil {
				w.logger.Error("failed to commit job", "job_id", jobID, "error", err)
				continue
			}
			w.jobsHandled++
		}
	}
}

// setTenantTx applies this worker's RLS tenant for the duration of tx.
func (w *Worker) setTenantTx(ctx context.Context, tx adapter.DBTransaction) error {
	tenantID := w.tenantID
	if tenantID == "" {
		tenantID = "default"
	}
	_, err := tx.Exec(ctx, "SELECT broker.broker_set_tenant($1)", tenantID)
	return err
}

// fetchNextJobTx fetches the next job from the queue within a transaction,
// claiming it with a lease that must be presented back to broker_run.
func (w *Worker) fetchNextJobTx(ctx context.Context, tx adapter.DBTransaction) (jobID int64, leaseToken string, err error) {
	var retval int
	var errmsg string
	var nullableJobID sql.NullInt64
	var nullableLeaseToken sql.NullString

	err = tx.QueryRow(ctx,
		"SELECT p_retval, p_errmsg, p_job_id, p_lease_token FROM broker.broker_get($1, $2, $3)",
		w.QueueNumber, w.InstanceID, w.leaseSeconds,
	).Scan(&retval, &errmsg, &nullableJobID, &nullableLeaseToken)

	if err != nil {
		return 0, "", fmt.Errorf("query error: %w", err)
	}

	if retval > 0 {
		return 0, "", fmt.Errorf("broker_get error: %s", errmsg)
	}

	if !nullableJobID.Valid {
		return 0, "", nil
	}

	return nullableJobID.Int64, nullableLeaseToken.String, nil
}

// runJobTx executes a leased job within a transaction. It only returns an
// error (triggering a rollback of the claim) on a genuine infra failure --
// job outcomes reported via p_job_status (requeued/completed/dead-lettered)
// are always committed.
func (w *Worker) runJobTx(ctx context.Context, tx adapter.DBTransaction, jobID int64, leaseToken string) (int, error) {
	w.logger.Debug("running job", "job_id", jobID)

	var retval int
	var errmsg string
	var jobStatus int

	err := tx.QueryRow(ctx,
		"SELECT p_retval, p_errmsg, p_job_status FROM broker.broker_run($1, $2)",
		jobID, leaseToken,
	).Scan(&retval, &errmsg, &jobStatus)

	if err != nil {
		return 0, fmt.Errorf("query error: %w", err)
	}

	if retval > 0 {
		return 0, fmt.Errorf("broker_run error: %s", errmsg)
	}

	w.logger.Debug("job finished", "job_id", jobID, "job_status", jobStatus)
	return jobStatus, nil
}

// fetchJobLabelsTx looks up the job_name/job_group of jobID for metric
// labeling. Best-effort: callers log and continue on error rather than
// failing the job over a metrics lookup.
func (w *Worker) fetchJobLabelsTx(ctx context.Context, tx adapter.DBTransaction, jobID int64) (jobName, jobGroup string, err error) {
	err = tx.QueryRow(ctx,
		"SELECT job_name, job_group FROM broker.broker_jobs WHERE id_broker_jobs = $1",
		jobID,
	).Scan(&jobName, &jobGroup)
	if err != nil {
		return "", "", fmt.Errorf("query error: %w", err)
	}
	return jobName, jobGroup, nil
}

// recordJobMetric routes a finished job attempt to the appropriate
// Prometheus counter/histogram based on the p_job_status broker_run
// reported (0=requeued, 2=completed, 3=dead-lettered).
func (w *Worker) recordJobMetric(jobStatus int, jobGroup, jobName string, duration time.Duration) {
	switch jobStatus {
	case 2:
		w.metrics.RecordJobCompleted(w.databaseName, jobGroup, jobName, duration)
	case 3:
		w.metrics.RecordJobFailed(w.databaseName, jobGroup, jobName, duration)
	case 0:
		w.metrics.RecordJobRequeued(w.databaseName, jobGroup, jobName, duration)
	}
}

// updateActivity updates the last activity timestamp
func (w *Worker) updateActivity() {
	w.mu.Lock()
	w.lastActivity = time.Now()
	w.mu.Unlock()
}

// GetStats returns worker statistics
func (w *Worker) GetStats() (lastActivity time.Time, jobsHandled int64, running bool) {
	w.mu.RLock()
	defer w.mu.RUnlock()
	return w.lastActivity, w.jobsHandled, w.running
}

// recoverPanic recovers from panics in the worker
func (w *Worker) recoverPanic() {
	if r := recover(); r != nil {
		w.logger.Error("worker panic recovered", "panic", r, "stack", string(debug.Stack()))
	}
}