mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-10-02 14:56:38 +08:00
refactor(repo): consolidate openflare-server to root and move subprojects to internal/apps
- Merge all files inside openflare-server to the repository root directory. - Relocate agent, relay, and flared subprojects from internal/ to internal/apps/. - Combine docker-compose files and update build context paths to root. - Update GitHub workflows and Dockerfiles to refer to new directories and package names. - Rewrite Go package imports across all files. - Resolve database renew test race condition and clean up docs.
This commit is contained in:
ryan
@@ -0,0 +1,14 @@
|
||||
// Copyright 2025 linux.do
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Package task 定义异步任务类型与调度常量
|
||||
package task
|
||||
|
||||
// 任务队列名称
|
||||
const (
|
||||
QueueDefault = "default"
|
||||
)
|
||||
|
||||
// DefaultMaxRetry 任务默认最大重试次数
|
||||
const DefaultMaxRetry = 3
|
||||
@@ -0,0 +1,17 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package task
|
||||
|
||||
const (
|
||||
errUnknownTaskType = "未知的任务类型: %s"
|
||||
errCreateTaskExecutionFailed = "创建任务执行记录失败: %w"
|
||||
errTaskEnqueueFailed = "任务入队失败: %w"
|
||||
errTaskExecutionNotFound = "任务执行记录不存在: %w"
|
||||
errRetryOnlyFailedTask = "只有失败的任务才能重试，当前状态: %s"
|
||||
errTaskNotRetryable = "该任务不支持重试"
|
||||
errTaskMaxRetryExceeded = "已达到最大重试次数 %d"
|
||||
errCreateRetryExecutionFailed = "创建重试任务执行记录失败: %w"
|
||||
errRetryTaskEnqueueFailed = "重试任务入队失败: %w"
|
||||
errUnregisteredTaskHandler = "未注册的任务处理器: %s"
|
||||
)
|
||||
@@ -0,0 +1,479 @@
|
||||
// Copyright 2025 linux.do
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package task
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/Rain-kl/Wavelet/internal/db/idgen"
|
||||
"github.com/Rain-kl/Wavelet/internal/model"
|
||||
"github.com/Rain-kl/Wavelet/pkg/logger"
|
||||
otel_trace "github.com/Rain-kl/Wavelet/pkg/trace"
|
||||
"github.com/hibiken/asynq"
|
||||
"go.opentelemetry.io/otel"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/codes"
|
||||
"go.opentelemetry.io/otel/propagation"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
// handlerRegistry 已注册的任务处理器
|
||||
var handlerRegistry = make(map[string]TaskHandler)
|
||||
|
||||
// CompletedHandler is called when a task execution completes.
|
||||
type CompletedHandler func(ctx context.Context, execution *model.TaskExecution, result *TaskResult, execErr error)
|
||||
|
||||
var taskCompletedHandlers []CompletedHandler
|
||||
|
||||
// OnTaskCompleted registers a handler for task completion events.
|
||||
// Handlers must be registered during application bootstrap before processing tasks.
|
||||
func OnTaskCompleted(handler CompletedHandler) {
|
||||
taskCompletedHandlers = append(taskCompletedHandlers, handler)
|
||||
}
|
||||
|
||||
// RegisterHandler 注册任务处理器
|
||||
// 传入任务类型标识（对应 constants.go 中的 AsynqTask 常量）和 TaskHandler 实现
|
||||
func RegisterHandler(asynqTaskType string, handler TaskHandler) {
|
||||
handlerRegistry[asynqTaskType] = handler
|
||||
}
|
||||
|
||||
// getHandler 获取已注册的处理器
|
||||
func getHandler(asynqTaskType string) (TaskHandler, bool) {
|
||||
h, ok := handlerRegistry[asynqTaskType]
|
||||
return h, ok
|
||||
}
|
||||
|
||||
// ValidateAndNormalizePayload 校验并标准化任务参数。
|
||||
// 如果 Handler 实现了 PayloadValidator，调用其 ValidatePayload 方法；
|
||||
// 否则直接返回原始 payload。
|
||||
func ValidateAndNormalizePayload(asynqTaskType string, payload []byte) ([]byte, error) {
|
||||
handler, ok := getHandler(asynqTaskType)
|
||||
if !ok {
|
||||
return payload, nil
|
||||
}
|
||||
if validator, ok := handler.(PayloadValidator); ok {
|
||||
return validator.ValidatePayload(payload)
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
// contextKey 用于 context 存取 taskID
|
||||
type contextKey string
|
||||
|
||||
const taskIDKey contextKey = "task_execution_task_id"
|
||||
const traceEnvelopeVersion = 1
|
||||
|
||||
type traceEnvelope struct {
|
||||
WaveletTraceEnvelope bool `json:"_wavelet_trace_envelope"`
|
||||
Version int `json:"version"`
|
||||
TraceContext map[string]string `json:"trace_context,omitempty"`
|
||||
Payload []byte `json:"payload"`
|
||||
}
|
||||
|
||||
// withTaskID 将 taskID 注入 context
|
||||
func withTaskID(ctx context.Context, taskID string) context.Context {
|
||||
return context.WithValue(ctx, taskIDKey, taskID)
|
||||
}
|
||||
|
||||
// GetTaskID 从 context 中获取 taskID
|
||||
func GetTaskID(ctx context.Context) string {
|
||||
if v, ok := ctx.Value(taskIDKey).(string); ok {
|
||||
return v
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// IsFinalAttempt 判断当前任务执行是否为最后一次重试尝试（若再次失败即为最终失败）
|
||||
func IsFinalAttempt(ctx context.Context) bool {
|
||||
retryCount, hasRetryCount := asynq.GetRetryCount(ctx)
|
||||
maxRetry, hasMaxRetry := asynq.GetMaxRetry(ctx)
|
||||
if !hasRetryCount || !hasMaxRetry {
|
||||
return true
|
||||
}
|
||||
return retryCount >= maxRetry
|
||||
}
|
||||
|
||||
// AppendLog 追加日志到任务执行记录
|
||||
// 在 TaskHandler.Execute 中调用，日志会自动追加到 TaskExecution.Log 字段
|
||||
func AppendLog(ctx context.Context, format string, args ...interface{}) {
|
||||
taskID := GetTaskID(ctx)
|
||||
if taskID == "" {
|
||||
// 上下文中没有 taskID，降级到普通日志
|
||||
logger.InfoF(ctx, format, args...)
|
||||
return
|
||||
}
|
||||
|
||||
logLine := fmt.Sprintf(format, args...)
|
||||
if err := model.AppendTaskExecutionLog(ctx, taskID, logLine); err != nil {
|
||||
logger.ErrorF(ctx, "[TaskExecutor] 追加任务日志失败 taskID=%s: %v", taskID, err)
|
||||
}
|
||||
}
|
||||
|
||||
// DispatchTask 下发任务（创建 TaskExecution 记录 → 入队 Asynq）
|
||||
func DispatchTask(ctx context.Context, taskType string, payload []byte, triggeredBy string) (string, error) {
|
||||
meta := GetTaskMeta(taskType)
|
||||
if meta == nil {
|
||||
return "", fmt.Errorf(errUnknownTaskType, taskType)
|
||||
}
|
||||
|
||||
// 生成唯一的 TaskID
|
||||
taskID := generateTaskID(taskType, triggeredBy)
|
||||
|
||||
// 创建任务执行记录
|
||||
execution := &model.TaskExecution{
|
||||
TaskID: taskID,
|
||||
TaskType: meta.AsynqTask,
|
||||
TaskName: meta.Name,
|
||||
Status: model.TaskExecutionStatusPending,
|
||||
Retryable: meta.Retryable,
|
||||
MaxRetry: meta.MaxRetry,
|
||||
RetryCount: 0,
|
||||
Payload: string(payload),
|
||||
TriggeredBy: triggeredBy,
|
||||
}
|
||||
|
||||
if err := model.CreateTaskExecution(ctx, execution); err != nil {
|
||||
return "", fmt.Errorf(errCreateTaskExecutionFailed, err)
|
||||
}
|
||||
|
||||
// 入队 Asynq
|
||||
taskInfo := asynq.NewTask(meta.AsynqTask, injectTaskTraceContext(ctx, payload))
|
||||
if _, err := AsynqClient.Enqueue(
|
||||
taskInfo,
|
||||
asynq.TaskID(taskID),
|
||||
asynq.MaxRetry(meta.MaxRetry),
|
||||
asynq.Queue(meta.Queue),
|
||||
); err != nil {
|
||||
// 入队失败，更新执行记录状态
|
||||
execution.Status = model.TaskExecutionStatusFailed
|
||||
execution.ErrorMessage = fmt.Sprintf("入队失败: %v", err)
|
||||
now := time.Now()
|
||||
execution.StartedAt = &now
|
||||
execution.FinishedAt = &now
|
||||
_ = model.UpdateTaskExecution(ctx, execution)
|
||||
return "", fmt.Errorf(errTaskEnqueueFailed, err)
|
||||
}
|
||||
|
||||
if err := model.AppendTaskExecutionLog(ctx, taskID, fmt.Sprintf("[系统] 任务已成功入队，等待调度执行 (队列: %s, 最大重试次数: %d)", meta.Queue, meta.MaxRetry)); err != nil {
|
||||
logger.ErrorF(ctx, "[TaskExecutor] 追加入队日志失败 taskID=%s: %v", taskID, err)
|
||||
}
|
||||
|
||||
return taskID, nil
|
||||
}
|
||||
|
||||
// RetryTask 重试失败的任务
|
||||
func RetryTask(ctx context.Context, id uint64) (string, error) {
|
||||
execution, err := model.GetTaskExecutionByID(ctx, id)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf(errTaskExecutionNotFound, err)
|
||||
}
|
||||
|
||||
if execution.Status != model.TaskExecutionStatusFailed {
|
||||
return "", fmt.Errorf(errRetryOnlyFailedTask, execution.Status)
|
||||
}
|
||||
|
||||
if !execution.Retryable {
|
||||
return "", errors.New(errTaskNotRetryable)
|
||||
}
|
||||
|
||||
// 生成新的 TaskID
|
||||
newTaskID := generateRetryTaskID(execution.TaskID, execution.RetryCount+1)
|
||||
|
||||
// 创建新的执行记录
|
||||
newExecution := &model.TaskExecution{
|
||||
TaskID: newTaskID,
|
||||
TaskType: execution.TaskType,
|
||||
TaskName: execution.TaskName,
|
||||
Status: model.TaskExecutionStatusPending,
|
||||
Retryable: execution.Retryable,
|
||||
MaxRetry: execution.MaxRetry,
|
||||
RetryCount: execution.RetryCount + 1,
|
||||
Payload: execution.Payload,
|
||||
TriggeredBy: "retry",
|
||||
}
|
||||
|
||||
if err := model.CreateTaskExecution(ctx, newExecution); err != nil {
|
||||
return "", fmt.Errorf(errCreateRetryExecutionFailed, err)
|
||||
}
|
||||
|
||||
meta := GetTaskMeta(execution.TaskType)
|
||||
queueName := QueueDefault
|
||||
if meta != nil {
|
||||
queueName = meta.Queue
|
||||
}
|
||||
|
||||
// 入队 Asynq
|
||||
taskInfo := asynq.NewTask(execution.TaskType, injectTaskTraceContext(ctx, []byte(execution.Payload)))
|
||||
if _, err := AsynqClient.Enqueue(
|
||||
taskInfo,
|
||||
asynq.TaskID(newTaskID),
|
||||
asynq.MaxRetry(execution.MaxRetry),
|
||||
asynq.Queue(queueName),
|
||||
); err != nil {
|
||||
newExecution.Status = model.TaskExecutionStatusFailed
|
||||
newExecution.ErrorMessage = fmt.Sprintf("重试入队失败: %v", err)
|
||||
now := time.Now()
|
||||
newExecution.StartedAt = &now
|
||||
newExecution.FinishedAt = &now
|
||||
_ = model.UpdateTaskExecution(ctx, newExecution)
|
||||
return "", fmt.Errorf(errRetryTaskEnqueueFailed, err)
|
||||
}
|
||||
|
||||
if err := model.AppendTaskExecutionLog(ctx, newTaskID, fmt.Sprintf("[系统] 手动触发重试，已重新创建任务并入队 (原任务ID: %s, 重试次数: %d/%d)", execution.TaskID, execution.RetryCount+1, execution.MaxRetry)); err != nil {
|
||||
logger.ErrorF(ctx, "[TaskExecutor] 追加重试日志失败 taskID=%s: %v", newTaskID, err)
|
||||
}
|
||||
|
||||
return newTaskID, nil
|
||||
}
|
||||
|
||||
// ProcessTask Asynq 实际调用的统一处理函数
|
||||
// Worker 注册时统一使用此函数，内部自动分发到对应的 TaskHandler
|
||||
func ProcessTask(ctx context.Context, t *asynq.Task) error {
|
||||
taskPayload := t.Payload()
|
||||
ctx, taskPayload, hasRemoteTraceContext := extractTaskTraceContext(ctx, taskPayload)
|
||||
|
||||
// 初始化 Trace
|
||||
ctx, span := otel_trace.Start(ctx, "TaskProcess_"+t.Type(), trace.WithSpanKind(trace.SpanKindConsumer))
|
||||
defer span.End()
|
||||
|
||||
// 添加任务信息到 Span
|
||||
span.SetAttributes(
|
||||
attribute.String("task.type", t.Type()),
|
||||
attribute.Int("task.payload_size", len(taskPayload)),
|
||||
attribute.Bool("task.trace_context_propagated", hasRemoteTraceContext),
|
||||
attribute.String("task.id", t.ResultWriter().TaskID()),
|
||||
)
|
||||
|
||||
taskID := t.ResultWriter().TaskID()
|
||||
|
||||
// 注入 taskID 到 context
|
||||
ctx = withTaskID(ctx, taskID)
|
||||
|
||||
// 查找处理器
|
||||
handler, ok := getHandler(t.Type())
|
||||
if !ok {
|
||||
err := fmt.Errorf(errUnregisteredTaskHandler, t.Type())
|
||||
logger.ErrorF(ctx, "[TaskExecutor] %v", err)
|
||||
span.SetStatus(codes.Error, err.Error())
|
||||
return err
|
||||
}
|
||||
|
||||
// 加载或动态创建执行记录
|
||||
now := time.Now()
|
||||
execution, err := getOrCreateTaskExecution(ctx, taskID, t, taskPayload, now)
|
||||
if err == nil {
|
||||
updateExecutionOnStart(ctx, execution, now)
|
||||
}
|
||||
|
||||
if execution != nil {
|
||||
AppendLog(ctx, "[系统] 开始执行异步任务 [名称: %s, 类型: %s]，重试次数: %d/%d",
|
||||
execution.TaskName, t.Type(), execution.RetryCount, execution.MaxRetry)
|
||||
} else {
|
||||
AppendLog(ctx, "[系统] 开始执行异步任务 [类型: %s]", t.Type())
|
||||
}
|
||||
|
||||
// 开始计时
|
||||
start := time.Now()
|
||||
|
||||
// 执行业务逻辑
|
||||
result, execErr := handler.Execute(ctx, taskPayload)
|
||||
|
||||
// 计算耗时并归档记录
|
||||
duration := time.Since(start)
|
||||
finishTime := time.Now()
|
||||
|
||||
completeTaskExecution(ctx, execution, t, duration, finishTime, result, execErr, span)
|
||||
|
||||
if execution == nil && execErr != nil {
|
||||
span.SetStatus(codes.Error, execErr.Error())
|
||||
return execErr
|
||||
}
|
||||
|
||||
return execErr
|
||||
}
|
||||
|
||||
func injectTaskTraceContext(ctx context.Context, payload []byte) []byte {
|
||||
carrier := propagation.MapCarrier{}
|
||||
otel.GetTextMapPropagator().Inject(ctx, carrier)
|
||||
if len(carrier) == 0 {
|
||||
return payload
|
||||
}
|
||||
|
||||
envelope := traceEnvelope{
|
||||
WaveletTraceEnvelope: true,
|
||||
Version: traceEnvelopeVersion,
|
||||
TraceContext: map[string]string(carrier),
|
||||
Payload: payload,
|
||||
}
|
||||
data, err := json.Marshal(envelope)
|
||||
if err != nil {
|
||||
logger.ErrorF(ctx, "[TaskExecutor] 序列化任务 Trace 上下文失败: %v", err)
|
||||
return payload
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
func extractTaskTraceContext(ctx context.Context, payload []byte) (context.Context, []byte, bool) {
|
||||
var envelope traceEnvelope
|
||||
if err := json.Unmarshal(payload, &envelope); err != nil {
|
||||
return ctx, payload, false
|
||||
}
|
||||
if !envelope.WaveletTraceEnvelope || envelope.Version != traceEnvelopeVersion {
|
||||
return ctx, payload, false
|
||||
}
|
||||
if len(envelope.TraceContext) == 0 {
|
||||
return ctx, envelope.Payload, false
|
||||
}
|
||||
|
||||
extractedCtx := otel.GetTextMapPropagator().Extract(ctx, propagation.MapCarrier(envelope.TraceContext))
|
||||
return extractedCtx, envelope.Payload, true
|
||||
}
|
||||
|
||||
func updateExecutionOnStart(ctx context.Context, execution *model.TaskExecution, now time.Time) {
|
||||
if execution == nil {
|
||||
return
|
||||
}
|
||||
dirty := false
|
||||
if retryCount, hasRetry := asynq.GetRetryCount(ctx); hasRetry && execution.RetryCount != retryCount {
|
||||
execution.RetryCount = retryCount
|
||||
dirty = true
|
||||
}
|
||||
if execution.Status != model.TaskExecutionStatusRunning {
|
||||
execution.Status = model.TaskExecutionStatusRunning
|
||||
execution.StartedAt = &now
|
||||
dirty = true
|
||||
}
|
||||
if dirty {
|
||||
if updateErr := model.UpdateTaskExecution(ctx, execution); updateErr != nil {
|
||||
logger.ErrorF(ctx, "[TaskExecutor] 更新执行状态失败 taskID=%s: %v", execution.TaskID, updateErr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// getOrCreateTaskExecution 获取已有的任务执行记录，如果不存在则针对已知任务类型动态创建记录
|
||||
func getOrCreateTaskExecution(ctx context.Context, taskID string, t *asynq.Task, payload []byte, now time.Time) (*model.TaskExecution, error) {
|
||||
execution, err := model.GetTaskExecutionByTaskID(ctx, taskID)
|
||||
if err == nil {
|
||||
return execution, nil
|
||||
}
|
||||
|
||||
meta := GetTaskMetaByAsynqTask(t.Type())
|
||||
if meta == nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
execution = &model.TaskExecution{
|
||||
TaskID: taskID,
|
||||
TaskType: meta.AsynqTask,
|
||||
TaskName: meta.Name,
|
||||
Status: model.TaskExecutionStatusRunning,
|
||||
Retryable: meta.Retryable,
|
||||
MaxRetry: meta.MaxRetry,
|
||||
RetryCount: 0,
|
||||
Payload: string(payload),
|
||||
TriggeredBy: "schedule",
|
||||
StartedAt: &now,
|
||||
}
|
||||
|
||||
if createErr := model.CreateTaskExecution(ctx, execution); createErr != nil {
|
||||
logger.ErrorF(ctx, "[TaskExecutor] 动态创建执行记录失败 taskID=%s: %v", taskID, createErr)
|
||||
return nil, createErr
|
||||
}
|
||||
|
||||
return execution, nil
|
||||
}
|
||||
|
||||
// completeTaskExecution 完成并更新任务执行记录的状态和执行结果
|
||||
func completeTaskExecution(ctx context.Context, execution *model.TaskExecution, t *asynq.Task, duration time.Duration, finishTime time.Time, result *TaskResult, execErr error, span trace.Span) {
|
||||
if execution == nil {
|
||||
return
|
||||
}
|
||||
|
||||
execution.Duration = duration.Milliseconds()
|
||||
execution.FinishedAt = &finishTime
|
||||
|
||||
if execErr != nil {
|
||||
handleFailedTask(ctx, execution, t, duration, execErr, span)
|
||||
} else {
|
||||
handleSuccessfulTask(ctx, execution, t, duration, result)
|
||||
}
|
||||
|
||||
if err := model.UpdateTaskExecution(ctx, execution); err != nil {
|
||||
logger.ErrorF(ctx, "[TaskExecutor] 更新执行记录失败 taskID=%s: %v", execution.TaskID, err)
|
||||
}
|
||||
if shouldFlushTaskExecutionLog(ctx, execErr) {
|
||||
if err := model.FlushTaskExecutionLog(ctx, execution.TaskID); err != nil {
|
||||
logger.ErrorF(ctx, "[TaskExecutor] 持久化任务日志失败 taskID=%s: %v", execution.TaskID, err)
|
||||
}
|
||||
}
|
||||
|
||||
notifyTaskCompleted(ctx, execution, result, execErr)
|
||||
}
|
||||
|
||||
func notifyTaskCompleted(ctx context.Context, execution *model.TaskExecution, result *TaskResult, execErr error) {
|
||||
if len(taskCompletedHandlers) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
asyncCtx := context.WithoutCancel(ctx)
|
||||
for _, handler := range taskCompletedHandlers {
|
||||
go handler(asyncCtx, execution, result, execErr)
|
||||
}
|
||||
}
|
||||
|
||||
func shouldFlushTaskExecutionLog(ctx context.Context, execErr error) bool {
|
||||
if execErr == nil {
|
||||
return true
|
||||
}
|
||||
|
||||
retryCount, hasRetryCount := asynq.GetRetryCount(ctx)
|
||||
maxRetry, hasMaxRetry := asynq.GetMaxRetry(ctx)
|
||||
if !hasRetryCount || !hasMaxRetry {
|
||||
return true
|
||||
}
|
||||
return retryCount >= maxRetry
|
||||
}
|
||||
|
||||
func handleFailedTask(ctx context.Context, execution *model.TaskExecution, t *asynq.Task, duration time.Duration, execErr error, span trace.Span) {
|
||||
execution.Status = model.TaskExecutionStatusFailed
|
||||
execution.ErrorMessage = execErr.Error()
|
||||
logger.ErrorF(ctx, "[TaskExecutor] 任务处理失败 Type: %s TaskID: %s Duration: %d ms Error: %v", t.Type(), execution.TaskID, duration.Milliseconds(), execErr)
|
||||
span.SetStatus(codes.Error, execErr.Error())
|
||||
span.RecordError(execErr)
|
||||
|
||||
AppendLog(ctx, "[系统] 任务执行失败，耗时: %d ms，错误原因: %v", duration.Milliseconds(), execErr)
|
||||
}
|
||||
|
||||
func handleSuccessfulTask(ctx context.Context, execution *model.TaskExecution, t *asynq.Task, duration time.Duration, result *TaskResult) {
|
||||
execution.Status = model.TaskExecutionStatusSucceeded
|
||||
execution.ErrorMessage = "" // 清除历史重试失败遗留的错误信息
|
||||
if result != nil {
|
||||
execution.Result = result.Message
|
||||
if result.Detail != "" {
|
||||
execution.Result = fmt.Sprintf("%s\n%s", result.Message, result.Detail)
|
||||
}
|
||||
}
|
||||
logger.InfoF(ctx, "[TaskExecutor] 任务处理完成 Type: %s TaskID: %s Duration: %d ms", t.Type(), execution.TaskID, duration.Milliseconds())
|
||||
|
||||
resultMsg := "成功"
|
||||
if result != nil {
|
||||
resultMsg = result.Message
|
||||
}
|
||||
AppendLog(ctx, "[系统] 任务执行成功，耗时: %d ms，执行结果: %s", duration.Milliseconds(), resultMsg)
|
||||
}
|
||||
|
||||
// generateTaskID 生成任务 ID
|
||||
func generateTaskID(taskType string, triggeredBy string) string {
|
||||
return fmt.Sprintf("%s_%s_%d", triggeredBy, taskType, idgen.NextUint64ID())
|
||||
}
|
||||
|
||||
// generateRetryTaskID 生成重试任务 ID
|
||||
func generateRetryTaskID(originalTaskID string, retryCount int) string {
|
||||
return fmt.Sprintf("retry_%d_%s", retryCount, originalTaskID)
|
||||
}
|
||||
@@ -0,0 +1,401 @@
|
||||
// Copyright 2025 linux.do
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package task
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Rain-kl/Wavelet/internal/model"
|
||||
"github.com/Rain-kl/Wavelet/internal/testhelper"
|
||||
"github.com/hibiken/asynq"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.opentelemetry.io/otel"
|
||||
"go.opentelemetry.io/otel/propagation"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
// mockHandler 用于测试的模拟任务处理器
|
||||
type mockHandler struct {
|
||||
executeFunc func(ctx context.Context, payload []byte) (*TaskResult, error)
|
||||
}
|
||||
|
||||
func (h *mockHandler) Execute(ctx context.Context, payload []byte) (*TaskResult, error) {
|
||||
if h.executeFunc != nil {
|
||||
return h.executeFunc(ctx, payload)
|
||||
}
|
||||
return &TaskResult{Message: "mock success"}, nil
|
||||
}
|
||||
|
||||
// successHandler 返回成功的处理器
|
||||
func successHandler() *mockHandler {
|
||||
return &mockHandler{
|
||||
executeFunc: func(ctx context.Context, payload []byte) (*TaskResult, error) {
|
||||
AppendLog(ctx, "执行成功，处理了 %d 条数据", 100)
|
||||
return &TaskResult{Message: "处理完成，共 100 条"}, nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// failHandler 返回失败的处理器
|
||||
func failHandler() *mockHandler {
|
||||
return &mockHandler{
|
||||
executeFunc: func(ctx context.Context, payload []byte) (*TaskResult, error) {
|
||||
AppendLog(ctx, "开始执行任务")
|
||||
return nil, fmt.Errorf("模拟执行失败: 数据库连接超时")
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
const testTaskType = "test:mock_task"
|
||||
|
||||
func setupTest(t *testing.T) func() {
|
||||
_, mr, cleanup := testhelper.SetupTestEnvironment(t)
|
||||
AsynqClient = asynq.NewClient(asynq.RedisClientOpt{
|
||||
Addr: mr.Addr(),
|
||||
})
|
||||
// 注册测试用 handler
|
||||
RegisterHandler(testTaskType, successHandler())
|
||||
return func() {
|
||||
if AsynqClient != nil {
|
||||
_ = AsynqClient.Close()
|
||||
AsynqClient = nil
|
||||
}
|
||||
cleanup()
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterAndGetHandler(t *testing.T) {
|
||||
_ = testTaskType
|
||||
cleanup := setupTest(t)
|
||||
defer cleanup()
|
||||
|
||||
// 验证 handler 已注册
|
||||
h, ok := getHandler(testTaskType)
|
||||
assert.True(t, ok, "handler should be registered")
|
||||
assert.NotNil(t, h)
|
||||
|
||||
// 未注册的 handler
|
||||
_, ok = getHandler("nonexistent")
|
||||
assert.False(t, ok, "non-existent handler should return false")
|
||||
}
|
||||
|
||||
func TestGetTaskIDFromContext(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
// 空 context
|
||||
taskID := GetTaskID(ctx)
|
||||
assert.Equal(t, "", taskID)
|
||||
|
||||
// 注入 taskID
|
||||
ctx = withTaskID(ctx, "test_task_123")
|
||||
taskID = GetTaskID(ctx)
|
||||
assert.Equal(t, "test_task_123", taskID)
|
||||
}
|
||||
|
||||
func TestAppendLogWithoutTaskID(t *testing.T) {
|
||||
cleanup := setupTest(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
// 没有 taskID 的 context，应降级到普通日志，不报错
|
||||
AppendLog(ctx, "这条日志应该降级处理，不会报错")
|
||||
}
|
||||
|
||||
func TestAppendLogWithTaskID(t *testing.T) {
|
||||
cleanup := setupTest(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
// 先创建一条执行记录
|
||||
execution := &model.TaskExecution{
|
||||
TaskID: "log_test_001",
|
||||
TaskType: testTaskType,
|
||||
TaskName: "测试任务",
|
||||
Status: model.TaskExecutionStatusRunning,
|
||||
TriggeredBy: "manual",
|
||||
}
|
||||
err := model.CreateTaskExecution(ctx, execution)
|
||||
require.NoError(t, err)
|
||||
|
||||
// 注入 taskID 并追加日志
|
||||
ctx = withTaskID(ctx, "log_test_001")
|
||||
AppendLog(ctx, "第一条日志")
|
||||
AppendLog(ctx, "处理了 %d 条数据", 50)
|
||||
|
||||
// 验证日志
|
||||
found, err := model.GetTaskExecutionByTaskID(ctx, "log_test_001")
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, found.Log, "第一条日志")
|
||||
assert.Contains(t, found.Log, "处理了 50 条数据")
|
||||
}
|
||||
|
||||
func TestTaskTraceContextEnvelope(t *testing.T) {
|
||||
payload := []byte(`{"hello":"wavelet"}`)
|
||||
traceID := "4bf92f3577b34da6a3ce929d0e0e4736"
|
||||
parentCtx := otel.GetTextMapPropagator().Extract(
|
||||
context.Background(),
|
||||
propagation.MapCarrier{
|
||||
"traceparent": "00-" + traceID + "-00f067aa0ba902b7-01",
|
||||
},
|
||||
)
|
||||
|
||||
wrappedPayload := injectTaskTraceContext(parentCtx, payload)
|
||||
require.NotEqual(t, string(payload), string(wrappedPayload))
|
||||
|
||||
gotCtx, gotPayload, ok := extractTaskTraceContext(context.Background(), wrappedPayload)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, payload, gotPayload)
|
||||
assert.Equal(t, traceID, trace.SpanContextFromContext(gotCtx).TraceID().String())
|
||||
}
|
||||
|
||||
func TestTaskTraceContextEnvelopeKeepsLegacyPayload(t *testing.T) {
|
||||
payload := []byte(`{"legacy":true}`)
|
||||
|
||||
gotCtx, gotPayload, ok := extractTaskTraceContext(context.Background(), payload)
|
||||
require.False(t, ok)
|
||||
assert.Equal(t, context.Background(), gotCtx)
|
||||
assert.Equal(t, payload, gotPayload)
|
||||
}
|
||||
|
||||
func TestTaskTraceContextEnvelopeSkipsEmptyContext(t *testing.T) {
|
||||
payload := []byte(`{"background":true}`)
|
||||
|
||||
wrappedPayload := injectTaskTraceContext(context.Background(), payload)
|
||||
assert.Equal(t, payload, wrappedPayload)
|
||||
}
|
||||
|
||||
func TestProcessTaskSuccess(t *testing.T) {
|
||||
cleanup := setupTest(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
// 注册成功 handler
|
||||
RegisterHandler(testTaskType, successHandler())
|
||||
|
||||
// 创建执行记录
|
||||
execution := &model.TaskExecution{
|
||||
TaskID: "process_success_001",
|
||||
TaskType: testTaskType,
|
||||
TaskName: "测试任务",
|
||||
Status: model.TaskExecutionStatusPending,
|
||||
Retryable: true,
|
||||
MaxRetry: 3,
|
||||
TriggeredBy: "manual",
|
||||
}
|
||||
err := model.CreateTaskExecution(ctx, execution)
|
||||
require.NoError(t, err)
|
||||
|
||||
// 通过 asynq 的 Task 不能直接设置 taskID，ProcessTask 通过 t.ResultWriter().TaskID() 获取
|
||||
// 但 asynq.Task 在没有经过 asynq server 的情况下 ResultWriter 可能为 nil
|
||||
// 我们需要在 ProcessTask 内部改用 taskID 注入的方式测试
|
||||
// 为了测试 ProcessTask，我们直接模拟调用 handler
|
||||
|
||||
// 直接通过 handler 测试
|
||||
handler, ok := getHandler(testTaskType)
|
||||
require.True(t, ok)
|
||||
|
||||
ctx = withTaskID(ctx, "process_success_001")
|
||||
result, err := handler.Execute(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "处理完成，共 100 条", result.Message)
|
||||
|
||||
// 验证日志被追加
|
||||
found, err := model.GetTaskExecutionByTaskID(ctx, "process_success_001")
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, found.Log, "执行成功，处理了 100 条数据")
|
||||
}
|
||||
|
||||
func TestProcessTaskFailure(t *testing.T) {
|
||||
cleanup := setupTest(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
// 注册失败 handler
|
||||
RegisterHandler(testTaskType, failHandler())
|
||||
|
||||
// 创建执行记录
|
||||
execution := &model.TaskExecution{
|
||||
TaskID: "process_fail_001",
|
||||
TaskType: testTaskType,
|
||||
TaskName: "测试任务",
|
||||
Status: model.TaskExecutionStatusPending,
|
||||
Retryable: true,
|
||||
MaxRetry: 3,
|
||||
TriggeredBy: "manual",
|
||||
}
|
||||
err := model.CreateTaskExecution(ctx, execution)
|
||||
require.NoError(t, err)
|
||||
|
||||
// 直接调用 handler
|
||||
handler, ok := getHandler(testTaskType)
|
||||
require.True(t, ok)
|
||||
|
||||
ctx = withTaskID(ctx, "process_fail_001")
|
||||
_, err = handler.Execute(ctx, nil)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "模拟执行失败")
|
||||
|
||||
// 验证日志
|
||||
found, err := model.GetTaskExecutionByTaskID(ctx, "process_fail_001")
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, found.Log, "开始执行任务")
|
||||
}
|
||||
|
||||
func TestCompleteTaskExecutionFlushesLog(t *testing.T) {
|
||||
cleanup := setupTest(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
execution := &model.TaskExecution{
|
||||
TaskID: "complete_flush_001",
|
||||
TaskType: testTaskType,
|
||||
TaskName: "测试任务",
|
||||
Status: model.TaskExecutionStatusRunning,
|
||||
TriggeredBy: "manual",
|
||||
}
|
||||
err := model.CreateTaskExecution(ctx, execution)
|
||||
require.NoError(t, err)
|
||||
|
||||
ctx = withTaskID(ctx, execution.TaskID)
|
||||
AppendLog(ctx, "任务执行中的日志")
|
||||
|
||||
finishTime := time.Now()
|
||||
completeTaskExecution(
|
||||
ctx,
|
||||
execution,
|
||||
asynq.NewTask(testTaskType, nil),
|
||||
100*time.Millisecond,
|
||||
finishTime,
|
||||
&TaskResult{Message: "处理完成"},
|
||||
nil,
|
||||
trace.SpanFromContext(ctx),
|
||||
)
|
||||
|
||||
found, err := model.GetTaskExecutionByTaskID(ctx, execution.TaskID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, model.TaskExecutionStatusSucceeded, found.Status)
|
||||
assert.Contains(t, found.Log, "任务执行中的日志")
|
||||
assert.Contains(t, found.Log, "任务执行成功")
|
||||
}
|
||||
|
||||
func TestRetryTask(t *testing.T) {
|
||||
cleanup := setupTest(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
// 创建一条失败的执行记录（可重试）
|
||||
now := time.Now()
|
||||
execution := &model.TaskExecution{
|
||||
TaskID: "retry_test_001",
|
||||
TaskType: testTaskType,
|
||||
TaskName: "测试任务",
|
||||
Status: model.TaskExecutionStatusFailed,
|
||||
Retryable: true,
|
||||
MaxRetry: 3,
|
||||
RetryCount: 0,
|
||||
ErrorMessage: "首次执行失败",
|
||||
StartedAt: &now,
|
||||
FinishedAt: &now,
|
||||
Duration: 100,
|
||||
TriggeredBy: "manual",
|
||||
}
|
||||
err := model.CreateTaskExecution(ctx, execution)
|
||||
require.NoError(t, err)
|
||||
|
||||
// 重试
|
||||
newTaskID, err := RetryTask(ctx, execution.ID)
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, newTaskID)
|
||||
assert.Contains(t, newTaskID, "retry_1_")
|
||||
|
||||
// 验证新记录
|
||||
newExecution, err := model.GetTaskExecutionByTaskID(ctx, newTaskID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, model.TaskExecutionStatusPending, newExecution.Status)
|
||||
assert.Equal(t, 1, newExecution.RetryCount)
|
||||
assert.Equal(t, "retry", newExecution.TriggeredBy)
|
||||
assert.Equal(t, execution.TaskType, newExecution.TaskType)
|
||||
assert.True(t, newExecution.Retryable)
|
||||
|
||||
// 原记录不变
|
||||
original, err := model.GetTaskExecutionByID(ctx, execution.ID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, model.TaskExecutionStatusFailed, original.Status)
|
||||
assert.Equal(t, 0, original.RetryCount)
|
||||
}
|
||||
|
||||
func TestRetryTaskNotFailed(t *testing.T) {
|
||||
cleanup := setupTest(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
// 创建一条成功的记录
|
||||
execution := &model.TaskExecution{
|
||||
TaskID: "retry_not_failed_001",
|
||||
TaskType: testTaskType,
|
||||
TaskName: "测试任务",
|
||||
Status: model.TaskExecutionStatusSucceeded,
|
||||
Retryable: true,
|
||||
MaxRetry: 3,
|
||||
TriggeredBy: "manual",
|
||||
}
|
||||
err := model.CreateTaskExecution(ctx, execution)
|
||||
require.NoError(t, err)
|
||||
|
||||
// 尝试重试成功的任务
|
||||
_, err = RetryTask(ctx, execution.ID)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "只有失败的任务才能重试")
|
||||
}
|
||||
|
||||
func TestRetryTaskNotRetryable(t *testing.T) {
|
||||
cleanup := setupTest(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
execution := &model.TaskExecution{
|
||||
TaskID: "retry_not_allowed_001",
|
||||
TaskType: testTaskType,
|
||||
TaskName: "测试任务",
|
||||
Status: model.TaskExecutionStatusFailed,
|
||||
Retryable: false,
|
||||
MaxRetry: 0,
|
||||
TriggeredBy: "manual",
|
||||
}
|
||||
err := model.CreateTaskExecution(ctx, execution)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = RetryTask(ctx, execution.ID)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "不支持重试")
|
||||
}
|
||||
|
||||
func TestRetryTaskNonExistent(t *testing.T) {
|
||||
cleanup := setupTest(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := RetryTask(ctx, 99999999)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "不存在")
|
||||
}
|
||||
|
||||
func TestGenerateTaskID(t *testing.T) {
|
||||
id1 := generateTaskID("test_type", "manual")
|
||||
id2 := generateTaskID("test_type", "manual")
|
||||
|
||||
// 两个 ID 应不同（包含 Snowflake ID）
|
||||
assert.NotEqual(t, id1, id2)
|
||||
assert.Contains(t, id1, "manual_test_type_")
|
||||
}
|
||||
|
||||
func TestGenerateRetryTaskID(t *testing.T) {
|
||||
id := generateRetryTaskID("original_task_123", 2)
|
||||
assert.Equal(t, "retry_2_original_task_123", id)
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
// Copyright 2025 linux.do
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package task
|
||||
|
||||
import "context"
|
||||
|
||||
// TaskResult 任务执行结果
|
||||
//
|
||||
//nolint:revive // TaskResult 保留完整名称以避免与通用 Result 混淆
|
||||
type TaskResult struct {
|
||||
Message string // 结果摘要，如 "共清理 120 个文件，耗时 3.2s"
|
||||
Detail string // 可选的详细结果 JSON
|
||||
}
|
||||
|
||||
// PayloadValidator 可选接口，带参数的任务 Handler 应实现此接口。
|
||||
// 框架在 Admin 下发时自动调用，完成参数校验和标准化（如 Trim 空白）。
|
||||
// 无参数的任务无需实现，框架会直接透传 payload。
|
||||
type PayloadValidator interface {
|
||||
ValidatePayload(payload []byte) ([]byte, error)
|
||||
}
|
||||
|
||||
// TaskHandler 异步任务处理器接口
|
||||
// 所有异步任务必须实现此接口，框架将自动管理任务执行记录的创建、状态流转和日志写入。
|
||||
//
|
||||
// 开发者只需实现 Execute 方法编写业务逻辑，在方法内通过 task.AppendLog(ctx, ...) 追加执行日志。
|
||||
// 任务的创建、状态更新、错误记录、重试计数全部由框架透明处理。
|
||||
//
|
||||
//nolint:revive // TaskHandler 保留完整名称以避免与通用 Handler 混淆
|
||||
type TaskHandler interface {
|
||||
// Execute 执行任务业务逻辑
|
||||
// - ctx: 已注入 Trace Span 和 taskID 的上下文
|
||||
// - payload: 调度时传入的原始参数（可为 nil）
|
||||
// - 返回 TaskResult 描述执行结果，或 error 表示执行失败
|
||||
Execute(ctx context.Context, payload []byte) (*TaskResult, error)
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
// Copyright 2025 linux.do
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Package handlers 注册异步任务处理器
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"github.com/Rain-kl/Wavelet/internal/apps/admin/push"
|
||||
"github.com/Rain-kl/Wavelet/internal/apps/openflare"
|
||||
"github.com/Rain-kl/Wavelet/internal/apps/upload"
|
||||
"github.com/Rain-kl/Wavelet/internal/apps/user"
|
||||
"github.com/Rain-kl/Wavelet/internal/task"
|
||||
)
|
||||
|
||||
// Register registers all built-in task handlers and their metadata.
|
||||
func Register() {
|
||||
task.RegisterHandler(upload.StorageMigrationTask, &upload.MigrationHandler{})
|
||||
task.RegisterTaskMeta(upload.StorageMigrationMeta)
|
||||
|
||||
// system cleanup
|
||||
task.RegisterHandler(upload.SystemCleanupTask, &upload.SystemCleanupHandler{})
|
||||
task.RegisterTaskMeta(upload.SystemCleanupMeta)
|
||||
|
||||
// upload
|
||||
task.RegisterHandler(upload.WarmImageCacheTask, &upload.WarmImageCacheHandler{})
|
||||
task.RegisterTaskMeta(upload.WarmImageCacheMeta)
|
||||
|
||||
task.RegisterHandler(upload.RebuildUploadStatsTask, &upload.RebuildUploadStatsHandler{})
|
||||
task.RegisterTaskMeta(upload.RebuildUploadStatsMeta)
|
||||
|
||||
// user
|
||||
task.RegisterHandler(user.SendEmailTask, &user.SendEmailHandler{})
|
||||
task.RegisterTaskMeta(user.SendEmailMeta)
|
||||
|
||||
// push
|
||||
task.RegisterHandler(push.SendNotificationTask, &push.PushHandler{})
|
||||
task.RegisterTaskMeta(push.SendNotificationMeta)
|
||||
|
||||
// openflare
|
||||
task.RegisterHandler(openflare.SSLRenewTask, &openflare.SSLRenewHandler{})
|
||||
task.RegisterTaskMeta(openflare.SSLRenewMeta)
|
||||
|
||||
task.RegisterHandler(openflare.DatabaseAutoCleanupTask, &openflare.DatabaseAutoCleanupHandler{})
|
||||
task.RegisterTaskMeta(openflare.DatabaseAutoCleanupMeta)
|
||||
|
||||
task.RegisterHandler(openflare.WAFIPGroupSyncTask, &openflare.WAFIPGroupSyncHandler{})
|
||||
task.RegisterTaskMeta(openflare.WAFIPGroupSyncMeta)
|
||||
|
||||
task.RegisterHandler(openflare.UptimeKumaSyncTask, &openflare.UptimeKumaSyncHandler{})
|
||||
task.RegisterTaskMeta(openflare.UptimeKumaSyncMeta)
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package task
|
||||
|
||||
import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
// TaskParam 任务参数定义
|
||||
//
|
||||
//nolint:revive // TaskParam 保留完整名称以避免与通用 Param 混淆
|
||||
type TaskParam struct {
|
||||
Name string `json:"name"` // 参数键名
|
||||
Label string `json:"label"` // 显示名称
|
||||
Type string `json:"type"` // 类型：string, text, number, boolean
|
||||
Required bool `json:"required"` // 是否必填
|
||||
Placeholder string `json:"placeholder"` // 占位符
|
||||
Description string `json:"description"` // 描述
|
||||
}
|
||||
|
||||
// TaskMeta 任务元数据
|
||||
//
|
||||
//nolint:revive // TaskMeta 保留完整名称以避免与通用 Meta 混淆
|
||||
type TaskMeta struct {
|
||||
Type string `json:"type"`
|
||||
AsynqTask string `json:"asynq_task"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
SupportsTime bool `json:"supports_time"`
|
||||
MaxRetry int `json:"max_retry"`
|
||||
Queue string `json:"queue"`
|
||||
Retryable bool `json:"retryable"` // 是否支持手动重试
|
||||
Params []TaskParam `json:"params,omitempty"`
|
||||
}
|
||||
|
||||
var (
|
||||
dispatchableTasksMutex sync.RWMutex
|
||||
dispatchableTasks []TaskMeta
|
||||
)
|
||||
|
||||
// RegisterTaskMeta 注册任务元数据到全局列表
|
||||
func RegisterTaskMeta(meta TaskMeta) {
|
||||
dispatchableTasksMutex.Lock()
|
||||
defer dispatchableTasksMutex.Unlock()
|
||||
for _, t := range dispatchableTasks {
|
||||
if t.Type == meta.Type {
|
||||
return
|
||||
}
|
||||
}
|
||||
dispatchableTasks = append(dispatchableTasks, meta)
|
||||
}
|
||||
|
||||
// GetDispatchableTasks 获取所有已注册的元数据列表（返回副本以避免并发并发读写冲突）
|
||||
func GetDispatchableTasks() []TaskMeta {
|
||||
dispatchableTasksMutex.RLock()
|
||||
defer dispatchableTasksMutex.RUnlock()
|
||||
|
||||
metas := make([]TaskMeta, len(dispatchableTasks))
|
||||
copy(metas, dispatchableTasks)
|
||||
return metas
|
||||
}
|
||||
|
||||
// GetTaskMeta 根据任务类型获取元数据
|
||||
func GetTaskMeta(taskType string) *TaskMeta {
|
||||
dispatchableTasksMutex.RLock()
|
||||
defer dispatchableTasksMutex.RUnlock()
|
||||
for _, t := range dispatchableTasks {
|
||||
if t.Type == taskType {
|
||||
copied := t
|
||||
return &copied
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetTaskMetaByAsynqTask 根据 Asynq 任务名称获取元数据
|
||||
func GetTaskMetaByAsynqTask(asynqTask string) *TaskMeta {
|
||||
dispatchableTasksMutex.RLock()
|
||||
defer dispatchableTasksMutex.RUnlock()
|
||||
for _, t := range dispatchableTasks {
|
||||
if t.AsynqTask == asynqTask {
|
||||
copied := t
|
||||
return &copied
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetRegisteredAsynqTasks 返回所有已注册的 Asynq 任务名称，以便动态注册路由
|
||||
func GetRegisteredAsynqTasks() []string {
|
||||
keys := make([]string, 0, len(handlerRegistry))
|
||||
for k := range handlerRegistry {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package task_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Rain-kl/Wavelet/internal/task"
|
||||
taskhandlers "github.com/Rain-kl/Wavelet/internal/task/handlers"
|
||||
)
|
||||
|
||||
func TestDuplicateTaskMeta(t *testing.T) {
|
||||
// Call Register twice to simulate being imported by multiple packages (routers, worker, etc.)
|
||||
taskhandlers.Register()
|
||||
taskhandlers.Register()
|
||||
|
||||
metas := task.GetDispatchableTasks()
|
||||
|
||||
// Check if we have duplicates by checking if a Type appears more than once
|
||||
seen := make(map[string]int)
|
||||
for _, m := range metas {
|
||||
seen[m.Type]++
|
||||
}
|
||||
|
||||
for taskType, count := range seen {
|
||||
if count > 1 {
|
||||
t.Errorf("Task type %q registered %d times, expected at most 1", taskType, count)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Package scheduler 提供定时任务调度功能
|
||||
package scheduler
|
||||
|
||||
const (
|
||||
errLoadLocationFailed = "failed to load location: %w"
|
||||
)
|
||||
@@ -0,0 +1,144 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os/signal"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/Rain-kl/Wavelet/internal/bootstrap"
|
||||
"github.com/Rain-kl/Wavelet/internal/model"
|
||||
"github.com/Rain-kl/Wavelet/internal/task"
|
||||
"github.com/Rain-kl/Wavelet/pkg/logger"
|
||||
|
||||
"github.com/hibiken/asynq"
|
||||
)
|
||||
|
||||
var (
|
||||
activeScheduler *asynq.Scheduler
|
||||
schedulerMutex sync.Mutex
|
||||
quitChan chan struct{}
|
||||
schedulerOnce sync.Once
|
||||
)
|
||||
|
||||
// GetAsynqClient 获取全局 AsynqClient
|
||||
func GetAsynqClient() *asynq.Client {
|
||||
return task.AsynqClient
|
||||
}
|
||||
|
||||
// StartScheduler 启动调度器 (该函数阻塞，直到调度器退出)
|
||||
func StartScheduler() error {
|
||||
bootstrap.RegisterScheduler()
|
||||
|
||||
var err error
|
||||
schedulerOnce.Do(func() {
|
||||
quitChan = make(chan struct{})
|
||||
done := quitChan
|
||||
|
||||
// 初始化并运行首次调度
|
||||
if err = ReloadScheduler(); err != nil {
|
||||
err = fmt.Errorf("initial reload failed: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
signalCtx, stopSignals := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stopSignals()
|
||||
|
||||
if waitForStop(done, signalCtx.Done()) {
|
||||
StopScheduler()
|
||||
}
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// StopScheduler 停止调度服务并解除 StartScheduler 阻塞
|
||||
func StopScheduler() {
|
||||
schedulerMutex.Lock()
|
||||
defer schedulerMutex.Unlock()
|
||||
|
||||
if activeScheduler != nil {
|
||||
activeScheduler.Shutdown()
|
||||
activeScheduler = nil
|
||||
}
|
||||
|
||||
if quitChan != nil {
|
||||
close(quitChan)
|
||||
quitChan = nil
|
||||
}
|
||||
}
|
||||
|
||||
// ReloadScheduler 重载调度器配置 (线程安全)
|
||||
func ReloadScheduler() error {
|
||||
schedulerMutex.Lock()
|
||||
defer schedulerMutex.Unlock()
|
||||
|
||||
// 1. 如果有运行中的调度器，先关闭它
|
||||
if activeScheduler != nil {
|
||||
activeScheduler.Shutdown()
|
||||
activeScheduler = nil
|
||||
}
|
||||
|
||||
// 2. 从数据库载入启用的定时任务配置
|
||||
schedules, err := model.ListActiveSchedules(context.Background())
|
||||
if err != nil {
|
||||
return fmt.Errorf("load schedules from db failed: %w", err)
|
||||
}
|
||||
|
||||
location, err := time.LoadLocation("Asia/Shanghai")
|
||||
if err != nil {
|
||||
return fmt.Errorf(errLoadLocationFailed, err)
|
||||
}
|
||||
|
||||
// 3. 实例化新的调度器
|
||||
newScheduler := asynq.NewScheduler(
|
||||
task.RedisOpt,
|
||||
&asynq.SchedulerOpts{
|
||||
Location: location,
|
||||
},
|
||||
)
|
||||
|
||||
// 4. 遍历并注册任务
|
||||
for _, s := range schedules {
|
||||
meta := task.GetTaskMeta(s.TaskType)
|
||||
if meta == nil {
|
||||
continue // 忽略排程配置中无效的任务类型
|
||||
}
|
||||
|
||||
// 构造 Asynq 载荷。定时任务使用对应 Meta 中的 Asynq 标识，同时将数据库中保存的 json 作为参数
|
||||
t := asynq.NewTask(meta.AsynqTask, []byte(s.Payload))
|
||||
|
||||
if _, err := newScheduler.Register(
|
||||
s.Cron,
|
||||
t,
|
||||
asynq.MaxRetry(meta.MaxRetry),
|
||||
asynq.Queue(meta.Queue),
|
||||
); err != nil {
|
||||
// 定时任务配置可能有误（如 Cron 格式不被 Asynq 识别），记录日志并跳过
|
||||
logger.ErrorF(context.Background(), "[Scheduler] 注册定时任务失败 id=%d name=%s: %v", s.ID, s.Name, err)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// 5. 启动并替换全局调度器。进程信号由 StartScheduler 统一处理。
|
||||
if err := newScheduler.Start(); err != nil {
|
||||
return fmt.Errorf("start scheduler failed: %w", err)
|
||||
}
|
||||
activeScheduler = newScheduler
|
||||
|
||||
logger.InfoF(context.Background(), "[Scheduler] 成功重新加载定时任务，共注册 %d 个活动任务", len(schedules))
|
||||
return nil
|
||||
}
|
||||
|
||||
func waitForStop(done <-chan struct{}, signals <-chan struct{}) bool {
|
||||
select {
|
||||
case <-done:
|
||||
return false
|
||||
case <-signals:
|
||||
return true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package scheduler
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestWaitForStop(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
closeDone bool
|
||||
closeSignal bool
|
||||
wantSignal bool
|
||||
}{
|
||||
{
|
||||
name: "explicit stop",
|
||||
closeDone: true,
|
||||
},
|
||||
{
|
||||
name: "process signal",
|
||||
closeSignal: true,
|
||||
wantSignal: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
done := make(chan struct{})
|
||||
signals := make(chan struct{})
|
||||
if tt.closeDone {
|
||||
close(done)
|
||||
}
|
||||
if tt.closeSignal {
|
||||
close(signals)
|
||||
}
|
||||
|
||||
if got := waitForStop(done, signals); got != tt.wantSignal {
|
||||
t.Errorf("waitForStop() = %t, want %t", got, tt.wantSignal)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
// Copyright 2025 linux.do
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package task
|
||||
|
||||
import (
|
||||
"github.com/Rain-kl/Wavelet/internal/config"
|
||||
"github.com/hibiken/asynq"
|
||||
)
|
||||
|
||||
// RedisOpt asynq Redis 连接配置（兼容 Standalone/Sentinel/Cluster）
|
||||
var RedisOpt asynq.RedisConnOpt
|
||||
|
||||
// AsynqClient asynq 客户端，用于任务入队
|
||||
var AsynqClient *asynq.Client
|
||||
|
||||
func init() {
|
||||
RedisOpt = NewRedisConnOpt()
|
||||
AsynqClient = asynq.NewClient(RedisOpt)
|
||||
}
|
||||
|
||||
// NewRedisConnOpt 根据配置返回对应的 asynq Redis 连接选项
|
||||
func NewRedisConnOpt() asynq.RedisConnOpt {
|
||||
cfg := config.Config.Redis
|
||||
addrs := cfg.Addrs
|
||||
|
||||
if cfg.ClusterMode {
|
||||
return asynq.RedisClusterClientOpt{
|
||||
Addrs: addrs,
|
||||
Username: cfg.Username,
|
||||
Password: cfg.Password,
|
||||
}
|
||||
}
|
||||
|
||||
if cfg.MasterName != "" {
|
||||
return asynq.RedisFailoverClientOpt{
|
||||
MasterName: cfg.MasterName,
|
||||
SentinelAddrs: addrs,
|
||||
Username: cfg.Username,
|
||||
Password: cfg.Password,
|
||||
DB: cfg.DB,
|
||||
}
|
||||
}
|
||||
|
||||
addr := "localhost:6379"
|
||||
if len(addrs) > 0 {
|
||||
addr = addrs[0]
|
||||
}
|
||||
return asynq.RedisClientOpt{
|
||||
Addr: addr,
|
||||
Username: cfg.Username,
|
||||
Password: cfg.Password,
|
||||
DB: cfg.DB,
|
||||
PoolSize: cfg.PoolSize,
|
||||
}
|
||||
}
|
||||
|
||||
// PrefixedQueue 返回带前缀的队列名，用于 Cluster 模式隔离
|
||||
func PrefixedQueue(queue string) string {
|
||||
prefix := config.Config.Redis.KeyPrefix
|
||||
if prefix == "" {
|
||||
return queue
|
||||
}
|
||||
return prefix + queue
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
// Copyright 2025 linux.do
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Package worker 提供 Asynq 任务处理服务器与中间件
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/hibiken/asynq"
|
||||
)
|
||||
|
||||
// taskLoggingMiddleware 任务处理中间件
|
||||
// 注意：OTel Span 创建、日志记录、TaskExecution 状态管理
|
||||
// 已由 task.ProcessTask 统一处理，此中间件保留用于未来扩展（如限流、监控等）
|
||||
func taskLoggingMiddleware(h asynq.Handler) asynq.Handler {
|
||||
return asynq.HandlerFunc(func(ctx context.Context, t *asynq.Task) error {
|
||||
return h.ProcessTask(ctx, t)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
// Copyright 2025 linux.do
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package worker
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/Rain-kl/Wavelet/internal/bootstrap"
|
||||
"github.com/Rain-kl/Wavelet/internal/config"
|
||||
"github.com/Rain-kl/Wavelet/internal/task"
|
||||
"github.com/hibiken/asynq"
|
||||
)
|
||||
|
||||
// workerShutdownTimeout Worker 优雅关闭超时时间
|
||||
const workerShutdownTimeout = 3 * time.Minute
|
||||
|
||||
// StartWorker 启动任务处理服务器
|
||||
func StartWorker() error {
|
||||
bootstrap.RegisterWorker()
|
||||
asynqServer := asynq.NewServer(
|
||||
task.RedisOpt,
|
||||
asynq.Config{
|
||||
Concurrency: config.Config.Worker.Concurrency,
|
||||
ShutdownTimeout: workerShutdownTimeout,
|
||||
Queues: buildQueuesFromConfig(),
|
||||
StrictPriority: config.Config.Worker.StrictPriority,
|
||||
},
|
||||
)
|
||||
|
||||
// 注册 Asynq 任务路由
|
||||
mux := asynq.NewServeMux()
|
||||
mux.Use(taskLoggingMiddleware)
|
||||
|
||||
// 统一使用 task.ProcessTask 处理所有任务类型
|
||||
// 框架内部自动分发到对应的 TaskHandler 实现
|
||||
// 动态注册所有已注册的任务处理器路由，框架内部自动分发到对应的 TaskHandler 实现
|
||||
for _, taskName := range task.GetRegisteredAsynqTasks() {
|
||||
mux.HandleFunc(taskName, task.ProcessTask)
|
||||
}
|
||||
|
||||
// 启动服务器
|
||||
return asynqServer.Run(mux)
|
||||
}
|
||||
|
||||
// buildQueuesFromConfig 从配置构建队列映射
|
||||
func buildQueuesFromConfig() map[string]int {
|
||||
queues := make(map[string]int)
|
||||
|
||||
// 从配置读取队列
|
||||
if len(config.Config.Worker.Queues) > 0 {
|
||||
for _, q := range config.Config.Worker.Queues {
|
||||
if q.Name != "" && q.Priority > 0 {
|
||||
queues[q.Name] = q.Priority
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 如果配置为空，使用默认队列
|
||||
if len(queues) == 0 {
|
||||
queues = map[string]int{
|
||||
task.QueueDefault: 1,
|
||||
}
|
||||
}
|
||||
|
||||
return queues
|
||||
}
|
||||
Reference in New Issue
Block a user