support rate limit

Author: sp-yduck

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/gorilla/websocket v1.5.0
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.8.4
+	golang.org/x/time v0.3.0
 )
 
 require (

go.sum

@@ -8,6 +8,8 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4=
+golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

rest/client.go

@@ -10,20 +10,28 @@ import (
 	"strings"
 	"time"
 
+	"golang.org/x/time/rate"
+
 	"github.com/sp-yduck/proxmox-go/api"
 )
 
 const (
 	defaultUserAgent = "sp-yduck/proxmox-go"
+	defaultQPS       = 20
 )
 
 type RESTClient struct {
-	endpoint    string
-	httpClient  *http.Client
+	// proxmox rest api endpoint
+	endpoint string
+
+	httpClient *http.Client
+
 	tokenid     string
 	token       string
 	session     *api.Session
 	credentials *TicketRequest
+
+	rateLimiter *rate.Limiter
 }
 
 type TicketRequest struct {
@@ -41,8 +49,9 @@ type ClientOption func(*RESTClient)
 
 func NewRESTClient(baseUrl string, opts ...ClientOption) (*RESTClient, error) {
 	client := &RESTClient{
-		endpoint:   complementURL(baseUrl),
-		httpClient: &http.Client{},
+		endpoint:    complementURL(baseUrl),
+		httpClient:  &http.Client{},
+		rateLimiter: rate.NewLimiter(rate.Every(1*time.Second), defaultQPS),
 	}
 
 	for _, option := range opts {
@@ -98,6 +107,16 @@ func WithAPIToken(tokenid, secret string) ClientOption {
 	}
 }
 
+func WithQPS(qps int) ClientOption {
+	return func(c *RESTClient) {
+		c.rateLimiter = rate.NewLimiter(rate.Every(1*time.Second), qps)
+	}
+}
+
+func (c *RESTClient) SetMaxQPS(qps int) {
+	c.rateLimiter = rate.NewLimiter(rate.Every(1*time.Second), qps)
+}
+
 func (c *RESTClient) Do(ctx context.Context, httpMethod, urlPath string, req, v interface{}) error {
 	endpoint := c.endpoint + urlPath
 
@@ -117,6 +136,9 @@ func (c *RESTClient) Do(ctx context.Context, httpMethod, urlPath string, req, v
 	httpReq.Header = c.makeAuthHeaders()
 	httpReq.Header.Add("Content-Type", "application/json")
 
+	if err := c.rateLimiter.Wait(ctx); err != nil {
+		return err
+	}
 	httpRsp, err := c.httpClient.Do(httpReq)
 	if err != nil {
 		return err

rest/client_test.go

@@ -0,0 +1,17 @@
+package rest
+
+import (
+	"context"
+	"fmt"
+)
+
+func (s *TestSuite) TestRateLimit() {
+	s.restclient.SetMaxQPS(10)
+	for i := 0; i < 13; i++ {
+		_, err := s.restclient.GetVersion(context.TODO())
+		if err != nil {
+			s.T().Errorf("failed to ger version: %v", err)
+		}
+		fmt.Println(i)
+	}
+}