check context not nil, and update test and docs

Author: JohnRoesler

errors.go

@@ -39,6 +39,7 @@ var (
 	ErrWeeklyJobMinutesSeconds       = fmt.Errorf("gocron: WeeklyJob: atTimes minutes and seconds must be between 0 and 59 inclusive")
 	ErrPanicRecovered                = fmt.Errorf("gocron: panic recovered")
 	ErrWithClockNil                  = fmt.Errorf("gocron: WithClock: clock must not be nil")
+	ErrWithContextNil                = fmt.Errorf("gocron: WithContext: context must not be nil")
 	ErrWithDistributedElectorNil     = fmt.Errorf("gocron: WithDistributedElector: elector must not be nil")
 	ErrWithDistributedLockerNil      = fmt.Errorf("gocron: WithDistributedLocker: locker must not be nil")
 	ErrWithDistributedJobLockerNil   = fmt.Errorf("gocron: WithDistributedJobLocker: locker must not be nil")

example_test.go

@@ -354,6 +354,21 @@ func ExampleNewScheduler() {
 	fmt.Println(s.Jobs())
 }
 
+func ExampleNewTask() {
+	s, _ := gocron.NewScheduler()
+	defer func() { _ = s.Shutdown() }()
+
+	_, _ = s.NewJob(
+		gocron.DurationJob(time.Second),
+		gocron.NewTask(
+			func(ctx context.Context) {
+				// gocron will pass in a context to the job and will cancel on shutdown.
+				// this allows you to listen for and handle cancellation within your job.
+			},
+		),
+	)
+}
+
 func ExampleOneTimeJob() {
 	s, _ := gocron.NewScheduler()
 	defer func() { _ = s.Shutdown() }()
@@ -598,6 +613,27 @@ func ExampleWithClock() {
 	// one, 2
 }
 
+func ExampleWithContext() {
+	s, _ := gocron.NewScheduler()
+	defer func() { _ = s.Shutdown() }()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	_, _ = s.NewJob(
+		gocron.DurationJob(
+			time.Second,
+		),
+		gocron.NewTask(
+			func(ctx context.Context) {
+				// gocron will pass in a context to the job and will cancel on shutdown.
+				// this allows you to listen for and handle cancellation within your job.
+			},
+		),
+		gocron.WithContext(ctx),
+	)
+}
+
 func ExampleWithDisabledDistributedJobLocker() {
 	// var _ gocron.Locker = (*myLocker)(nil)
 	//

job.go

@@ -83,6 +83,9 @@ type task struct {
 type Task func() task
 
 // NewTask provides the job's task function and parameters.
+// If you set the first argument of your Task func to be a context.Context,
+// gocron will pass in a context to the job and will cancel on shutdown.
+// This allows you to listen for and handle cancellation within your job.
 func NewTask(function any, parameters ...any) Task {
 	return func() task {
 		return task{
@@ -705,8 +708,14 @@ func WithIdentifier(id uuid.UUID) JobOption {
 }
 
 // WithContext sets the parent context for the job
+// If you set the first argument of your Task func to be a context.Context,
+// gocron will pass in a context to the job and will cancel on shutdown.
+// This allows you to listen for and handle cancellation within your job.
 func WithContext(ctx context.Context) JobOption {
 	return func(j *internalJob, _ time.Time) error {
+		if ctx == nil {
+			return ErrWithContextNil
+		}
 		j.parentCtx = ctx
 		return nil
 	}

scheduler.go

@@ -21,6 +21,9 @@ type Scheduler interface {
 	// NewJob creates a new job in the Scheduler. The job is scheduled per the provided
 	// definition when the Scheduler is started. If the Scheduler is already running
 	// the job will be scheduled when the Scheduler is started.
+	// If you set the first argument of your Task func to be a context.Context,
+	// gocron will pass in a context to the job and will cancel on shutdown.
+	// This allows you to listen for and handle cancellation within your job.
 	NewJob(JobDefinition, Task, ...JobOption) (Job, error)
 	// RemoveByTags removes all jobs that have at least one of the provided tags.
 	RemoveByTags(...string)

scheduler_test.go

@@ -394,6 +394,77 @@ func TestScheduler_StopLongRunningJobs(t *testing.T) {
 		require.NoError(t, s.StopJobs())
 		time.Sleep(100 * time.Millisecond)
 	})
+	t.Run("start, run job, stop jobs before job is completed - manual context cancel", func(t *testing.T) {
+		s := newTestScheduler(t,
+			WithStopTimeout(50*time.Millisecond),
+		)
+
+		ctx, cancel := context.WithCancel(context.Background())
+
+		_, 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")
+					}
+				}, ctx,
+			),
+			WithStartAt(
+				WithStartImmediately(),
+			),
+			WithSingletonMode(LimitModeReschedule),
+		)
+		require.NoError(t, err)
+
+		s.Start()
+
+		time.Sleep(20 * time.Millisecond)
+		// the running job is canceled, no unexpected timeout error
+		cancel()
+		require.NoError(t, s.StopJobs())
+		time.Sleep(100 * time.Millisecond)
+	})
+	t.Run("start, run job, stop jobs before job is completed - manual context cancel WithContext", func(t *testing.T) {
+		s := newTestScheduler(t,
+			WithStopTimeout(50*time.Millisecond),
+		)
+
+		ctx, cancel := context.WithCancel(context.Background())
+
+		_, 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),
+			WithContext(ctx),
+		)
+		require.NoError(t, err)
+
+		s.Start()
+
+		time.Sleep(20 * time.Millisecond)
+		// the running job is canceled, no unexpected timeout error
+		cancel()
+		require.NoError(t, s.StopJobs())
+		time.Sleep(100 * time.Millisecond)
+	})
 }
 
 func TestScheduler_Shutdown(t *testing.T) {
@@ -576,6 +647,12 @@ func TestScheduler_NewJobErrors(t *testing.T) {
 			nil,
 			ErrCronJobInvalid,
 		},
+		{
+			"context nil",
+			DurationJob(time.Second),
+			[]JobOption{WithContext(nil)}, //nolint:staticcheck
+			ErrWithContextNil,
+		},
 		{
 			"duration job time interval is zero",
 			DurationJob(0 * time.Second),