Merge pull request #2 from dl1998/wait-implementation

Extend implementation with Wait method
dl1998 · Feb 22, 2024 · 2f8e6b7 · 2f8e6b7
2 parents 9ae4233 + c23c74e
commit 2f8e6b7
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/examples/basic/main.go b/examples/basic/main.go
@@ -29,7 +29,7 @@ func main() {
 
 	task := ScheduleDemoTask(*newScheduler)
 
-	time.Sleep(10 * time.Second)
+	task.Wait()
 
 	fmt.Println()
 	fmt.Println(task)

diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
@@ -177,7 +177,7 @@ func (scheduler *Scheduler) ScheduleTask(name string, startTime *time.Time, dura
 
 		ticker := time.NewTicker(interval)
 		defer ticker.Stop()
-		defer scheduler.removeTask(scheduledTask)
+		defer scheduler.StopTask(scheduledTask)
 
 		for {
 			select {
@@ -200,7 +200,10 @@ func (scheduler *Scheduler) ScheduleTask(name string, startTime *time.Time, dura
 
 // StopTask encapsulates task stopping sequence.
 func (scheduler *Scheduler) StopTask(task *Task) {
-	close(task.stopSignal)
+	if task.stopSignal != nil {
+		close(task.stopSignal)
+		task.stopSignal = nil
+	}
 	scheduler.removeTask(task)
 }
 
@@ -246,3 +249,9 @@ func (task *Task) SetToContext(name string, value interface{}) {
 func (task *Task) RemoveFromContext(name string) {
 	delete(task.context, name)
 }
+
+func (task *Task) Wait() {
+	if task.stopSignal != nil {
+		<-task.stopSignal
+	}
+}