query-orchestration/internal/backgroundtask/runner.go

package backgroundtask

import (
	"context"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"
)

// Runner executes a work function periodically or when signaled.
// It provides graceful shutdown, work coalescing, and statistics tracking.
type Runner struct {
	config       map[string]any // Configuration passed to work function
	logger       *slog.Logger   // Logger for output
	workFunc     WorkFunc       // The function to execute
	wakeCh       chan struct{}  // Channel to signal immediate work
	stopCh       chan struct{}  // Channel to signal shutdown
	doneCh       chan struct{}  // Channel closed when goroutine exits
	ticker       *time.Ticker   // Periodic timer
	mu           sync.Mutex     // Protects started flag
	started      bool           // Whether Run has been called
	stopped      atomic.Bool    // Whether shutdown has been called
	working      atomic.Bool    // Whether work is currently executing
	pendingWork  atomic.Bool    // Whether work was requested during execution
	interval     time.Duration  // Time between periodic executions
	stats        Stats          // Execution statistics
	shutdownOnce sync.Once      // Ensures shutdown runs once
}

// Initialize creates a new Runner with the provided configuration, logger, and work function.
// Options can be provided to customize the runner's behavior.
func Initialize(config map[string]any, logger *slog.Logger, workFunc WorkFunc, opts ...Option) (*Runner, error) {
	if config == nil {
		return nil, ErrNilConfig
	}
	if workFunc == nil {
		return nil, ErrNilWorkFunc
	}
	if logger == nil {
		logger = slog.Default()
	}

	r := &Runner{
		config:   config,
		logger:   logger,
		workFunc: workFunc,
		wakeCh:   make(chan struct{}, 1),
		stopCh:   make(chan struct{}),
		doneCh:   make(chan struct{}),
		interval: DefaultInterval,
	}

	// Apply options
	for _, opt := range opts {
		if err := opt(r); err != nil {
			return nil, err
		}
	}

	return r, nil
}

// Stats returns a pointer to the current statistics.
// This method is thread-safe and can be called concurrently.
func (r *Runner) Stats() *Stats {
	return &r.stats
}

// Run starts the background goroutine that executes the work function.
// This method is idempotent - calling it multiple times has no effect if already running.
// Returns ErrAlreadyRunning if called after shutdown.
func (r *Runner) Run() error {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.stopped.Load() {
		return ErrStopped
	}

	if r.started {
		// Idempotent - already running
		return nil
	}

	r.started = true
	r.ticker = time.NewTicker(r.interval)

	// Start the worker goroutine
	go r.worker()

	r.logger.Info("background runner started",
		"interval", r.interval,
		"config_type", r.config != nil,
	)

	return nil
}

// Signal requests immediate execution of the work function.
// If the runner is already executing work, the signal is coalesced and
// another execution will occur after the current one completes.
// Returns ErrStopped if the runner has been shut down.
func (r *Runner) Signal() error {
	if r.stopped.Load() {
		return ErrStopped
	}

	// Non-blocking send - if channel is full, work is already pending
	select {
	case r.wakeCh <- struct{}{}:
		r.logger.Debug("work signaled")
	default:
		// Channel full - work already pending
		// Set pendingWork flag to ensure we run again after current work
		r.pendingWork.Store(true)
		r.logger.Debug("work signal coalesced")
	}

	return nil
}

// Shutdown gracefully stops the runner, waiting for any in-progress work to complete.
// The provided context controls the shutdown timeout. If the context has no deadline,
// a default timeout of MaxShutdownTimeout (60 seconds) is applied.
// This method is idempotent - multiple calls will not cause issues.
func (r *Runner) Shutdown(ctx context.Context) error {
	var err error

	r.shutdownOnce.Do(func() {
		r.logger.Info("shutting down background runner")

		// Mark as stopped to prevent new signals
		r.stopped.Store(true)

		// Stop the ticker
		if r.ticker != nil {
			r.ticker.Stop()
		}

		// Signal the worker to stop
		close(r.stopCh)

		// Apply default timeout if context has no deadline
		if _, hasDeadline := ctx.Deadline(); !hasDeadline {
			var cancel context.CancelFunc
			ctx, cancel = context.WithTimeout(ctx, MaxShutdownTimeout)
			defer cancel()
		}

		// Wait for worker to finish or timeout
		select {
		case <-r.doneCh:
			r.logger.Info("background runner shutdown complete")
		case <-ctx.Done():
			r.logger.Error("background runner shutdown timeout")
			err = ErrShutdownTimeout
		}
	})

	return err
}

// worker is the main goroutine that executes the work function.
func (r *Runner) worker() {
	defer close(r.doneCh)
	defer r.ticker.Stop()

	r.logger.Debug("worker started")

	for {
		select {
		case <-r.wakeCh:
			r.logger.Debug("worker woken by signal")
			r.tryWork("signal")

		case <-r.ticker.C:
			r.logger.Debug("worker woken by ticker")
			r.tryWork("ticker")

		case <-r.stopCh:
			r.logger.Debug("worker stopping")
			return
		}
	}
}

// tryWork attempts to execute the work function, handling coalescing and pending work.
func (r *Runner) tryWork(trigger string) {
	// Check if already working
	if !r.working.CompareAndSwap(false, true) {
		// Already working - set pending flag
		r.pendingWork.Store(true)
		r.logger.Debug("work already in progress, pending flag set", "trigger", trigger)
		return
	}

	// Execute work
	r.doWork(trigger)

	// Clear working flag
	r.working.Store(false)

	// Check if there's pending work
	if r.pendingWork.CompareAndSwap(true, false) {
		r.logger.Debug("processing pending work")
		// Recursively process pending work (will only recurse once due to working flag)
		r.tryWork("pending")
	}
}

// doWork executes the actual work function with proper context and stats tracking.
func (r *Runner) doWork(trigger string) {
	// Create context that will be canceled on shutdown
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Create a done channel for this work execution
	done := make(chan struct{})
	defer close(done)

	// Monitor for shutdown
	go func() {
		select {
		case <-r.stopCh:
			cancel()
		case <-done:
			// Work completed
		}
	}()

	// Update stats
	r.stats.enterWork()
	defer r.stats.exitWork()

	startTime := time.Now()
	r.stats.setLastStart(startTime)

	r.logger.Debug("executing work", "trigger", trigger)

	// Execute the work function
	err := r.workFunc(ctx, r.logger, r.config)

	// Update stats
	r.stats.setLastFinish(time.Now())
	r.stats.incrementWorkCycles()
	if err != nil {
		r.stats.setLastErr(err)
		r.logger.Error("work function error", "error", err, "duration", time.Since(startTime))
	} else {
		r.stats.setLastErr(nil)
		r.logger.Debug("work completed", "duration", time.Since(startTime))
	}
}