github.com/wfusion/gofusion@v1.1.14/common/infra/watermill/message/router/middleware/timeout.go (about)

     1  package middleware
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/wfusion/gofusion/common/infra/watermill/message"
     8  )
     9  
    10  // Timeout makes the handler cancel the incoming message's context after a specified time.
    11  // Any timeout-sensitive functionality of the handler should listen on msg.Context().Done() to know when to fail.
    12  func Timeout(timeout time.Duration) func(message.HandlerFunc) message.HandlerFunc {
    13  	return func(h message.HandlerFunc) message.HandlerFunc {
    14  		return func(msg *message.Message) ([]*message.Message, error) {
    15  			ctx, cancel := context.WithTimeout(msg.Context(), timeout)
    16  			defer func() {
    17  				cancel()
    18  			}()
    19  
    20  			msg.SetContext(ctx)
    21  			return h(msg)
    22  		}
    23  	}
    24  }