Add pre and post enqueue callback options for Scheduler

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- `PreEnqueueFunc`, `PostEnqueueFunc` is added in `Scheduler` and deprecated `EnqueueErrorHandler` (PR: https://github.com/hibiken/asynq/pull/476)
+
 ## [0.23.0] - 2022-04-11
 
 ### Added

scheduler.go

@@ -33,6 +33,8 @@ type Scheduler struct {
 	location        *time.Location
 	done            chan struct{}
 	wg              sync.WaitGroup
+	preEnqueueFunc  func(task *Task, opts []Option)
+	postEnqueueFunc func(info *TaskInfo, err error)
 	errHandler      func(task *Task, opts []Option, err error)
 
 	// guards idmap
@@ -75,6 +77,8 @@ func NewScheduler(r RedisConnOpt, opts *SchedulerOpts) *Scheduler {
 		cron:            cron.New(cron.WithLocation(loc)),
 		location:        loc,
 		done:            make(chan struct{}),
+		preEnqueueFunc:  opts.PreEnqueueFunc,
+		postEnqueueFunc: opts.PostEnqueueFunc,
 		errHandler:      opts.EnqueueErrorHandler,
 		idmap:           make(map[string]cron.EntryID),
 	}
@@ -105,6 +109,15 @@ type SchedulerOpts struct {
 	// If unset, the UTC time zone (time.UTC) is used.
 	Location *time.Location
 
+	// PreEnqueueFunc, if provided, is called before a task gets enqueued by Scheduler.
+	// The callback function should return quickly to not block the current thread.
+	PreEnqueueFunc func(task *Task, opts []Option)
+
+	// PostEnqueueFunc, if provided, is called after a task gets enqueued by Scheduler.
+	// The callback function should return quickly to not block the current thread.
+	PostEnqueueFunc func(info *TaskInfo, err error)
+
+	// Deprecated: Use PostEnqueueFunc instead
 	// EnqueueErrorHandler gets called when scheduler cannot enqueue a registered task
 	// due to an error.
 	EnqueueErrorHandler func(task *Task, opts []Option, err error)
@@ -120,11 +133,19 @@ type enqueueJob struct {
 	logger          *log.Logger
 	client          *Client
 	rdb             *rdb.RDB
+	preEnqueueFunc  func(task *Task, opts []Option)
+	postEnqueueFunc func(info *TaskInfo, err error)
 	errHandler      func(task *Task, opts []Option, err error)
 }
 
 func (j *enqueueJob) Run() {
+	if j.preEnqueueFunc != nil {
+		j.preEnqueueFunc(j.task, j.opts)
+	}
 	info, err := j.client.Enqueue(j.task, j.opts...)
+	if j.postEnqueueFunc != nil {
+		j.postEnqueueFunc(info, err)
+	}
 	if err != nil {
 		j.logger.Errorf("scheduler could not enqueue a task %+v: %v", j.task, err)
 		if j.errHandler != nil {
@@ -155,6 +176,8 @@ func (s *Scheduler) Register(cronspec string, task *Task, opts ...Option) (entry
 		client:          s.client,
 		rdb:             s.rdb,
 		logger:          s.logger,
+		preEnqueueFunc:  s.preEnqueueFunc,
+		postEnqueueFunc: s.postEnqueueFunc,
 		errHandler:      s.errHandler,
 	}
 	cronID, err := s.cron.AddJob(cronspec, job)

scheduler_test.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
+
 	"github.com/hibiken/asynq/internal/base"
 	"github.com/hibiken/asynq/internal/testutil"
 )
@@ -154,3 +155,56 @@ func TestSchedulerUnregister(t *testing.T) {
 		}
 	}
 }
+
+func TestSchedulerPostAndPreEnqueueHandler(t *testing.T) {
+	var (
+		preMu       sync.Mutex
+		preCounter  int
+		postMu      sync.Mutex
+		postCounter int
+	)
+	preHandler := func(task *Task, opts []Option) {
+		preMu.Lock()
+		preCounter++
+		preMu.Unlock()
+	}
+	postHandler := func(info *TaskInfo, err error) {
+		postMu.Lock()
+		postCounter++
+		postMu.Unlock()
+	}
+
+	// Connect to non-existent redis instance to simulate a redis server being down.
+	scheduler := NewScheduler(
+		getRedisConnOpt(t),
+		&SchedulerOpts{
+			PreEnqueueFunc:  preHandler,
+			PostEnqueueFunc: postHandler,
+		},
+	)
+
+	task := NewTask("test", nil)
+
+	if _, err := scheduler.Register("@every 3s", task); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := scheduler.Start(); err != nil {
+		t.Fatal(err)
+	}
+	// Scheduler should attempt to enqueue the task three times (every 3s).
+	time.Sleep(10 * time.Second)
+	scheduler.Shutdown()
+
+	preMu.Lock()
+	if preCounter != 3 {
+		t.Errorf("PreEnqueueFunc was called %d times, want 3", preCounter)
+	}
+	preMu.Unlock()
+
+	postMu.Lock()
+	if postCounter != 3 {
+		t.Errorf("PostEnqueueFunc was called %d times, want 3", postCounter)
+	}
+	postMu.Unlock()
+}