mirror of
https://github.com/binwiederhier/ntfy.git
synced 2026-11-10 14:44:04 +08:00
Read the period end of a daily denial under the counter lock
This commit is contained in:
1 parent
2b46b90df9
commit
3fef91d342
2 files changed
+29
-1
No files matched your search
+2
-1
@@ -125,11 +125,12 @@ func (s *Service) Allow(key string, rule Rule, n int64) Decision {
|
||||
c := cs.counterNoLock(key, start)
|
||||
total := c.base[i] + c.unflushed[i]
|
||||
if rule.Daily > 0 && total+n > rule.Daily {
|
||||
end := time.Unix(c.periodStart+daySeconds, 0) // Read under the lock: another call may roll the counter over
|
||||
cs.mu.Unlock()
|
||||
if rule.Burst > 0 {
|
||||
s.refundTokens(key, rule.Kind, n)
|
||||
}
|
||||
return Decision{RetryAfter: time.Unix(c.periodStart+daySeconds, 0).Sub(dbNow)}
|
||||
return Decision{RetryAfter: end.Sub(dbNow)}
|
||||
}
|
||||
c.unflushed[i] += n
|
||||
if rule.Threshold == 0 || total+n >= rule.Threshold {
|
||||
|
||||
@@ -841,3 +841,30 @@ func TestService_Persist_LeftoverAfterRollover(t *testing.T) {
|
||||
}, usageRowOf(t, store, "b"))
|
||||
})
|
||||
}
|
||||
|
||||
func TestService_Allow_DenyRacesRollover(t *testing.T) {
|
||||
// A daily denial computes its retry time from the counter it decided on, while other calls
|
||||
// roll the counter over; run with -race
|
||||
var ticks atomic.Int64
|
||||
start := time.Date(2026, 10, 25, 12, 0, 0, 0, time.UTC)
|
||||
s := newService(&Config{}, func() time.Time {
|
||||
return start.Add(time.Duration(ticks.Add(1)) * time.Hour)
|
||||
})
|
||||
rule := Rule{
|
||||
Kind: KindMessages,
|
||||
Daily: 1,
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
for g := 0; g < 8; g++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for i := 0; i < 3000; i++ {
|
||||
if d := s.Allow("k", rule, 1); !d.OK {
|
||||
require.Greater(t, d.RetryAfter, time.Duration(0))
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
Reference in new issue
Block a user