From a757eff276403ce9a93f255a9826e356ca3b20b2 Mon Sep 17 00:00:00 2001
From: Guang Jiong Lou <7991675+27149chen@users.noreply.github.com>
Date: Sat, 8 Feb 2025 00:28:32 +0800
Subject: [PATCH] fix scheduler restart (#825)

---
 scheduler.go      | 20 +++++++++++++++----
 scheduler_test.go | 49 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/scheduler.go b/scheduler.go
index 7b3e72b..fd8fac2 100644
--- a/scheduler.go
+++ b/scheduler.go
@@ -238,11 +238,8 @@ func (s *scheduler) stopScheduler() {
 	for _, j := range s.jobs {
 		j.stop()
 	}
-	for id, j := range s.jobs {
+	for _, j := range s.jobs {
 		<-j.ctx.Done()
-
-		j.ctx, j.cancel = context.WithCancel(s.shutdownCtx)
-		s.jobs[id] = j
 	}
 	var err error
 	if s.started {
@@ -254,6 +251,21 @@ func (s *scheduler) stopScheduler() {
 			err = ErrStopExecutorTimedOut
 		}
 	}
+	for id, j := range s.jobs {
+		oldCtx := j.ctx
+		if j.parentCtx == nil {
+			j.parentCtx = s.shutdownCtx
+		}
+		j.ctx, j.cancel = context.WithCancel(j.parentCtx)
+
+		// also replace the old context with the new one in the parameters
+		if len(j.parameters) > 0 && j.parameters[0] == oldCtx {
+			j.parameters[0] = j.ctx
+		}
+
+		s.jobs[id] = j
+	}
+
 	s.stopErrCh <- err
 	s.started = false
 	s.logger.Debug("gocron: scheduler stopped")
diff --git a/scheduler_test.go b/scheduler_test.go
index b9c591b..4fc7549 100644
--- a/scheduler_test.go
+++ b/scheduler_test.go
@@ -467,6 +467,55 @@ func TestScheduler_StopLongRunningJobs(t *testing.T) {
 	})
 }
 
+func TestScheduler_StopAndStartLongRunningJobs(t *testing.T) {
+	t.Run("start, run job, stop jobs before job is completed", func(t *testing.T) {
+		s := newTestScheduler(t,
+			WithStopTimeout(50*time.Millisecond),
+		)
+
+		restart := false
+		restartP := &restart
+
+		_, err := s.NewJob(
+			DurationJob(
+				50*time.Millisecond,
+			),
+			NewTask(
+				func(ctx context.Context) {
+					select {
+					case <-ctx.Done():
+						if *restartP {
+							t.Fatal("job should not been canceled after restart")
+						}
+					case <-time.After(100 * time.Millisecond):
+						if !*restartP {
+							t.Fatal("job can not been canceled")
+						}
+
+					}
+				},
+			),
+			WithStartAt(
+				WithStartImmediately(),
+			),
+			WithSingletonMode(LimitModeReschedule),
+		)
+		require.NoError(t, err)
+
+		s.Start()
+
+		time.Sleep(20 * time.Millisecond)
+		// the running job is canceled, no unexpected timeout error
+		require.NoError(t, s.StopJobs())
+
+		*restartP = true
+
+		s.Start()
+
+		time.Sleep(200 * time.Millisecond)
+	})
+}
+
 func TestScheduler_Shutdown(t *testing.T) {
 	defer verifyNoGoroutineLeaks(t)