forked from pothtonswer/discordmuffin
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:
parent
7bb0965a6f
commit
013faa1da4
2 changed files with 44 additions and 5 deletions
41
ratelimit.go
41
ratelimit.go
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
2
wsapi.go
2
wsapi.go
|
@ -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
|
||||||
|
|
Loading…
Reference in a new issue