diff --git a/ratelimit.go b/ratelimit.go index 876e98a..223c0d0 100644 --- a/ratelimit.go +++ b/ratelimit.go @@ -3,17 +3,26 @@ package discordgo import ( "net/http" "strconv" + "strings" "sync" "sync/atomic" "time" ) +// customRateLimit holds information for defining a custom rate limit +type customRateLimit struct { + suffix string + requests int + reset time.Duration +} + // RateLimiter holds all ratelimit buckets type RateLimiter struct { sync.Mutex - global *int64 - buckets map[string]*Bucket - globalRateLimit time.Duration + global *int64 + buckets map[string]*Bucket + globalRateLimit time.Duration + customRateLimits []*customRateLimit } // NewRatelimiter returns a new RateLimiter @@ -22,6 +31,13 @@ func NewRatelimiter() *RateLimiter { return &RateLimiter{ buckets: make(map[string]*Bucket), global: new(int64), + customRateLimits: []*customRateLimit{ + &customRateLimit{ + suffix: "//reactions//", + requests: 1, + reset: 200 * time.Millisecond, + }, + }, } } @@ -40,6 +56,14 @@ func (r *RateLimiter) getBucket(key string) *Bucket { global: r.global, } + // Check if there is a custom ratelimit set for this bucket ID. + for _, rl := range r.customRateLimits { + if strings.HasSuffix(b.Key, rl.suffix) { + b.customRateLimit = rl + break + } + } + r.buckets[key] = b return b } @@ -76,13 +100,28 @@ type Bucket struct { limit int reset time.Time global *int64 + + lastReset time.Time + customRateLimit *customRateLimit } // 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() + + // Check if the bucket uses a custom ratelimiter + if rl := b.customRateLimit; rl != nil { + if time.Now().Sub(b.lastReset) >= rl.reset { + b.remaining = rl.requests - 1 + b.lastReset = time.Now() + } + if b.remaining < 1 { + b.reset = time.Now().Add(rl.reset) + } + return nil + } + if headers == nil { return nil } diff --git a/wsapi.go b/wsapi.go index d050ffd..baec234 100644 --- a/wsapi.go +++ b/wsapi.go @@ -199,7 +199,7 @@ type helloOp struct { Trace []string `json:"_trace"` } -// Number of heartbeat intervals to wait until forcing a connection restart. +// FailedHeartbeatAcks is the Number of heartbeat intervals to wait until forcing a connection restart. const FailedHeartbeatAcks time.Duration = 5 * time.Millisecond // heartbeat sends regular heartbeats to Discord so it knows the client