bring in khatru and eventstore.
This commit is contained in:
@@ -0,0 +1,42 @@
|
||||
package policies
|
||||
|
||||
import (
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/puzpuzpuz/xsync/v3"
|
||||
)
|
||||
|
||||
func startRateLimitSystem[K comparable](
|
||||
tokensPerInterval int,
|
||||
interval time.Duration,
|
||||
maxTokens int,
|
||||
) func(key K) (ratelimited bool) {
|
||||
negativeBuckets := xsync.NewMapOf[K, *atomic.Int32]()
|
||||
maxTokensInt32 := int32(maxTokens)
|
||||
|
||||
go func() {
|
||||
for {
|
||||
time.Sleep(interval)
|
||||
for key, bucket := range negativeBuckets.Range {
|
||||
newv := bucket.Add(int32(-tokensPerInterval))
|
||||
if newv <= 0 {
|
||||
negativeBuckets.Delete(key)
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return func(key K) bool {
|
||||
nb, _ := negativeBuckets.LoadOrStore(key, &atomic.Int32{})
|
||||
|
||||
if nb.Load() < maxTokensInt32 {
|
||||
nb.Add(1)
|
||||
// rate limit not reached yet
|
||||
return false
|
||||
}
|
||||
|
||||
// rate limit reached
|
||||
return true
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user