mirror of
https://git.sr.ht/~phw/scotty
synced 2025-06-01 19:38:34 +02:00
Rework ratelimit code
Simplify variables and avoid potential error if retry header reading fails
This commit is contained in:
parent
7542657925
commit
4ad89d287d
1 changed files with 10 additions and 11 deletions
|
@ -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
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue