1
0
Fork 0
mirror of https://github.com/diamondburned/arikawa.git synced 2025-02-12 22:43:32 +00:00

Integrated rate limiter

This commit is contained in:
diamondburned 2020-01-08 10:43:15 -08:00
parent e4cd4f9b69
commit af87f116cd
3 changed files with 13 additions and 18 deletions

View file

@ -3,6 +3,7 @@ package api
import ( import (
"net/http" "net/http"
"github.com/diamondburned/arikawa/api/rate"
"github.com/diamondburned/arikawa/httputil" "github.com/diamondburned/arikawa/httputil"
) )
@ -20,13 +21,16 @@ var UserAgent = "DiscordBot (https://github.com/diamondburned/arikawa, v0.0.1)"
type Client struct { type Client struct {
httputil.Client httputil.Client
Limiter *rate.Limiter
Token string Token string
} }
func NewClient(token string) *Client { func NewClient(token string) *Client {
cli := &Client{ cli := &Client{
Client: httputil.NewClient(), Client: httputil.NewClient(),
Token: token, Limiter: rate.NewLimiter(),
Token: token,
} }
tw := httputil.NewTransportWrapper() tw := httputil.NewTransportWrapper()
@ -41,7 +45,11 @@ func NewClient(token string) *Client {
r.Header.Set("X-RateLimit-Precision", "millisecond") r.Header.Set("X-RateLimit-Precision", "millisecond")
return nil // Rate limit stuff
return cli.Limiter.Acquire(r.Context(), r.URL.Path)
}
tw.Post = func(r *http.Response) error {
return cli.Limiter.Release(r.Request.URL.Path, r.Header)
} }
cli.Client.Transport = tw cli.Client.Transport = tw

View file

@ -61,6 +61,8 @@ func NewLimiter() *Limiter {
} }
func (l *Limiter) getBucket(path string, store bool) *bucket { func (l *Limiter) getBucket(path string, store bool) *bucket {
path = ParseBucketKey(path)
bc, ok := l.buckets.Load(path) bc, ok := l.buckets.Load(path)
if !ok && !store { if !ok && !store {
return nil return nil

View file

@ -1,15 +0,0 @@
package httputil
import "fmt"
// URL extends the normal URL and allows for a general string.
type URL struct {
Base string
URL string
}
func URLf(base string, v ...interface{}) URL {
return URL{base, fmt.Sprintf(base, v...)}
}
func (url URL) String() string { return url.URL }