Files
waf-platform/EdgeHttpDNS/internal/nodes/task_manager.go
2026-02-27 10:35:22 +08:00

132 lines
3.8 KiB
Go

package nodes
import (
"fmt"
"log"
"time"
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
teaconst "github.com/TeaOSLab/EdgeHttpDNS/internal/const"
"github.com/TeaOSLab/EdgeHttpDNS/internal/rpc"
)
type TaskManager struct {
quitCh <-chan struct{}
ticker *time.Ticker
version int64
snapshotManager *SnapshotManager
}
func NewTaskManager(quitCh <-chan struct{}, snapshotManager *SnapshotManager) *TaskManager {
return &TaskManager{
quitCh: quitCh,
ticker: time.NewTicker(20 * time.Second),
version: 0,
snapshotManager: snapshotManager,
}
}
func (m *TaskManager) Start() {
defer m.ticker.Stop()
m.processTasks()
for {
select {
case <-m.ticker.C:
m.processTasks()
case <-m.quitCh:
return
}
}
}
func (m *TaskManager) processTasks() {
rpcClient, err := rpc.SharedRPC()
if err != nil {
log.Println("[HTTPDNS_NODE][task]rpc unavailable:", err.Error())
return
}
resp, err := rpcClient.NodeTaskRPC.FindNodeTasks(rpcClient.Context(), &pb.FindNodeTasksRequest{
Version: m.version,
})
if err != nil {
log.Println("[HTTPDNS_NODE][task]fetch tasks failed:", err.Error())
return
}
for _, task := range resp.GetNodeTasks() {
ok, errorMessage := m.handleTask(task)
_, reportErr := rpcClient.NodeTaskRPC.ReportNodeTaskDone(rpcClient.Context(), &pb.ReportNodeTaskDoneRequest{
NodeTaskId: task.GetId(),
IsOk: ok,
Error: errorMessage,
})
if reportErr != nil {
log.Println("[HTTPDNS_NODE][task]report task result failed:", reportErr.Error())
}
if task.GetVersion() > m.version {
m.version = task.GetVersion()
}
}
}
func (m *TaskManager) handleTask(task *pb.NodeTask) (bool, string) {
taskType := task.GetType()
requestID := fmt.Sprintf("task-%d", task.GetId())
switch taskType {
case teaconst.TaskTypeHTTPDNSConfigChanged:
if m.snapshotManager != nil {
if err := m.snapshotManager.RefreshNow(taskType); err != nil {
reportRuntimeLog("error", "config", "task-manager", "refresh snapshot failed: "+err.Error(), requestID)
return false, err.Error()
}
}
reportRuntimeLog("info", "config", "task-manager", "HTTPDNS configuration updated", requestID)
return true, ""
case teaconst.TaskTypeHTTPDNSAppChanged:
if m.snapshotManager != nil {
if err := m.snapshotManager.RefreshNow(taskType); err != nil {
reportRuntimeLog("error", "app", "task-manager", "refresh snapshot failed: "+err.Error(), requestID)
return false, err.Error()
}
}
reportRuntimeLog("info", "app", "task-manager", "HTTPDNS app policy updated", requestID)
return true, ""
case teaconst.TaskTypeHTTPDNSDomainChanged:
if m.snapshotManager != nil {
if err := m.snapshotManager.RefreshNow(taskType); err != nil {
reportRuntimeLog("error", "domain", "task-manager", "refresh snapshot failed: "+err.Error(), requestID)
return false, err.Error()
}
}
reportRuntimeLog("info", "domain", "task-manager", "HTTPDNS domain binding updated", requestID)
return true, ""
case teaconst.TaskTypeHTTPDNSRuleChanged:
if m.snapshotManager != nil {
if err := m.snapshotManager.RefreshNow(taskType); err != nil {
reportRuntimeLog("error", "rule", "task-manager", "refresh snapshot failed: "+err.Error(), requestID)
return false, err.Error()
}
}
reportRuntimeLog("info", "rule", "task-manager", "HTTPDNS custom rule updated", requestID)
return true, ""
case teaconst.TaskTypeHTTPDNSTLSChanged:
if m.snapshotManager != nil {
if err := m.snapshotManager.RefreshNow(taskType); err != nil {
reportRuntimeLog("error", "tls", "task-manager", "refresh snapshot failed: "+err.Error(), requestID)
return false, err.Error()
}
}
reportRuntimeLog("info", "tls", "task-manager", "HTTPDNS TLS config updated", requestID)
return true, ""
default:
reportRuntimeLog("warning", "task", "task-manager", "unknown task type: "+taskType, requestID)
return true, ""
}
}