mirror of
https://github.com/oarkflow/mq.git
synced 2025-09-27 12:22:08 +08:00
feat: sig
This commit is contained in:
7
mq.go
7
mq.go
@@ -125,16 +125,13 @@ type RateLimiter struct {
|
||||
C chan struct{}
|
||||
}
|
||||
|
||||
// Modified RateLimiter: use blocking send to avoid discarding tokens.
|
||||
func NewRateLimiter(rate int, burst int) *RateLimiter {
|
||||
rl := &RateLimiter{C: make(chan struct{}, burst)}
|
||||
ticker := time.NewTicker(time.Second / time.Duration(rate))
|
||||
go func() {
|
||||
for range ticker.C {
|
||||
select {
|
||||
case rl.C <- struct{}{}:
|
||||
default:
|
||||
// bucket full; token discarded
|
||||
}
|
||||
rl.C <- struct{}{} // blocking send; tokens queue for deferred task processing
|
||||
}
|
||||
}()
|
||||
return rl
|
||||
|
26
options.go
26
options.go
@@ -130,6 +130,8 @@ func defaultOptions() *Options {
|
||||
maxMemoryLoad: 5000000,
|
||||
storage: NewMemoryTaskStorage(10 * time.Minute),
|
||||
logger: logger.NewDefaultLogger(),
|
||||
BrokerRateLimiter: NewRateLimiter(10, 5),
|
||||
ConsumerRateLimiter: NewRateLimiter(10, 5),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -254,3 +256,27 @@ func WithJitterPercent(val float64) Option {
|
||||
opts.jitterPercent = val
|
||||
}
|
||||
}
|
||||
|
||||
func WithBrokerRateLimiter(rate int, burst int) Option {
|
||||
return func(opts *Options) {
|
||||
opts.BrokerRateLimiter = NewRateLimiter(rate, burst)
|
||||
}
|
||||
}
|
||||
|
||||
func WithConsumerRateLimiter(rate int, burst int) Option {
|
||||
return func(opts *Options) {
|
||||
opts.ConsumerRateLimiter = NewRateLimiter(rate, burst)
|
||||
}
|
||||
}
|
||||
|
||||
func DisableBrokerRateLimit() Option {
|
||||
return func(opts *Options) {
|
||||
opts.BrokerRateLimiter = nil
|
||||
}
|
||||
}
|
||||
|
||||
func DisableConsumerRateLimit() Option {
|
||||
return func(opts *Options) {
|
||||
opts.ConsumerRateLimiter = nil
|
||||
}
|
||||
}
|
||||
|
Reference in New Issue
Block a user