Repository: go-co-op/gocron (go)

Issue:
[BUG] - A long running job can not be cancelled when run s.StopJobs()
### Describe the bug

A long running job can not be cancelled when run s.StopJobs(), and always return "timed out waiting for jobs to finish"

### To Reproduce

Steps to reproduce the behavior:

1. Create a Scheduler and add a long running job and set WithStartImmediately
2. start the scheduler
3. wait for a while before the job is completed, run s.StopJobs()
4. timeout error returns

### Version

v2.14.2

### Expected behavior

if the job has a ctx field, Scheduler should cancel it when running StopJobs() instead of waiting for it

### Additional context

<!---Add any other context about the problem here.--->


Pull request description:
cancel running jobs if possible when run Scheduler.StopJobs()
### What does this do?

when run Scheduler.StopJobs(), the scheduler will wait for the running jobs to be completed. If there are long running jobs, it will always return "timed out waiting for jobs to finish".

This pr try to cancel the running jobs if possible

### Which issue(s) does this PR fix/relate to?
Resolves #818


### List any changes that modify/break current functionality
You can now add a task function which accept ctx as the first param, but you do need pass a ctx actually, gocron will handle it properly
```
	j, err := s.NewJob(
		gocron.DurationJob(
			10*time.Second,
		),
		gocron.NewTask(
			func(ctx context.Context) {
				// do things
			},
		),
	)
```
but if you have you own ctx, fell free to pass it
```
	j, err := s.NewJob(
		gocron.DurationJob(
			10*time.Second,
		),
		gocron.NewTask(
			func(ctx context.Context) {
				// do things
			},
                         ctx,
		),
	)
```
or you can also pass it as a job option to control the whole thing
```
	j, err := s.NewJob(
		gocron.DurationJob(
			10*time.Second,
		),
		gocron.NewTask(
			func(ctx context.Context) {
				// do things
			},
		),
                gocron.WithContext(ctx)
	)
```


### Have you included tests for your changes?
yes

### Did you document any new/modified functionality?

- [ ] Updated `example_test.go`
- [ ] Updated `README.md`

### Notes


Patch:
diff --git a/job.go b/job.go
index 700a0b6..85a1857 100644
--- a/job.go
+++ b/job.go
@@ -18,11 +18,12 @@ import (
 // internalJob stores the information needed by the scheduler
 // to manage scheduling, starting and stopping the job
 type internalJob struct {
-	ctx    context.Context
-	cancel context.CancelFunc
-	id     uuid.UUID
-	name   string
-	tags   []string
+	ctx       context.Context
+	parentCtx context.Context
+	cancel    context.CancelFunc
+	id        uuid.UUID
+	name      string
+	tags      []string
 	jobSchedule
 
 	// as some jobs may queue up, it's possible to
@@ -703,6 +704,14 @@ func WithIdentifier(id uuid.UUID) JobOption {
 	}
 }
 
+// WithContext sets the parent context for the job
+func WithContext(ctx context.Context) JobOption {
+	return func(j *internalJob, _ time.Time) error {
+		j.parentCtx = ctx
+		return nil
+	}
+}
+
 // -----------------------------------------------
 // -----------------------------------------------
 // ------------- Job Event Listeners -------------
diff --git a/scheduler.go b/scheduler.go
index 825323e..8789f50 100644
--- a/scheduler.go
+++ b/scheduler.go
@@ -648,8 +648,6 @@ func (s *scheduler) addOrUpdateJob(id uuid.UUID, definition JobDefinition, taskW
 		j.id = id
 	}
 
-	j.ctx, j.cancel = context.WithCancel(s.shutdownCtx)
-
 	if taskWrapper == nil {
 		return nil, ErrNewJobTaskNil
 	}
@@ -664,10 +662,6 @@ func (s *scheduler) addOrUpdateJob(id uuid.UUID, definition JobDefinition, taskW
 		return nil, ErrNewJobTaskNotFunc
 	}
 
-	if err := s.verifyParameterType(taskFunc, tsk); err != nil {
-		return nil, err
-	}
-
 	j.name = runtime.FuncForPC(taskFunc.Pointer()).Name()
 	j.function = tsk.function
 	j.parameters = tsk.parameters
@@ -686,6 +680,28 @@ func (s *scheduler) addOrUpdateJob(id uuid.UUID, definition JobDefinition, taskW
 		}
 	}
 
+	if j.parentCtx == nil {
+		j.parentCtx = s.shutdownCtx
+	}
+	j.ctx, j.cancel = context.WithCancel(j.parentCtx)
+
+	if !taskFunc.IsZero() && taskFunc.Type().NumIn() > 0 {
+		// if the first parameter is a context.Context and params have no context.Context, add current ctx to the params
+		if taskFunc.Type().In(0) == reflect.TypeOf((*context.Context)(nil)).Elem() {
+			if len(tsk.parameters) == 0 {
+				tsk.parameters = []any{j.ctx}
+				j.parameters = []any{j.ctx}
+			} else if _, ok := tsk.parameters[0].(context.Context); !ok {
+				tsk.parameters = append([]any{j.ctx}, tsk.parameters...)
+				j.parameters = append([]any{j.ctx}, j.parameters...)
+			}
+		}
+	}
+
+	if err := s.verifyParameterType(taskFunc, tsk); err != nil {
+		return nil, err
+	}
+
 	if err := definition.setup(&j, s.location, s.exec.clock.Now()); err != nil {
 		return nil, err
 	}


Test patch:
diff --git a/scheduler_test.go b/scheduler_test.go
index 8d46d31..3c47472 100644
--- a/scheduler_test.go
+++ b/scheduler_test.go
@@ -361,6 +361,41 @@ func TestScheduler_StopTimeout(t *testing.T) {
 	}
 }
 
+func TestScheduler_StopLongRunningJobs(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),
+		)
+
+		_, err := s.NewJob(
+			DurationJob(
+				50*time.Millisecond,
+			),
+			NewTask(
+				func(ctx context.Context) {
+					select {
+					case <-ctx.Done():
+					case <-time.After(100 * time.Millisecond):
+						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())
+		time.Sleep(100 * time.Millisecond)
+	})
+}
+
 func TestScheduler_Shutdown(t *testing.T) {
 	defer verifyNoGoroutineLeaks(t)