Rework context values handling

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
+}

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")
+		})
+	})
+}

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 ""
-}

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)
@@ -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)
+		}
+	}
+}