Timeout middleware: context is not restored

[turip]

If you take a look at the Timeout middleware:

func Timeout(timeout time.Duration) func(message.HandlerFunc) message.HandlerFunc {
	return func(h message.HandlerFunc) message.HandlerFunc {
		return func(msg *message.Message) ([]*message.Message, error) {
			ctx, cancel := context.WithTimeout(msg.Context(), timeout)
			defer func() {
				cancel()
			}()

			msg.SetContext(ctx)
			return h(msg)
		}
	}
}

It is visible that the original context is never set to the original value if a timeout occurs. However the Retry middleware uses the message's context to ensure that the MaxElapsedTime is supported:

                 if r.MaxElapsedTime > 0 {
			var cancel func()
			ctx, cancel = context.WithTimeout(ctx, r.MaxElapsedTime)
			defer cancel()
		}

This means that if I have this chain Retry(Timeout(handler)), and a timeout occurs then the Retry will not be effective.

I would recommend to save the original context in Timeout and restore that in the defer function(). Or is this the expected behavior?