package worker import ( "context" "database/sql" "fmt" "runtime/debug" "sync" "time" "git.warky.dev/wdevs/pgsql-broker/pkg/broker/adapter" ) // 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 } // 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 } // 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, } } // 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 { tx.Rollback() w.logger.Error("failed to set tenant", "error", err) return } jobID, leaseToken, err := w.fetchNextJobTx(ctx, tx) if err != nil { tx.Rollback() w.logger.Error("failed to fetch job", "error", err) return } if jobID <= 0 { tx.Rollback() // No job found, rollback return // No more jobs } // Run the job if err := w.runJobTx(ctx, tx, jobID, leaseToken); err != nil { tx.Rollback() // Rollback on genuine infra failure 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) (int64, string, 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) 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 fmt.Errorf("query error: %w", err) } if retval > 0 { return fmt.Errorf("broker_run error: %s", errmsg) } w.logger.Debug("job finished", "job_id", jobID, "job_status", jobStatus) return nil } // 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())) } }