ITWorx-Pulse-Public/internal/probe/scheduler.go

package probe

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"
)

type Executor interface {
	Execute(context.Context, Definition) (Result, error)
}

type RetryableError struct{ Err error }

func (e RetryableError) Error() string {
	if e.Err == nil {
		return "retryable probe error"
	}
	return e.Err.Error()
}
func (e RetryableError) Unwrap() error { return e.Err }

type SchedulerConfig struct {
	MaxConcurrent  int
	MaxAttempts    int
	AttemptTimeout time.Duration
	RetryBackoff   time.Duration
	Now            func() time.Time
}

func (c SchedulerConfig) withDefaults() SchedulerConfig {
	if c.MaxConcurrent == 0 {
		c.MaxConcurrent = 16
	}
	if c.MaxAttempts == 0 {
		c.MaxAttempts = 2
	}
	if c.AttemptTimeout == 0 {
		c.AttemptTimeout = 10 * time.Second
	}
	if c.RetryBackoff == 0 {
		c.RetryBackoff = 100 * time.Millisecond
	}
	if c.Now == nil {
		c.Now = func() time.Time { return time.Now().UTC() }
	}
	return c
}

func (c SchedulerConfig) Validate() error {
	if c.MaxConcurrent < 1 || c.MaxConcurrent > 64 || c.MaxAttempts < 1 || c.MaxAttempts > 3 || c.AttemptTimeout <= 0 || c.AttemptTimeout > 2*time.Minute || c.RetryBackoff < 0 || c.RetryBackoff > time.Minute {
		return errors.New("probe scheduler configuration is outside safe bounds")
	}
	return nil
}

type Metrics struct {
	Runs           int64     `json:"runs"`
	Completed      int64     `json:"completed"`
	Failed         int64     `json:"failed"`
	TimedOut       int64     `json:"timedOut"`
	Retried        int64     `json:"retried"`
	SkippedOverlap int64     `json:"skippedOverlap"`
	Active         int64     `json:"active"`
	LastRunAt      time.Time `json:"lastRunAt"`
}

type RunReport struct {
	Results []Result `json:"results"`
	Metrics Metrics  `json:"metrics"`
}

type Scheduler struct {
	executor Executor
	config   SchedulerConfig
	sem      chan struct{}
	mu       sync.Mutex
	inflight map[string]context.CancelFunc
	closed   bool
	wg       sync.WaitGroup
	metrics  Metrics
}

func NewScheduler(executor Executor, config SchedulerConfig) (*Scheduler, error) {
	config = config.withDefaults()
	if executor == nil {
		return nil, errors.New("probe executor is required")
	}
	if err := config.Validate(); err != nil {
		return nil, err
	}
	return &Scheduler{executor: executor, config: config, sem: make(chan struct{}, config.MaxConcurrent), inflight: make(map[string]context.CancelFunc)}, nil
}

func (s *Scheduler) Run(ctx context.Context, definitions []Definition) (RunReport, error) {
	if s == nil {
		return RunReport{}, errors.New("probe scheduler is nil")
	}
	if ctx == nil {
		return RunReport{}, errors.New("probe scheduler context is nil")
	}
	if err := ctx.Err(); err != nil {
		return RunReport{}, err
	}
	items := append([]Definition(nil), definitions...)
	sort.SliceStable(items, func(i, j int) bool { return items[i].ID < items[j].ID })
	for _, definition := range items {
		if !definition.Enabled || definition.ArchivedAt != nil {
			continue
		}
		if err := definition.Validate(); err != nil {
			return RunReport{}, fmt.Errorf("validate probe %s: %w", definition.ID, err)
		}
	}
	results := make(chan Result, len(items))
	for _, definition := range items {
		if definition.Enabled && definition.ArchivedAt == nil {
			s.begin(definition, ctx, results)
		}
	}
	s.wg.Wait()
	close(results)
	report := RunReport{Results: make([]Result, 0, len(results))}
	for result := range results {
		report.Results = append(report.Results, result)
	}
	sort.SliceStable(report.Results, func(i, j int) bool { return report.Results[i].ProbeID < report.Results[j].ProbeID })
	s.mu.Lock()
	report.Metrics = s.metrics
	s.mu.Unlock()
	return report, nil
}
func (s *Scheduler) begin(definition Definition, parent context.Context, results chan<- Result) bool {
	probeID := definition.ID
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return false
	}
	if _, exists := s.inflight[probeID]; exists {
		s.metrics.SkippedOverlap++
		s.mu.Unlock()
		return false
	}
	child, cancel := context.WithCancel(parent)
	s.inflight[probeID] = cancel
	s.wg.Add(1)
	s.metrics.Runs++
	s.metrics.Active++
	s.mu.Unlock()
	go func() {
		defer s.wg.Done()
		defer func() { s.mu.Lock(); delete(s.inflight, probeID); s.metrics.Active--; s.mu.Unlock() }()
		select {
		case s.sem <- struct{}{}:
		case <-child.Done():
			result := s.failureResult(probeID, child.Err(), 0)
			s.recordFailure(result)
			results <- result
			return
		}
		defer func() { <-s.sem }()
		result := s.executeDefinition(child, definition)
		results <- result
	}()
	return true
}

func (s *Scheduler) executeDefinition(ctx context.Context, definition Definition) Result {
	started := s.config.Now().UTC()
	var lastErr error
	var attempts int
	for attempts = 1; attempts <= s.config.MaxAttempts; attempts++ {
		attemptCtx, cancel := context.WithTimeout(ctx, s.config.AttemptTimeout)
		result, err := s.executor.Execute(attemptCtx, definition)
		deadline := errors.Is(attemptCtx.Err(), context.DeadlineExceeded)
		cancel()
		if err == nil {
			result = normalizeResult(result, definition.ID, attempts, started, s.config.Now)
			s.recordSuccess(result, false)
			return result
		}
		lastErr = err
		if errors.Is(ctx.Err(), context.Canceled) || errors.Is(ctx.Err(), context.DeadlineExceeded) {
			break
		}
		if attempts >= s.config.MaxAttempts || (!deadline && !isRetryable(err)) {
			break
		}
		s.mu.Lock()
		s.metrics.Retried++
		s.mu.Unlock()
		if !waitBackoff(ctx, s.config.RetryBackoff) {
			break
		}
	}
	result := s.failureResult(definition.ID, lastErr, attempts)
	if errors.Is(ctx.Err(), context.Canceled) || errors.Is(ctx.Err(), context.DeadlineExceeded) {
		result.ErrorClass = "canceled"
	}
	s.recordFailure(result)
	return result
}

func (s *Scheduler) failureResult(probeID string, err error, attempts int) Result {
	result := Result{ProbeID: probeID, State: "unknown", Attempts: attempts, ErrorClass: "execution_error", ErrorMessage: boundedError(err)}
	if errors.Is(err, context.DeadlineExceeded) {
		result.ErrorClass = "timeout"
	}
	if errors.Is(err, context.Canceled) {
		result.ErrorClass = "canceled"
	}
	return normalizeResult(result, probeID, attempts, s.config.Now(), s.config.Now)
}

func (s *Scheduler) recordSuccess(result Result, _ bool) {
	s.mu.Lock()
	s.metrics.Completed++
	s.metrics.LastRunAt = result.CompletedAt
	s.mu.Unlock()
}
func (s *Scheduler) recordFailure(result Result) {
	s.mu.Lock()
	s.metrics.Failed++
	if result.ErrorClass == "timeout" {
		s.metrics.TimedOut++
	}
	s.metrics.LastRunAt = result.CompletedAt
	s.mu.Unlock()
}

func (s *Scheduler) Shutdown(ctx context.Context) error {
	if s == nil {
		return nil
	}
	if ctx == nil {
		ctx = context.Background()
	}
	s.mu.Lock()
	if !s.closed {
		s.closed = true
		for _, cancel := range s.inflight {
			cancel()
		}
	}
	s.mu.Unlock()
	done := make(chan struct{})
	go func() { s.wg.Wait(); close(done) }()
	select {
	case <-done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (s *Scheduler) Metrics() Metrics { s.mu.Lock(); defer s.mu.Unlock(); return s.metrics }

func normalizeResult(result Result, probeID string, attempts int, started time.Time, now func() time.Time) Result {
	result.ProbeID = probeID
	if result.State != "up" && result.State != "degraded" && result.State != "down" && result.State != "unknown" {
		result.State = "unknown"
		if result.ErrorClass == "" {
			result.ErrorClass = "invalid_result"
		}
	}
	result.Attempts = attempts
	if result.ObservedAt.IsZero() {
		result.ObservedAt = started.UTC()
	}
	if result.CompletedAt.IsZero() {
		result.CompletedAt = now().UTC()
	}
	result.ErrorMessage = boundText(result.ErrorMessage, 256)
	return result
}

func isRetryable(err error) bool { var retryable RetryableError; return errors.As(err, &retryable) }
func waitBackoff(ctx context.Context, delay time.Duration) bool {
	if delay <= 0 {
		return true
	}
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-timer.C:
		return true
	case <-ctx.Done():
		return false
	}
}
func boundedError(err error) string {
	if err == nil {
		return ""
	}
	return boundText(err.Error(), 256)
}
func boundText(value string, max int) string {
	value = strings.TrimSpace(value)
	if len(value) > max {
		return value[:max]
	}
	return value
}