forked from pothtonswer/discordmuffin
56b19073d3
* Added ratelimiter Handles the new ratelimit headers - X-RateLimit-Remaining - X-RateLimit-Reset - X-RateLimit-Global * Pad only reset time with a second * Moved ratelimiter out of internal package * Change for loop, move global ratelimit check inside sleep check * Moved ratelimiter locking to getBucket * Added global bucket * Changed how bucket id's are done Now each request function will need to specify the bucket id if the endpoint contains minor variables * Allow empty bucketID in request * Remove some uneeded Endpoint* function calls * Added test for global ratelimits * Fixed a silly little mistake causing incorrect ratelimits * Update test comments, Fixed treating a endpoint as 2 in ratelimiting * Use date header from discord instead of relying on sys time sync * Update all REST functions to use RequestWithBucketID * Embed mutex into bucket * Added webhook and reaction buckets
112 lines
2.8 KiB
Go
112 lines
2.8 KiB
Go
package discordgo
|
|
|
|
import (
|
|
"net/http"
|
|
"strconv"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// This test takes ~2 seconds to run
|
|
func TestRatelimitReset(t *testing.T) {
|
|
rl := NewRatelimiter()
|
|
|
|
sendReq := func(endpoint string) {
|
|
bucket := rl.LockBucket(endpoint)
|
|
|
|
headers := http.Header(make(map[string][]string))
|
|
|
|
headers.Set("X-RateLimit-Remaining", "0")
|
|
// Reset for approx 2 seconds from now
|
|
headers.Set("X-RateLimit-Reset", strconv.FormatInt(time.Now().Add(time.Second*2).Unix(), 10))
|
|
headers.Set("Date", time.Now().Format(time.RFC850))
|
|
|
|
err := bucket.Release(headers)
|
|
if err != nil {
|
|
t.Errorf("Release returned error: %v", err)
|
|
}
|
|
}
|
|
|
|
sent := time.Now()
|
|
sendReq("/guilds/99/channels")
|
|
sendReq("/guilds/55/channels")
|
|
sendReq("/guilds/66/channels")
|
|
|
|
sendReq("/guilds/99/channels")
|
|
sendReq("/guilds/55/channels")
|
|
sendReq("/guilds/66/channels")
|
|
|
|
// We hit the same endpoint 2 times, so we should only be ratelimited 2 second
|
|
// And always less than 4 seconds (unless you're on a stoneage computer or using swap or something...)
|
|
if time.Since(sent) >= time.Second && time.Since(sent) < time.Second*4 {
|
|
t.Log("OK", time.Since(sent))
|
|
} else {
|
|
t.Error("Did not ratelimit correctly, got:", time.Since(sent))
|
|
}
|
|
}
|
|
|
|
// This test takes ~1 seconds to run
|
|
func TestRatelimitGlobal(t *testing.T) {
|
|
rl := NewRatelimiter()
|
|
|
|
sendReq := func(endpoint string) {
|
|
bucket := rl.LockBucket(endpoint)
|
|
|
|
headers := http.Header(make(map[string][]string))
|
|
|
|
headers.Set("X-RateLimit-Global", "1")
|
|
// Reset for approx 1 seconds from now
|
|
headers.Set("Retry-After", "1000")
|
|
|
|
err := bucket.Release(headers)
|
|
if err != nil {
|
|
t.Errorf("Release returned error: %v", err)
|
|
}
|
|
}
|
|
|
|
sent := time.Now()
|
|
|
|
// This should trigger a global ratelimit
|
|
sendReq("/guilds/99/channels")
|
|
time.Sleep(time.Millisecond * 100)
|
|
|
|
// This shouldn't go through in less than 1 second
|
|
sendReq("/guilds/55/channels")
|
|
|
|
if time.Since(sent) >= time.Second && time.Since(sent) < time.Second*2 {
|
|
t.Log("OK", time.Since(sent))
|
|
} else {
|
|
t.Error("Did not ratelimit correctly, got:", time.Since(sent))
|
|
}
|
|
}
|
|
|
|
func BenchmarkRatelimitSingleEndpoint(b *testing.B) {
|
|
rl := NewRatelimiter()
|
|
for i := 0; i < b.N; i++ {
|
|
sendBenchReq("/guilds/99/channels", rl)
|
|
}
|
|
}
|
|
|
|
func BenchmarkRatelimitParallelMultiEndpoints(b *testing.B) {
|
|
rl := NewRatelimiter()
|
|
b.RunParallel(func(pb *testing.PB) {
|
|
i := 0
|
|
for pb.Next() {
|
|
sendBenchReq("/guilds/"+strconv.Itoa(i)+"/channels", rl)
|
|
i++
|
|
}
|
|
})
|
|
}
|
|
|
|
// Does not actually send requests, but locks the bucket and releases it with made-up headers
|
|
func sendBenchReq(endpoint string, rl *RateLimiter) {
|
|
bucket := rl.LockBucket(endpoint)
|
|
|
|
headers := http.Header(make(map[string][]string))
|
|
|
|
headers.Set("X-RateLimit-Remaining", "10")
|
|
headers.Set("X-RateLimit-Reset", strconv.FormatInt(time.Now().Unix(), 10))
|
|
headers.Set("Date", time.Now().Format(time.RFC850))
|
|
|
|
bucket.Release(headers)
|
|
}
|