Rework ratelimit code

Simplify variables and avoid potential error if retry header reading fails
This commit is contained in:
Philipp Wolfer 2025-05-24 16:47:13 +02:00
parent 7542657925
commit 4ad89d287d
No known key found for this signature in database
GPG key ID: 8FDF744D4919943B

View file

@ -1,5 +1,5 @@
/* /*
Copyright © 2023 Philipp Wolfer <phw@uploadedlobster.com> Copyright © 2023-2025 Philipp Wolfer <phw@uploadedlobster.com>
Scotty is free software: you can redistribute it and/or modify it under the 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 terms of the GNU General Public License as published by the Free Software
@ -25,9 +25,9 @@ import (
) )
const ( const (
RetryCount = 5 RetryCount = 5
DefaultRateLimitWaitSeconds = 5 DefaultRateLimitWait = 5 * time.Second
MaxWaitTimeSeconds = 60 MaxWaitTime = 60 * time.Second
) )
// Implements rate HTTP header based limiting for resty. // Implements rate HTTP header based limiting for resty.
@ -47,16 +47,15 @@ func EnableHTTPHeaderRateLimit(client *resty.Client, resetInHeader string) {
return code == http.StatusTooManyRequests || code >= http.StatusInternalServerError return code == http.StatusTooManyRequests || code >= http.StatusInternalServerError
}, },
) )
client.SetRetryMaxWaitTime(time.Duration(MaxWaitTimeSeconds * time.Second)) client.SetRetryMaxWaitTime(MaxWaitTime)
client.SetRetryAfter(func(client *resty.Client, resp *resty.Response) (time.Duration, error) { client.SetRetryAfter(func(client *resty.Client, resp *resty.Response) (time.Duration, error) {
var err error retryAfter := DefaultRateLimitWait
var retryAfter int = DefaultRateLimitWaitSeconds
if resp.StatusCode() == http.StatusTooManyRequests { if resp.StatusCode() == http.StatusTooManyRequests {
retryAfter, err = strconv.Atoi(resp.Header().Get(resetInHeader)) retryAfterHeader, err := strconv.Atoi(resp.Header().Get(resetInHeader))
if err != nil { if err == nil {
retryAfter = DefaultRateLimitWaitSeconds retryAfter = time.Duration(retryAfterHeader) * time.Second
} }
} }
return time.Duration(retryAfter * int(time.Second)), err return retryAfter, nil
}) })
} }