@@ -156,6 +156,13 @@ func WithCleanContext(h GenericHandlerFunc) GenericHandlerFunc {
156156return func (parent context.Context ,tb Deps ,args json.RawMessage ) (ret json.RawMessage ,err error ) {
157157child ,childCancel := context .WithCancel (context .Background ())
158158defer childCancel ()
159+ // Ensure that the child context has the same deadline as the parent
160+ // context.
161+ if deadline ,ok := parent .Deadline ();ok {
162+ deadlineCtx ,deadlineCancel := context .WithDeadline (child ,deadline )
163+ defer deadlineCancel ()
164+ child = deadlineCtx
165+ }
159166// Ensure that cancellation propagates from the parent context to the child context.
160167go func () {
161168select {
@@ -165,13 +172,6 @@ func WithCleanContext(h GenericHandlerFunc) GenericHandlerFunc {
165172childCancel ()
166173}
167174}()
168- // Also ensure that the child context has the same deadline as the parent
169- // context.
170- if deadline ,ok := parent .Deadline ();ok {
171- deadlineCtx ,deadlineCancel := context .WithDeadline (child ,deadline )
172- defer deadlineCancel ()
173- child = deadlineCtx
174- }
175175return h (child ,tb ,args )
176176}
177177}