diff --git a/_examples/rated.go b/_examples/rated.go index e37c31c..f8bb179 100644 --- a/_examples/rated.go +++ b/_examples/rated.go @@ -90,15 +90,26 @@ func init() { pre := "[Rate5] " go func() { + var lastcount = 0 + var count = 0 for { select { case msg := <-rd: - println(pre + "Limit: " + msg) + fmt.Printf("%s Limit: %s \n", pre, msg) + count++ case msg := <-rrd: - println(pre + "RegLimit: " + msg) + fmt.Printf("%s RegLimit: %s \n", pre, msg) + count++ case msg := <-crd: - println(pre + "CmdLimit: " + msg) + fmt.Printf("%s CmdLimit: %s \n", pre, msg) + count++ default: + if count - lastcount >= 25 { + lastcount = count + fmt.Println("Rater: ", Rater.GetGrandTotalRated()) + fmt.Println("RegRater: ", RegRater.GetGrandTotalRated()) + fmt.Println("CmdRater: ", CmdRater.GetGrandTotalRated()) + } time.Sleep(time.Duration(10) * time.Millisecond) } } diff --git a/models.go b/models.go index 547a82d..32101c1 100644 --- a/models.go +++ b/models.go @@ -1,4 +1,4 @@ -package ratelimit +package rate5 import ( "sync" @@ -28,11 +28,12 @@ type Limiter struct { // Ruleset is the actual ratelimitting model Ruleset Policy /* Debug mode (toggled here) enables debug messages - delivered through a channel. See: DebugChannel() */ + delivered through a channel. See: DebugChannel() */ Debug bool + count int known map[interface{}]int - mu *sync.Mutex + mu *sync.RWMutex } // Policy defines the mechanics of our ratelimiter @@ -40,10 +41,10 @@ type Policy struct { // Window defines the duration in seconds that we should keep track of ratelimit triggers Window int /* Burst is the amount of times that Check will not trigger a limit - within the duration defined by Window */ + within the duration defined by Window */ Burst int /* Strict mode punishes triggers of the ratelimit - by increasing the amount of time they have to wait - every time they trigger the limitter */ + by increasing the amount of time they have to wait + every time they trigger the limitter */ Strict bool } diff --git a/ratelimiter.go b/ratelimiter.go index 671a15b..e947f7a 100644 --- a/ratelimiter.go +++ b/ratelimiter.go @@ -1,4 +1,4 @@ -package ratelimit +package rate5 import ( "fmt" @@ -37,7 +37,7 @@ func newLimiter(window int, burst int, strict bool) *Limiter { } q.Patrons = cache.New(time.Duration(q.Ruleset.Window)*time.Second, 5*time.Second) q.known = make(map[interface{}]int) - q.mu = &sync.Mutex{} + q.mu = &sync.RWMutex{} return q } @@ -69,6 +69,7 @@ func (q *Limiter) Check(from Identity) bool { if count > q.Ruleset.Burst { if !q.Ruleset.Strict { q.debugPrint("ratelimit (limited): ", count, " ", src) + q.increment() return true } q.mu.Lock() @@ -83,6 +84,7 @@ func (q *Limiter) Check(from Identity) bool { println("Rate5: " + err.Error()) } q.debugPrint("ratelimit (strictly limited): ", count, " ", src) + q.increment() return true } return false @@ -97,6 +99,19 @@ func (q *Limiter) Peek(from Identity) bool { return false } +func (q *Limiter) increment() { + q.mu.Lock() + defer q.mu.Unlock() + q.count++ +} + +// GetGrandTotalRated returns the historic total amount of times we have ever reported something as ratelimited +func (q *Limiter) GetGrandTotalRated() int { + q.mu.RLock() + defer q.mu.RUnlock() + return q.count +} + func (q *Limiter) debugPrint(a ...interface{}) { if q.Debug { debugChannel <- fmt.Sprint(a...)