-
Notifications
You must be signed in to change notification settings - Fork 116
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(robot): handle ratelimiting with constant backoff
Add a constant backoff in case we are being rate limited by the Robot API. This is done through an opaque wrapping of the robot.Client interface, so callers do not need to do this manually.
- Loading branch information
1 parent
8bb131f
commit 76a0366
Showing
6 changed files
with
191 additions
and
17 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,93 @@ | ||
package robot | ||
|
||
import ( | ||
"fmt" | ||
"strings" | ||
"time" | ||
|
||
hrobotmodels "github.com/syself/hrobot-go/models" | ||
) | ||
|
||
type rateLimitClient struct { | ||
robotClient Client | ||
|
||
waitTime time.Duration | ||
exceeded bool | ||
lastChecked time.Time | ||
} | ||
|
||
func WithRateLimit(rateLimitWaitTime time.Duration, robotClient Client) Client { | ||
return &rateLimitClient{ | ||
robotClient: robotClient, | ||
|
||
waitTime: rateLimitWaitTime, | ||
} | ||
} | ||
|
||
func (c *rateLimitClient) ServerGet(id int) (*hrobotmodels.Server, error) { | ||
if c.isExceeded() { | ||
return nil, c.getRateLimitError() | ||
} | ||
|
||
server, err := c.robotClient.ServerGet(id) | ||
c.handleError(err) | ||
return server, err | ||
} | ||
|
||
func (c *rateLimitClient) ServerGetList() ([]hrobotmodels.Server, error) { | ||
if c.isExceeded() { | ||
return nil, c.getRateLimitError() | ||
} | ||
|
||
servers, err := c.robotClient.ServerGetList() | ||
c.handleError(err) | ||
return servers, err | ||
} | ||
|
||
func (c *rateLimitClient) ResetGet(id int) (*hrobotmodels.Reset, error) { | ||
if c.isExceeded() { | ||
return nil, c.getRateLimitError() | ||
} | ||
|
||
reset, err := c.robotClient.ResetGet(id) | ||
c.handleError(err) | ||
return reset, err | ||
} | ||
|
||
func (c *rateLimitClient) set() { | ||
c.exceeded = true | ||
c.lastChecked = time.Now() | ||
} | ||
|
||
func (c *rateLimitClient) isExceeded() bool { | ||
if !c.exceeded { | ||
return false | ||
} | ||
|
||
if time.Now().Before(c.lastChecked.Add(c.waitTime)) { | ||
return true | ||
} | ||
// Waiting time is over. Should try again | ||
c.exceeded = false | ||
c.lastChecked = time.Time{} | ||
return false | ||
} | ||
|
||
func (c *rateLimitClient) handleError(err error) { | ||
if err == nil { | ||
return | ||
} | ||
|
||
if hrobotmodels.IsError(err, hrobotmodels.ErrorCodeRateLimitExceeded) || strings.Contains(err.Error(), "server responded with status code 403") { | ||
c.set() | ||
} | ||
} | ||
|
||
func (c *rateLimitClient) getRateLimitError() error { | ||
if !c.isExceeded() { | ||
return nil | ||
} | ||
|
||
nextPossibleCall := c.lastChecked.Add(c.waitTime) | ||
return fmt.Errorf("rate limit exceeded, next try at %q", nextPossibleCall.String()) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,68 @@ | ||
package robot | ||
|
||
import ( | ||
"strings" | ||
"testing" | ||
"time" | ||
|
||
"github.com/stretchr/testify/assert" | ||
hrobotmodels "github.com/syself/hrobot-go/models" | ||
|
||
"github.com/hetznercloud/hcloud-cloud-controller-manager/internal/mocks" | ||
) | ||
|
||
func TestRateLimit(t *testing.T) { | ||
mock := mocks.RobotClient{} | ||
mock.On("ServerGetList").Return([]hrobotmodels.Server{}, nil).Once() | ||
|
||
client := WithRateLimit(5*time.Minute, &mock) | ||
|
||
servers, err := client.ServerGetList() | ||
assert.NoError(t, err) | ||
assert.Len(t, servers, 0) | ||
mock.AssertNumberOfCalls(t, "ServerGetList", 1) | ||
|
||
mock.On("ServerGetList").Return(nil, hrobotmodels.Error{Code: hrobotmodels.ErrorCodeRateLimitExceeded, Message: "Rate limit exceeded"}).Once() | ||
_, err = client.ServerGetList() | ||
assert.Error(t, err) | ||
mock.AssertNumberOfCalls(t, "ServerGetList", 2) | ||
|
||
// No further call should be made | ||
_, err = client.ServerGetList() | ||
assert.Error(t, err) | ||
mock.AssertNumberOfCalls(t, "ServerGetList", 2) | ||
} | ||
|
||
func TestRateLimitIsExceeded(t *testing.T) { | ||
client := rateLimitClient{ | ||
waitTime: 5 * time.Minute, | ||
exceeded: true, | ||
lastChecked: time.Now(), | ||
} | ||
// Just exceeded | ||
assert.True(t, client.isExceeded()) | ||
|
||
// Exceeded longer than wait time ago | ||
client.lastChecked = time.Now().Add(-6 * time.Minute) | ||
assert.False(t, client.isExceeded()) | ||
|
||
// Not exceeded ever | ||
client.exceeded = false | ||
client.lastChecked = time.Time{} | ||
assert.False(t, client.isExceeded()) | ||
} | ||
|
||
func TestRateLimitGetRateLimitError(t *testing.T) { | ||
client := rateLimitClient{ | ||
waitTime: 5 * time.Minute, | ||
} | ||
err := client.getRateLimitError() | ||
assert.NoError(t, err) | ||
|
||
client.exceeded = true | ||
client.lastChecked = time.Now() | ||
|
||
err = client.getRateLimitError() | ||
assert.Error(t, err) | ||
assert.True(t, strings.HasPrefix(err.Error(), "rate limit exceeded, next try at ")) | ||
} |