feat: add inspector

Author: zhufuyi

pkg/sasynq/README.md

@@ -18,6 +18,13 @@ func NewClient(cfg RedisConfig) *Client {
 	}
 }
 
+// NewFromClient creates a new producer client from an existing asynq.Client.
+func NewFromClient(c *asynq.Client) *Client {
+	return &Client{
+		Client: c,
+	}
+}
+
 // Enqueue enqueues the given task to a queue.
 func (c *Client) Enqueue(task *asynq.Task, opts ...asynq.Option) (*asynq.TaskInfo, error) {
 	return c.Client.Enqueue(task, opts...)
@@ -57,3 +64,14 @@ func (c *Client) EnqueueAt(t time.Time, typeName string, payload any, opts ...as
 	info, err := c.Client.Enqueue(task, opts...)
 	return task, info, err
 }
+
+// EnqueueUnique enqueues a task with unique in the queue for a specified duration.
+func (c *Client) EnqueueUnique(keepTime time.Duration, typeName string, payload any, opts ...asynq.Option) (*asynq.Task, *asynq.TaskInfo, error) {
+	task, err := NewTask(typeName, payload)
+	if err != nil {
+		return nil, nil, err
+	}
+	opts = append(opts, asynq.Unique(keepTime))
+	info, err := c.Client.Enqueue(task, opts...)
+	return task, info, err
+}

pkg/sasynq/client.go

@@ -1,49 +1,105 @@
 package sasynq
 
 import (
+	"fmt"
 	"log"
+	"math/rand"
 	"testing"
 	"time"
+
+	"github.com/hibiken/asynq"
 )
 
 func runProducer(client *Client) error {
 	userPayload1 := &EmailPayload{UserID: 101, Message: "Critical Update"}
-	_, info, err := client.EnqueueNow(TypeEmailSend, userPayload1, WithQueue("critical"), WithRetry(5))
+	_, info1, err := client.EnqueueNow(TypeEmailSend, userPayload1, WithQueue("critical"), WithMaxRetry(5), asynq.Retention(60*time.Second))
 	if err != nil {
 		return err
 	}
-	log.Printf("enqueued task: type=%s, id=%s queue=%s", TypeEmailSend, info.ID, info.Queue)
+	log.Printf("enqueued task: type=%s, id=%s queue=%s", TypeEmailSend, info1.ID, info1.Queue)
 
 	userPayload2 := &SMSPayload{UserID: 202, Message: "Weekly Newsletter"}
-	_, info, err = client.EnqueueIn(time.Second*5, TypeSMSSend, userPayload2, WithQueue("default"), WithRetry(3))
+	_, info2, err := client.EnqueueIn(time.Second*5, TypeSMSSend, userPayload2, WithQueue("default"), WithMaxRetry(3), asynq.Retention(60*time.Second))
 	if err != nil {
 		return err
 	}
-	log.Printf("enqueued task: type=%s, id=%s queue=%s", TypeSMSSend, info.ID, info.Queue)
+	log.Printf("enqueued task: type=%s, id=%s queue=%s", TypeSMSSend, info2.ID, info2.Queue)
+	cancelTask("default", info2.ID, true) // cancel task will succeed
 
 	userPayload3 := &MsgNotificationPayload{UserID: 303, Message: "Promotional Offer"}
-	_, info, err = client.EnqueueAt(time.Now().Add(time.Second*10), TypeMsgNotification, userPayload3, WithQueue("low"), WithRetry(1))
+	_, info3, err := client.EnqueueAt(time.Now().Add(time.Second*10), TypeMsgNotification, userPayload3, WithQueue("low"), WithMaxRetry(1), asynq.Retention(60*time.Second))
 	if err != nil {
 		return err
 	}
-	log.Printf("enqueued task: type=%s, id=%s queue=%s", TypeMsgNotification, info.ID, info.Queue)
+	log.Printf("enqueued task: type=%s, id=%s queue=%s", TypeMsgNotification, info3.ID, info3.Queue)
+	cancelTask("low", info3.ID, true) // cancel task will succeed
+
+	userPayload4 := &UniqueTaskPayload{UserID: 404, Message: "unique task"}
+	_, info4, err := client.EnqueueUnique(time.Minute, TypeUniqueTask, userPayload4, WithQueue("default"), WithMaxRetry(2))
+	if err != nil {
+		return err
+	}
+	log.Printf("enqueued task: type=%s, id=%s queue=%s", TypeUniqueTask, info4.ID, info4.Queue)
+	_, _, err = client.EnqueueUnique(time.Minute, TypeUniqueTask, userPayload4, WithQueue("default"), WithMaxRetry(2))
+	if err != nil {
+		log.Printf("triggered duplicate task error:%v", err)
+	}
 
 	// Equivalent EnqueueNow function
-	userPayload4 := &EmailPayload{UserID: 404, Message: "Important Notification"}
-	task, err := NewTask(TypeEmailSend, userPayload4)
+	userPayload5 := &EmailPayload{UserID: 505, Message: "Important Notification"}
+	task, err := NewTask(TypeEmailSend, userPayload5)
 	if err != nil {
 		return err
 	}
-	info, err = client.Enqueue(task, WithQueue("low"), WithRetry(1), WithDeadline(time.Now().Add(time.Second*15)), WithUniqueID("custom-id-xxxx-xxxx"))
+	info5, err := client.Enqueue(task, WithQueue("low"), WithMaxRetry(3), WithDeadline(time.Now().Add(time.Second*15)), WithTaskID(fmt.Sprintf("unique-%d", rand.Int63n(1e10))), asynq.Retention(60*time.Second))
 	if err != nil {
 		return err
 	}
-	log.Printf("enqueued task: type=%s, id=%s queue=%s", TypeEmailSend, info.ID, info.Queue)
+	log.Printf("enqueued task: type=%s, id=%s queue=%s", TypeEmailSend, info5.ID, info5.Queue)
 
 	return nil
 }
 
-func TestNewClient(t *testing.T) {
+func cancelTask(queue string, taskID string, isScheduled bool) {
+	fmt.Println()
+	defer fmt.Println()
+	time.Sleep(time.Second)
+
+	inspector := NewInspector(getRedisConfig())
+
+	info, err := inspector.GetTaskInfo(queue, taskID)
+	if err != nil {
+		log.Printf("get task info failed: %s, queue=%s, taskID=%s", err, queue, taskID)
+		return
+	}
+	log.Printf("task status: type=%s, id=%s queue=%s, status=%s", info.Type, info.ID, info.Queue, info.State.String())
+	if info.State == asynq.TaskStateCompleted {
+		return
+	}
+
+	time.Sleep(time.Millisecond * 100)
+	if isScheduled {
+		err = inspector.CancelTask(queue, info.ID)
+	} else {
+		err = inspector.CancelTask("", info.ID) // queue is empty string for non-scheduled tasks
+	}
+
+	if err != nil {
+		log.Printf("cancel task failed: %s, queue=%s, taskID=%s", err, queue, info.ID)
+		return
+	}
+	log.Printf("cancel task succeeded: type=%s, id=%s queue=%s", info.Type, info.ID, info.Queue)
+
+	time.Sleep(time.Millisecond * 100)
+	info2, err := inspector.GetTaskInfo(queue, info.ID)
+	if err != nil {
+		log.Printf("get task info after cancel failed: %s, queue=%s, taskID=%s", err, queue, info.ID)
+		return
+	}
+	log.Printf("get task status after cancel: type=%s, id=%s queue=%s, status=%s", info2.Type, info2.ID, info2.Queue, info2.State.String())
+}
+
+func TestProducer(t *testing.T) {
 	client := NewClient(getRedisConfig())
 
 	err := runProducer(client)

pkg/sasynq/client_test.go

@@ -1,8 +1,9 @@
 package sasynq
 
 import (
+	"crypto/tls"
+
 	"github.com/hibiken/asynq"
-	"go.uber.org/zap"
 )
 
 // RedisMode defines the Redis connection mode.
@@ -35,11 +36,13 @@ type RedisConfig struct {
 	Username string `yaml:"username"`
 	Password string `yaml:"password"`
 	DB       int    `yaml:"db"`
+
+	TLSConfig *tls.Config `yaml:"tlsConfig"`
 }
 
 // GetAsynqRedisConnOpt converts RedisConfig to asynq's RedisConnOpt interface.
 // This is the core of the high-availability switching logic.
-func (c *RedisConfig) GetAsynqRedisConnOpt() asynq.RedisConnOpt {
+func (c RedisConfig) GetAsynqRedisConnOpt() asynq.RedisConnOpt {
 	switch c.Mode {
 	case RedisModeSentinel:
 		return asynq.RedisFailoverClientOpt{
@@ -48,21 +51,24 @@ func (c *RedisConfig) GetAsynqRedisConnOpt() asynq.RedisConnOpt {
 			Username:      c.Username,
 			Password:      c.Password,
 			DB:            c.DB,
+			TLSConfig:     c.TLSConfig,
 		}
 	case RedisModeCluster:
 		return asynq.RedisClusterClientOpt{
-			Addrs:    c.ClusterAddrs,
-			Username: c.Username,
-			Password: c.Password,
+			Addrs:     c.ClusterAddrs,
+			Username:  c.Username,
+			Password:  c.Password,
+			TLSConfig: c.TLSConfig,
 		}
 	case RedisModeSingle:
 		fallthrough
 	default:
 		return asynq.RedisClientOpt{
-			Addr:     c.Addr,
-			Username: c.Username,
-			Password: c.Password,
-			DB:       c.DB,
+			Addr:      c.Addr,
+			Username:  c.Username,
+			Password:  c.Password,
+			DB:        c.DB,
+			TLSConfig: c.TLSConfig,
 		}
 	}
 }
@@ -73,13 +79,9 @@ type ServerConfig struct {
 }
 
 // DefaultServerConfig returns a default server configuration.
-func DefaultServerConfig(l ...*zap.Logger) ServerConfig {
-	var zapLogger *zap.Logger
-	if len(l) == 0 {
-		zapLogger = defaultLogger
-	} else {
-		zapLogger = l[0]
-	}
+func DefaultServerConfig(opts ...LoggerOption) ServerConfig {
+	o := defaultLoggerOptions()
+	o.apply(opts...)
 
 	cfg := &asynq.Config{
 		Concurrency: 10,
@@ -88,7 +90,7 @@ func DefaultServerConfig(l ...*zap.Logger) ServerConfig {
 			"default":  3,
 			"low":      1,
 		},
-		Logger: NewZapLogger(zapLogger),
+		Logger: NewZapLogger(o.logger, o.zapSkip),
 	}
 
 	return ServerConfig{cfg}

pkg/sasynq/config.go

@@ -0,0 +1,38 @@
+package sasynq
+
+import (
+	"github.com/hibiken/asynq"
+)
+
+// Inspector provides access to the Redis backend used by asynq.
+type Inspector struct {
+	*asynq.Inspector
+}
+
+// NewInspector creates a new Inspector instance.
+// Note: A new Redis connection will be created, in actual use, only once
+func NewInspector(cfg RedisConfig) *Inspector {
+	return &Inspector{asynq.NewInspector(cfg.GetAsynqRedisConnOpt())}
+}
+
+// CancelTask cancels the processing of a task.
+func (i *Inspector) CancelTask(queue string, taskID string) error {
+	if queue == "" {
+		return i.Inspector.CancelProcessing(taskID)
+	}
+
+	return i.Inspector.DeleteTask(queue, taskID)
+}
+
+// GetTaskInfo returns information about a task.
+func (i *Inspector) GetTaskInfo(queue string, taskID string) (*asynq.TaskInfo, error) {
+	return i.Inspector.GetTaskInfo(queue, taskID)
+}
+
+// Close closes the inspector.
+func (i *Inspector) Close() error {
+	if i.Inspector == nil {
+		return nil
+	}
+	return i.Inspector.Close()
+}

pkg/sasynq/inspector.go

@@ -21,7 +21,8 @@ type LoggerOption func(*loggerOptions)
 
 type loggerOptions struct {
 	logger    *zap.Logger
-	maxLength int
+	zapSkip   int // default is 2
+	maxLength int // default is 300
 }
 
 func (o *loggerOptions) apply(opts ...LoggerOption) {
@@ -34,6 +35,7 @@ func defaultLoggerOptions() *loggerOptions {
 	return &loggerOptions{
 		logger:    defaultLogger,
 		maxLength: defaultMaxLength,
+		zapSkip:   2,
 	}
 }
 
@@ -55,6 +57,15 @@ func WithMaxLength(l int) LoggerOption {
 	}
 }
 
+// WithZapSkip sets the number of callers to skip when logging.
+func WithZapSkip(s int) LoggerOption {
+	return func(o *loggerOptions) {
+		if s >= 0 {
+			o.zapSkip = s
+		}
+	}
+}
+
 // LoggingMiddleware logs information about each processed task.
 func LoggingMiddleware(opts ...LoggerOption) func(next asynq.Handler) asynq.Handler {
 	o := defaultLoggerOptions()
@@ -116,9 +127,8 @@ type ZapLogger struct {
 	zLog *zap.Logger
 }
 
-func NewZapLogger(l *zap.Logger) asynq.Logger {
-	//zLog := l.WithOptions(zap.AddCallerSkip(5)).With(zap.String("component", "asynq"))
-	zLog := l.WithOptions(zap.AddCallerSkip(2)).With(zap.String("component", "asynq"))
+func NewZapLogger(l *zap.Logger, skip int) asynq.Logger {
+	zLog := l.WithOptions(zap.AddCallerSkip(skip)).With(zap.String("asynq", "true"))
 	return &ZapLogger{
 		zLog: zLog,
 	}

pkg/sasynq/logger_middleware.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 
 	"github.com/hibiken/asynq"
-	"go.uber.org/zap"
 )
 
 // SchedulerOption set options.
@@ -43,12 +42,11 @@ func WithSchedulerLogLevel(level asynq.LogLevel) SchedulerOption {
 }
 
 // WithSchedulerLogger sets the logger for the scheduler.
-func WithSchedulerLogger(l *zap.Logger) SchedulerOption {
+func WithSchedulerLogger(opts ...LoggerOption) SchedulerOption {
+	opt := defaultLoggerOptions()
+	opt.apply(opts...)
 	return func(o *schedulerOptions) {
-		if l == nil {
-			return
-		}
-		o.logger = NewZapLogger(l)
+		o.logger = NewZapLogger(opt.logger, opt.zapSkip)
 	}
 }
 
@@ -94,7 +92,7 @@ func (s *Scheduler) RegisterTask(cronSpec string, typeName string, payload any,
 	return s.Scheduler.Register(cronSpec, task, opts...)
 }
 
-// Unregister removes a periodic task.
+// Unregister removes a periodic task, cancel task execution.
 func (s *Scheduler) Unregister(entryID string) error {
 	return s.Scheduler.Unregister(entryID)
 }