5
0
mirror of https://github.com/cwinfo/matterbridge.git synced 2024-11-10 09:50:27 +00:00
matterbridge/vendor/github.com/bwmarrin/discordgo/ratelimit.go
2017-01-28 22:45:32 +01:00

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
}