mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-09-29 22:06:38 +08:00
ryan
297 lines
8.8 KiB
Go
297 lines
8.8 KiB
Go
/*
|
||
Copyright 2025-2026 linux.do
|
||
|
||
Licensed under the Apache License, Version 2.0 (the "License");
|
||
you may not use this file except in compliance with the License.
|
||
You may obtain a copy of the License at
|
||
|
||
http://www.apache.org/licenses/LICENSE-2.0
|
||
|
||
Unless required by applicable law or agreed to in writing, software
|
||
distributed under the License is distributed on an "AS IS" BASIS,
|
||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
See the License for the specific language governing permissions and
|
||
limitations under the License.
|
||
*/
|
||
|
||
package task
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"time"
|
||
|
||
"github.com/hibiken/asynq"
|
||
"github.com/linux-do/credit/internal/db/idgen"
|
||
"github.com/linux-do/credit/internal/logger"
|
||
"github.com/linux-do/credit/internal/model"
|
||
"github.com/linux-do/credit/internal/otel_trace"
|
||
"go.opentelemetry.io/otel/attribute"
|
||
"go.opentelemetry.io/otel/codes"
|
||
"go.opentelemetry.io/otel/trace"
|
||
)
|
||
|
||
// handlerRegistry 已注册的任务处理器
|
||
var handlerRegistry = make(map[string]TaskHandler)
|
||
|
||
// 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
|
||
}
|
||
|
||
// contextKey 用于 context 存取 taskID
|
||
type contextKey string
|
||
|
||
const taskIDKey contextKey = "task_execution_task_id"
|
||
|
||
// 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 ""
|
||
}
|
||
|
||
// 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("未知的任务类型: %s", 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("创建任务执行记录失败: %w", err)
|
||
}
|
||
|
||
// 入队 Asynq
|
||
taskInfo := asynq.NewTask(meta.AsynqTask, 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("任务入队失败: %w", 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("任务执行记录不存在: %w", err)
|
||
}
|
||
|
||
if execution.Status != model.TaskExecutionStatusFailed {
|
||
return "", fmt.Errorf("只有失败的任务才能重试，当前状态: %s", execution.Status)
|
||
}
|
||
|
||
if !execution.Retryable {
|
||
return "", fmt.Errorf("该任务不支持重试")
|
||
}
|
||
|
||
if execution.RetryCount >= execution.MaxRetry {
|
||
return "", fmt.Errorf("已达到最大重试次数 %d", execution.MaxRetry)
|
||
}
|
||
|
||
// 生成新的 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("创建重试任务执行记录失败: %w", err)
|
||
}
|
||
|
||
// 入队 Asynq
|
||
taskInfo := asynq.NewTask(execution.TaskType, []byte(execution.Payload))
|
||
if _, err := AsynqClient.Enqueue(
|
||
taskInfo,
|
||
asynq.TaskID(newTaskID),
|
||
asynq.MaxRetry(execution.MaxRetry),
|
||
asynq.Queue(PrefixedQueue(QueueDefault)),
|
||
); 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("重试任务入队失败: %w", err)
|
||
}
|
||
|
||
return newTaskID, nil
|
||
}
|
||
|
||
// ProcessTask Asynq 实际调用的统一处理函数
|
||
// Worker 注册时统一使用此函数，内部自动分发到对应的 TaskHandler
|
||
func ProcessTask(ctx context.Context, t *asynq.Task) error {
|
||
// 初始化 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(t.Payload())),
|
||
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("未注册的任务处理器: %s", t.Type())
|
||
logger.ErrorF(ctx, "[TaskExecutor] %v", err)
|
||
span.SetStatus(codes.Error, err.Error())
|
||
return err
|
||
}
|
||
|
||
// 从数据库加载执行记录
|
||
execution, err := model.GetTaskExecutionByTaskID(ctx, taskID)
|
||
if err != nil {
|
||
logger.ErrorF(ctx, "[TaskExecutor] 查询执行记录失败 taskID=%s: %v", taskID, err)
|
||
// 执行记录不存在，仍然执行任务但不记录状态
|
||
_, execErr := handler.Execute(ctx, t.Payload())
|
||
if execErr != nil {
|
||
span.SetStatus(codes.Error, execErr.Error())
|
||
return execErr
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// 更新状态为 running
|
||
now := time.Now()
|
||
execution.Status = model.TaskExecutionStatusRunning
|
||
execution.StartedAt = &now
|
||
if err := model.UpdateTaskExecution(ctx, execution); err != nil {
|
||
logger.ErrorF(ctx, "[TaskExecutor] 更新执行状态失败 taskID=%s: %v", taskID, err)
|
||
}
|
||
|
||
// 开始计时
|
||
start := time.Now()
|
||
|
||
// 执行业务逻辑
|
||
result, execErr := handler.Execute(ctx, t.Payload())
|
||
|
||
// 计算耗时
|
||
duration := time.Since(start)
|
||
finishTime := time.Now()
|
||
execution.Duration = duration.Milliseconds()
|
||
execution.FinishedAt = &finishTime
|
||
|
||
if execErr != nil {
|
||
// 执行失败
|
||
execution.Status = model.TaskExecutionStatusFailed
|
||
execution.ErrorMessage = execErr.Error()
|
||
|
||
logger.ErrorF(ctx,
|
||
"[TaskExecutor] 任务处理失败 Type: %s TaskID: %s Duration: %d ms Error: %v",
|
||
t.Type(), taskID, duration.Milliseconds(), execErr,
|
||
)
|
||
|
||
span.SetStatus(codes.Error, execErr.Error())
|
||
span.RecordError(execErr)
|
||
} else {
|
||
// 执行成功
|
||
execution.Status = model.TaskExecutionStatusSucceeded
|
||
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(), taskID, duration.Milliseconds(),
|
||
)
|
||
}
|
||
|
||
// 更新执行记录
|
||
if err := model.UpdateTaskExecution(ctx, execution); err != nil {
|
||
logger.ErrorF(ctx, "[TaskExecutor] 更新执行记录失败 taskID=%s: %v", taskID, err)
|
||
}
|
||
|
||
return execErr
|
||
}
|
||
|
||
// generateTaskID 生成任务 ID
|
||
func generateTaskID(taskType string, triggeredBy string) string {
|
||
uniqueID := idgen.NextUint64ID()
|
||
return fmt.Sprintf("%s_%s_%d", triggeredBy, taskType, uniqueID)
|
||
}
|
||
|
||
// generateRetryTaskID 生成重试任务 ID
|
||
func generateRetryTaskID(originalTaskID string, retryCount int) string {
|
||
return fmt.Sprintf("retry_%d_%s", retryCount, originalTaskID)
|
||
}
|