fix(core): delegate driver and task extension lookups to root context and improve test reliability

backend/plugins/drivers/driver_asynq_worker/executor.go

@@ -154,7 +154,7 @@ func DispatchTask(ctx context.Context, taskType string, payload []byte, triggere
 
 	// 入队 Asynq
 	taskInfo := asynq.NewTask(meta.AsynqTask, injectTaskTraceContext(ctx, payload))
-	if _, err := AsynqClient.Enqueue(
+	if _, err := GetAsynqClient().Enqueue(
 		taskInfo,
 		asynq.TaskID(taskID),
 		asynq.MaxRetry(meta.MaxRetry),
@@ -220,7 +220,7 @@ func RetryTask(ctx context.Context, id uint64) (string, error) {
 
 	// 入队 Asynq
 	taskInfo := asynq.NewTask(execution.TaskType, injectTaskTraceContext(ctx, []byte(execution.Payload)))
-	if _, err := AsynqClient.Enqueue(
+	if _, err := GetAsynqClient().Enqueue(
 		taskInfo,
 		asynq.TaskID(newTaskID),
 		asynq.MaxRetry(execution.MaxRetry),

backend/plugins/drivers/driver_asynq_worker/plugin.go

@@ -13,10 +13,10 @@ import (
 	"time"
 
 	"github.com/hibiken/asynq"
+	"github.com/redis/go-redis/v9"
 
 	"Wavelet/core"
 	"Wavelet/core/contracts"
-	"Wavelet/pkg/config"
 )
 
 const (
@@ -90,7 +90,6 @@ type Plugin struct {
 // 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},
@@ -139,6 +138,8 @@ func (p *Plugin) Apply(ctx *core.Context) error {
 
 	ctx.OnDispose(func() error {
 		SetActiveTaskExtension(nil)
+		SetRedisClient(nil)
+		ResetAsynqClient()
 		shutdownCtx, cancel := context.WithTimeout(context.Background(), p.shutdownTimeout)
 		defer cancel()
 		return p.Stop(shutdownCtx)
@@ -173,29 +174,21 @@ func (p *Plugin) Start(_ context.Context) error {
 		}
 	}
 
-	for _, taskName := range GetRegisteredAsynqTasks() {
-		mux.HandleFunc(taskName, ProcessTask)
+	opt := p.redisOpt
+	if opt == nil {
+		opt = NewRedisConnOpt()
+	}
+	RedisOpt = opt
+
+	if getRedisClient() == nil {
+		if mk, ok := opt.(interface{ MakeRedisClient() interface{} }); ok {
+			if client, ok := mk.MakeRedisClient().(redis.UniversalClient); ok {
+				SetRedisClient(client)
+			}
+		}
 	}
 
 	if p.server == nil {
-		opt := p.redisOpt
-		if opt == nil {
-			if RedisOpt != nil {
-				opt = RedisOpt
-			} else {
-				redisCfg := config.Config.Redis
-				addr := "127.0.0.1:6379"
-				if len(redisCfg.Addrs) > 0 && redisCfg.Addrs[0] != "" {
-					addr = redisCfg.Addrs[0]
-				}
-				opt = asynq.RedisClientOpt{
-					Addr:     addr,
-					Username: redisCfg.Username,
-					Password: redisCfg.Password,
-					DB:       redisCfg.DB,
-				}
-			}
-		}
 		p.server = asynq.NewServer(
 			opt,
 			asynq.Config{
@@ -276,15 +269,17 @@ func toAsynqHandler(h any) (asynq.Handler, 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())
+			c, payload, _ := extractTaskTraceContext(c, t.Payload())
+			return fn(c, 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())
+		return asynq.HandlerFunc(func(c context.Context, t *asynq.Task) error {
+			_, payload, _ := extractTaskTraceContext(c, t.Payload())
+			return fn(payload)
 		}), nil
 	case func() error:
 		return asynq.HandlerFunc(func(_ context.Context, _ *asynq.Task) error {

backend/plugins/drivers/driver_asynq_worker/utils.go

@@ -4,6 +4,8 @@
 package driver_asynq_worker
 
 import (
+	"sync"
+
 	"github.com/hibiken/asynq"
 	"github.com/redis/go-redis/v9"
 	"github.com/redis/go-redis/v9/maintnotifications"
@@ -54,9 +56,37 @@ var RedisOpt asynq.RedisConnOpt
 // AsynqClient asynq 客户端，用于任务入队
 var AsynqClient *asynq.Client
 
-func init() {
-	RedisOpt = NewRedisConnOpt()
-	AsynqClient = asynq.NewClient(RedisOpt)
+var asynqClientMu sync.RWMutex
+
+// GetAsynqClient returns the active Asynq client, dynamically creating or refreshing it if needed.
+func GetAsynqClient() *asynq.Client {
+	asynqClientMu.RLock()
+	if AsynqClient != nil {
+		asynqClientMu.RUnlock()
+		return AsynqClient
+	}
+	asynqClientMu.RUnlock()
+
+	asynqClientMu.Lock()
+	defer asynqClientMu.Unlock()
+	if AsynqClient != nil {
+		return AsynqClient
+	}
+
+	opt := NewRedisConnOpt()
+	RedisOpt = opt
+	AsynqClient = asynq.NewClient(opt)
+	return AsynqClient
+}
+
+// ResetAsynqClient resets the Asynq client so it will be re-created with current config.
+func ResetAsynqClient() {
+	asynqClientMu.Lock()
+	defer asynqClientMu.Unlock()
+	if AsynqClient != nil {
+		_ = AsynqClient.Close()
+		AsynqClient = nil
+	}
 }
 
 // NewRedisConnOpt 根据配置返回对应的 asynq Redis 连接选项