kayos/Rate5
1
2
Fork 0
mirror of https://github.com/yunginnanet/Rate5 synced 2024-06-29 18:40:48 +00:00
Code Issues Releases Wiki Activity

Refactor (perf): Trade slightly higher execution time for less memory allocation

Browse Source
This commit is contained in:
kayos@tcp.direct 2024-02-06 23:46:42 -08:00
parent 6c81593f2a
commit 598aa983e8
Signed by: kayos
GPG Key ID: 4B841471B4BEE979
1 changed files with 20 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
ratelimiter.go
View File

@ -106,7 +106,7 @@ func (q *Limiter) getHitsPtr(src string) *int64 {
return newPtr return newPtr
} }
func (q *Limiter) strictLogic(src string, count int64) { func (q *Limiter) strictLogic(src string, count *atomic.Int64) {
knownHits := q.getHitsPtr(src) knownHits := q.getHitsPtr(src)
atomic.AddInt64(knownHits, 1) atomic.AddInt64(knownHits, 1)
var extwindow int64 var extwindow int64
@ -132,26 +132,31 @@ func (q *Limiter) CheckStringer(from fmt.Stringer) bool {
// Check checks and increments an Identities UniqueKey() output against a list of cached strings to determine and raise it's ratelimitting status. // Check checks and increments an Identities UniqueKey() output against a list of cached strings to determine and raise it's ratelimitting status.
func (q *Limiter) Check(from Identity) (limited bool) { func (q *Limiter) Check(from Identity) (limited bool) {
var aval any
var count int64 var count int64
var err error var ok bool
src := from.UniqueKey() aval, ok = q.Patrons.Get(from.UniqueKey())
count, err = q.Patrons.IncrementInt64(src, 1) switch {
if err != nil { case !ok:
// IncrementInt64 should only error if the value is not an int64, so we can assume it's a new key. q.debugPrintf("ratelimit %s (new) ", from.UniqueKey())
q.debugPrintf("ratelimit %s (new) ", src) aval = &atomic.Int64{}
aval.(*atomic.Int64).Store(1)
// We can't reproduce this throwing an error, we can only assume that the key is new. // We can't reproduce this throwing an error, we can only assume that the key is new.
_ = q.Patrons.Add(src, int64(1), time.Duration(q.Ruleset.Window)*time.Second) _ = q.Patrons.Add(from.UniqueKey(), aval, time.Duration(q.Ruleset.Window)*time.Second)
return false
}
if count < q.Ruleset.Burst {
return false return false
case ok && aval != nil:
count = aval.(*atomic.Int64).Add(1)
_ = q.Patrons.Replace(from.UniqueKey(), aval, time.Duration(q.Ruleset.Window)*time.Second)
if count < q.Ruleset.Burst {
return false
}
} }
if q.Ruleset.Strict { if q.Ruleset.Strict {
q.strictLogic(src, count) q.strictLogic(from.UniqueKey(), aval.(*atomic.Int64))
} else { return true
q.debugPrintf("ratelimit %s: last count %d. time: %s",
src, count, time.Duration(q.Ruleset.Window)*time.Second)
} }
q.debugPrintf("ratelimit %s: last count %d. time: %s",
from.UniqueKey(), count, time.Duration(q.Ruleset.Window)*time.Second)
return true return true
} }