From 0fad386f1579bdf2b4788c57a66da392c7a616ab Mon Sep 17 00:00:00 2001 From: Jason Song Date: Tue, 12 May 2026 15:11:45 +0800 Subject: [PATCH 1/5] feat: support DeriveContext --- inner_job.go | 9 ++++++++- job_option.go | 10 ++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/inner_job.go b/inner_job.go index 9163aee..014634a 100644 --- a/inner_job.go +++ b/inner_job.go @@ -27,6 +27,7 @@ type innerJob struct { key string spec string before BeforeFunc + deriveContext DeriveContext run RunFunc after AfterFunc retryTimes int @@ -96,7 +97,13 @@ func (j *innerJob) Run() { task.BeginAt = &beginAt for i := 0; i < j.retryTimes; i++ { - task.Return = safeRun(ctx, j.run) + if j.deriveContext != nil { + ctx, cancel := j.deriveContext(ctx, task) + task.Return = safeRun(ctx, j.run) + cancel() + } else { + task.Return = safeRun(ctx, j.run) + } atomic.AddInt64(&j.statistics.TotalRun, 1) if i > 0 { atomic.AddInt64(&j.statistics.RetriedRun, 1) diff --git a/job_option.go b/job_option.go index 0447ca8..e51f2f3 100644 --- a/job_option.go +++ b/job_option.go @@ -20,6 +20,9 @@ type AfterFunc func(task Task) // RetryInterval indicates how long should delay before retrying when run failed `triedTimes` times. type RetryInterval func(triedTimes int) time.Duration +// DeriveContext indicates how to derive a new context for Run with the given Task, +type DeriveContext func(ctx context.Context, task Task) (context.Context, context.CancelFunc) + // WithBeforeFunc specifies what to do before Run. func WithBeforeFunc(before BeforeFunc) JobOption { return func(job *innerJob) { @@ -63,3 +66,10 @@ func WithGroup(group Group) JobOption { job.group = group } } + +// WithDeriveContext indicates how to derive a new context for Run with the given Task. +func WithDeriveContext(deriveContext DeriveContext) JobOption { + return func(job *innerJob) { + job.deriveContext = deriveContext + } +} From 9a7286ab2505b28b1b1e5efada8f054dc928ca7b Mon Sep 17 00:00:00 2001 From: Jason Song Date: Tue, 12 May 2026 15:21:55 +0800 Subject: [PATCH 2/5] feat: support WithBeforeContextFunc & WithAfterContextFunc --- inner_job.go | 33 +++++++++++++++++++++------------ job_option.go | 26 +++++++++++++++++++++++++- 2 files changed, 46 insertions(+), 13 deletions(-) diff --git a/inner_job.go b/inner_job.go index 014634a..95f1651 100644 --- a/inner_job.go +++ b/inner_job.go @@ -26,9 +26,11 @@ type innerJob struct { entryGetter entryGetter key string spec string - before BeforeFunc deriveContext DeriveContext + ctxBefore BeforeContextFunc + before BeforeFunc run RunFunc + ctxAfter AfterContextFunc after AfterFunc retryTimes int retryInterval RetryInterval @@ -76,9 +78,20 @@ func (j *innerJob) Run() { ctx, cancel := context.WithDeadline(context.WithValue(parentCtx, keyContextTask, task), nextAt) defer cancel() - if j.before != nil && j.before(task) { - task.Skipped = true - atomic.AddInt64(&j.statistics.SkippedTask, 1) + if j.deriveContext != nil { + ctx = j.deriveContext(ctx, task) + } + + if j.ctxBefore != nil { + if j.ctxBefore(ctx, task) { + task.Skipped = true + atomic.AddInt64(&j.statistics.SkippedTask, 1) + } + } else if j.before != nil { + if j.before(task) { + task.Skipped = true + atomic.AddInt64(&j.statistics.SkippedTask, 1) + } } if !task.Skipped { @@ -97,13 +110,7 @@ func (j *innerJob) Run() { task.BeginAt = &beginAt for i := 0; i < j.retryTimes; i++ { - if j.deriveContext != nil { - ctx, cancel := j.deriveContext(ctx, task) - task.Return = safeRun(ctx, j.run) - cancel() - } else { - task.Return = safeRun(ctx, j.run) - } + task.Return = safeRun(ctx, j.run) atomic.AddInt64(&j.statistics.TotalRun, 1) if i > 0 { atomic.AddInt64(&j.statistics.RetriedRun, 1) @@ -135,7 +142,9 @@ func (j *innerJob) Run() { } } - if j.after != nil { + if j.ctxAfter != nil { + j.ctxAfter(ctx, task) + } else if j.after != nil { j.after(task) } diff --git a/job_option.go b/job_option.go index e51f2f3..269baa8 100644 --- a/job_option.go +++ b/job_option.go @@ -9,34 +9,58 @@ import ( type JobOption func(job *innerJob) // BeforeFunc represents the function could be called before Run. +// Deprecated: Use BeforeContextFunc instead, it will be ignored if BeforeContextFunc is set. type BeforeFunc func(task Task) (skip bool) +// BeforeContextFunc represents the function could be called before Run with the given Task. +type BeforeContextFunc func(ctx context.Context, task Task) (skip bool) + // RunFunc represents the function could be called by a cron. type RunFunc func(ctx context.Context) error // AfterFunc represents the function could be called after Run. +// Deprecated: Use AfterContextFunc instead, it will be ignored if AfterContextFunc is set. type AfterFunc func(task Task) +// AfterContextFunc represents the function could be called after Run with the given Task. +type AfterContextFunc func(ctx context.Context, task Task) + // RetryInterval indicates how long should delay before retrying when run failed `triedTimes` times. type RetryInterval func(triedTimes int) time.Duration // DeriveContext indicates how to derive a new context for Run with the given Task, -type DeriveContext func(ctx context.Context, task Task) (context.Context, context.CancelFunc) +type DeriveContext func(ctx context.Context, task Task) context.Context // WithBeforeFunc specifies what to do before Run. +// Deprecated: Use WithBeforeContextFunc instead. func WithBeforeFunc(before BeforeFunc) JobOption { return func(job *innerJob) { job.before = before } } +// WithBeforeContextFunc specifies what to do before Run with the given Task. +func WithBeforeContextFunc(before BeforeContextFunc) JobOption { + return func(job *innerJob) { + job.ctxBefore = before + } +} + // WithAfterFunc specifies what to do after Run. +// Deprecated: Use WithAfterContextFunc instead. func WithAfterFunc(after AfterFunc) JobOption { return func(job *innerJob) { job.after = after } } +// WithAfterContextFunc specifies what to do after Run with the given Task. +func WithAfterContextFunc(after AfterContextFunc) JobOption { + return func(job *innerJob) { + job.ctxAfter = after + } +} + // WithRetryTimes specifies max times to retry, // retryTimes will be set as 1 if it is less than 1. func WithRetryTimes(retryTimes int) JobOption { From 46a24329b70ff0f8620c73c7a2fa0039bcc564b9 Mon Sep 17 00:00:00 2001 From: Jason Song Date: Tue, 12 May 2026 15:37:55 +0800 Subject: [PATCH 3/5] test: add more cases --- inner_job_test.go | 80 +++++++++++++++++++++++++++++++++ job_option.go | 2 +- job_option_test.go | 108 +++++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 186 insertions(+), 4 deletions(-) diff --git a/inner_job_test.go b/inner_job_test.go index 6029722..a00a370 100644 --- a/inner_job_test.go +++ b/inner_job_test.go @@ -90,8 +90,11 @@ func Test_innerJob_Run(t *testing.T) { entryGetter entryGetter key string spec string + deriveContext DeriveContext + ctxBefore BeforeContextFunc before BeforeFunc run RunFunc + ctxAfter AfterContextFunc after AfterFunc retryTimes int retryInterval RetryInterval @@ -329,6 +332,80 @@ func Test_innerJob_Run(t *testing.T) { RetriedRun: 0, }, }, + { + name: "ctx_before_skip", + fields: fields{ + cron: NewCron(WithAtomic(atomic)), + entryID: 1, + entryGetter: mockEntryGetter, + ctxBefore: func(ctx context.Context, task Task) (skip bool) { + return true + }, + run: func(ctx context.Context) error { + return nil + }, + ctxAfter: func(ctx context.Context, task Task) { + if !task.Skipped { + t.Fatal("expected task to be skipped") + } + }, + retryTimes: 1, + }, + statistics: Statistics{ + TotalTask: 1, + SkippedTask: 1, + }, + }, + { + name: "ctx_before_no_skip", + fields: fields{ + cron: NewCron(WithAtomic(atomic)), + entryID: 1, + entryGetter: mockEntryGetter, + ctxBefore: func(ctx context.Context, task Task) (skip bool) { + return false + }, + run: func(ctx context.Context) error { + return nil + }, + ctxAfter: func(ctx context.Context, task Task) { + if task.Return != nil { + t.Fatal(task.Return) + } + }, + retryTimes: 1, + }, + statistics: Statistics{ + TotalTask: 1, + PassedTask: 1, + TotalRun: 1, + PassedRun: 1, + }, + }, + { + name: "derive_context", + fields: fields{ + cron: NewCron(WithAtomic(atomic)), + entryID: 1, + entryGetter: mockEntryGetter, + deriveContext: func(ctx context.Context, task Task) context.Context { + return context.WithValue(ctx, "test_key", "test_value") + }, + run: func(ctx context.Context) error { + if ctx.Value("test_key") != "test_value" { + return errors.New("context value not found") + } + return nil + }, + retryTimes: 1, + }, + statistics: Statistics{ + TotalTask: 1, + PassedTask: 1, + TotalRun: 1, + PassedRun: 1, + }, + }, { name: "panic by runtime", fields: fields{ @@ -370,8 +447,11 @@ func Test_innerJob_Run(t *testing.T) { entryGetter: tt.fields.entryGetter, key: tt.fields.key, spec: tt.fields.spec, + deriveContext: tt.fields.deriveContext, + ctxBefore: tt.fields.ctxBefore, before: tt.fields.before, run: tt.fields.run, + ctxAfter: tt.fields.ctxAfter, after: tt.fields.after, retryTimes: tt.fields.retryTimes, retryInterval: tt.fields.retryInterval, diff --git a/job_option.go b/job_option.go index 269baa8..5a164de 100644 --- a/job_option.go +++ b/job_option.go @@ -28,7 +28,7 @@ type AfterContextFunc func(ctx context.Context, task Task) // RetryInterval indicates how long should delay before retrying when run failed `triedTimes` times. type RetryInterval func(triedTimes int) time.Duration -// DeriveContext indicates how to derive a new context for Run with the given Task, +// DeriveContext indicates how to derive a new context for Run with the given Task. type DeriveContext func(ctx context.Context, task Task) context.Context // WithBeforeFunc specifies what to do before Run. diff --git a/job_option_test.go b/job_option_test.go index b4bcfd4..a5c0cb9 100644 --- a/job_option_test.go +++ b/job_option_test.go @@ -1,15 +1,14 @@ package dcron import ( + "context" "fmt" "testing" "time" ) func TestWithAfterFunc(t *testing.T) { - after := func(task Task) { - - } + after := func(task Task) {} type args struct { after AfterFunc @@ -164,3 +163,106 @@ func TestWithNoMutex(t *testing.T) { }) } } + +func TestWithBeforeContextFunc(t *testing.T) { + before := func(ctx context.Context, task Task) (skip bool) { + return false + } + + type args struct { + before BeforeContextFunc + } + tests := []struct { + name string + args args + check func(t *testing.T, option JobOption) + }{ + { + name: "regular", + args: args{ + before: before, + }, + check: func(t *testing.T, option JobOption) { + j := &innerJob{} + option(j) + if fmt.Sprintf("%p", j.ctxBefore) != fmt.Sprintf("%p", before) { + t.Fatal() + } + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := WithBeforeContextFunc(tt.args.before) + tt.check(t, got) + }) + } +} + +func TestWithAfterContextFunc(t *testing.T) { + after := func(ctx context.Context, task Task) {} + + type args struct { + after AfterContextFunc + } + tests := []struct { + name string + args args + check func(t *testing.T, option JobOption) + }{ + { + name: "regular", + args: args{ + after: after, + }, + check: func(t *testing.T, option JobOption) { + j := &innerJob{} + option(j) + if fmt.Sprintf("%p", j.ctxAfter) != fmt.Sprintf("%p", after) { + t.Fatal() + } + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := WithAfterContextFunc(tt.args.after) + tt.check(t, got) + }) + } +} + +func TestWithDeriveContext(t *testing.T) { + deriveContext := func(ctx context.Context, task Task) context.Context { + return ctx + } + + type args struct { + deriveContext DeriveContext + } + tests := []struct { + name string + args args + check func(t *testing.T, option JobOption) + }{ + { + name: "regular", + args: args{ + deriveContext: deriveContext, + }, + check: func(t *testing.T, option JobOption) { + j := &innerJob{} + option(j) + if fmt.Sprintf("%p", j.deriveContext) != fmt.Sprintf("%p", deriveContext) { + t.Fatal() + } + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := WithDeriveContext(tt.args.deriveContext) + tt.check(t, got) + }) + } +} From 93e372f310abff4007313d71ead447abc8575f5f Mon Sep 17 00:00:00 2001 From: Jason Song Date: Tue, 12 May 2026 15:41:23 +0800 Subject: [PATCH 4/5] docs: comment for WithDeriveContext --- job_option.go | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/job_option.go b/job_option.go index 5a164de..4868029 100644 --- a/job_option.go +++ b/job_option.go @@ -28,7 +28,7 @@ type AfterContextFunc func(ctx context.Context, task Task) // RetryInterval indicates how long should delay before retrying when run failed `triedTimes` times. type RetryInterval func(triedTimes int) time.Duration -// DeriveContext indicates how to derive a new context for Run with the given Task. +// DeriveContext indicates how to derive a new context from the job's base context and the current Task. type DeriveContext func(ctx context.Context, task Task) context.Context // WithBeforeFunc specifies what to do before Run. @@ -91,7 +91,10 @@ func WithGroup(group Group) JobOption { } } -// WithDeriveContext indicates how to derive a new context for Run with the given Task. +// WithDeriveContext specifies how to derive a new context for the entire job execution, including +// before/after hooks, Run, and retry logic. The returned context must derive from the provided ctx +// to preserve the deadline, cancellation signal, and the embedded Task value (accessible via TaskFromContext). +// Returning a detached context (e.g., context.Background()) will break deadline enforcement and retry timeout logic. func WithDeriveContext(deriveContext DeriveContext) JobOption { return func(job *innerJob) { job.deriveContext = deriveContext From bde3e8e682012f73e745f84663e60c073d5c5cb2 Mon Sep 17 00:00:00 2001 From: Jason Song Date: Tue, 12 May 2026 15:42:55 +0800 Subject: [PATCH 5/5] test: fix lint --- inner_job_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/inner_job_test.go b/inner_job_test.go index a00a370..7078a11 100644 --- a/inner_job_test.go +++ b/inner_job_test.go @@ -389,10 +389,10 @@ func Test_innerJob_Run(t *testing.T) { entryID: 1, entryGetter: mockEntryGetter, deriveContext: func(ctx context.Context, task Task) context.Context { - return context.WithValue(ctx, "test_key", "test_value") + return context.WithValue(ctx, ctxKey("test_key"), "test_value") }, run: func(ctx context.Context) error { - if ctx.Value("test_key") != "test_value" { + if ctx.Value(ctxKey("test_key")) != "test_value" { return errors.New("context value not found") } return nil