Skip to content

Commit af649c6

Browse files
vladhssvladhss
authored
Merge pull request #67 from OpenVPN/rate-limit-updates
Integrate with rate limit response headers from API
2 parents dc969f9 + cef87ba commit af649c6

8 files changed

+94
-51
lines changed

‎cloudconnexa/cloudconnexa.go‎

Lines changed: 51 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"fmt"
88
"io"
99
"net/http"
10+
"strconv"
1011
"strings"
1112
"time"
1213

@@ -21,9 +22,10 @@ const (
2122
type Client struct {
2223
client *http.Client
2324

24-
BaseURL string
25-
Token string
26-
RateLimiter *rate.Limiter
25+
BaseURL string
26+
Token string
27+
ReadRateLimiter *rate.Limiter
28+
UpdateRateLimiter *rate.Limiter
2729

2830
UserAgent string
2931

@@ -115,11 +117,12 @@ func NewClient(baseURL, clientID, clientSecret string) (*Client, error) {
115117
}
116118

117119
c := &Client{
118-
client: httpClient,
119-
BaseURL: baseURL,
120-
Token: credentials.AccessToken,
121-
UserAgent: userAgent,
122-
RateLimiter: rate.NewLimiter(rate.Every(1*time.Second), 5),
120+
client: httpClient,
121+
BaseURL: baseURL,
122+
Token: credentials.AccessToken,
123+
UserAgent: userAgent,
124+
ReadRateLimiter: rate.NewLimiter(rate.Every(1*time.Second), 1),
125+
UpdateRateLimiter: rate.NewLimiter(rate.Every(4*time.Second), 1),
123126
}
124127
c.common.client = c
125128
c.HostConnectors = (*HostConnectorsService)(&c.common)
@@ -158,7 +161,13 @@ func (c *Client) setCommonHeaders(req *http.Request) {
158161
// DoRequest executes an HTTP request with authentication and rate limiting.
159162
// It automatically adds the Bearer token, sets headers, and handles errors.
160163
func (c *Client) DoRequest(req *http.Request) ([]byte, error) {
161-
err := c.RateLimiter.Wait(context.Background())
164+
var rateLimiter *rate.Limiter
165+
if req.Method == "GET" {
166+
rateLimiter = c.ReadRateLimiter
167+
} else {
168+
rateLimiter = c.UpdateRateLimiter
169+
}
170+
err := rateLimiter.Wait(context.Background())
162171
if err != nil {
163172
return nil, err
164173
}
@@ -185,9 +194,42 @@ func (c *Client) DoRequest(req *http.Request) ([]byte, error) {
185194
return nil, &ErrClientResponse{status: res.StatusCode, body: string(body)}
186195
}
187196

197+
err = c.AssignLimits(res, rateLimiter)
198+
if err != nil {
199+
return nil, err
200+
}
201+
188202
return body, nil
189203
}
190204

205+
// AssignLimits adjusts the rate limiter according to values received in response headers from the API
206+
func (c *Client) AssignLimits(res *http.Response, rateLimiter *rate.Limiter) error {
207+
rateHeader := res.Header.Get("X-RateLimit-Replenish-Rate")
208+
timeHeader := res.Header.Get("X-RateLimit-Replenish-Time")
209+
remainingHeader := res.Header.Get("X-RateLimit-Remaining")
210+
211+
if rateHeader != "" && timeHeader != "" && remainingHeader != "" {
212+
rateValue, err := strconv.Atoi(rateHeader)
213+
if err != nil {
214+
return err
215+
}
216+
timeValue, err := strconv.Atoi(timeHeader)
217+
if err != nil {
218+
return err
219+
}
220+
remainingValue, err := strconv.Atoi(remainingHeader)
221+
if err != nil {
222+
return err
223+
}
224+
if remainingValue <= 0 {
225+
remainingValue = 1
226+
}
227+
rateLimiter.SetLimit(rate.Every(time.Duration(timeValue * 1_000_000_000 / rateValue)))
228+
rateLimiter.SetBurst(remainingValue)
229+
}
230+
return nil
231+
}
232+
191233
// GetV1Url returns the base URL for CloudConnexa API v1 endpoints.
192234
func (c *Client) GetV1Url() string {
193235
return c.BaseURL + "/api/v1"

‎cloudconnexa/cloudconnexa_test.go‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -79,10 +79,11 @@ func TestDoRequest(t *testing.T) {
7979
defer server.Close()
8080

8181
client := &Client{
82-
client: server.Client(),
83-
BaseURL: server.URL,
84-
Token: "mock-access-token",
85-
RateLimiter: rate.NewLimiter(rate.Every(1), 5),
82+
client: server.Client(),
83+
BaseURL: server.URL,
84+
Token: "mock-access-token",
85+
ReadRateLimiter: rate.NewLimiter(rate.Every(1), 5),
86+
UpdateRateLimiter: rate.NewLimiter(rate.Every(1), 5),
8687
}
8788

8889
tests := []struct {

‎cloudconnexa/devices_test.go‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,11 @@ import (
1313
// createTestClient creates a test client with the given server
1414
func createTestClient(server *httptest.Server) *Client {
1515
client := &Client{
16-
client: server.Client(),
17-
BaseURL: server.URL,
18-
Token: "test-token",
19-
RateLimiter: rate.NewLimiter(rate.Every(1*time.Second), 5),
16+
client: server.Client(),
17+
BaseURL: server.URL,
18+
Token: "test-token",
19+
ReadRateLimiter: rate.NewLimiter(rate.Every(1), 5),
20+
UpdateRateLimiter: rate.NewLimiter(rate.Every(1), 5),
2021
}
2122
client.Devices = (*DevicesService)(&service{client: client})
2223
client.Sessions = (*SessionsService)(&service{client: client})
@@ -280,7 +281,8 @@ func TestDevicesService_ListByUserID(t *testing.T) {
280281

281282
func TestDevicesService_List_InvalidSize(t *testing.T) {
282283
client := &Client{
283-
RateLimiter: rate.NewLimiter(rate.Every(1*time.Second), 5),
284+
ReadRateLimiter: rate.NewLimiter(rate.Every(1), 5),
285+
UpdateRateLimiter: rate.NewLimiter(rate.Every(1), 5),
284286
}
285287
client.Devices = (*DevicesService)(&service{client: client})
286288

‎cloudconnexa/dns_records_direct_test.go‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,21 +2,20 @@ package cloudconnexa
22

33
import (
44
"encoding/json"
5+
"golang.org/x/time/rate"
56
"net/http"
67
"net/http/httptest"
78
"testing"
8-
"time"
9-
10-
"golang.org/x/time/rate"
119
)
1210

1311
// createTestDNSClient creates a test client with the given server for DNS testing
1412
func createTestDNSClient(server *httptest.Server) *Client {
1513
client := &Client{
16-
client: server.Client(),
17-
BaseURL: server.URL,
18-
Token: "test-token",
19-
RateLimiter: rate.NewLimiter(rate.Every(1*time.Second), 5),
14+
client: server.Client(),
15+
BaseURL: server.URL,
16+
Token: "test-token",
17+
ReadRateLimiter: rate.NewLimiter(rate.Every(1), 5),
18+
UpdateRateLimiter: rate.NewLimiter(rate.Every(1), 5),
2019
}
2120
client.DNSRecords = (*DNSRecordsService)(&service{client: client})
2221
return client

‎cloudconnexa/host_ip_services_dto_test.go‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,21 +2,20 @@ package cloudconnexa
22

33
import (
44
"encoding/json"
5+
"golang.org/x/time/rate"
56
"net/http"
67
"net/http/httptest"
78
"testing"
8-
"time"
9-
10-
"golang.org/x/time/rate"
119
)
1210

1311
// createTestHostIPServicesClient creates a test client with the given server for host IP services testing
1412
func createTestHostIPServicesClient(server *httptest.Server) *Client {
1513
client := &Client{
16-
client: server.Client(),
17-
BaseURL: server.URL,
18-
Token: "test-token",
19-
RateLimiter: rate.NewLimiter(rate.Every(1*time.Second), 5),
14+
client: server.Client(),
15+
BaseURL: server.URL,
16+
Token: "test-token",
17+
ReadRateLimiter: rate.NewLimiter(rate.Every(1), 5),
18+
UpdateRateLimiter: rate.NewLimiter(rate.Every(1), 5),
2019
}
2120
client.HostIPServices = (*HostIPServicesService)(&service{client: client})
2221
return client

‎cloudconnexa/network_connectors_ipsec_test.go‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,21 +2,20 @@ package cloudconnexa
22

33
import (
44
"encoding/json"
5+
"golang.org/x/time/rate"
56
"net/http"
67
"net/http/httptest"
78
"testing"
8-
"time"
9-
10-
"golang.org/x/time/rate"
119
)
1210

1311
// createTestIPsecClient creates a test client with the given server for IPsec testing
1412
func createTestIPsecClient(server *httptest.Server) *Client {
1513
client := &Client{
16-
client: server.Client(),
17-
BaseURL: server.URL,
18-
Token: "test-token",
19-
RateLimiter: rate.NewLimiter(rate.Every(1*time.Second), 5),
14+
client: server.Client(),
15+
BaseURL: server.URL,
16+
Token: "test-token",
17+
ReadRateLimiter: rate.NewLimiter(rate.Every(1), 5),
18+
UpdateRateLimiter: rate.NewLimiter(rate.Every(1), 5),
2019
}
2120
client.NetworkConnectors = (*NetworkConnectorsService)(&service{client: client})
2221
return client

‎cloudconnexa/sessions_test.go‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,11 @@ import (
1313
// createTestSessionsClient creates a test client with the given server for sessions testing
1414
func createTestSessionsClient(server *httptest.Server) *Client {
1515
client := &Client{
16-
client: server.Client(),
17-
BaseURL: server.URL,
18-
Token: "test-token",
19-
RateLimiter: rate.NewLimiter(rate.Every(1*time.Second), 5),
16+
client: server.Client(),
17+
BaseURL: server.URL,
18+
Token: "test-token",
19+
ReadRateLimiter: rate.NewLimiter(rate.Every(1), 5),
20+
UpdateRateLimiter: rate.NewLimiter(rate.Every(1), 5),
2021
}
2122
client.Sessions = (*SessionsService)(&service{client: client})
2223
return client
@@ -176,7 +177,8 @@ func TestSessionsService_ListByDateRange(t *testing.T) {
176177

177178
func TestSessionsService_List_InvalidSize(t *testing.T) {
178179
client := &Client{
179-
RateLimiter: rate.NewLimiter(rate.Every(1*time.Second), 5),
180+
ReadRateLimiter: rate.NewLimiter(rate.Every(1), 5),
181+
UpdateRateLimiter: rate.NewLimiter(rate.Every(1), 5),
180182
}
181183
client.Sessions = (*SessionsService)(&service{client: client})
182184

‎cloudconnexa/user_groups_direct_test.go‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,21 +2,20 @@ package cloudconnexa
22

33
import (
44
"encoding/json"
5+
"golang.org/x/time/rate"
56
"net/http"
67
"net/http/httptest"
78
"testing"
8-
"time"
9-
10-
"golang.org/x/time/rate"
119
)
1210

1311
// createTestUserGroupsClient creates a test client with the given server for user groups testing
1412
func createTestUserGroupsClient(server *httptest.Server) *Client {
1513
client := &Client{
16-
client: server.Client(),
17-
BaseURL: server.URL,
18-
Token: "test-token",
19-
RateLimiter: rate.NewLimiter(rate.Every(1*time.Second), 5),
14+
client: server.Client(),
15+
BaseURL: server.URL,
16+
Token: "test-token",
17+
ReadRateLimiter: rate.NewLimiter(rate.Every(1), 5),
18+
UpdateRateLimiter: rate.NewLimiter(rate.Every(1), 5),
2019
}
2120
client.UserGroups = (*UserGroupsService)(&service{client: client})
2221
return client

0 commit comments

Comments
 (0)