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

43 lines
773 B
Go
Raw Permalink Normal View History

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)
2024-09-25 02:29:00 +00:00
for key, bucket := range negativeBuckets.Range {
newv := bucket.Add(int32(-tokensPerInterval))
if newv <= 0 {
negativeBuckets.Delete(key)
}
2024-09-25 02:29:00 +00:00
}
}
}()
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
}
}