From 5af9bbfec9e4ca72189349c4f3ca8b18fa1ed5e0 Mon Sep 17 00:00:00 2001
From: Maycon Santos <mlsmaycon@gmail.com>
Date: Sat, 2 Mar 2024 01:02:22 +0100
Subject: [PATCH] check for closed channel and adjust tests

---
 management/server/scheduler.go      | 14 +++++++++-----
 management/server/scheduler_test.go | 20 +++++++++++++++-----
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/management/server/scheduler.go b/management/server/scheduler.go
index c6a5da539..356348056 100644
--- a/management/server/scheduler.go
+++ b/management/server/scheduler.go
@@ -93,7 +93,14 @@ func (wm *DefaultScheduler) Schedule(in time.Duration, ID string, job func() (ne
 		for {
 			select {
 			case <-ticker.C:
-				log.Debugf("time to do a scheduled job %s", ID)
+				select {
+				case <-cancel:
+					log.Debugf("scheduled job %s was canceled, stop timer", ID)
+					ticker.Stop()
+					return
+				default:
+					log.Debugf("time to do a scheduled job %s", ID)
+				}
 				runIn, reschedule := job()
 				if !reschedule {
 					wm.mu.Lock()
@@ -108,11 +115,8 @@ func (wm *DefaultScheduler) Schedule(in time.Duration, ID string, job func() (ne
 					ticker.Reset(runIn)
 				}
 			case <-cancel:
+				log.Debugf("job %s was canceled, stopping timer", ID)
 				ticker.Stop()
-				wm.mu.Lock()
-				defer wm.mu.Unlock()
-				delete(wm.jobs, ID)
-				log.Debugf("stopped scheduled job %s ", ID)
 				return
 			}
 		}
diff --git a/management/server/scheduler_test.go b/management/server/scheduler_test.go
index 0c0cef99b..4b2c2e30d 100644
--- a/management/server/scheduler_test.go
+++ b/management/server/scheduler_test.go
@@ -2,11 +2,12 @@ package server
 
 import (
 	"fmt"
-	"github.com/stretchr/testify/assert"
 	"math/rand"
 	"sync"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestScheduler_Performance(t *testing.T) {
@@ -36,15 +37,24 @@ func TestScheduler_Cancel(t *testing.T) {
 	jobID1 := "test-scheduler-job-1"
 	jobID2 := "test-scheduler-job-2"
 	scheduler := NewDefaultScheduler()
-	scheduler.Schedule(2*time.Second, jobID1, func() (nextRunIn time.Duration, reschedule bool) {
-		return 0, false
+	tChan := make(chan struct{})
+	p := []string{jobID1, jobID2}
+	scheduler.Schedule(2*time.Millisecond, jobID1, func() (nextRunIn time.Duration, reschedule bool) {
+		tt := p[0]
+		<-tChan
+		t.Logf("job %s", tt)
+		return 2 * time.Millisecond, true
 	})
-	scheduler.Schedule(2*time.Second, jobID2, func() (nextRunIn time.Duration, reschedule bool) {
-		return 0, false
+	scheduler.Schedule(2*time.Millisecond, jobID2, func() (nextRunIn time.Duration, reschedule bool) {
+		return 2 * time.Millisecond, true
 	})
 
+	time.Sleep(4 * time.Millisecond)
 	assert.Len(t, scheduler.jobs, 2)
 	scheduler.Cancel([]string{jobID1})
+	close(tChan)
+	p = []string{}
+	time.Sleep(4 * time.Millisecond)
 	assert.Len(t, scheduler.jobs, 1)
 	assert.NotNil(t, scheduler.jobs[jobID2])
 }