1
0
Fork 0
mirror of https://github.com/binwiederhier/ntfy.git synced 2025-05-20 13:58:21 +02:00

Making RateLimiter and FixedLimiter, so they can both work with LimitWriter

This commit is contained in:
Philipp Heckel 2022-01-12 17:03:28 -05:00
parent f6b9ebb693
commit c76e55a1c8
7 changed files with 127 additions and 67 deletions

View file

@ -24,7 +24,7 @@ type visitor struct {
config *Config
ip string
requests *rate.Limiter
subscriptions *util.Limiter
subscriptions util.Limiter
emails *rate.Limiter
seen time.Time
mu sync.Mutex
@ -35,7 +35,7 @@ func newVisitor(conf *Config, ip string) *visitor {
config: conf,
ip: ip,
requests: rate.NewLimiter(rate.Every(conf.VisitorRequestLimitReplenish), conf.VisitorRequestLimitBurst),
subscriptions: util.NewLimiter(int64(conf.VisitorSubscriptionLimit)),
subscriptions: util.NewFixedLimiter(int64(conf.VisitorSubscriptionLimit)),
emails: rate.NewLimiter(rate.Every(conf.VisitorEmailLimitReplenish), conf.VisitorEmailLimitBurst),
seen: time.Now(),
}
@ -62,7 +62,7 @@ func (v *visitor) EmailAllowed() error {
func (v *visitor) SubscriptionAllowed() error {
v.mu.Lock()
defer v.mu.Unlock()
if err := v.subscriptions.Add(1); err != nil {
if err := v.subscriptions.Allow(1); err != nil {
return errVisitorLimitReached
}
return nil
@ -71,7 +71,7 @@ func (v *visitor) SubscriptionAllowed() error {
func (v *visitor) RemoveSubscription() {
v.mu.Lock()
defer v.mu.Unlock()
v.subscriptions.Sub(1)
v.subscriptions.Allow(-1)
}
func (v *visitor) Keepalive() {