[management] scheduler cancel all jobs (#4158)

management/server/scheduler.go

@@ -11,14 +11,17 @@ import (
 // Scheduler is an interface which implementations can schedule and cancel jobs
 type Scheduler interface {
 	Cancel(ctx context.Context, IDs []string)
+	CancelAll(ctx context.Context)
 	Schedule(ctx context.Context, in time.Duration, ID string, job func() (nextRunIn time.Duration, reschedule bool))
 	IsSchedulerRunning(ID string) bool
 }
 
 // MockScheduler is a mock implementation of  Scheduler
 type MockScheduler struct {
-	CancelFunc   func(ctx context.Context, IDs []string)
-	ScheduleFunc func(ctx context.Context, in time.Duration, ID string, job func() (nextRunIn time.Duration, reschedule bool))
+	CancelFunc             func(ctx context.Context, IDs []string)
+	CancelAllFunc          func(ctx context.Context)
+	ScheduleFunc           func(ctx context.Context, in time.Duration, ID string, job func() (nextRunIn time.Duration, reschedule bool))
+	IsSchedulerRunningFunc func(ID string) bool
 }
 
 // Cancel mocks the Cancel function of the Scheduler interface
@@ -30,6 +33,15 @@ func (mock *MockScheduler) Cancel(ctx context.Context, IDs []string) {
 	log.WithContext(ctx).Warnf("MockScheduler doesn't have Cancel function defined ")
 }
 
+// CancelAll mocks the CancelAll function of the Scheduler interface
+func (mock *MockScheduler) CancelAll(ctx context.Context) {
+	if mock.CancelAllFunc != nil {
+		mock.CancelAllFunc(ctx)
+		return
+	}
+	log.WithContext(ctx).Warnf("MockScheduler doesn't have CancelAll function defined ")
+}
+
 // Schedule mocks the Schedule function of the Scheduler interface
 func (mock *MockScheduler) Schedule(ctx context.Context, in time.Duration, ID string, job func() (nextRunIn time.Duration, reschedule bool)) {
 	if mock.ScheduleFunc != nil {
@@ -40,7 +52,9 @@ func (mock *MockScheduler) Schedule(ctx context.Context, in time.Duration, ID st
 }
 
 func (mock *MockScheduler) IsSchedulerRunning(ID string) bool {
-	// MockScheduler does not implement IsSchedulerRunning, so we return false
+	if mock.IsSchedulerRunningFunc != nil {
+		return mock.IsSchedulerRunningFunc(ID)
+	}
 	log.Warnf("MockScheduler doesn't have IsSchedulerRunning function defined")
 	return false
 }
@@ -52,6 +66,15 @@ type DefaultScheduler struct {
 	mu   *sync.Mutex
 }
 
+func (wm *DefaultScheduler) CancelAll(ctx context.Context) {
+	wm.mu.Lock()
+	defer wm.mu.Unlock()
+
+	for id := range wm.jobs {
+		wm.cancel(ctx, id)
+	}
+}
+
 // NewDefaultScheduler creates an instance of a DefaultScheduler
 func NewDefaultScheduler() *DefaultScheduler {
 	return &DefaultScheduler{

management/server/scheduler_test.go

@@ -75,6 +75,38 @@ func TestScheduler_Cancel(t *testing.T) {
 	assert.NotNil(t, scheduler.jobs[jobID2])
 }
 
+func TestScheduler_CancelAll(t *testing.T) {
+	jobID1 := "test-scheduler-job-1"
+	jobID2 := "test-scheduler-job-2"
+	scheduler := NewDefaultScheduler()
+	tChan := make(chan struct{})
+	p := []string{jobID1, jobID2}
+	scheduletime := 2 * time.Millisecond
+	sleepTime := 4 * time.Millisecond
+	if runtime.GOOS == "windows" {
+		// sleep and ticker are slower on windows see https://github.com/golang/go/issues/44343
+		sleepTime = 20 * time.Millisecond
+	}
+
+	scheduler.Schedule(context.Background(), scheduletime, jobID1, func() (nextRunIn time.Duration, reschedule bool) {
+		tt := p[0]
+		<-tChan
+		t.Logf("job %s", tt)
+		return scheduletime, true
+	})
+	scheduler.Schedule(context.Background(), scheduletime, jobID2, func() (nextRunIn time.Duration, reschedule bool) {
+		return scheduletime, true
+	})
+
+	time.Sleep(sleepTime)
+	assert.Len(t, scheduler.jobs, 2)
+	scheduler.CancelAll(context.Background())
+	close(tChan)
+	p = []string{}
+	time.Sleep(sleepTime)
+	assert.Len(t, scheduler.jobs, 0)
+}
+
 func TestScheduler_Schedule(t *testing.T) {
 	jobID := "test-scheduler-job-1"
 	scheduler := NewDefaultScheduler()