From 0f5cb49b4c713fc0ff4a09ccf6c321c366ad2f26 Mon Sep 17 00:00:00 2001 From: Philipp Wolfer Date: Fri, 24 Nov 2023 00:22:36 +0100 Subject: [PATCH] Refactored common rate limit code into separate module --- internal/backends/funkwhale/client.go | 24 +---------- internal/backends/listenbrainz/client.go | 32 +++------------ internal/backends/spotify/client.go | 32 ++++----------- internal/ratelimit/httpheader.go | 52 ++++++++++++++++++++++++ 4 files changed, 69 insertions(+), 71 deletions(-) create mode 100644 internal/ratelimit/httpheader.go diff --git a/internal/backends/funkwhale/client.go b/internal/backends/funkwhale/client.go index 4080ac7..ebd049c 100644 --- a/internal/backends/funkwhale/client.go +++ b/internal/backends/funkwhale/client.go @@ -23,15 +23,13 @@ package funkwhale import ( "errors" - "net/http" "strconv" - "time" "github.com/go-resty/resty/v2" + "go.uploadedlobster.com/scotty/internal/ratelimit" ) const MaxItemsPerGet = 50 -const DefaultRateLimitWaitSeconds = 5 type Client struct { HttpClient *resty.Client @@ -46,25 +44,7 @@ func NewClient(serverUrl string, token string) Client { client.SetHeader("Accept", "application/json") // Handle rate limiting (see https://docs.funkwhale.audio/developer/api/rate-limit.html) - client.SetRetryCount(5) - client.AddRetryCondition( - func(r *resty.Response, err error) bool { - code := r.StatusCode() - return code == http.StatusTooManyRequests || code >= http.StatusInternalServerError - }, - ) - client.SetRetryMaxWaitTime(time.Duration(1 * time.Minute)) - client.SetRetryAfter(func(client *resty.Client, resp *resty.Response) (time.Duration, error) { - var err error - var retryAfter int = DefaultRateLimitWaitSeconds - if resp.StatusCode() == http.StatusTooManyRequests { - retryAfter, err = strconv.Atoi(resp.Header().Get("Retry-After")) - if err != nil { - retryAfter = DefaultRateLimitWaitSeconds - } - } - return time.Duration(retryAfter * int(time.Second)), err - }) + ratelimit.EnableHTTPHeaderRateLimit(client, "Retry-After") return Client{ HttpClient: client, diff --git a/internal/backends/listenbrainz/client.go b/internal/backends/listenbrainz/client.go index c90a637..da59144 100644 --- a/internal/backends/listenbrainz/client.go +++ b/internal/backends/listenbrainz/client.go @@ -23,20 +23,18 @@ package listenbrainz import ( "errors" - "net/http" "strconv" "time" "github.com/go-resty/resty/v2" + "go.uploadedlobster.com/scotty/internal/ratelimit" ) -const listenBrainzBaseURL = "https://api.listenbrainz.org/1/" - const ( - DefaultItemsPerGet = 25 - MaxItemsPerGet = 1000 - MaxListensPerRequest = 1000 - DefaultRateLimitWaitSeconds = 5 + listenBrainzBaseURL = "https://api.listenbrainz.org/1/" + DefaultItemsPerGet = 25 + MaxItemsPerGet = 1000 + MaxListensPerRequest = 1000 ) type Client struct { @@ -52,25 +50,7 @@ func NewClient(token string) Client { client.SetHeader("Accept", "application/json") // Handle rate limiting (see https://listenbrainz.readthedocs.io/en/latest/users/api/index.html#rate-limiting) - client.SetRetryCount(5) - client.AddRetryCondition( - func(r *resty.Response, err error) bool { - code := r.StatusCode() - return code == http.StatusTooManyRequests || code >= http.StatusInternalServerError - }, - ) - client.SetRetryMaxWaitTime(time.Duration(1 * time.Minute)) - client.SetRetryAfter(func(client *resty.Client, resp *resty.Response) (time.Duration, error) { - var err error - var retryAfter int = DefaultRateLimitWaitSeconds - if resp.StatusCode() == http.StatusTooManyRequests { - retryAfter, err = strconv.Atoi(resp.Header().Get("X-RateLimit-Reset-In")) - if err != nil { - retryAfter = DefaultRateLimitWaitSeconds - } - } - return time.Duration(retryAfter * int(time.Second)), err - }) + ratelimit.EnableHTTPHeaderRateLimit(client, "X-RateLimit-Reset-In") return Client{ HttpClient: client, diff --git a/internal/backends/spotify/client.go b/internal/backends/spotify/client.go index 93fa17f..2ec12b0 100644 --- a/internal/backends/spotify/client.go +++ b/internal/backends/spotify/client.go @@ -25,17 +25,18 @@ package spotify import ( "context" "errors" - "net/http" "strconv" "time" "github.com/go-resty/resty/v2" + "go.uploadedlobster.com/scotty/internal/ratelimit" "golang.org/x/oauth2" ) -const baseURL = "https://api.spotify.com/v1/" -const MaxItemsPerGet = 50 -const DefaultRateLimitWaitSeconds = 5 +const ( + baseURL = "https://api.spotify.com/v1/" + MaxItemsPerGet = 50 +) type Client struct { HttpClient *resty.Client @@ -47,25 +48,10 @@ func NewClient(token oauth2.TokenSource) Client { client := resty.NewWithClient(httpClient) client.SetBaseURL(baseURL) client.SetHeader("Accept", "application/json") - client.SetRetryCount(5) - client.AddRetryCondition( - func(r *resty.Response, err error) bool { - code := r.StatusCode() - return code == http.StatusTooManyRequests || code >= http.StatusInternalServerError - }, - ) - client.SetRetryMaxWaitTime(time.Duration(1 * time.Minute)) - client.SetRetryAfter(func(client *resty.Client, resp *resty.Response) (time.Duration, error) { - var err error - var retryAfter int = DefaultRateLimitWaitSeconds - if resp.StatusCode() == http.StatusTooManyRequests { - retryAfter, err = strconv.Atoi(resp.Header().Get("Retry-After")) - if err != nil { - retryAfter = DefaultRateLimitWaitSeconds - } - } - return time.Duration(retryAfter * int(time.Second)), err - }) + + // Handle rate limiting (see https://developer.spotify.com/documentation/web-api/concepts/rate-limits) + ratelimit.EnableHTTPHeaderRateLimit(client, "Retry-After") + return Client{ HttpClient: client, } diff --git a/internal/ratelimit/httpheader.go b/internal/ratelimit/httpheader.go new file mode 100644 index 0000000..91a05b5 --- /dev/null +++ b/internal/ratelimit/httpheader.go @@ -0,0 +1,52 @@ +/* +Copyright © 2023 Philipp Wolfer + +Scotty is free software: you can redistribute it and/or modify it under the +terms of the GNU General Public License as published by the Free Software +Foundation, either version 3 of the License, or (at your option) any later version. + +Scotty is distributed in the hope that it will be useful, but WITHOUT ANY +WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR +A PARTICULAR PURPOSE. See the GNU General Public License for more details. + +You should have received a copy of the GNU General Public License along with +Scotty. If not, see . +*/ + +package ratelimit + +import ( + "net/http" + "strconv" + "time" + + "github.com/go-resty/resty/v2" +) + +const ( + RetryCount = 5 + DefaultRateLimitWaitSeconds = 5 + MaxWaitTimeSeconds = 60 +) + +func EnableHTTPHeaderRateLimit(client *resty.Client, resetInHeader string) { + client.SetRetryCount(RetryCount) + client.AddRetryCondition( + func(r *resty.Response, err error) bool { + code := r.StatusCode() + return code == http.StatusTooManyRequests || code >= http.StatusInternalServerError + }, + ) + client.SetRetryMaxWaitTime(time.Duration(MaxWaitTimeSeconds * time.Second)) + client.SetRetryAfter(func(client *resty.Client, resp *resty.Response) (time.Duration, error) { + var err error + var retryAfter int = DefaultRateLimitWaitSeconds + if resp.StatusCode() == http.StatusTooManyRequests { + retryAfter, err = strconv.Atoi(resp.Header().Get(resetInHeader)) + if err != nil { + retryAfter = DefaultRateLimitWaitSeconds + } + } + return time.Duration(retryAfter * int(time.Second)), err + }) +}