happyDomain/internal/app/testscheduler.go

// This file is part of the happyDomain (R) project.
// Copyright (c) 2020-2026 happyDomain
// Authors: Pierre-Olivier Mercier, et al.
//
// This program is offered under a commercial and under the AGPL license.
// For commercial licensing, contact us at <contact@happydomain.org>.
//
// For AGPL licensing:
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

package app

import (
	"container/heap"
	"context"
	"fmt"
	"log"
	"runtime"
	"sync"
	"time"

	"git.happydns.org/happyDomain/internal/storage"
	"git.happydns.org/happyDomain/internal/usecase/testresult"
	"git.happydns.org/happyDomain/model"
)

const (
	SchedulerCheckInterval     = 1 * time.Minute // How often to check for due tests
	SchedulerCleanupInterval   = 24 * time.Hour  // How often to clean up old executions
	SchedulerDiscoveryInterval = 1 * time.Hour   // How often to auto-discover new targets
	TestExecutionTimeout       = 5 * time.Minute // Max time for a single test
	MaxRetries                 = 3               // Max retry attempts for failed tests
)

// Priority levels for test execution queue
const (
	PriorityOnDemand  = iota // On-demand tests (highest priority)
	PriorityOverdue          // Overdue scheduled tests
	PriorityScheduled        // Regular scheduled tests
)

// testScheduler manages background test execution
type testScheduler struct {
	cfg              *happydns.Options
	store            storage.Storage
	pluginUsecase    happydns.TestPluginUsecase
	resultUsecase    *testresult.TestResultUsecase
	scheduleUsecase  *testresult.TestScheduleUsecase
	stop             chan struct{}   // closed to stop the main Run loop
	stopWorkers      chan struct{}   // closed to stop all workers simultaneously
	runNowChan       chan *queueItem // on-demand items routed through the main loop
	workAvail        chan struct{}   // non-blocking signals that queue has new work
	queue            *priorityQueue
	activeExecutions map[string]*activeExecution
	workers          []*worker
	mu               sync.RWMutex
	wg               sync.WaitGroup
	runtimeEnabled   bool
	running          bool
}

// activeExecution tracks a running test execution
type activeExecution struct {
	execution *happydns.TestExecution
	cancel    context.CancelFunc
	startTime time.Time
}

// queueItem represents a test execution request in the queue
type queueItem struct {
	schedule  *happydns.TestSchedule
	execution *happydns.TestExecution
	priority  int
	queuedAt  time.Time
	retries   int
}

// --- container/heap implementation for priorityQueue ---

// priorityHeap is the underlying heap, ordered by priority then arrival time.
type priorityHeap []*queueItem

func (h priorityHeap) Len() int { return len(h) }
func (h priorityHeap) Less(i, j int) bool {
	if h[i].priority != h[j].priority {
		return h[i].priority < h[j].priority
	}
	return h[i].queuedAt.Before(h[j].queuedAt)
}
func (h priorityHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *priorityHeap) Push(x any)   { *h = append(*h, x.(*queueItem)) }
func (h *priorityHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	old[n-1] = nil // avoid memory leak
	*h = old[:n-1]
	return x
}

// priorityQueue is a thread-safe min-heap of queueItems.
type priorityQueue struct {
	h  priorityHeap
	mu sync.Mutex
}

func newPriorityQueue() *priorityQueue {
	pq := &priorityQueue{}
	heap.Init(&pq.h)
	return pq
}

// Push adds an item to the queue.
func (q *priorityQueue) Push(item *queueItem) {
	q.mu.Lock()
	defer q.mu.Unlock()
	heap.Push(&q.h, item)
}

// Pop removes and returns the highest-priority item, or nil if empty.
func (q *priorityQueue) Pop() *queueItem {
	q.mu.Lock()
	defer q.mu.Unlock()
	if q.h.Len() == 0 {
		return nil
	}
	return heap.Pop(&q.h).(*queueItem)
}

// Len returns the queue length.
func (q *priorityQueue) Len() int {
	q.mu.Lock()
	defer q.mu.Unlock()
	return q.h.Len()
}

// worker processes tests from the queue
type worker struct {
	id        int
	scheduler *testScheduler
}

// disabledScheduler is a no-op implementation used when scheduler is disabled
type disabledScheduler struct{}

func (d *disabledScheduler) Run()   {}
func (d *disabledScheduler) Close() {}

// TriggerOnDemandTest returns an error indicating the scheduler is disabled
func (d *disabledScheduler) TriggerOnDemandTest(pluginName string, targetType happydns.TestScopeType, targetId happydns.Identifier, userId happydns.Identifier, options happydns.PluginOptions) (happydns.Identifier, error) {
	return happydns.Identifier{}, fmt.Errorf("test scheduler is disabled in configuration")
}

// GetSchedulerStatus returns a status indicating the scheduler is disabled
func (d *disabledScheduler) GetSchedulerStatus() happydns.SchedulerStatus {
	return happydns.SchedulerStatus{
		ConfigEnabled:  false,
		RuntimeEnabled: false,
		Running:        false,
	}
}

// SetEnabled returns an error since the scheduler is disabled in configuration
func (d *disabledScheduler) SetEnabled(enabled bool) error {
	return fmt.Errorf("scheduler is disabled in configuration, cannot enable at runtime")
}

// RescheduleUpcomingTests returns an error since the scheduler is disabled
func (d *disabledScheduler) RescheduleUpcomingTests() (int, error) {
	return 0, fmt.Errorf("test scheduler is disabled in configuration")
}

// newTestScheduler creates a new test scheduler
func newTestScheduler(
	cfg *happydns.Options,
	store storage.Storage,
	pluginUsecase happydns.TestPluginUsecase,
) *testScheduler {
	numWorkers := cfg.TestWorkers
	if numWorkers <= 0 {
		numWorkers = runtime.NumCPU()
	}

	scheduler := &testScheduler{
		cfg:              cfg,
		store:            store,
		pluginUsecase:    pluginUsecase,
		resultUsecase:    testresult.NewTestResultUsecase(store, cfg),
		scheduleUsecase:  testresult.NewTestScheduleUsecase(store, cfg, store, pluginUsecase),
		stop:             make(chan struct{}),
		stopWorkers:      make(chan struct{}),
		runNowChan:       make(chan *queueItem, 100),
		workAvail:        make(chan struct{}, numWorkers),
		queue:            newPriorityQueue(),
		activeExecutions: make(map[string]*activeExecution),
		workers:          make([]*worker, numWorkers),
		runtimeEnabled:   true,
	}

	for i := 0; i < numWorkers; i++ {
		scheduler.workers[i] = &worker{
			id:        i,
			scheduler: scheduler,
		}
	}

	return scheduler
}

// enqueue pushes an item to the priority queue and wakes one idle worker.
func (s *testScheduler) enqueue(item *queueItem) {
	s.queue.Push(item)
	select {
	case s.workAvail <- struct{}{}:
	default:
		// All workers are already busy or already notified; they will drain
		// the queue on their own after finishing the current item.
	}
}

// Close stops the scheduler and waits for all workers to finish.
func (s *testScheduler) Close() {
	log.Println("Stopping test scheduler...")

	// Unblock the main Run loop.
	close(s.stop)

	// Unblock all workers simultaneously.
	close(s.stopWorkers)

	// Cancel all active test executions.
	s.mu.Lock()
	for _, exec := range s.activeExecutions {
		exec.cancel()
	}
	s.mu.Unlock()

	// Wait for all workers to finish their current item.
	s.wg.Wait()

	log.Println("Test scheduler stopped")
}

// Run starts the scheduler main loop. It must not be called more than once.
func (s *testScheduler) Run() {
	s.mu.Lock()
	s.running = true
	s.mu.Unlock()

	defer func() {
		s.mu.Lock()
		s.running = false
		s.mu.Unlock()
	}()

	log.Printf("Starting test scheduler with %d workers...\n", len(s.workers))

	// Reschedule overdue tests before starting workers so that tests missed
	// during a server suspend or shutdown are spread into the near future
	// instead of all firing at once.
	if n, err := s.scheduleUsecase.RescheduleOverdueTests(); err != nil {
		log.Printf("Warning: failed to reschedule overdue tests: %v\n", err)
	} else if n > 0 {
		log.Printf("Rescheduled %d overdue test(s) into the near future\n", n)
	}

	// Start workers
	for _, w := range s.workers {
		s.wg.Add(1)
		go w.run(&s.wg)
	}

	// Main scheduling loop
	checkTicker := time.NewTicker(SchedulerCheckInterval)
	cleanupTicker := time.NewTicker(SchedulerCleanupInterval)
	discoveryTicker := time.NewTicker(SchedulerDiscoveryInterval)
	defer checkTicker.Stop()
	defer cleanupTicker.Stop()
	defer discoveryTicker.Stop()

	// Initial discovery: create default schedules for all existing targets
	if err := s.scheduleUsecase.DiscoverAndEnsureSchedules(); err != nil {
		log.Printf("Warning: schedule discovery encountered errors: %v\n", err)
	}
	// Initial check
	s.checkSchedules()

	for {
		select {
		case <-checkTicker.C:
			s.checkSchedules()

		case <-cleanupTicker.C:
			s.cleanup()

		case <-discoveryTicker.C:
			if err := s.scheduleUsecase.DiscoverAndEnsureSchedules(); err != nil {
				log.Printf("Warning: schedule discovery encountered errors: %v\n", err)
			}

		case item := <-s.runNowChan:
			s.enqueue(item)

		case <-s.stop:
			return
		}
	}
}

// checkSchedules checks for due tests and queues them
func (s *testScheduler) checkSchedules() {
	s.mu.RLock()
	enabled := s.runtimeEnabled
	s.mu.RUnlock()
	if !enabled {
		return
	}

	dueSchedules, err := s.scheduleUsecase.ListDueSchedules()
	if err != nil {
		log.Printf("Error listing due schedules: %v\n", err)
		return
	}

	now := time.Now()
	for _, schedule := range dueSchedules {
		// Determine priority based on how overdue the test is
		priority := PriorityScheduled
		if schedule.NextRun.Add(schedule.Interval).Before(now) {
			priority = PriorityOverdue
		}

		// Create execution record
		execution := &happydns.TestExecution{
			ScheduleId: &schedule.Id,
			PluginName: schedule.PluginName,
			OwnerId:    schedule.OwnerId,
			TargetType: schedule.TargetType,
			TargetId:   schedule.TargetId,
			Status:     happydns.TestExecutionPending,
			StartedAt:  time.Now(),
			Options:    schedule.Options,
		}

		if err := s.resultUsecase.CreateTestExecution(execution); err != nil {
			log.Printf("Error creating execution for schedule %s: %v\n", schedule.Id.String(), err)
			continue
		}

		s.enqueue(&queueItem{
			schedule:  schedule,
			execution: execution,
			priority:  priority,
			queuedAt:  now,
			retries:   0,
		})
	}

	// Mark scheduler run
	if err := s.store.TestSchedulerRun(); err != nil {
		log.Printf("Error marking scheduler run: %v\n", err)
	}
}

// TriggerOnDemandTest triggers an immediate test execution.
// It creates the execution record synchronously (so the caller gets an ID back)
// and then routes the item through runNowChan so the main loop controls
// all queue insertions.
func (s *testScheduler) TriggerOnDemandTest(pluginName string, targetType happydns.TestScopeType, targetId happydns.Identifier, ownerId happydns.Identifier, options happydns.PluginOptions) (happydns.Identifier, error) {
	schedule := &happydns.TestSchedule{
		PluginName: pluginName,
		OwnerId:    ownerId,
		TargetType: targetType,
		TargetId:   targetId,
		Interval:   0, // On-demand, no interval
		Enabled:    true,
		Options:    options,
	}

	execution := &happydns.TestExecution{
		ScheduleId: nil,
		PluginName: pluginName,
		OwnerId:    ownerId,
		TargetType: targetType,
		TargetId:   targetId,
		Status:     happydns.TestExecutionPending,
		StartedAt:  time.Now(),
		Options:    options,
	}

	if err := s.resultUsecase.CreateTestExecution(execution); err != nil {
		return happydns.Identifier{}, err
	}

	item := &queueItem{
		schedule:  schedule,
		execution: execution,
		priority:  PriorityOnDemand,
		queuedAt:  time.Now(),
		retries:   0,
	}

	// Route through the main loop when possible; fall back to direct enqueue
	// if the channel is full so that the caller never blocks.
	select {
	case s.runNowChan <- item:
	default:
		s.enqueue(item)
	}

	return execution.Id, nil
}

// GetSchedulerStatus returns a snapshot of the current scheduler state
func (s *testScheduler) GetSchedulerStatus() happydns.SchedulerStatus {
	s.mu.RLock()
	activeCount := len(s.activeExecutions)
	running := s.running
	runtimeEnabled := s.runtimeEnabled
	s.mu.RUnlock()

	nextSchedules, _ := s.scheduleUsecase.ListUpcomingSchedules(20)

	return happydns.SchedulerStatus{
		ConfigEnabled:  !s.cfg.DisableScheduler,
		RuntimeEnabled: runtimeEnabled,
		Running:        running,
		WorkerCount:    len(s.workers),
		QueueSize:      s.queue.Len(),
		ActiveCount:    activeCount,
		NextSchedules:  nextSchedules,
	}
}

// SetEnabled enables or disables the scheduler at runtime
func (s *testScheduler) SetEnabled(enabled bool) error {
	s.mu.Lock()
	wasEnabled := s.runtimeEnabled
	s.runtimeEnabled = enabled
	s.mu.Unlock()

	if enabled && !wasEnabled {
		// Spread out any overdue tests to avoid a thundering herd, then
		// immediately enqueue whatever is now due.
		if n, err := s.scheduleUsecase.RescheduleOverdueTests(); err != nil {
			log.Printf("Warning: failed to reschedule overdue tests on re-enable: %v\n", err)
		} else if n > 0 {
			log.Printf("Rescheduled %d overdue test(s) after scheduler re-enable\n", n)
		}
		s.checkSchedules()
	}

	return nil
}

// RescheduleUpcomingTests randomizes the next run time of all enabled schedules
// within their respective intervals, delegating to the schedule usecase.
func (s *testScheduler) RescheduleUpcomingTests() (int, error) {
	return s.scheduleUsecase.RescheduleUpcomingTests()
}

// cleanup removes old execution records and expired test results
func (s *testScheduler) cleanup() {
	log.Println("Running scheduler cleanup...")

	// Delete completed/failed execution records older than 7 days
	if err := s.resultUsecase.DeleteCompletedExecutions(7 * 24 * time.Hour); err != nil {
		log.Printf("Error cleaning up old executions: %v\n", err)
	}

	// Delete test results older than the configured retention period
	if err := s.resultUsecase.CleanupOldResults(); err != nil {
		log.Printf("Error cleaning up old test results: %v\n", err)
	}

	log.Println("Scheduler cleanup complete")
}

// run is the worker's main loop. It drains the queue eagerly and waits for a
// workAvail signal when idle, rather than sleeping on a fixed timer.
func (w *worker) run(wg *sync.WaitGroup) {
	defer wg.Done()

	log.Printf("Worker %d started\n", w.id)

	for {
		// Drain: try to grab work before blocking.
		if item := w.scheduler.queue.Pop(); item != nil {
			w.executeTest(item)
			continue
		}

		// Queue is empty; wait for new work or a stop signal.
		select {
		case <-w.scheduler.workAvail:
			// Loop back to attempt a Pop.
		case <-w.scheduler.stopWorkers:
			log.Printf("Worker %d stopped\n", w.id)
			return
		}
	}
}

// executeTest runs a test plugin and stores the result
func (w *worker) executeTest(item *queueItem) {
	ctx, cancel := context.WithTimeout(context.Background(), TestExecutionTimeout)
	defer cancel()

	execution := item.execution
	schedule := item.schedule

	// Always update schedule NextRun after execution, whether it succeeds or fails.
	// This prevents the schedule from being re-queued on the next tick if the test fails.
	if item.execution.ScheduleId != nil {
		defer func() {
			if err := w.scheduler.scheduleUsecase.UpdateScheduleAfterRun(*item.execution.ScheduleId); err != nil {
				log.Printf("Worker %d: Error updating schedule after run: %v\n", w.id, err)
			}
		}()
	}

	// Mark execution as running
	execution.Status = happydns.TestExecutionRunning
	if err := w.scheduler.resultUsecase.UpdateTestExecution(execution); err != nil {
		log.Printf("Worker %d: Error updating execution status: %v\n", w.id, err)
		_ = w.scheduler.resultUsecase.FailTestExecution(execution.Id, err.Error())
		return
	}

	// Track active execution
	w.scheduler.mu.Lock()
	w.scheduler.activeExecutions[execution.Id.String()] = &activeExecution{
		execution: execution,
		cancel:    cancel,
		startTime: time.Now(),
	}
	w.scheduler.mu.Unlock()

	defer func() {
		w.scheduler.mu.Lock()
		delete(w.scheduler.activeExecutions, execution.Id.String())
		w.scheduler.mu.Unlock()
	}()

	// Get the plugin
	plugin, err := w.scheduler.pluginUsecase.GetTestPlugin(schedule.PluginName)
	if err != nil {
		errMsg := fmt.Sprintf("plugin not found: %s - %v", schedule.PluginName, err)
		log.Printf("Worker %d: %s\n", w.id, errMsg)
		_ = w.scheduler.resultUsecase.FailTestExecution(execution.Id, errMsg)
		return
	}

	// Merge options: global defaults < user opts < domain/service opts < schedule opts
	mergedOptions, err := w.scheduler.scheduleUsecase.PrepareTestOptions(schedule)
	if err != nil {
		// Non-fatal: PrepareTestOptions already falls back to schedule-only options
		log.Printf("Worker %d: warning, could not prepare plugin options for %s: %v\n", w.id, schedule.PluginName, err)
	}

	// Prepare metadata
	meta := make(map[string]string)
	meta["target_type"] = schedule.TargetType.String()
	meta["target_id"] = schedule.TargetId.String()

	// Run the test
	startTime := time.Now()
	resultChan := make(chan *happydns.PluginResult, 1)
	errorChan := make(chan error, 1)

	go func() {
		defer func() {
			if r := recover(); r != nil {
				errorChan <- fmt.Errorf("plugin panicked: %v", r)
			}
		}()
		result, err := plugin.RunTest(mergedOptions, meta)
		if err != nil {
			errorChan <- err
		} else {
			resultChan <- result
		}
	}()

	// Wait for result or timeout
	var pluginResult *happydns.PluginResult
	var testErr error

	select {
	case pluginResult = <-resultChan:
		// Test completed successfully
	case testErr = <-errorChan:
		// Test returned an error
	case <-ctx.Done():
		// Timeout
		testErr = fmt.Errorf("test execution timeout after %v", TestExecutionTimeout)
	}

	duration := time.Since(startTime)

	// Store the result
	result := &happydns.TestResult{
		PluginName:    schedule.PluginName,
		TestType:      schedule.TargetType,
		TargetId:      schedule.TargetId,
		OwnerId:       schedule.OwnerId,
		ExecutedAt:    time.Now(),
		ScheduledTest: item.execution.ScheduleId != nil,
		Options:       schedule.Options,
		Duration:      duration,
	}

	if testErr != nil {
		result.Status = happydns.PluginResultStatusKO
		result.StatusLine = "Test execution failed"
		result.Error = testErr.Error()
	} else if pluginResult != nil {
		result.Status = pluginResult.Status
		result.StatusLine = pluginResult.StatusLine
		result.Report = pluginResult.Report
	} else {
		result.Status = happydns.PluginResultStatusKO
		result.StatusLine = "Unknown error"
		result.Error = "No result or error returned from plugin"
	}

	// Save the result
	if err := w.scheduler.resultUsecase.CreateTestResult(result); err != nil {
		log.Printf("Worker %d: Error saving test result: %v\n", w.id, err)
		_ = w.scheduler.resultUsecase.FailTestExecution(execution.Id, err.Error())
		return
	}

	// Complete the execution
	if err := w.scheduler.resultUsecase.CompleteTestExecution(execution.Id, result.Id); err != nil {
		log.Printf("Worker %d: Error completing execution: %v\n", w.id, err)
		return
	}

	log.Printf("Worker %d: Completed test %s for target %s (status: %d, duration: %v)\n",
		w.id, schedule.PluginName, schedule.TargetId.String(), result.Status, duration)
}