Hardcoded reactions ratelimit (#398)

* add custom ratelimits

* check for nil ratelimiter

* Don't expose custom ratelimits to Session

* attempt to fix race conditions

* use defer instead

* Slightly improved ratelimiter

You shouldn't need to change the ratelimiters ratelimits while its
running so I removed the functions SetCustomRatelimit and
RemoveCustomRatelimit.
This commit is contained in:
Necroforger 2017-07-29 10:53:02 -04:00 committed by Chris Rhodes
parent 7bb0965a6f
commit 013faa1da4
2 changed files with 44 additions and 5 deletions

View file

@ -3,17 +3,26 @@ package discordgo
import ( import (
"net/http" "net/http"
"strconv" "strconv"
"strings"
"sync" "sync"
"sync/atomic" "sync/atomic"
"time" "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 // RateLimiter holds all ratelimit buckets
type RateLimiter struct { type RateLimiter struct {
sync.Mutex sync.Mutex
global *int64 global *int64
buckets map[string]*Bucket buckets map[string]*Bucket
globalRateLimit time.Duration globalRateLimit time.Duration
customRateLimits []*customRateLimit
} }
// NewRatelimiter returns a new RateLimiter // NewRatelimiter returns a new RateLimiter
@ -22,6 +31,13 @@ func NewRatelimiter() *RateLimiter {
return &RateLimiter{ return &RateLimiter{
buckets: make(map[string]*Bucket), buckets: make(map[string]*Bucket),
global: new(int64), 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, 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 r.buckets[key] = b
return b return b
} }
@ -76,13 +100,28 @@ type Bucket struct {
limit int limit int
reset time.Time reset time.Time
global *int64 global *int64
lastReset time.Time
customRateLimit *customRateLimit
} }
// Release unlocks the bucket and reads the headers to update the buckets ratelimit info // 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. // and locks up the whole thing in case if there's a global ratelimit.
func (b *Bucket) Release(headers http.Header) error { func (b *Bucket) Release(headers http.Header) error {
defer b.Unlock() 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 { if headers == nil {
return nil return nil
} }

View file

@ -199,7 +199,7 @@ type helloOp struct {
Trace []string `json:"_trace"` 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 const FailedHeartbeatAcks time.Duration = 5 * time.Millisecond
// heartbeat sends regular heartbeats to Discord so it knows the client // heartbeat sends regular heartbeats to Discord so it knows the client