From 0fe202cc2a413491e7e9b3fb29421a3af1f39cc4 Mon Sep 17 00:00:00 2001 From: Roman Timofeev Date: Mon, 5 May 2025 17:11:55 +0300 Subject: [PATCH 1/2] Add WithDeadlineFunc/WithTimeout middlewares --- middleware.go | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/middleware.go b/middleware.go index f018831..7e852c6 100644 --- a/middleware.go +++ b/middleware.go @@ -248,3 +248,20 @@ func WithMetrics(app string) MiddlewareFunc { } } } + +// WithTimeout returns a MiddlewareFunc that wraps a Func with a context timeout. +func WithTimeout(timeout time.Duration) MiddlewareFunc { + return WithDeadlineFunc(func(ctx context.Context) time.Time { return time.Now().Add(timeout) }) +} + +// WithDeadlineFunc returns a MiddlewareFunc that wraps a Func with a context deadline. +// The deadline is determined by calling the provided function df. +func WithDeadlineFunc(df func(context.Context) time.Time) MiddlewareFunc { + return func(next Func) Func { + return func(ctx context.Context) error { + ctx, cancel := context.WithDeadline(ctx, df(ctx)) + defer cancel() + return next(ctx) + } + } +} From 178973adcd8b65727f7018aae372315081ed3b5f Mon Sep 17 00:00:00 2001 From: Roman Timofeev Date: Mon, 5 May 2025 17:20:45 +0300 Subject: [PATCH 2/2] add ContextValue wrapper and use empty struct as context key --- context.go | 22 ++++++++++++++ context_test.go | 81 +++++++++++++++++++++++++++++++++++++++++++++++++ cron.go | 41 ++++++++----------------- middleware.go | 14 ++++----- 4 files changed, 122 insertions(+), 36 deletions(-) create mode 100644 context.go create mode 100644 context_test.go diff --git a/context.go b/context.go new file mode 100644 index 0000000..91b0778 --- /dev/null +++ b/context.go @@ -0,0 +1,22 @@ +package cron + +import ( + "context" +) + +type ctxKey struct{} + +type ContextValue[K ~struct{}, T any] struct{} + +func NewContextValue[K ~struct{}, T any]() *ContextValue[K, T] { + return &ContextValue[K, T]{} +} + +func (p *ContextValue[K, T]) WithValue(ctx context.Context, v T) context.Context { + return context.WithValue(ctx, K{}, v) +} + +func (p *ContextValue[K, T]) FromContext(ctx context.Context) T { + v, _ := ctx.Value(K{}).(T) + return v +} diff --git a/context_test.go b/context_test.go new file mode 100644 index 0000000..85073e2 --- /dev/null +++ b/context_test.go @@ -0,0 +1,81 @@ +package cron + +import ( + "context" + "testing" + + . "github.com/smartystreets/goconvey/convey" +) + +type ( + keyString struct{} + keyInt struct{} + keyBool struct{} + keyPtr struct{} + keyStructA struct{} + keyStructB struct{} +) + +type sampleStruct struct { + ID int + Name string +} + +func TestCtxValue(t *testing.T) { + Convey("ContextValue should correctly store and retrieve values", t, func() { + Convey("string value", func() { + cv := NewContextValue[keyString, string]() + ctx := cv.WithValue(context.Background(), "hello") + So(cv.FromContext(ctx), ShouldEqual, "hello") + }) + + Convey("int value", func() { + cv := NewContextValue[keyInt, int]() + ctx := cv.WithValue(context.Background(), 123) + So(cv.FromContext(ctx), ShouldEqual, 123) + }) + + Convey("bool value", func() { + cv := NewContextValue[keyBool, bool]() + ctx := cv.WithValue(context.Background(), true) + So(cv.FromContext(ctx), ShouldBeTrue) + }) + + Convey("pointer value", func() { + cv := NewContextValue[keyPtr, *int]() + val := 42 + ctx := cv.WithValue(context.Background(), &val) + got := cv.FromContext(ctx) + So(got, ShouldNotBeNil) + So(*got, ShouldEqual, 42) + }) + + Convey("struct value", func() { + cv := NewContextValue[keyStructA, sampleStruct]() + value := sampleStruct{ID: 1, Name: "test"} + ctx := cv.WithValue(context.Background(), value) + So(cv.FromContext(ctx), ShouldResemble, value) + }) + + Convey("absent value returns zero", func() { + cv := NewContextValue[keyString, string]() + So(cv.FromContext(context.Background()), ShouldEqual, "") + }) + + Convey("different keys do not conflict", func() { + cvA := NewContextValue[keyStructA, string]() + cvB := NewContextValue[keyStructB, string]() + ctx := cvA.WithValue(context.Background(), "valueA") + So(cvA.FromContext(ctx), ShouldEqual, "valueA") + So(cvB.FromContext(ctx), ShouldEqual, "") + }) + + Convey("value can be overwritten", func() { + cv := NewContextValue[keyString, string]() + ctx := context.Background() + ctx = cv.WithValue(ctx, "first") + ctx = cv.WithValue(ctx, "second") + So(cv.FromContext(ctx), ShouldEqual, "second") + }) + }) +} diff --git a/cron.go b/cron.go index 4c9f0c6..836913c 100644 --- a/cron.go +++ b/cron.go @@ -12,15 +12,22 @@ import ( ) const ( - maintenanceKey contextKey = "maintenance" - nameKey contextKey = "name" - stateIdle cronState = "idle" stateDisabled cronState = "disabled" stateRunning cronState = "running" stateSkipped cronState = "skipped" ) +type ( + maintenanceCtxKey ctxKey + nameCtxKey ctxKey +) + +var ( + maintenance = NewContextValue[maintenanceCtxKey, bool]() + name = NewContextValue[nameCtxKey, string]() +) + var ( ErrSkipped = errors.New("skipped") ErrNotFound = errors.New("job not found") @@ -147,8 +154,8 @@ func (cm *Manager) Run(ctx context.Context) error { } // set context - ctx = NewNameContext(ctx, j.name) - ctx = NewMaintenanceContext(ctx, j.isMaintenance) + ctx = name.WithValue(ctx, j.name) + ctx = maintenance.WithValue(ctx, j.isMaintenance) // invoke main func with middleware cm.updateState(idx, stateRunning, nil) @@ -240,27 +247,3 @@ func newJob(name string, schedule Schedule, fn Func, isMaintenance bool) job { }, } } - -func NewMaintenanceContext(ctx context.Context, isMaintenance bool) context.Context { - return context.WithValue(ctx, maintenanceKey, isMaintenance) -} - -func MaintenanceFromContext(ctx context.Context) bool { - if r, ok := ctx.Value(maintenanceKey).(bool); ok { - return r - } - - return false -} - -func NewNameContext(ctx context.Context, name string) context.Context { - return context.WithValue(ctx, nameKey, name) -} - -func NameFromContext(ctx context.Context) string { - if v, ok := ctx.Value(nameKey).(string); ok { - return v - } - - return "" -} diff --git a/middleware.go b/middleware.go index 7e852c6..0a595c5 100644 --- a/middleware.go +++ b/middleware.go @@ -33,11 +33,11 @@ func WithLogger(pf LogPrintf, managerName string) MiddlewareFunc { pf("cron job %s job=%s duration=%v err=%q manager=%s maintenance=%v", state, - NameFromContext(ctx), + name.FromContext(ctx), time.Since(start), errMsg, managerName, - MaintenanceFromContext(ctx), + maintenance.FromContext(ctx), ) return err } @@ -57,7 +57,7 @@ func WithSLog(lg Logger) MiddlewareFunc { start := time.Now() err := next(ctx) - d, name := time.Since(start), NameFromContext(ctx) + d, name := time.Since(start), name.FromContext(ctx) switch { case errors.Is(err, ErrSkipped): lg.Print(ctx, "cron job skipped", "job", name, "duration", d) @@ -90,7 +90,7 @@ func WithSentry() MiddlewareFunc { if err != nil { sentryHub := sentry.CurrentHub().Clone() sentryHub.WithScope(func(scope *sentry.Scope) { - scope.SetTag("cron", NameFromContext(ctx)) + scope.SetTag("cron", name.FromContext(ctx)) }) sentryHub.CaptureException(err) } @@ -149,7 +149,7 @@ func WithSkipActive() MiddlewareFunc { return func(next Func) Func { return func(ctx context.Context) error { - name := NameFromContext(ctx) + name := name.FromContext(ctx) // check for running function mu.Lock() @@ -184,7 +184,7 @@ func WithMaintenance(p LogPrintf) MiddlewareFunc { return func(next Func) Func { return func(ctx context.Context) error { - name, isMaintenance := NameFromContext(ctx), MaintenanceFromContext(ctx) + name, isMaintenance := name.FromContext(ctx), maintenance.FromContext(ctx) if isMaintenance { pf("cron getting maintenance lock=%v", name) mutex.Lock() @@ -232,7 +232,7 @@ func WithMetrics(app string) MiddlewareFunc { return func(next Func) Func { return func(ctx context.Context) error { - name, start, state := NameFromContext(ctx), time.Now(), "ok" + name, start, state := name.FromContext(ctx), time.Now(), "ok" statActive.WithLabelValues(app, name).Inc() err := next(ctx)