aiops/backend/pkg/siteprobe/worker.go

package siteprobe

import (
	"context"
	"log"
	"sync"
	"time"

	"ezops/model"

	"gorm.io/gorm"
)

// Worker 站点拨测后台采集器 // Logan.Li
type Worker struct {
	db     *gorm.DB
	stopCh chan struct{}
	wg     sync.WaitGroup
}

// NewWorker 创建拨测 worker
func NewWorker(db *gorm.DB) *Worker {
	return &Worker{db: db, stopCh: make(chan struct{})}
}

// Start 启动拨测循环
func (w *Worker) Start(intervalSec int) {
	if intervalSec < 5 {
		intervalSec = 10
	}
	w.wg.Add(1)
	go func() {
		defer w.wg.Done()
		ticker := time.NewTicker(time.Duration(intervalSec) * time.Second)
		defer ticker.Stop()
		log.Printf("[SiteProbe] 拨测引擎启动，扫描间隔 %ds", intervalSec)
		w.runOnce()
		for {
			select {
			case <-ticker.C:
				w.runOnce()
			case <-w.stopCh:
				log.Printf("[SiteProbe] 拨测引擎停止")
				return
			}
		}
	}()
}

// Stop 停止拨测
func (w *Worker) Stop() {
	close(w.stopCh)
	w.wg.Wait()
}

// RunTaskNow 立即执行单个任务拨测 // Logan.Li
func (w *Worker) RunTaskNow(taskID uint) error {
	var task model.SiteTask
	if err := w.db.First(&task, taskID).Error; err != nil {
		return err
	}
	w.probeTask(task)
	return nil
}

func (w *Worker) runOnce() {
	var tasks []model.SiteTask
	w.db.Where("status = ?", model.SiteTaskStatusRunning).Find(&tasks)
	now := time.Now()
	sem := make(chan struct{}, 8) // 最多 8 路并发拨测，避免串行阻塞 // Logan.Li
	var wg sync.WaitGroup
	for _, task := range tasks {
		task.Frequency = NormalizeProbeFrequency(string(task.TaskType), task.Frequency)
		if task.LastCheckAt != nil && now.Sub(*task.LastCheckAt) < time.Duration(task.Frequency)*time.Second {
			continue
		}
		wg.Add(1)
		go func(t model.SiteTask) {
			defer wg.Done()
			sem <- struct{}{}
			defer func() { <-sem }()
			w.probeTask(t)
		}(task)
	}
	wg.Wait()
}

func (w *Worker) probeTask(task model.SiteTask) {
	if task.TaskType == "" {
		task.TaskType = model.SiteTaskType(inferTaskType(task))
	}
	outcome := Execute(context.Background(), task)
	now := time.Now()

	result := model.SiteResult{
		TaskID:           task.ID,
		Status:           outcome.Status,
		HTTPStatusCode:   outcome.HTTPStatus,
		ResponseTime:     outcome.ResponseTime,
		TotalTime:        outcome.TotalTime,
		RemoteIP:         outcome.RemoteIP,
		ErrorMessage:     outcome.ErrorMessage,
		SSLCertExpiry:    outcome.SSLCertExpiry,
		SSLCertIssuer:    outcome.SSLCertIssuer,
		SSLCertSubject:   outcome.SSLCertSubject,
		SSLDaysRemaining: outcome.SSLDaysRemaining,
		CheckedAt:        now,
	}
	w.db.Create(&result)

	lastResult := "fail"
	if outcome.Status == model.SiteResultSuccess {
		lastResult = "success"
	}
	avg := outcome.ResponseTime
	if task.AvgResponse > 0 {
		avg = (task.AvgResponse + outcome.ResponseTime) / 2
	}
	w.db.Model(&model.SiteTask{}).Where("id = ?", task.ID).Updates(map[string]interface{}{
		"task_type":     task.TaskType,
		"last_result":   lastResult,
		"last_check_at": now,
		"avg_response":  avg,
	})
	// 拨测恢复后即时关闭探测类告警，与统计栏 down_count 同步 // Logan.Li
	ResolveRecoveredProbeAlerts(w.db, task.ID, lastResult)
}