feat(queue): decorate handler context before process events

Author: cmilesio

fake_queue.go

@@ -75,6 +75,8 @@ func (f *FakeQueue) WithContext(ctx context.Context) queueRuntime {
 	return &clone
 }
 
+func (f *FakeQueue) setHandlerContextDecorator(func(context.Context) context.Context) {}
+
 // Dispatch records a typed job payload in-memory using the fake default queue.
 // @group Testing
 //

observability.go

@@ -869,15 +869,20 @@ func (q *observedQueue) Register(jobType string, handler Handler) {
 		runtime.Register(jobType, handler)
 		return
 	}
-	runtime.Register(jobType, wrapObservedHandler(q.observer, q.driver, "", jobType, handler))
+	runtime.Register(jobType, wrapObservedHandler(q.observer, q.driver, "", jobType, nil, handler))
 }
 
 func (q *observedQueue) Driver() Driver {
 	return q.driver
 }
 
-func wrapObservedHandler(observer Observer, driver Driver, queueName string, jobType string, handler Handler) Handler {
+func wrapObservedHandler(observer Observer, driver Driver, queueName string, jobType string, ctxDecorator func(context.Context) context.Context, handler Handler) Handler {
 	return func(ctx context.Context, job Job) error {
+		if ctxDecorator != nil {
+			if decorated := ctxDecorator(ctx); decorated != nil {
+				ctx = decorated
+			}
+		}
 		opts := job.jobOptions()
 		effectiveQueue := queueName
 		if effectiveQueue == "" {

observability_branches_test.go

@@ -60,6 +60,15 @@ func (r *observerRecorder) Observe(_ context.Context, event Event) {
 	r.events = append(r.events, event)
 }
 
+type observerContextRecorder struct {
+	values []string
+}
+
+func (r *observerContextRecorder) Observe(ctx context.Context, _ Event) {
+	value, _ := ctx.Value("decorated").(string)
+	r.values = append(r.values, value)
+}
+
 func TestChannelObserver_DropIfFullAndBlockingSend(t *testing.T) {
 	ch := make(chan Event, 1)
 	ch <- Event{Kind: EventEnqueueAccepted}
@@ -116,7 +125,7 @@ func TestObservedQueue_DispatchClassifiesErrors(t *testing.T) {
 func TestWrapObservedHandler_EmitsRetriedAndArchived(t *testing.T) {
 	t.Run("retry path", func(t *testing.T) {
 		recorder := &observerRecorder{}
-		h := wrapObservedHandler(recorder, DriverSync, "", "job:retry", func(context.Context, Job) error {
+		h := wrapObservedHandler(recorder, DriverSync, "", "job:retry", nil, func(context.Context, Job) error {
 			return errors.New("boom")
 		})
 
@@ -134,7 +143,7 @@ func TestWrapObservedHandler_EmitsRetriedAndArchived(t *testing.T) {
 
 	t.Run("archive path", func(t *testing.T) {
 		recorder := &observerRecorder{}
-		h := wrapObservedHandler(recorder, DriverSync, "", "job:archive", func(context.Context, Job) error {
+		h := wrapObservedHandler(recorder, DriverSync, "", "job:archive", nil, func(context.Context, Job) error {
 			return errors.New("boom")
 		})
 
@@ -151,6 +160,30 @@ func TestWrapObservedHandler_EmitsRetriedAndArchived(t *testing.T) {
 	})
 }
 
+func TestWrapObservedHandler_DecoratesObserverContext(t *testing.T) {
+	recorder := &observerContextRecorder{}
+	h := wrapObservedHandler(recorder, DriverSync, "", "job:decorated", func(ctx context.Context) context.Context {
+		return context.WithValue(ctx, "decorated", "jobs")
+	}, func(ctx context.Context, _ Job) error {
+		if got, _ := ctx.Value("decorated").(string); got != "jobs" {
+			t.Fatalf("handler ctx value = %q, want jobs", got)
+		}
+		return nil
+	})
+
+	if err := h(context.Background(), NewJob("job:decorated").OnQueue("default")); err != nil {
+		t.Fatalf("wrapped handler returned error: %v", err)
+	}
+	if len(recorder.values) != 2 {
+		t.Fatalf("expected 2 observed events, got %d", len(recorder.values))
+	}
+	for i, got := range recorder.values {
+		if got != "jobs" {
+			t.Fatalf("observer ctx value at index %d = %q, want jobs", i, got)
+		}
+	}
+}
+
 func TestObservedQueue_WrapperMethods(t *testing.T) {
 	recorder := &observerRecorder{}
 	inner := &queueBackendStub{

queue.go

@@ -43,6 +43,9 @@ type queueRuntime interface {
 	// Ready checks backend readiness for dispatch/worker operation.
 	// @group Driver Integration
 	Ready(ctx context.Context) error
+
+	// setHandlerContextDecorator decorates handler execution context at registration time.
+	setHandlerContextDecorator(func(context.Context) context.Context)
 }
 
 // WorkerpoolConfig configures the in-memory workerpool q.
@@ -180,6 +183,7 @@ type queueCommon struct {
 	cfg    Config
 	driver Driver
 	ctx    context.Context
+	handlerContextDecorator func(context.Context) context.Context
 }
 
 type nativeQueueRuntime struct {
@@ -229,6 +233,13 @@ func (q *queueCommon) WithContext(ctx context.Context) *queueCommon {
 	return &clone
 }
 
+func (q *queueCommon) setHandlerContextDecorator(fn func(context.Context) context.Context) {
+	if q == nil {
+		return
+	}
+	q.handlerContextDecorator = fn
+}
+
 func (q *queueCommon) Dispatch(job any) error {
 	dispatchJob, err := q.jobFromAny(job)
 	if err != nil {
@@ -259,6 +270,20 @@ func (q *externalQueueRuntime) WithContext(ctx context.Context) queueRuntime {
 	return &clone
 }
 
+func (q *nativeQueueRuntime) setHandlerContextDecorator(fn func(context.Context) context.Context) {
+	if q == nil {
+		return
+	}
+	q.common.setHandlerContextDecorator(fn)
+}
+
+func (q *externalQueueRuntime) setHandlerContextDecorator(fn func(context.Context) context.Context) {
+	if q == nil {
+		return
+	}
+	q.common.setHandlerContextDecorator(fn)
+}
+
 func (q *nativeQueueRuntime) BusRegister(jobType string, handler busruntime.Handler) {
 	if handler == nil {
 		q.Register(jobType, nil)
@@ -510,7 +535,7 @@ func (q *queueCommon) wrapRegisteredHandler(jobType string, handler Handler) Han
 	if q.cfg.Driver == DriverRedis {
 		return handler
 	}
-	return wrapObservedHandler(q.cfg.Observer, q.cfg.Driver, "", jobType, handler)
+	return wrapObservedHandler(q.cfg.Observer, q.cfg.Driver, "", jobType, q.handlerContextDecorator, handler)
 }
 
 func (q *queueCommon) dispatchBusJob(ctx context.Context, jobType string, payload []byte, opts busruntime.JobOptions) error {

runtime.go

@@ -101,6 +101,7 @@ type Option func(*runtimeOptions)
 type runtimeOptions struct {
 	busOpts []bus.Option
 	workers int
+	handlerContextDecorator func(context.Context) context.Context
 }
 
 func (o *runtimeOptions) apply(opts []Option) {
@@ -207,6 +208,15 @@ func WithWorkers(count int) Option {
 	}
 }
 
+// WithHandlerContextDecorator decorates queue handler execution context before
+// process lifecycle events and handler execution run.
+// @group Queue
+func WithHandlerContextDecorator(fn func(context.Context) context.Context) Option {
+	return func(o *runtimeOptions) {
+		o.handlerContextDecorator = fn
+	}
+}
+
 // Queue is the high-level user-facing queue API.
 // It composes the queue runtime with the internal orchestration engine.
 // @group Queue
@@ -230,6 +240,9 @@ func newQueueFromRuntime(q queueRuntime, opts ...Option) (*Queue, error) {
 	if ro.workers > 0 && q != nil {
 		q = q.Workers(ro.workers)
 	}
+	if ro.handlerContextDecorator != nil && q != nil {
+		q.setHandlerContextDecorator(ro.handlerContextDecorator)
+	}
 	b, err := bus.New(q, ro.busOpts...)
 	if err != nil {
 		return nil, err