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

     1  package middleware
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/wfusion/gofusion/common/infra/watermill/message"
     7  )
     8  
     9  // Throttle provides a middleware that limits the amount of messages processed per unit of time.
    10  // This may be done e.g. to prevent excessive load caused by running a handler on a long queue of unprocessed messages.
    11  type Throttle struct {
    12  	ticker *time.Ticker
    13  }
    14  
    15  // NewThrottle creates a new Throttle middleware.
    16  // Example duration and count: NewThrottle(10, time.Second) for 10 messages per second
    17  func NewThrottle(count int64, duration time.Duration) *Throttle {
    18  	return &Throttle{
    19  		ticker: time.NewTicker(duration / time.Duration(count)),
    20  	}
    21  }
    22  
    23  // Middleware returns the Throttle middleware.
    24  func (t Throttle) Middleware(h message.HandlerFunc) message.HandlerFunc {
    25  	return func(message *message.Message) ([]*message.Message, error) {
    26  		// throttle is shared by multiple handlers, which will wait for their "tick"
    27  		<-t.ticker.C
    28  
    29  		return h(message)
    30  	}
    31  }