mirror of
https://github.com/42wim/matterbridge.git
synced 2024-11-24 11:42:03 -08:00
158 lines
3.6 KiB
Go
158 lines
3.6 KiB
Go
|
package discordgo
|
||
|
|
||
|
import (
|
||
|
"net/http"
|
||
|
"strconv"
|
||
|
"sync"
|
||
|
"time"
|
||
|
)
|
||
|
|
||
|
// RateLimiter holds all ratelimit buckets
|
||
|
type RateLimiter struct {
|
||
|
sync.Mutex
|
||
|
global *Bucket
|
||
|
buckets map[string]*Bucket
|
||
|
globalRateLimit time.Duration
|
||
|
}
|
||
|
|
||
|
// NewRatelimiter returns a new RateLimiter
|
||
|
func NewRatelimiter() *RateLimiter {
|
||
|
|
||
|
return &RateLimiter{
|
||
|
buckets: make(map[string]*Bucket),
|
||
|
global: &Bucket{Key: "global"},
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// getBucket retrieves or creates a bucket
|
||
|
func (r *RateLimiter) getBucket(key string) *Bucket {
|
||
|
r.Lock()
|
||
|
defer r.Unlock()
|
||
|
|
||
|
if bucket, ok := r.buckets[key]; ok {
|
||
|
return bucket
|
||
|
}
|
||
|
|
||
|
b := &Bucket{
|
||
|
remaining: 1,
|
||
|
Key: key,
|
||
|
global: r.global,
|
||
|
}
|
||
|
|
||
|
r.buckets[key] = b
|
||
|
return b
|
||
|
}
|
||
|
|
||
|
// LockBucket Locks until a request can be made
|
||
|
func (r *RateLimiter) LockBucket(bucketID string) *Bucket {
|
||
|
|
||
|
b := r.getBucket(bucketID)
|
||
|
|
||
|
b.Lock()
|
||
|
|
||
|
// If we ran out of calls and the reset time is still ahead of us
|
||
|
// then we need to take it easy and relax a little
|
||
|
if b.remaining < 1 && b.reset.After(time.Now()) {
|
||
|
time.Sleep(b.reset.Sub(time.Now()))
|
||
|
|
||
|
}
|
||
|
|
||
|
// Check for global ratelimits
|
||
|
r.global.Lock()
|
||
|
r.global.Unlock()
|
||
|
|
||
|
b.remaining--
|
||
|
return b
|
||
|
}
|
||
|
|
||
|
// Bucket represents a ratelimit bucket, each bucket gets ratelimited individually (-global ratelimits)
|
||
|
type Bucket struct {
|
||
|
sync.Mutex
|
||
|
Key string
|
||
|
remaining int
|
||
|
limit int
|
||
|
reset time.Time
|
||
|
global *Bucket
|
||
|
}
|
||
|
|
||
|
// Release unlocks the bucket and reads the headers to update the buckets ratelimit info
|
||
|
// and locks up the whole thing in case if there's a global ratelimit.
|
||
|
func (b *Bucket) Release(headers http.Header) error {
|
||
|
|
||
|
defer b.Unlock()
|
||
|
if headers == nil {
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
remaining := headers.Get("X-RateLimit-Remaining")
|
||
|
reset := headers.Get("X-RateLimit-Reset")
|
||
|
global := headers.Get("X-RateLimit-Global")
|
||
|
retryAfter := headers.Get("Retry-After")
|
||
|
|
||
|
// If it's global just keep the main ratelimit mutex locked
|
||
|
if global != "" {
|
||
|
parsedAfter, err := strconv.Atoi(retryAfter)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
|
||
|
// Lock it in a new goroutine so that this isn't a blocking call
|
||
|
go func() {
|
||
|
// Make sure if several requests were waiting we don't sleep for n * retry-after
|
||
|
// where n is the amount of requests that were going on
|
||
|
sleepTo := time.Now().Add(time.Duration(parsedAfter) * time.Millisecond)
|
||
|
|
||
|
b.global.Lock()
|
||
|
|
||
|
sleepDuration := sleepTo.Sub(time.Now())
|
||
|
if sleepDuration > 0 {
|
||
|
time.Sleep(sleepDuration)
|
||
|
}
|
||
|
|
||
|
b.global.Unlock()
|
||
|
}()
|
||
|
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
// Update reset time if either retry after or reset headers are present
|
||
|
// Prefer retryafter because it's more accurate with time sync and whatnot
|
||
|
if retryAfter != "" {
|
||
|
parsedAfter, err := strconv.ParseInt(retryAfter, 10, 64)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
b.reset = time.Now().Add(time.Duration(parsedAfter) * time.Millisecond)
|
||
|
|
||
|
} else if reset != "" {
|
||
|
// Calculate the reset time by using the date header returned from discord
|
||
|
discordTime, err := http.ParseTime(headers.Get("Date"))
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
|
||
|
unix, err := strconv.ParseInt(reset, 10, 64)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
|
||
|
// Calculate the time until reset and add it to the current local time
|
||
|
// some extra time is added because without it i still encountered 429's.
|
||
|
// The added amount is the lowest amount that gave no 429's
|
||
|
// in 1k requests
|
||
|
delta := time.Unix(unix, 0).Sub(discordTime) + time.Millisecond*250
|
||
|
b.reset = time.Now().Add(delta)
|
||
|
}
|
||
|
|
||
|
// Udpate remaining if header is present
|
||
|
if remaining != "" {
|
||
|
parsedRemaining, err := strconv.ParseInt(remaining, 10, 32)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
b.remaining = int(parsedRemaining)
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
}
|