refactor(layout): consolidate backend codebase into backend/ package and clean root directory

- Moved cmd/, core/, plugins/, pkg/, downstream/, and main.go into backend/ directory
- Batch updated all Go source files to import github.com/Rain-kl/Wavelet/backend/...
- Updated Makefile, scripts/swagger.sh, architecture guards, and platform skills
- Passed all quality gates (100% tests, 0 lint issues, clean build)
This commit is contained in:
ryan
2026-08-28 12:56:02 +08:00
parent 33b38f8687
commit 43dc97e48c
319 changed files with 912 additions and 1031 deletions
@@ -0,0 +1,9 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package driver_asynq_cron provides Asynq cron schedule driver plugin for Cordis.
package driver_asynq_cron
const (
errLoadLocationFailed = "failed to load location: %w"
)
@@ -0,0 +1,20 @@
-- +goose Up
CREATE TABLE IF NOT EXISTS w_schedules (
id BIGINT PRIMARY KEY,
name VARCHAR(128) NOT NULL,
task_type VARCHAR(64) NOT NULL,
cron VARCHAR(64) NOT NULL,
payload TEXT,
is_active BOOLEAN NOT NULL DEFAULT TRUE,
created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS idx_w_schedules_is_active ON w_schedules (is_active);
-- Seed initial cleanup task
INSERT INTO w_schedules (id, name, task_type, cron, payload, is_active, created_at, updated_at)
VALUES (1, '系统定期垃圾清理', 'system_cleanup', '0 3 * * *', '{}', TRUE, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
ON CONFLICT (id) DO NOTHING;
-- +goose Down
DROP TABLE IF EXISTS w_schedules;
@@ -0,0 +1,228 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package driver_asynq_cron provides the Asynq cron schedule driver plugin for Cordis.
package driver_asynq_cron
import (
"context"
"embed"
"encoding/json"
"fmt"
"sync"
"time"
"github.com/hibiken/asynq"
"github.com/Rain-kl/Wavelet/backend/core"
"github.com/Rain-kl/Wavelet/backend/plugins/drivers/driver_asynq_worker"
)
//go:embed migrations/*.sql
var cronMigrations embed.FS
// Option configures the Asynq cron scheduler driver plugin.
type Option func(*Plugin)
// WithRedisOpt sets the Redis connection options for Asynq scheduler.
func WithRedisOpt(opt asynq.RedisConnOpt) Option {
return func(p *Plugin) {
p.redisOpt = opt
}
}
// WithLocation sets the timezone location for the cron scheduler.
func WithLocation(loc *time.Location) Option {
return func(p *Plugin) {
p.location = loc
}
}
// WithSchedulerOpts sets custom Asynq SchedulerOpts.
func WithSchedulerOpts(opts *asynq.SchedulerOpts) Option {
return func(p *Plugin) {
p.schedulerOpts = opts
}
}
// WithScheduler injects a pre-configured Asynq Scheduler instance.
func WithScheduler(sched *asynq.Scheduler) Option {
return func(p *Plugin) {
p.scheduler = sched
}
}
// Plugin implements core.Plugin and core.Driver for Asynq Cron Scheduler.
type Plugin struct {
mu sync.RWMutex
redisOpt asynq.RedisConnOpt
location *time.Location
schedulerOpts *asynq.SchedulerOpts
scheduler *asynq.Scheduler
running bool
coreCtx *core.Context
}
// New creates a new Asynq Cron Scheduler driver plugin.
func New(opts ...Option) *Plugin {
p := &Plugin{
redisOpt: driver_asynq_worker.RedisOpt,
location: time.Local,
}
for _, opt := range opts {
if opt != nil {
opt(p)
}
}
return p
}
// Name returns the unique plugin identifier.
func (p *Plugin) Name() string {
return "driver_asynq_cron"
}
// Apply mounts the Asynq Cron Scheduler driver into the micro-kernel Context.
func (p *Plugin) Apply(ctx *core.Context) error {
p.mu.Lock()
p.coreCtx = ctx
p.mu.Unlock()
// Register migrations for w_schedules table
ctx.Migrations().Register("driver_asynq_cron", cronMigrations)
ctx.OnDispose(func() error {
return p.Stop(context.Background())
})
return ctx.RegisterDriver(p)
}
// Type returns DriverTypeScheduler.
func (p *Plugin) Type() core.DriverType {
return core.DriverTypeScheduler
}
// Start registers scheduled cron entries from Context and boots the Asynq scheduler.
func (p *Plugin) Start(_ context.Context) error {
p.mu.Lock()
defer p.mu.Unlock()
if p.running {
return nil
}
if p.scheduler == nil {
opts := p.schedulerOpts
if opts == nil {
opts = &asynq.SchedulerOpts{
Location: p.location,
}
} else if opts.Location == nil && p.location != nil {
opts.Location = p.location
}
p.scheduler = asynq.NewScheduler(p.redisOpt, opts)
}
if p.coreCtx != nil && p.coreCtx.Schedules() != nil {
for _, sd := range p.coreCtx.Schedules().Schedules() {
payloadBytes, err := encodePayload(sd.Payload)
if err != nil {
return fmt.Errorf("driver_asynq_cron: encode payload for task %q failed: %w", sd.TaskType, err)
}
task := asynq.NewTask(sd.TaskType, payloadBytes)
asynqOpts := buildAsynqOptions(sd.Options)
if _, err := p.scheduler.Register(sd.Spec, task, asynqOpts...); err != nil {
return fmt.Errorf("driver_asynq_cron: register schedule %q for task %q failed: %w", sd.Spec, sd.TaskType, err)
}
}
}
if err := p.scheduler.Start(); err != nil {
return fmt.Errorf("driver_asynq_cron: start scheduler failed: %w", err)
}
p.running = true
return nil
}
// Stop gracefully shuts down the Asynq cron scheduler.
func (p *Plugin) Stop(_ context.Context) error {
p.mu.Lock()
defer p.mu.Unlock()
if !p.running {
return nil
}
p.running = false
if p.scheduler != nil {
p.scheduler.Shutdown()
p.scheduler = nil
}
return nil
}
// IsRunning returns whether the scheduler is running.
func (p *Plugin) IsRunning() bool {
p.mu.RLock()
defer p.mu.RUnlock()
return p.running
}
// Scheduler returns the underlying Asynq scheduler instance.
func (p *Plugin) Scheduler() *asynq.Scheduler {
p.mu.RLock()
defer p.mu.RUnlock()
return p.scheduler
}
func encodePayload(payload any) ([]byte, error) {
if payload == nil {
return nil, nil
}
switch val := payload.(type) {
case []byte:
return val, nil
case string:
return []byte(val), nil
default:
return json.Marshal(val)
}
}
func buildAsynqOptions(opts map[string]any) []asynq.Option {
if len(opts) == 0 {
return nil
}
var res []asynq.Option
if q, ok := opts["queue"].(string); ok && q != "" {
res = append(res, asynq.Queue(q))
}
if retry, ok := opts["retry"].(int); ok {
res = append(res, asynq.MaxRetry(retry))
} else if maxRetry, ok := opts["max_retry"].(int); ok {
res = append(res, asynq.MaxRetry(maxRetry))
}
if timeout, ok := opts["timeout"].(time.Duration); ok {
res = append(res, asynq.Timeout(timeout))
}
if retention, ok := opts["retention"].(time.Duration); ok {
res = append(res, asynq.Retention(retention))
}
return res
}
@@ -0,0 +1,37 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_cron
import (
"context"
"time"
db "github.com/Rain-kl/Wavelet/backend/plugins/infra/database"
)
// Schedule 定时任务配置表
type Schedule struct {
ID uint64 `json:"id,string" gorm:"primaryKey"`
Name string `json:"name" gorm:"size:128;not null"`
TaskType string `json:"task_type" gorm:"size:64;not null"`
Cron string `json:"cron" gorm:"size:64;not null"`
Payload string `json:"payload" gorm:"type:text"`
IsActive bool `json:"is_active" gorm:"not null;default:true"`
CreatedAt time.Time `json:"created_at" gorm:"autoCreateTime;index"`
UpdatedAt time.Time `json:"updated_at" gorm:"autoUpdateTime"`
}
// TableName 表名
func (Schedule) TableName() string {
return "w_schedules"
}
// ListActiveSchedules 查询所有已启用的定时任务配置
func ListActiveSchedules(ctx context.Context) ([]Schedule, error) {
var schedules []Schedule
if err := db.DB(ctx).Where("is_active = ?", true).Find(&schedules).Error; err != nil {
return nil, err
}
return schedules, nil
}
@@ -0,0 +1,140 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_cron
import (
"context"
"fmt"
"os/signal"
"sync"
"syscall"
"time"
"github.com/hibiken/asynq"
"github.com/Rain-kl/Wavelet/backend/pkg/logger"
"github.com/Rain-kl/Wavelet/backend/plugins/drivers/driver_asynq_worker"
)
var (
activeScheduler *asynq.Scheduler
schedulerMutex sync.Mutex
quitChan chan struct{}
schedulerOnce sync.Once
)
// GetAsynqClient 获取全局 AsynqClient
func GetAsynqClient() *asynq.Client {
return driver_asynq_worker.AsynqClient
}
// StartScheduler 启动调度器 (该函数阻塞,直到调度器退出)
func StartScheduler() error {
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 := 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(
driver_asynq_worker.RedisOpt,
&asynq.SchedulerOpts{
Location: location,
},
)
// 4. 遍历并注册任务
for _, s := range schedules {
meta := driver_asynq_worker.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 driver_asynq_cron
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,13 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
// 任务队列名称
const (
QueueDefault = "default"
)
// DefaultMaxRetry 任务默认最大重试次数
const DefaultMaxRetry = 3
@@ -0,0 +1,17 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
const (
errUnknownTaskType = "未知的任务类型: %s"
errCreateTaskExecutionFailed = "创建任务执行记录失败: %w"
errTaskEnqueueFailed = "任务入队失败: %w"
errTaskExecutionNotFound = "任务执行记录不存在: %w"
errRetryOnlyFailedTask = "只有失败的任务才能重试,当前状态: %s"
errTaskNotRetryable = "该任务不支持重试"
errTaskMaxRetryExceeded = "已达到最大重试次数 %d"
errCreateRetryExecutionFailed = "创建重试任务执行记录失败: %w"
errRetryTaskEnqueueFailed = "重试任务入队失败: %w"
errUnregisteredTaskHandler = "未注册的任务处理器: %s"
)
@@ -0,0 +1,500 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
import (
"context"
"encoding/json"
"errors"
"fmt"
"sync"
"time"
"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"
"github.com/Rain-kl/Wavelet/backend/pkg/idgen"
"github.com/Rain-kl/Wavelet/backend/pkg/logger"
otel_trace "github.com/Rain-kl/Wavelet/backend/pkg/trace"
"github.com/Rain-kl/Wavelet/backend/pkg/util"
)
var (
handlerRegistryMutex sync.RWMutex
handlerRegistry = make(map[string]TaskHandler)
completedHandlersMutex sync.RWMutex
taskCompletedHandlers []CompletedHandler
)
// CompletedHandler is called when a task execution completes.
type CompletedHandler func(ctx context.Context, execution *TaskExecution, result *TaskResult, execErr error)
// OnTaskCompleted registers a handler for task completion events.
// Handlers must be registered during application bootstrap before processing tasks.
func OnTaskCompleted(handler CompletedHandler) {
completedHandlersMutex.Lock()
defer completedHandlersMutex.Unlock()
taskCompletedHandlers = append(taskCompletedHandlers, handler)
}
// RegisterHandler 注册任务处理器
// 传入任务类型标识(对应 AsynqTask 常量)和 TaskHandler 实现
func RegisterHandler(asynqTaskType string, handler TaskHandler) {
handlerRegistryMutex.Lock()
defer handlerRegistryMutex.Unlock()
handlerRegistry[asynqTaskType] = handler
}
// getHandler 获取已注册的处理器
func getHandler(asynqTaskType string) (TaskHandler, bool) {
handlerRegistryMutex.RLock()
defer handlerRegistryMutex.RUnlock()
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 := 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 := &TaskExecution{
TaskID: taskID,
TaskType: meta.AsynqTask,
TaskName: meta.Name,
Status: TaskExecutionStatusPending,
Retryable: meta.Retryable,
MaxRetry: meta.MaxRetry,
RetryCount: 0,
Payload: string(payload),
TriggeredBy: triggeredBy,
}
if err := 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 = TaskExecutionStatusFailed
execution.ErrorMessage = fmt.Sprintf("入队失败: %v", err)
now := time.Now()
execution.StartedAt = &now
execution.FinishedAt = &now
_ = updateTaskExecution(ctx, execution)
return "", fmt.Errorf(errTaskEnqueueFailed, err)
}
if err := 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 := getTaskExecutionByID(ctx, id)
if err != nil {
return "", fmt.Errorf(errTaskExecutionNotFound, err)
}
if execution.Status != TaskExecutionStatusFailed {
return "", fmt.Errorf(errRetryOnlyFailedTask, execution.Status)
}
if !execution.Retryable {
return "", errors.New(errTaskNotRetryable)
}
// 生成新的 TaskID
newTaskID := generateRetryTaskID(execution.TaskID, execution.RetryCount+1)
// 创建新的执行记录
newExecution := &TaskExecution{
TaskID: newTaskID,
TaskType: execution.TaskType,
TaskName: execution.TaskName,
Status: TaskExecutionStatusPending,
Retryable: execution.Retryable,
MaxRetry: execution.MaxRetry,
RetryCount: execution.RetryCount + 1,
Payload: execution.Payload,
TriggeredBy: "retry",
}
if err := 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 = TaskExecutionStatusFailed
newExecution.ErrorMessage = fmt.Sprintf("重试入队失败: %v", err)
now := time.Now()
newExecution.StartedAt = &now
newExecution.FinishedAt = &now
_ = updateTaskExecution(ctx, newExecution)
return "", fmt.Errorf(errRetryTaskEnqueueFailed, err)
}
if err := 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 *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 != TaskExecutionStatusRunning {
execution.Status = TaskExecutionStatusRunning
execution.StartedAt = &now
dirty = true
}
if dirty {
if updateErr := 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) (*TaskExecution, error) {
execution, err := getTaskExecutionByTaskID(ctx, taskID)
if err == nil {
return execution, nil
}
meta := GetTaskMetaByAsynqTask(t.Type())
if meta == nil {
return nil, err
}
execution = &TaskExecution{
TaskID: taskID,
TaskType: meta.AsynqTask,
TaskName: meta.Name,
Status: TaskExecutionStatusRunning,
Retryable: meta.Retryable,
MaxRetry: meta.MaxRetry,
RetryCount: 0,
Payload: string(payload),
TriggeredBy: "schedule",
StartedAt: &now,
}
if createErr := 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 *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 := updateTaskExecution(ctx, execution); err != nil {
logger.ErrorF(ctx, "[TaskExecutor] 更新执行记录失败 taskID=%s: %v", execution.TaskID, err)
}
if shouldFlushTaskExecutionLog(ctx, execErr) {
if err := 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 *TaskExecution, result *TaskResult, execErr error) {
completedHandlersMutex.RLock()
defer completedHandlersMutex.RUnlock()
if len(taskCompletedHandlers) == 0 {
return
}
asyncCtx := context.WithoutCancel(ctx)
for _, handler := range taskCompletedHandlers {
h := handler
util.Go(func() {
h(asyncCtx, execution, result, execErr)
})
}
}
func shouldFlushTaskExecutionLog(ctx context.Context, execErr error) bool {
if isTerminalTaskExecutionError(execErr) {
return true
}
retryCount, hasRetryCount := asynq.GetRetryCount(ctx)
maxRetry, hasMaxRetry := asynq.GetMaxRetry(ctx)
if !hasRetryCount || !hasMaxRetry {
return true
}
return retryCount >= maxRetry
}
func isTerminalTaskExecutionError(execErr error) bool {
return execErr == nil || errors.Is(execErr, asynq.SkipRetry)
}
func handleFailedTask(ctx context.Context, execution *TaskExecution, t *asynq.Task, duration time.Duration, execErr error, span trace.Span) {
execution.Status = 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 *TaskExecution, t *asynq.Task, duration time.Duration, result *TaskResult) {
execution.Status = 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,402 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
import (
"context"
"errors"
"fmt"
"testing"
"time"
"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"
"github.com/Rain-kl/Wavelet/backend/pkg/testhelper"
)
// 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 := &TaskExecution{
TaskID: "log_test_001",
TaskType: testTaskType,
TaskName: "测试任务",
Status: TaskExecutionStatusRunning,
TriggeredBy: "manual",
}
err := CreateTaskExecution(ctx, execution)
require.NoError(t, err)
// 注入 taskID 并追加日志
ctx = withTaskID(ctx, "log_test_001")
AppendLog(ctx, "第一条日志")
AppendLog(ctx, "处理了 %d 条数据", 50)
// 验证日志
found, err := 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 := &TaskExecution{
TaskID: "process_success_001",
TaskType: testTaskType,
TaskName: "测试任务",
Status: TaskExecutionStatusPending,
Retryable: true,
MaxRetry: 3,
TriggeredBy: "manual",
}
err := CreateTaskExecution(ctx, execution)
require.NoError(t, err)
// 直接通过 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 := 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 := &TaskExecution{
TaskID: "process_fail_001",
TaskType: testTaskType,
TaskName: "测试任务",
Status: TaskExecutionStatusPending,
Retryable: true,
MaxRetry: 3,
TriggeredBy: "manual",
}
err := 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 := 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 := &TaskExecution{
TaskID: "complete_flush_001",
TaskType: testTaskType,
TaskName: "测试任务",
Status: TaskExecutionStatusRunning,
TriggeredBy: "manual",
}
err := 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 := GetTaskExecutionByTaskID(ctx, execution.TaskID)
require.NoError(t, err)
assert.Equal(t, TaskExecutionStatusSucceeded, found.Status)
assert.Contains(t, found.Log, "任务执行中的日志")
assert.Contains(t, found.Log, "任务执行成功")
}
func TestPermanentErrorIsTerminalForLogFlush(t *testing.T) {
assert.True(t, isTerminalTaskExecutionError(PermanentError("配置无效")))
assert.False(t, isTerminalTaskExecutionError(errors.New("temporary failure")))
}
func TestRetryTask(t *testing.T) {
cleanup := setupTest(t)
defer cleanup()
ctx := context.Background()
// 创建一条失败的执行记录(可重试)
now := time.Now()
execution := &TaskExecution{
TaskID: "retry_test_001",
TaskType: testTaskType,
TaskName: "测试任务",
Status: TaskExecutionStatusFailed,
Retryable: true,
MaxRetry: 3,
RetryCount: 0,
ErrorMessage: "首次执行失败",
StartedAt: &now,
FinishedAt: &now,
Duration: 100,
TriggeredBy: "manual",
}
err := 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 := GetTaskExecutionByTaskID(ctx, newTaskID)
require.NoError(t, err)
assert.Equal(t, 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 := GetTaskExecutionByID(ctx, execution.ID)
require.NoError(t, err)
assert.Equal(t, TaskExecutionStatusFailed, original.Status)
assert.Equal(t, 0, original.RetryCount)
}
func TestRetryTaskNotFailed(t *testing.T) {
cleanup := setupTest(t)
defer cleanup()
ctx := context.Background()
// 创建一条成功的记录
execution := &TaskExecution{
TaskID: "retry_not_failed_001",
TaskType: testTaskType,
TaskName: "测试任务",
Status: TaskExecutionStatusSucceeded,
Retryable: true,
MaxRetry: 3,
TriggeredBy: "manual",
}
err := 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 := &TaskExecution{
TaskID: "retry_not_allowed_001",
TaskType: testTaskType,
TaskName: "测试任务",
Status: TaskExecutionStatusFailed,
Retryable: false,
MaxRetry: 0,
TriggeredBy: "manual",
}
err := 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 driver_asynq_worker
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 方法编写业务逻辑,在方法内通过 driver_asynq_worker.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,100 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
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 {
handlerRegistryMutex.RLock()
defer handlerRegistryMutex.RUnlock()
keys := make([]string, 0, len(handlerRegistry))
for k := range handlerRegistry {
keys = append(keys, k)
}
return keys
}
@@ -0,0 +1,98 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
import (
"sync"
"testing"
)
func resetDispatchableTasks() {
dispatchableTasksMutex.Lock()
defer dispatchableTasksMutex.Unlock()
dispatchableTasks = nil
}
func TestRegisterTaskMeta_ThreadSafe(t *testing.T) {
resetDispatchableTasks()
defer resetDispatchableTasks()
var wg sync.WaitGroup
workers := 20
iterations := 100
for i := 0; i < workers; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
for j := 0; j < iterations; j++ {
RegisterTaskMeta(TaskMeta{
Type: "task_type_a",
AsynqTask: "asynq_task_a",
Name: "Task A",
})
RegisterTaskMeta(TaskMeta{
Type: "task_type_b",
AsynqTask: "asynq_task_b",
Name: "Task B",
})
}
}(i)
}
wg.Wait()
tasks := GetDispatchableTasks()
if len(tasks) != 2 {
t.Fatalf("expected 2 unique tasks, got %d", len(tasks))
}
}
func TestGetTaskMeta_Isolation(t *testing.T) {
resetDispatchableTasks()
defer resetDispatchableTasks()
RegisterTaskMeta(TaskMeta{
Type: "task_isolated",
AsynqTask: "asynq_isolated",
Name: "Original Name",
})
meta := GetTaskMeta("task_isolated")
if meta == nil {
t.Fatal("expected task to be found")
}
// Modify returned struct
meta.Name = "Modified Name"
// Fetch again and verify original data is unaffected
meta2 := GetTaskMeta("task_isolated")
if meta2.Name != "Original Name" {
t.Fatalf("expected name to be 'Original Name', got %s", meta2.Name)
}
}
func TestGetTaskMetaByAsynqTask_Isolation(t *testing.T) {
resetDispatchableTasks()
defer resetDispatchableTasks()
RegisterTaskMeta(TaskMeta{
Type: "task_isolated",
AsynqTask: "asynq_isolated",
Name: "Original Name",
})
meta := GetTaskMetaByAsynqTask("asynq_isolated")
if meta == nil {
t.Fatal("expected task to be found")
}
meta.Name = "Modified Name"
meta2 := GetTaskMetaByAsynqTask("asynq_isolated")
if meta2.Name != "Original Name" {
t.Fatalf("expected name to be 'Original Name', got %s", meta2.Name)
}
}
@@ -0,0 +1,28 @@
-- +goose Up
CREATE TABLE IF NOT EXISTS w_task_executions (
id BIGINT PRIMARY KEY,
task_id VARCHAR(128) NOT NULL UNIQUE,
task_type VARCHAR(64) NOT NULL,
task_name VARCHAR(128),
status VARCHAR(32) NOT NULL,
retryable BOOLEAN NOT NULL DEFAULT FALSE,
max_retry INTEGER NOT NULL DEFAULT 0,
retry_count INTEGER NOT NULL DEFAULT 0,
log TEXT,
error_message TEXT,
result TEXT,
started_at TIMESTAMPTZ,
finished_at TIMESTAMPTZ,
duration BIGINT,
payload TEXT,
triggered_by VARCHAR(32) NOT NULL DEFAULT 'system',
created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS idx_w_task_executions_task_type ON w_task_executions (task_type);
CREATE INDEX IF NOT EXISTS idx_w_task_executions_status ON w_task_executions (status);
CREATE INDEX IF NOT EXISTS idx_w_task_executions_started_at ON w_task_executions (started_at);
CREATE INDEX IF NOT EXISTS idx_w_task_executions_created_at ON w_task_executions (created_at);
-- +goose Down
DROP TABLE IF EXISTS w_task_executions;
@@ -0,0 +1,35 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
import (
"strings"
"github.com/hibiken/asynq"
)
const defaultPermanentErrorMessage = "任务无法继续执行"
type permanentTaskError struct {
message string
}
// PermanentError marks a safe domain message as a non-retryable task failure.
// It intentionally accepts no underlying error so Error never exposes provider,
// URL, header, response-body, or other sensitive implementation details.
func PermanentError(message string) error {
message = strings.TrimSpace(message)
if message == "" {
message = defaultPermanentErrorMessage
}
return &permanentTaskError{message: message}
}
func (e *permanentTaskError) Error() string {
return e.message
}
func (e *permanentTaskError) Unwrap() error {
return asynq.SkipRetry
}
@@ -0,0 +1,37 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
import (
"errors"
"testing"
"github.com/hibiken/asynq"
)
func TestPermanentError(t *testing.T) {
err := PermanentError("自定义永久错误")
if err == nil {
t.Fatal("PermanentError returned nil")
}
if err.Error() != "自定义永久错误" {
t.Fatalf("Error() = %q, want %q", err.Error(), "自定义永久错误")
}
if !errors.Is(err, asynq.SkipRetry) {
t.Fatal("PermanentError does not unwrap to asynq.SkipRetry")
}
}
func TestPermanentErrorDefaultMessage(t *testing.T) {
err := PermanentError(" ")
if err == nil {
t.Fatal("PermanentError returned nil")
}
if err.Error() != defaultPermanentErrorMessage {
t.Fatalf("Error() = %q, want default %q", err.Error(), defaultPermanentErrorMessage)
}
if !errors.Is(err, asynq.SkipRetry) {
t.Fatal("PermanentError does not unwrap to asynq.SkipRetry")
}
}
@@ -0,0 +1,256 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package driver_asynq_worker provides the Asynq worker driver plugin for Cordis.
package driver_asynq_worker
import (
"context"
"embed"
"errors"
"fmt"
"sync"
"time"
"github.com/hibiken/asynq"
"github.com/Rain-kl/Wavelet/backend/core"
)
const (
defaultConcurrency = 10
defaultShutdownTimeout = 10 * time.Second
)
//go:embed migrations/*.sql
var workerMigrations embed.FS
// Option configures the Asynq worker driver plugin.
type Option func(*Plugin)
// WithRedisOpt sets the Redis connection options for Asynq.
func WithRedisOpt(opt asynq.RedisConnOpt) Option {
return func(p *Plugin) {
p.redisOpt = opt
}
}
// WithConcurrency sets the worker concurrency limit.
func WithConcurrency(concurrency int) Option {
return func(p *Plugin) {
p.concurrency = concurrency
}
}
// WithQueues sets the queue priorities mapping.
func WithQueues(queues map[string]int) Option {
return func(p *Plugin) {
p.queues = queues
}
}
// WithStrictPriority sets whether to process queues strictly in priority order.
func WithStrictPriority(strict bool) Option {
return func(p *Plugin) {
p.strictPriority = strict
}
}
// WithShutdownTimeout sets the timeout for graceful worker shutdown.
func WithShutdownTimeout(d time.Duration) Option {
return func(p *Plugin) {
p.shutdownTimeout = d
}
}
// WithServer injects an existing Asynq server instance.
func WithServer(srv *asynq.Server) Option {
return func(p *Plugin) {
p.server = srv
}
}
// Plugin implements core.Plugin and core.Driver for Asynq background worker server.
type Plugin struct {
mu sync.RWMutex
redisOpt asynq.RedisConnOpt
concurrency int
queues map[string]int
strictPriority bool
shutdownTimeout time.Duration
server *asynq.Server
mux *asynq.ServeMux
running bool
coreCtx *core.Context
}
// New creates a new Asynq Worker driver plugin.
func New(opts ...Option) *Plugin {
p := &Plugin{
redisOpt: RedisOpt,
concurrency: defaultConcurrency,
shutdownTimeout: defaultShutdownTimeout,
queues: map[string]int{"default": 1},
}
for _, opt := range opts {
if opt != nil {
opt(p)
}
}
return p
}
// Name returns the unique plugin identifier.
func (p *Plugin) Name() string {
return "driver_asynq_worker"
}
// Apply mounts the Asynq Worker driver into the micro-kernel Context.
func (p *Plugin) Apply(ctx *core.Context) error {
p.mu.Lock()
p.coreCtx = ctx
p.mu.Unlock()
// Register migrations for w_task_executions table
ctx.Migrations().Register("driver_asynq_worker", workerMigrations)
ctx.OnDispose(func() error {
shutdownCtx, cancel := context.WithTimeout(context.Background(), p.shutdownTimeout)
defer cancel()
return p.Stop(shutdownCtx)
})
return ctx.RegisterDriver(p)
}
// Type returns DriverTypeWorker.
func (p *Plugin) Type() core.DriverType {
return core.DriverTypeWorker
}
// Start boots the Asynq worker server and starts processing background tasks.
func (p *Plugin) Start(_ context.Context) error {
p.mu.Lock()
defer p.mu.Unlock()
if p.running {
return nil
}
mux := asynq.NewServeMux()
if p.coreCtx != nil && p.coreCtx.Tasks() != nil {
for _, td := range p.coreCtx.Tasks().Tasks() {
handler, err := toAsynqHandler(td.Handler)
if err != nil {
return fmt.Errorf("driver_asynq_worker: invalid handler for task pattern %q: %w", td.Pattern, err)
}
mux.Handle(td.Pattern, handler)
}
}
for _, taskName := range GetRegisteredAsynqTasks() {
mux.HandleFunc(taskName, ProcessTask)
}
if p.server == nil {
p.server = asynq.NewServer(
p.redisOpt,
asynq.Config{
Concurrency: p.concurrency,
Queues: p.queues,
StrictPriority: p.strictPriority,
ShutdownTimeout: p.shutdownTimeout,
},
)
}
if err := p.server.Start(mux); err != nil {
return fmt.Errorf("driver_asynq_worker: start server failed: %w", err)
}
p.mux = mux
p.running = true
return nil
}
// Stop gracefully shuts down the Asynq worker server.
func (p *Plugin) Stop(_ context.Context) error {
p.mu.Lock()
defer p.mu.Unlock()
if !p.running {
return nil
}
p.running = false
if p.server != nil {
p.server.Stop()
p.server.Shutdown()
p.server = nil
}
return nil
}
// IsRunning returns whether the worker server is running.
func (p *Plugin) IsRunning() bool {
p.mu.RLock()
defer p.mu.RUnlock()
return p.running
}
// Server returns the underlying Asynq server instance.
func (p *Plugin) Server() *asynq.Server {
p.mu.RLock()
defer p.mu.RUnlock()
return p.server
}
// Mux returns the underlying Asynq serve mux.
func (p *Plugin) Mux() *asynq.ServeMux {
p.mu.RLock()
defer p.mu.RUnlock()
return p.mux
}
func toAsynqHandler(h any) (asynq.Handler, error) {
if h == nil {
return nil, errors.New("nil handler")
}
switch fn := h.(type) {
case asynq.HandlerFunc:
return fn, nil
case asynq.Handler:
return fn, nil
case TaskHandler:
return asynq.HandlerFunc(func(c context.Context, t *asynq.Task) error {
RegisterHandler(t.Type(), fn)
return ProcessTask(c, t)
}), nil
case func(context.Context, *asynq.Task) error:
return asynq.HandlerFunc(fn), nil
case func(context.Context, []byte) error:
return asynq.HandlerFunc(func(c context.Context, t *asynq.Task) error {
return fn(c, t.Payload())
}), nil
case func(context.Context) error:
return asynq.HandlerFunc(func(c context.Context, _ *asynq.Task) error {
return fn(c)
}), nil
case func([]byte) error:
return asynq.HandlerFunc(func(_ context.Context, t *asynq.Task) error {
return fn(t.Payload())
}), nil
case func() error:
return asynq.HandlerFunc(func(_ context.Context, _ *asynq.Task) error {
return fn()
}), nil
default:
return nil, fmt.Errorf("unsupported task handler type: %T", h)
}
}
@@ -0,0 +1,126 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
import (
"context"
"errors"
"fmt"
"strings"
"time"
"github.com/redis/go-redis/v9"
"github.com/Rain-kl/Wavelet/backend/pkg/idgen"
cachepkg "github.com/Rain-kl/Wavelet/backend/plugins/infra/cache"
db "github.com/Rain-kl/Wavelet/backend/plugins/infra/database"
)
const (
taskExecutionLogRedisKeyPrefix = "task:execution:log:"
taskExecutionLogExpiration = 24 * time.Hour
taskExecutionLogMaxLines = 1000
)
func taskExecutionLogRedisKey(taskID string) string {
return cachepkg.PrefixedKey(taskExecutionLogRedisKeyPrefix + taskID)
}
func createTaskExecution(ctx context.Context, execution *TaskExecution) error {
if execution.ID == 0 {
execution.ID = idgen.NextUint64ID()
}
return db.DB(ctx).Create(execution).Error
}
func updateTaskExecution(ctx context.Context, execution *TaskExecution) error {
return db.DB(ctx).Omit("log").Save(execution).Error
}
func getTaskExecutionByID(ctx context.Context, id uint64) (*TaskExecution, error) {
var execution TaskExecution
if err := db.DB(ctx).Where("id = ?", id).First(&execution).Error; err != nil {
return nil, err
}
_ = loadTaskExecutionLog(ctx, &execution)
return &execution, nil
}
func getTaskExecutionByTaskID(ctx context.Context, taskID string) (*TaskExecution, error) {
var execution TaskExecution
if err := db.DB(ctx).Where("task_id = ?", taskID).First(&execution).Error; err != nil {
return nil, err
}
_ = loadTaskExecutionLog(ctx, &execution)
return &execution, nil
}
func appendTaskExecutionLog(ctx context.Context, taskID string, logLine string) error {
if cachepkg.Redis == nil {
return errors.New("redis client is not initialized")
}
now := time.Now().Format("15:04:05")
line := fmt.Sprintf("[%s] %s\n", now, logLine)
key := taskExecutionLogRedisKey(taskID)
_, err := cachepkg.Redis.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
pipe.RPush(ctx, key, line)
pipe.LTrim(ctx, key, -taskExecutionLogMaxLines, -1)
pipe.Expire(ctx, key, taskExecutionLogExpiration)
return nil
})
if err != nil {
return fmt.Errorf("append task execution log to redis: %w", err)
}
return nil
}
func flushTaskExecutionLog(ctx context.Context, taskID string) error {
if cachepkg.Redis == nil {
return errors.New("redis client is not initialized")
}
key := taskExecutionLogRedisKey(taskID)
logLines, err := cachepkg.Redis.LRange(ctx, key, 0, -1).Result()
if err != nil {
return fmt.Errorf("get task execution log from redis: %w", err)
}
if len(logLines) == 0 {
return nil
}
logText := strings.Join(logLines, "")
result := db.DB(ctx).Model(&TaskExecution{}).
Where("task_id = ?", taskID).
Update("log", logText)
if result.Error != nil {
return fmt.Errorf("persist task execution log: %w", result.Error)
}
if result.RowsAffected == 0 {
return fmt.Errorf("persist task execution log: task %q not found", taskID)
}
if err := cachepkg.Redis.Del(ctx, key).Err(); err != nil {
return fmt.Errorf("delete persisted task execution log from redis: %w", err)
}
return nil
}
func loadTaskExecutionLog(ctx context.Context, execution *TaskExecution) error {
if cachepkg.Redis == nil {
return nil
}
logLines, err := cachepkg.Redis.LRange(ctx, taskExecutionLogRedisKey(execution.TaskID), 0, -1).Result()
if err != nil {
return fmt.Errorf("get task execution log from redis: %w", err)
}
if len(logLines) == 0 {
return nil
}
execution.Log = strings.Join(logLines, "")
return nil
}
@@ -0,0 +1,155 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
import (
"context"
"errors"
"time"
"gorm.io/gorm"
db "github.com/Rain-kl/Wavelet/backend/plugins/infra/database"
)
// TaskExecutionStatus 任务执行状态
//
//nolint:revive
type TaskExecutionStatus string
// Task execution status constants.
const (
TaskExecutionStatusPending TaskExecutionStatus = "pending"
TaskExecutionStatusRunning TaskExecutionStatus = "running"
TaskExecutionStatusSucceeded TaskExecutionStatus = "succeeded"
TaskExecutionStatusFailed TaskExecutionStatus = "failed"
)
// TaskExecution 任务执行记录
//
//nolint:revive
type TaskExecution struct {
ID uint64 `json:"id,string" gorm:"primaryKey"`
TaskID string `json:"task_id" gorm:"size:128;uniqueIndex;not null"`
TaskType string `json:"task_type" gorm:"size:64;index;not null"`
TaskName string `json:"task_name" gorm:"size:128"`
Status TaskExecutionStatus `json:"status" gorm:"size:32;index;not null"`
Retryable bool `json:"retryable" gorm:"not null;default:false"`
MaxRetry int `json:"max_retry" gorm:"not null;default:0"`
RetryCount int `json:"retry_count" gorm:"not null;default:0"`
Log string `json:"log" gorm:"type:text"`
ErrorMessage string `json:"error_message" gorm:"type:text"`
Result string `json:"result" gorm:"type:text"`
StartedAt *time.Time `json:"started_at" gorm:"index"`
FinishedAt *time.Time `json:"finished_at"`
Duration int64 `json:"duration" gorm:"comment:耗时毫秒"`
Payload string `json:"payload" gorm:"type:text"`
TriggeredBy string `json:"triggered_by" gorm:"size:32;not null;default:system"`
CreatedAt time.Time `json:"created_at" gorm:"autoCreateTime;index"`
UpdatedAt time.Time `json:"updated_at" gorm:"autoUpdateTime"`
}
// TableName 表名
func (TaskExecution) TableName() string {
return "w_task_executions"
}
// CreateTaskExecution 创建任务执行记录
func CreateTaskExecution(ctx context.Context, exec *TaskExecution) error {
return db.DB(ctx).Create(exec).Error
}
// GetTaskExecutionByTaskID 根据 TaskID 查询执行记录
func GetTaskExecutionByTaskID(ctx context.Context, taskID string) (*TaskExecution, error) {
var exec TaskExecution
if err := db.DB(ctx).Where("task_id = ?", taskID).First(&exec).Error; err != nil {
return nil, err
}
_ = loadTaskExecutionLog(ctx, &exec)
return &exec, nil
}
// GetTaskExecutionByID 根据主键 ID 查询执行记录
func GetTaskExecutionByID(ctx context.Context, id uint64) (*TaskExecution, error) {
var exec TaskExecution
if err := db.DB(ctx).Where("id = ?", id).First(&exec).Error; err != nil {
return nil, err
}
_ = loadTaskExecutionLog(ctx, &exec)
return &exec, nil
}
// GetLatestTaskExecutionByTaskType 获取指定任务类型的最新执行记录
func GetLatestTaskExecutionByTaskType(ctx context.Context, taskType string) (*TaskExecution, bool, error) {
var exec TaskExecution
err := db.DB(ctx).Where("task_type = ?", taskType).Order("id DESC").First(&exec).Error
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, false, nil
}
return nil, false, err
}
return &exec, true, nil
}
// TaskExecutionCleanupStats describes task execution log cleanup results.
//
//nolint:revive
type TaskExecutionCleanupStats struct {
HighFrequencyDeleted int64
LowFrequencyDeleted int64
}
// CleanupTaskExecutionLogs removes finished task execution logs according to frequency-based retention.
func CleanupTaskExecutionLogs(ctx context.Context, now time.Time) (TaskExecutionCleanupStats, error) {
const (
frequencyWindowDays = 30
highFrequencyThreshold = frequencyWindowDays
)
frequencyWindowStart := now.AddDate(0, 0, -frequencyWindowDays)
highFrequencyCutoff := now.AddDate(0, 0, -3)
lowFrequencyCutoff := now.AddDate(0, 0, -30)
terminalStatuses := []TaskExecutionStatus{TaskExecutionStatusSucceeded, TaskExecutionStatusFailed}
var highFrequencyTaskTypes []string
if err := db.DB(ctx).
Model(&TaskExecution{}).
Select("task_type").
Where("created_at >= ?", frequencyWindowStart).
Group("task_type").
Having("COUNT(*) > ?", highFrequencyThreshold).
Pluck("task_type", &highFrequencyTaskTypes).Error; err != nil {
return TaskExecutionCleanupStats{}, err
}
var highFrequencyDeleted int64
if len(highFrequencyTaskTypes) > 0 {
highFrequencyResult := db.DB(ctx).
Where("status IN ?", terminalStatuses).
Where("created_at < ?", highFrequencyCutoff).
Where("task_type IN ?", highFrequencyTaskTypes).
Delete(&TaskExecution{})
if highFrequencyResult.Error != nil {
return TaskExecutionCleanupStats{}, highFrequencyResult.Error
}
highFrequencyDeleted = highFrequencyResult.RowsAffected
}
lowFrequencyQuery := db.DB(ctx).
Where("status IN ?", terminalStatuses).
Where("created_at < ?", lowFrequencyCutoff)
if len(highFrequencyTaskTypes) > 0 {
lowFrequencyQuery = lowFrequencyQuery.Where("task_type NOT IN ?", highFrequencyTaskTypes)
}
lowFrequencyResult := lowFrequencyQuery.Delete(&TaskExecution{})
if lowFrequencyResult.Error != nil {
return TaskExecutionCleanupStats{}, lowFrequencyResult.Error
}
return TaskExecutionCleanupStats{
HighFrequencyDeleted: highFrequencyDeleted,
LowFrequencyDeleted: lowFrequencyResult.RowsAffected,
}, nil
}
@@ -0,0 +1,115 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
import (
"github.com/hibiken/asynq"
"github.com/redis/go-redis/v9"
"github.com/redis/go-redis/v9/maintnotifications"
"github.com/Rain-kl/Wavelet/backend/pkg/config"
)
type redisClientConnOpt struct {
options redis.Options
}
func (opt redisClientConnOpt) MakeRedisClient() interface{} {
return redis.NewClient(&opt.options)
}
type redisClusterConnOpt struct {
options redis.ClusterOptions
}
func (opt redisClusterConnOpt) MakeRedisClient() interface{} {
return redis.NewClusterClient(&opt.options)
}
type redisFailoverConnOpt struct {
options redis.FailoverOptions
maintNotificationsEnabled bool
}
func (opt redisFailoverConnOpt) MakeRedisClient() interface{} {
client := redis.NewFailoverClient(&opt.options)
// go-redis v9.16 does not expose maintenance notification settings on
// FailoverOptions, so apply the configured mode before the client is used.
client.Options().MaintNotificationsConfig = maintNotificationsConfig(opt.maintNotificationsEnabled)
return client
}
func maintNotificationsConfig(enabled bool) *maintnotifications.Config {
mode := maintnotifications.ModeDisabled
if enabled {
mode = maintnotifications.ModeAuto
}
return &maintnotifications.Config{Mode: mode}
}
// 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 redisClusterConnOpt{
options: redis.ClusterOptions{
Addrs: addrs,
Username: cfg.Username,
Password: cfg.Password,
MaintNotificationsConfig: maintNotificationsConfig(cfg.MaintNotifications),
},
}
}
if cfg.MasterName != "" {
return redisFailoverConnOpt{
maintNotificationsEnabled: cfg.MaintNotifications,
options: redis.FailoverOptions{
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 redisClientConnOpt{
options: redis.Options{
Addr: addr,
Username: cfg.Username,
Password: cfg.Password,
DB: cfg.DB,
PoolSize: cfg.PoolSize,
MaintNotificationsConfig: maintNotificationsConfig(cfg.MaintNotifications),
},
}
}
// PrefixedQueue 返回带前缀的队列名,用于 Cluster 模式隔离
func PrefixedQueue(queue string) string {
prefix := config.Config.Redis.KeyPrefix
if prefix == "" {
return queue
}
return prefix + queue
}
@@ -0,0 +1,59 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_asynq_worker
import (
"testing"
"github.com/Rain-kl/Wavelet/backend/pkg/config"
"github.com/redis/go-redis/v9/maintnotifications"
)
func TestMaintNotificationsConfig(t *testing.T) {
tests := []struct {
name string
enabled bool
wantMode maintnotifications.Mode
}{
{
name: "disabled mode when flag is false",
enabled: false,
wantMode: maintnotifications.ModeDisabled,
},
{
name: "auto mode when flag is true",
enabled: true,
wantMode: maintnotifications.ModeAuto,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
cfg := maintNotificationsConfig(tt.enabled)
if cfg == nil {
t.Fatalf("maintNotificationsConfig(%v) = nil, want non-nil", tt.enabled)
}
if cfg.Mode != tt.wantMode {
t.Fatalf("maintNotificationsConfig(%v).Mode = %v, want %v", tt.enabled, cfg.Mode, tt.wantMode)
}
})
}
}
func TestPrefixedQueue(t *testing.T) {
oldPrefix := config.Config.Redis.KeyPrefix
defer func() {
config.Config.Redis.KeyPrefix = oldPrefix
}()
config.Config.Redis.KeyPrefix = "test:"
if got := PrefixedQueue("default"); got != "test:default" {
t.Fatalf("PrefixedQueue() = %q, want %q", got, "test:default")
}
config.Config.Redis.KeyPrefix = ""
if got := PrefixedQueue("default"); got != "default" {
t.Fatalf("PrefixedQueue() = %q, want %q", got, "default")
}
}
@@ -0,0 +1,123 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_http
import (
"context"
"errors"
"log"
"net"
"net/http"
"os"
"os/signal"
"strconv"
"syscall"
"time"
"github.com/Rain-kl/Wavelet/backend/pkg/config"
"github.com/Rain-kl/Wavelet/backend/pkg/trace"
"github.com/Rain-kl/Wavelet/backend/pkg/util"
"github.com/Rain-kl/Wavelet/backend/plugins/domain/auth"
"github.com/Rain-kl/Wavelet/backend/plugins/domain/risk_control"
"github.com/gin-contrib/sessions"
"github.com/gin-contrib/sessions/redis"
"github.com/gin-gonic/gin"
"go.opentelemetry.io/contrib/instrumentation/github.com/gin-gonic/gin/otelgin"
)
// BuildEngine 构建并初始化 Gin 路由引擎及全部中间件和路由
func BuildEngine() (*gin.Engine, error) {
// 运行模式
if config.Config.App.IsProduction() {
gin.SetMode(gin.ReleaseMode)
}
// 初始化路由
r := gin.New()
r.Use(gin.Recovery())
r.Use(corsMiddleware())
cfg := config.Config.Redis
addrs := cfg.Addrs
sessionAddr := "localhost:6379"
if len(addrs) > 0 {
sessionAddr = addrs[0]
}
sessionStore, err := redis.NewStoreWithDB(
cfg.MinIdleConn,
"tcp",
sessionAddr,
cfg.Username,
cfg.Password,
strconv.Itoa(cfg.DB),
[]byte(config.Config.App.SessionSecret),
)
if err != nil {
return nil, err
}
// 设置 Session Redis Key 前缀
if cfg.KeyPrefix != "" {
if err := redis.SetKeyPrefix(sessionStore, cfg.KeyPrefix+"session:"); err != nil {
log.Printf("[API] set session key prefix failed: %v\n", err)
}
}
sessionStore.Options(auth.GetSessionOptions(config.Config.App.SessionAge))
r.Use(sessions.Sessions(config.Config.App.SessionCookieName, sessionStore))
// 补充中间件
r.Use(otelgin.Middleware(config.Config.App.AppName), errorHandlerMiddleware(), loggerMiddleware(), risk_control.Middleware())
return r, nil
}
// Serve 启动 HTTP API 服务。onStarted 仅会在 HTTP 地址成功绑定后调用。
func Serve(onStarted func()) {
r, err := BuildEngine()
if err != nil {
log.Fatalf("[API] init session store failed: %v\n", err)
}
srv := &http.Server{
Addr: config.Config.App.Addr,
Handler: r,
ReadHeaderTimeout: 10 * time.Second,
}
listener, err := (&net.ListenConfig{}).Listen(context.Background(), "tcp", config.Config.App.Addr)
if err != nil {
log.Fatalf("[API] server failed to listen on %s: %v\n", config.Config.App.Addr, err)
}
if onStarted != nil {
onStarted()
}
util.Go(func() {
log.Printf("[API] server listening on %s\n", config.Config.App.Addr)
if err := srv.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
log.Fatalf("[API] server failed: %v\n", err)
}
})
quit := make(chan os.Signal, 1)
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
<-quit
shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Duration(config.Config.App.GracefulShutdownTimeout)*time.Second)
trace.Shutdown(shutdownCtx)
if err := srv.Shutdown(shutdownCtx); err != nil {
log.Printf("[API] server forced to shutdown: %v\n", err)
cancel()
os.Exit(1)
}
cancel()
log.Println("[API] server exited")
}
@@ -0,0 +1,111 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package driver_http 提供 HTTP 路由中间件与服务启动
package driver_http
import (
"context"
"net/http"
"strconv"
"strings"
"time"
"github.com/Rain-kl/Wavelet/backend/pkg/config"
"github.com/Rain-kl/Wavelet/backend/pkg/logger"
"github.com/Rain-kl/Wavelet/backend/pkg/response"
otel_trace "github.com/Rain-kl/Wavelet/backend/pkg/trace"
database "github.com/Rain-kl/Wavelet/backend/plugins/infra/database"
"github.com/gin-gonic/gin"
"go.opentelemetry.io/otel/codes"
"go.opentelemetry.io/otel/trace"
)
func loggerMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
// 初始化 Trace
ctx, span := otel_trace.Start(c.Request.Context(), "LoggerMiddleware")
defer span.End()
// 开始计时
start := time.Now()
// 记录请求路径和 Query
path := c.Request.URL.Path
raw := c.Request.URL.RawQuery
if raw != "" {
path = path + "?" + raw
}
// 执行请求
c.Next()
// 停止计时
end := time.Now()
latency := end.Sub(start)
// 打印日志
// 排除健康检查接口
healthPath := config.Config.App.APIPrefix + "/health"
if c.Request.URL.Path != healthPath {
logger.InfoF(
ctx,
"[LoggerMiddleware] %s %s\nStartTime: %s\nEndTime: %s\nLatency: %d\nClientIP: %s\nResponse: %d %d",
c.Request.Method,
path,
start.Format(time.RFC3339),
end.Format(time.RFC3339),
latency.Milliseconds(),
c.ClientIP(),
c.Writer.Status(),
c.Writer.Size(),
)
}
// 设置 Span 状态
if c.Writer.Status() >= http.StatusBadRequest {
span := trace.SpanFromContext(ctx)
span.SetStatus(codes.Error, strconv.Itoa(c.Writer.Status()))
}
}
}
func isOriginAllowed(ctx context.Context, origin string) bool {
var val string
if err := database.DB(ctx).Table("w_system_configs").Where("key = ?", "server_address").Pluck("value", &val).Error; err != nil || val == "" {
return false
}
allowedOrigins := strings.Split(val, ",")
for _, allowed := range allowedOrigins {
allowed = strings.TrimRight(strings.TrimSpace(allowed), "/")
if allowed != "" && strings.EqualFold(allowed, origin) {
return true
}
}
return false
}
func corsMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
origin := c.Request.Header.Get("Origin")
if origin != "" && isOriginAllowed(c.Request.Context(), origin) {
c.Writer.Header().Set("Access-Control-Allow-Origin", origin)
c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, accept, origin, Cache-Control, X-Requested-With, X-Access-Token, X-Cap-Token")
c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE, PATCH")
}
if c.Request.Method == "OPTIONS" {
c.AbortWithStatus(http.StatusNoContent)
return
}
c.Next()
}
}
// errorHandlerMiddleware 委托给 response.ErrorHandlerMiddleware,保持路由层单一入口。
func errorHandlerMiddleware() gin.HandlerFunc {
return response.ErrorHandlerMiddleware()
}
@@ -0,0 +1,100 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package driver_http
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/Rain-kl/Wavelet/backend/pkg/testhelper"
"github.com/gin-gonic/gin"
)
func TestCORSMiddleware(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
gin.SetMode(gin.TestMode)
clearConfigCache := func() {}
t.Run("missing server_address configuration returns no CORS headers", func(t *testing.T) {
clearConfigCache()
if err := dbConn.Table("w_system_configs").Where("key = ?", "server_address").Update("value", "").Error; err != nil {
t.Fatalf("failed to update config: %v", err)
}
r := gin.New()
r.Use(corsMiddleware())
r.GET("/test", func(c *gin.Context) {
c.String(http.StatusOK, "ok")
})
req, _ := http.NewRequest(http.MethodGet, "/test", nil)
req.Header.Set("Origin", "http://attacker.com")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Header().Get("Access-Control-Allow-Origin") != "" {
t.Errorf("expected no Access-Control-Allow-Origin, got %s", w.Header().Get("Access-Control-Allow-Origin"))
}
})
t.Run("configured server_address allows exact origin match", func(t *testing.T) {
if err := dbConn.Table("w_system_configs").Where("key = ?", "server_address").Update("value", "http://trusted.com").Error; err != nil {
t.Fatalf("failed to update config: %v", err)
}
r := gin.New()
r.Use(corsMiddleware())
r.GET("/test", func(c *gin.Context) {
c.String(http.StatusOK, "ok")
})
// Trusted origin
req, _ := http.NewRequest(http.MethodGet, "/test", nil)
req.Header.Set("Origin", "http://trusted.com")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Header().Get("Access-Control-Allow-Origin") != "http://trusted.com" {
t.Errorf("expected Access-Control-Allow-Origin http://trusted.com, got %s", w.Header().Get("Access-Control-Allow-Origin"))
}
if w.Header().Get("Access-Control-Allow-Credentials") != "true" {
t.Errorf("expected Access-Control-Allow-Credentials true, got %s", w.Header().Get("Access-Control-Allow-Credentials"))
}
// Untrusted origin
req, _ = http.NewRequest(http.MethodGet, "/test", nil)
req.Header.Set("Origin", "http://attacker.com")
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Header().Get("Access-Control-Allow-Origin") != "" {
t.Errorf("expected no Access-Control-Allow-Origin for attacker, got %s", w.Header().Get("Access-Control-Allow-Origin"))
}
})
t.Run("preflight OPTIONS request responds with 204", func(t *testing.T) {
if err := dbConn.Table("w_system_configs").Where("key = ?", "server_address").Update("value", "http://trusted.com").Error; err != nil {
t.Fatalf("failed to update config: %v", err)
}
r := gin.New()
r.Use(corsMiddleware())
req, _ := http.NewRequest(http.MethodOptions, "/test", nil)
req.Header.Set("Origin", "http://trusted.com")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusNoContent {
t.Errorf("expected status 204, got %d", w.Code)
}
if w.Header().Get("Access-Control-Allow-Methods") == "" {
t.Error("expected Access-Control-Allow-Methods header")
}
})
}
@@ -0,0 +1,249 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package driver_http provides the Gin HTTP web server driver plugin for Cordis.
package driver_http
import (
"context"
"errors"
"fmt"
"net"
"net/http"
"sync"
"time"
"github.com/gin-gonic/gin"
"github.com/Rain-kl/Wavelet/backend/core"
"github.com/Rain-kl/Wavelet/backend/pkg/util"
)
const (
defaultAddr = ":8080"
defaultReadHeaderTimeout = 10 * time.Second
defaultShutdownTimeout = 5 * time.Second
)
// Option configures the HTTP driver plugin.
type Option func(*Plugin)
// WithAddr sets the TCP address for the HTTP server to listen on.
func WithAddr(addr string) Option {
return func(p *Plugin) {
p.addr = addr
}
}
// WithEngine sets a pre-configured Gin engine for the HTTP server.
func WithEngine(engine *gin.Engine) Option {
return func(p *Plugin) {
p.engine = engine
}
}
// WithReadHeaderTimeout sets the ReadHeaderTimeout for http.Server.
func WithReadHeaderTimeout(d time.Duration) Option {
return func(p *Plugin) {
p.readHeaderTimeout = d
}
}
// WithShutdownTimeout sets the fallback timeout for graceful server shutdown.
func WithShutdownTimeout(d time.Duration) Option {
return func(p *Plugin) {
p.shutdownTimeout = d
}
}
// Plugin implements core.Plugin and core.Driver for Gin HTTP Web Server.
type Plugin struct {
mu sync.RWMutex
addr string
engine *gin.Engine
server *http.Server
listener net.Listener
running bool
readHeaderTimeout time.Duration
shutdownTimeout time.Duration
coreCtx *core.Context
}
// New creates a new Gin HTTP server driver plugin.
func New(opts ...Option) *Plugin {
p := &Plugin{
addr: defaultAddr,
readHeaderTimeout: defaultReadHeaderTimeout,
shutdownTimeout: defaultShutdownTimeout,
}
for _, opt := range opts {
if opt != nil {
opt(p)
}
}
return p
}
// Name returns the unique plugin identifier.
func (p *Plugin) Name() string {
return "driver_http"
}
// Apply mounts the HTTP driver plugin into the micro-kernel Context.
func (p *Plugin) Apply(ctx *core.Context) error {
p.mu.Lock()
p.coreCtx = ctx
p.mu.Unlock()
ctx.OnDispose(func() error {
shutdownCtx, cancel := context.WithTimeout(context.Background(), p.shutdownTimeout)
defer cancel()
return p.Stop(shutdownCtx)
})
return ctx.RegisterDriver(p)
}
// Type returns DriverTypeHTTP.
func (p *Plugin) Type() core.DriverType {
return core.DriverTypeHTTP
}
// Start boots the Gin HTTP server, binds routes collected from ctx.Router(), and starts listening.
func (p *Plugin) Start(ctx context.Context) error {
p.mu.Lock()
defer p.mu.Unlock()
if p.running {
return nil
}
if p.engine == nil {
p.engine = gin.New()
}
// Mount routes collected in Context RouterExtension
if p.coreCtx != nil && p.coreCtx.Router() != nil {
for _, rd := range p.coreCtx.Router().Routes() {
allHandlers := make([]gin.HandlerFunc, 0, len(rd.Middlewares)+len(rd.Handlers))
for _, m := range rd.Middlewares {
gh, err := toGinHandler(m)
if err != nil {
return fmt.Errorf("driver_http: invalid middleware for route %s %s: %w", rd.Method, rd.Path, err)
}
allHandlers = append(allHandlers, gh)
}
for _, h := range rd.Handlers {
gh, err := toGinHandler(h)
if err != nil {
return fmt.Errorf("driver_http: invalid handler for route %s %s: %w", rd.Method, rd.Path, err)
}
allHandlers = append(allHandlers, gh)
}
p.engine.Handle(rd.Method, rd.Path, allHandlers...)
}
}
p.server = &http.Server{
Addr: p.addr,
Handler: p.engine,
ReadHeaderTimeout: p.readHeaderTimeout,
}
lc := &net.ListenConfig{}
listener, err := lc.Listen(ctx, "tcp", p.addr)
if err != nil {
return fmt.Errorf("driver_http: listen on %s failed: %w", p.addr, err)
}
p.listener = listener
p.addr = listener.Addr().String()
p.running = true
srv := p.server
util.Go(func() {
if serveErr := srv.Serve(listener); serveErr != nil && !errors.Is(serveErr, http.ErrServerClosed) {
_ = serveErr
}
})
return nil
}
// Stop gracefully stops the HTTP server.
//
//nolint:contextcheck
func (p *Plugin) Stop(ctx context.Context) error {
p.mu.Lock()
defer p.mu.Unlock()
if !p.running {
return nil
}
p.running = false
var err error
if p.server != nil {
if ctx == nil {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(context.Background(), p.shutdownTimeout)
defer cancel()
}
err = p.server.Shutdown(ctx)
}
if p.listener != nil {
_ = p.listener.Close()
p.listener = nil
}
return err
}
// Addr returns the current listening address (or configured address if not yet started).
func (p *Plugin) Addr() string {
p.mu.RLock()
defer p.mu.RUnlock()
return p.addr
}
// Engine returns the underlying Gin engine.
func (p *Plugin) Engine() *gin.Engine {
p.mu.RLock()
defer p.mu.RUnlock()
return p.engine
}
// IsRunning returns whether the HTTP server is currently running.
func (p *Plugin) IsRunning() bool {
p.mu.RLock()
defer p.mu.RUnlock()
return p.running
}
func toGinHandler(h any) (gin.HandlerFunc, error) {
if h == nil {
return nil, errors.New("nil handler")
}
switch fn := h.(type) {
case gin.HandlerFunc:
return fn, nil
case func(*gin.Context):
return gin.HandlerFunc(fn), nil
case http.HandlerFunc:
return gin.WrapF(fn), nil
case func(http.ResponseWriter, *http.Request):
return gin.WrapF(fn), nil
case http.Handler:
return gin.WrapH(fn), nil
default:
return nil, fmt.Errorf("unsupported handler type: %T", h)
}
}
+291
View File
@@ -0,0 +1,291 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package drivers_test
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"sync/atomic"
"testing"
"time"
"github.com/alicebob/miniredis/v2"
"github.com/gin-gonic/gin"
"github.com/hibiken/asynq"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/Rain-kl/Wavelet/backend/core"
"github.com/Rain-kl/Wavelet/backend/core/extpoints"
"github.com/Rain-kl/Wavelet/backend/plugins/drivers/driver_asynq_cron"
"github.com/Rain-kl/Wavelet/backend/plugins/drivers/driver_asynq_worker"
"github.com/Rain-kl/Wavelet/backend/plugins/drivers/driver_http"
)
func init() {
gin.SetMode(gin.TestMode)
}
func TestHTTPDriverLifecycle(t *testing.T) {
ctx := core.NewContext(context.Background())
var globalMiddlewareCalled atomic.Bool
var groupMiddlewareCalled atomic.Bool
// Register global middleware
ctx.Router().Use(func(c *gin.Context) {
globalMiddlewareCalled.Store(true)
c.Next()
})
// Register standard gin handler
ctx.Router().GET("/ping", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"message": "pong"})
})
// Register route group with middleware
v1 := ctx.Router().Group("/api/v1", func(c *gin.Context) {
groupMiddlewareCalled.Store(true)
c.Next()
})
v1.POST("/echo", func(c *gin.Context) {
var req map[string]any
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, req)
})
// Register standard http.HandlerFunc
v1.GET("/legacy", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("legacy response"))
}))
// Create and apply HTTP driver with dynamic port
httpPlugin := driver_http.New(driver_http.WithAddr("127.0.0.1:0"))
require.Equal(t, "driver_http", httpPlugin.Name())
err := httpPlugin.Apply(ctx)
require.NoError(t, err)
// Verify driver registered in context
d, ok := ctx.Driver(core.DriverTypeHTTP)
require.True(t, ok)
require.Equal(t, core.DriverTypeHTTP, d.Type())
// Start HTTP Server
err = d.Start(context.Background())
require.NoError(t, err)
assert.True(t, httpPlugin.IsRunning())
addr := httpPlugin.Addr()
require.NotEmpty(t, addr)
// Verify GET /ping
resp, err := http.Get(fmt.Sprintf("http://%s/ping", addr))
require.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
body, err := io.ReadAll(resp.Body)
require.NoError(t, err)
assert.Contains(t, string(body), "pong")
assert.True(t, globalMiddlewareCalled.Load())
// Verify POST /api/v1/echo
echoPayload := []byte(`{"title":"test-echo","value":42}`)
respEcho, err := http.Post(fmt.Sprintf("http://%s/api/v1/echo", addr), "application/json", bytes.NewReader(echoPayload))
require.NoError(t, err)
defer respEcho.Body.Close()
assert.Equal(t, http.StatusOK, respEcho.StatusCode)
var echoResult map[string]any
err = json.NewDecoder(respEcho.Body).Decode(&echoResult)
require.NoError(t, err)
assert.Equal(t, "test-echo", echoResult["title"])
assert.Equal(t, float64(42), echoResult["value"])
assert.True(t, groupMiddlewareCalled.Load())
// Verify GET /api/v1/legacy
respLegacy, err := http.Get(fmt.Sprintf("http://%s/api/v1/legacy", addr))
require.NoError(t, err)
defer respLegacy.Body.Close()
assert.Equal(t, http.StatusOK, respLegacy.StatusCode)
bodyLegacy, err := io.ReadAll(respLegacy.Body)
require.NoError(t, err)
assert.Equal(t, "legacy response", string(bodyLegacy))
// Stop HTTP Server
stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
err = d.Stop(stopCtx)
require.NoError(t, err)
assert.False(t, httpPlugin.IsRunning())
// Idempotent Stop
err = d.Stop(stopCtx)
require.NoError(t, err)
}
func TestAsynqWorkerDriverLifecycle(t *testing.T) {
mr, err := miniredis.Run()
require.NoError(t, err)
defer mr.Close()
ctx := core.NewContext(context.Background())
var emailTaskProcessed atomic.Bool
var simpleTaskProcessed atomic.Bool
// Register tasks with various handler signatures
ctx.Tasks().Register("email:send", func(c context.Context, payload []byte) error {
var data map[string]string
if err := json.Unmarshal(payload, &data); err != nil {
return err
}
if data["to"] == "user@example.com" {
emailTaskProcessed.Store(true)
}
return nil
}, extpoints.WithTaskConcurrency(2))
ctx.Tasks().Register("maintenance:cleanup", func() error {
simpleTaskProcessed.Store(true)
return nil
})
workerPlugin := driver_asynq_worker.New(
driver_asynq_worker.WithRedisOpt(asynq.RedisClientOpt{Addr: mr.Addr()}),
driver_asynq_worker.WithConcurrency(2),
driver_asynq_worker.WithShutdownTimeout(1*time.Second),
)
require.Equal(t, "driver_asynq_worker", workerPlugin.Name())
err = workerPlugin.Apply(ctx)
require.NoError(t, err)
// Verify driver registered
d, ok := ctx.Driver(core.DriverTypeWorker)
require.True(t, ok)
require.Equal(t, core.DriverTypeWorker, d.Type())
// Start Worker Driver
err = d.Start(context.Background())
require.NoError(t, err)
assert.True(t, workerPlugin.IsRunning())
// Enqueue tasks using asynq.Client
client := asynq.NewClient(asynq.RedisClientOpt{Addr: mr.Addr()})
defer client.Close()
emailPayload, _ := json.Marshal(map[string]string{"to": "user@example.com"})
_, err = client.Enqueue(asynq.NewTask("email:send", emailPayload))
require.NoError(t, err)
_, err = client.Enqueue(asynq.NewTask("maintenance:cleanup", nil))
require.NoError(t, err)
// Wait for worker to consume tasks
require.Eventually(t, func() bool {
return emailTaskProcessed.Load() && simpleTaskProcessed.Load()
}, 3*time.Second, 50*time.Millisecond)
// Stop Worker Driver
stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
err = d.Stop(stopCtx)
require.NoError(t, err)
assert.False(t, workerPlugin.IsRunning())
// Idempotent Stop
err = d.Stop(stopCtx)
require.NoError(t, err)
}
func TestAsynqCronDriverLifecycle(t *testing.T) {
mr, err := miniredis.Run()
require.NoError(t, err)
defer mr.Close()
ctx := core.NewContext(context.Background())
// Register cron schedules
ctx.Schedules().RegisterCron("@every 1s", "sync:stats", map[string]string{"type": "daily"},
extpoints.WithScheduleOption("queue", "default"),
extpoints.WithScheduleOption("retry", 3),
)
ctx.Schedules().RegisterCron("0 0 * * *", "report:generate", "raw_payload")
cronPlugin := driver_asynq_cron.New(
driver_asynq_cron.WithRedisOpt(asynq.RedisClientOpt{Addr: mr.Addr()}),
driver_asynq_cron.WithLocation(time.UTC),
)
require.Equal(t, "driver_asynq_cron", cronPlugin.Name())
err = cronPlugin.Apply(ctx)
require.NoError(t, err)
// Verify driver registered
d, ok := ctx.Driver(core.DriverTypeScheduler)
require.True(t, ok)
require.Equal(t, core.DriverTypeScheduler, d.Type())
// Start Cron Scheduler Driver
err = d.Start(context.Background())
require.NoError(t, err)
assert.True(t, cronPlugin.IsRunning())
// Stop Cron Scheduler Driver
stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
err = d.Stop(stopCtx)
require.NoError(t, err)
assert.False(t, cronPlugin.IsRunning())
// Idempotent Stop
err = d.Stop(stopCtx)
require.NoError(t, err)
}
func TestMultipleDriversInContext(t *testing.T) {
mr, err := miniredis.Run()
require.NoError(t, err)
defer mr.Close()
ctx := core.NewContext(context.Background())
httpPlugin := driver_http.New(driver_http.WithAddr("127.0.0.1:0"))
workerPlugin := driver_asynq_worker.New(driver_asynq_worker.WithRedisOpt(asynq.RedisClientOpt{Addr: mr.Addr()}))
cronPlugin := driver_asynq_cron.New(driver_asynq_cron.WithRedisOpt(asynq.RedisClientOpt{Addr: mr.Addr()}))
require.NoError(t, httpPlugin.Apply(ctx))
require.NoError(t, workerPlugin.Apply(ctx))
require.NoError(t, cronPlugin.Apply(ctx))
drivers := ctx.Drivers()
assert.Len(t, drivers, 3)
_, ok := ctx.Driver(core.DriverTypeHTTP)
assert.True(t, ok)
_, ok = ctx.Driver(core.DriverTypeWorker)
assert.True(t, ok)
_, ok = ctx.Driver(core.DriverTypeScheduler)
assert.True(t, ok)
_, ok = ctx.Driver("non_existent")
assert.False(t, ok)
}