chore: Refactor task schedulers to use context.Context for lifecycle management

common/task/fifo_task_scheduler.go

@@ -21,6 +21,7 @@
 package task
 
 import (
+	"context"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -38,7 +39,8 @@ type fifoTaskSchedulerImpl struct {
 	options      *FIFOTaskSchedulerOptions
 	dispatcherWG sync.WaitGroup
 	taskCh       chan PriorityTask
-	shutdownCh   chan struct{}
+	ctx          context.Context
+	cancel       context.CancelFunc
 
 	processor Processor
 }
@@ -52,13 +54,15 @@ func NewFIFOTaskScheduler(
 	metricsClient metrics.Client,
 	options *FIFOTaskSchedulerOptions,
 ) Scheduler {
+	ctx, cancel := context.WithCancel(context.Background())
 	return &fifoTaskSchedulerImpl{
 		status:       common.DaemonStatusInitialized,
 		logger:       logger,
 		metricsScope: metricsClient.Scope(metrics.TaskSchedulerScope),
 		options:      options,
 		taskCh:       make(chan PriorityTask, options.QueueSize),
-		shutdownCh:   make(chan struct{}),
+		ctx:          ctx,
+		cancel:       cancel,
 		processor: NewParallelTaskProcessor(
 			logger,
 			metricsClient,
@@ -91,7 +95,7 @@ func (f *fifoTaskSchedulerImpl) Stop() {
 		return
 	}
 
-	close(f.shutdownCh)
+	f.cancel()
 
 	f.processor.Stop()
 
@@ -119,7 +123,7 @@ func (f *fifoTaskSchedulerImpl) Submit(task PriorityTask) error {
 			f.drainAndNackTasks()
 		}
 		return nil
-	case <-f.shutdownCh:
+	case <-f.ctx.Done():
 		return ErrTaskSchedulerClosed
 	}
 }
@@ -136,7 +140,7 @@ func (f *fifoTaskSchedulerImpl) TrySubmit(task PriorityTask) (bool, error) {
 			f.drainAndNackTasks()
 		}
 		return true, nil
-	case <-f.shutdownCh:
+	case <-f.ctx.Done():
 		return false, ErrTaskSchedulerClosed
 	default:
 		return false, nil
@@ -153,7 +157,7 @@ func (f *fifoTaskSchedulerImpl) dispatcher() {
 				f.logger.Error("failed to submit task to processor", tag.Error(err))
 				task.Nack(err)
 			}
-		case <-f.shutdownCh:
+		case <-f.ctx.Done():
 			return
 		}
 	}

common/task/weighted_round_robin_task_scheduler.go

@@ -21,6 +21,7 @@
 package task
 
 import (
+	"context"
 	"errors"
 	"sync"
 	"sync/atomic"
@@ -38,7 +39,8 @@ type weightedRoundRobinTaskSchedulerImpl[K comparable] struct {
 
 	status       int32
 	pool         *WeightedRoundRobinChannelPool[K, PriorityTask]
-	shutdownCh   chan struct{}
+	ctx          context.Context
+	cancel       context.CancelFunc
 	notifyCh     chan struct{}
 	dispatcherWG sync.WaitGroup
 	logger       log.Logger
@@ -67,6 +69,7 @@ func NewWeightedRoundRobinTaskScheduler[K comparable](
 	options *WeightedRoundRobinTaskSchedulerOptions[K],
 ) (Scheduler, error) {
 	metricsScope := metricsClient.Scope(metrics.TaskSchedulerScope)
+	ctx, cancel := context.WithCancel(context.Background())
 	scheduler := &weightedRoundRobinTaskSchedulerImpl[K]{
 		status: common.DaemonStatusInitialized,
 		pool: NewWeightedRoundRobinChannelPool[K, PriorityTask](
@@ -77,7 +80,8 @@ func NewWeightedRoundRobinTaskScheduler[K comparable](
 				BufferSize:              options.QueueSize,
 				IdleChannelTTLInSeconds: defaultIdleChannelTTLInSeconds,
 			}),
-		shutdownCh:   make(chan struct{}),
+		ctx:          ctx,
+		cancel:       cancel,
 		notifyCh:     make(chan struct{}, 1),
 		logger:       logger,
 		metricsScope: metricsScope,
@@ -105,7 +109,7 @@ func (w *weightedRoundRobinTaskSchedulerImpl[K]) Stop() {
 		return
 	}
 
-	close(w.shutdownCh)
+	w.cancel()
 
 	taskChs := w.pool.GetAllChannels()
 	for _, taskCh := range taskChs {
@@ -139,7 +143,7 @@ func (w *weightedRoundRobinTaskSchedulerImpl[K]) Submit(task PriorityTask) error
 			drainAndNackPriorityTask(taskCh)
 		}
 		return nil
-	case <-w.shutdownCh:
+	case <-w.ctx.Done():
 		return ErrTaskSchedulerClosed
 	}
 }
@@ -165,7 +169,7 @@ func (w *weightedRoundRobinTaskSchedulerImpl[K]) TrySubmit(
 			w.notifyDispatcher()
 		}
 		return true, nil
-	case <-w.shutdownCh:
+	case <-w.ctx.Done():
 		return false, ErrTaskSchedulerClosed
 	default:
 		return false, nil
@@ -179,7 +183,7 @@ func (w *weightedRoundRobinTaskSchedulerImpl[K]) dispatcher() {
 		select {
 		case <-w.notifyCh:
 			w.dispatchTasks()
-		case <-w.shutdownCh:
+		case <-w.ctx.Done():
 			return
 		}
 	}
@@ -198,7 +202,7 @@ func (w *weightedRoundRobinTaskSchedulerImpl[K]) dispatchTasks() {
 					w.logger.Error("fail to submit task to processor", tag.Error(err))
 					task.Nack(err)
 				}
-			case <-w.shutdownCh:
+			case <-w.ctx.Done():
 				return
 			default:
 			}

common/task/weighted_round_robin_task_scheduler_test.go

@@ -219,7 +219,7 @@ func (s *weightedRoundRobinTaskSchedulerSuite) TestDispatcher_SubmitWithNoError(
 	}()
 
 	taskWG.Wait()
-	close(s.scheduler.shutdownCh)
+	s.scheduler.cancel()
 
 	<-doneCh
 }
@@ -248,7 +248,7 @@ func (s *weightedRoundRobinTaskSchedulerSuite) TestDispatcher_FailToSubmit() {
 	}()
 
 	taskWG.Wait()
-	close(s.scheduler.shutdownCh)
+	s.scheduler.cancel()
 
 	<-doneCh
 }
@@ -372,9 +372,9 @@ func testSchedulerContract(
 	s.True(common.AwaitWaitGroup(&taskWG, 10*time.Second))
 	switch schedulerImpl := scheduler.(type) {
 	case *fifoTaskSchedulerImpl:
-		<-schedulerImpl.shutdownCh
+		<-schedulerImpl.ctx.Done()
 	case *weightedRoundRobinTaskSchedulerImpl[int]:
-		<-schedulerImpl.shutdownCh
+		<-schedulerImpl.ctx.Done()
 	default:
 		s.Fail("unknown task scheduler type")
 	}