well-goknown/vendor/github.com/fiatjaf/khatru/policies/helpers.go

44 lines
804 B
Go

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)
negativeBuckets.Range(func(key K, bucket *atomic.Int32) bool {
newv := bucket.Add(int32(-tokensPerInterval))
if newv <= 0 {
negativeBuckets.Delete(key)
}
return true
})
}
}()
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
}
}