Introducing RateLimiter to Control Request Rate (Issue #4) (#5)

* Introducing RateLimiter to Control Request Rate #4

* Updating README to describe how-to use rate limiting

* RateLimiting to RateLimit

Author: 0x19

README.md

@@ -24,7 +24,7 @@ import "github.com/txpull/sourcify-go"
 
 ### Creating a Client
 
-To interact with the Sourcify API, you need to create a client using the `NewClient` function. You can provide optional configuration options using `ClientOption` functions. For example, you can specify a custom base URL or a custom HTTP client and set retry configuration in case sourcify servers are temporairly unavailable.
+To interact with the Sourcify API, you need to create a client using the `NewClient` function. You can provide optional configuration options using `ClientOption` functions. For example, you can specify a custom base URL or a custom HTTP client, set retry configuration in case sourcify servers are temporairly unavailable and set rate limits in order to control the rate at which HTTP requests are sent to the sourcify servers.
 
 ```go
 client := sourcify.NewClient(
@@ -34,6 +34,7 @@ client := sourcify.NewClient(
 		sourcify.WithMaxRetries(3),
 		sourcify.WithDelay(2*time.Second),
 	),
+	sourcify.WithRateLimit(10, 1*time.Second),
 )
 ``` 
 

client.go

@@ -37,6 +37,7 @@ type Client struct {
 	BaseURL      string       // The base URL of the Sourcify API.
 	HTTPClient   *http.Client // The HTTP client to use for making requests.
 	RetryOptions RetryOptions // The retry options for the client.
+	RateLimiter  *RateLimiter // The rate limiter for the client.
 }
 
 // WithHTTPClient allows you to provide your own http.Client for the Sourcify client.
@@ -62,6 +63,13 @@ func WithRetryOptions(options ...RetryOption) ClientOption {
 	}
 }
 
+// WithRateLimit allows you to configure rate limits for the Sourcify client.
+func WithRateLimit(max int, duration time.Duration) ClientOption {
+	return func(c *Client) {
+		c.RateLimiter = NewRateLimiter(max, duration)
+	}
+}
+
 // NewClient initializes a new Sourcify client with optional configurations.
 // By default, it uses the Sourcify API's base URL (https://sourcify.dev/server),
 // the default http.Client, and no retry options.
@@ -150,6 +158,10 @@ func (c *Client) doRequestWithRetry(req *http.Request) (io.ReadCloser, int, erro
 	attempt := 0
 
 	for {
+		if c.RateLimiter != nil {
+			c.RateLimiter.Wait()
+		}
+
 		attempt++
 		resp, err := c.HTTPClient.Do(req)
 		if err != nil {

client_rate_limiter.go

@@ -0,0 +1,65 @@
+package sourcify
+
+import "time"
+
+// RateLimiter represents a rate limiter that controls the rate of actions using the token bucket algorithm.
+// It provides a mechanism to prevent an HTTP client from exceeding a certain rate of requests.
+// The Max field represents the maximum number of actions that can be performed per 'Duration'.
+// The Duration field represents the time duration for which 'Max' number of actions can be performed.
+// These fields together determine the capacity of the token bucket and the rate at which tokens are added to the bucket.
+// The bucket field is a channel that models the token bucket. A token is consumed from the bucket each time an action is taken.
+// The capacity of the bucket determines the maximum burstiness of the actions, while the rate at which tokens are added
+// to the bucket determines the sustainable average rate of actions.
+type RateLimiter struct {
+	// Max is the maximum number of actions that can be performed per 'Duration'.
+	Max int
+	// Duration is the time duration for which 'Max' number of actions can be performed.
+	Duration time.Duration
+	// bucket is a channel that models the token bucket. A token is consumed from the bucket each time an action is taken.
+	bucket chan struct{}
+}
+
+// NewRateLimiter creates a new rate limiter.
+// The rate limiter uses the token bucket algorithm to control the rate of actions.
+// It initially creates a bucket of capacity 'Max' and then adds a token to the bucket every 'Duration'.
+// It allows a maximum of 'Max' actions to be performed per 'Duration'.
+// If an action is attempted when the bucket is empty, the action blocks until a token is added to the bucket.
+// This blocking behaviour ensures that the rate of actions does not exceed the specified rate.
+//
+// Parameters:
+// max - The maximum number of actions that can be performed per 'duration'. It is the capacity of the token bucket.
+// duration - The time duration for which 'max' number of actions can be performed.
+//
+// Returns:
+// A pointer to the created RateLimiter.
+func NewRateLimiter(max int, duration time.Duration) *RateLimiter {
+	bucket := make(chan struct{}, max)
+
+	// Initially, the bucket is filled to its capacity.
+	for i := 0; i < max; i++ {
+		bucket <- struct{}{}
+	}
+
+	// A ticker is set up to add a token to the bucket every 'duration'.
+	// If the bucket is full, the addition of a new token blocks until there is room in the bucket.
+	// This ensures that the rate of actions doesn't exceed the specified rate.
+	go func() {
+		ticker := time.NewTicker(duration)
+		for range ticker.C {
+			bucket <- struct{}{}
+		}
+	}()
+
+	return &RateLimiter{
+		Max:      max,
+		Duration: duration,
+		bucket:   bucket,
+	}
+}
+
+// Wait is used to perform an action with rate limiting.
+// If the token bucket (i.e., 'bucket' field of RateLimiter) is empty, Wait blocks until a token is added to the bucket.
+// If a token is available in the bucket, Wait consumes the token and returns immediately, allowing the action to be performed.
+func (r *RateLimiter) Wait() {
+	<-r.bucket
+}

client_rate_limiter_test.go

@@ -0,0 +1,57 @@
+package sourcify
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewRateLimiter(t *testing.T) {
+	// Create a new rate limiter with max 2 actions per second
+	rateLimiter := NewRateLimiter(2, time.Second)
+
+	assert.NotNil(t, rateLimiter)
+	assert.Equal(t, 2, rateLimiter.Max)
+	assert.Equal(t, time.Second, rateLimiter.Duration)
+}
+
+func TestRateLimiter_Wait(t *testing.T) {
+	// Create a new rate limiter with max 1 action per second
+	rateLimiter := NewRateLimiter(1, time.Second)
+
+	// Record the start time
+	start := time.Now()
+
+	// Perform 3 actions
+	for i := 0; i < 3; i++ {
+		rateLimiter.Wait()
+	}
+
+	// Record the end time
+	end := time.Now()
+
+	// The duration between start and end should be at least 2 seconds,
+	// since the rate limiter allows only 1 action per second.
+	assert.GreaterOrEqual(t, end.Sub(start).Seconds(), 2.0)
+}
+
+func TestRateLimiter_Wait_Burst(t *testing.T) {
+	// Create a new rate limiter with max 5 actions per 100 milliseconds
+	rateLimiter := NewRateLimiter(5, 100*time.Millisecond)
+
+	// Record the start time
+	start := time.Now()
+
+	// Perform 5 actions, should be processed in a burst
+	for i := 0; i < 5; i++ {
+		rateLimiter.Wait()
+	}
+
+	// Record the end time
+	end := time.Now()
+
+	// The duration between start and end should be less than 100 milliseconds,
+	// since all the actions are processed in a burst.
+	assert.Less(t, end.Sub(start).Seconds(), 0.1)
+}

client_test.go

@@ -182,3 +182,31 @@ func TestDoRequestWithRetry_SuccessfulRetry(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, "Hello, world!", string(body))
 }
+
+func TestWithRateLimiting(t *testing.T) {
+	client := NewClient(WithRateLimit(10, 1*time.Second))
+
+	assert.NotNil(t, client.RateLimiter)
+	assert.Equal(t, 10, client.RateLimiter.Max)
+	assert.Equal(t, 1*time.Second, client.RateLimiter.Duration)
+}
+
+func TestRateLimiting(t *testing.T) {
+	handler := func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprint(w, "Hello, world!")
+	}
+	server := httptest.NewServer(http.HandlerFunc(handler))
+	defer server.Close()
+
+	client := NewClient(
+		WithBaseURL(server.URL),
+		WithRateLimit(1, 1*time.Second),
+	)
+
+	req, _ := http.NewRequest("GET", server.URL, nil)
+
+	// Perform first request - should pass
+	resp, _, err := client.doRequestWithRetry(req)
+	assert.NoError(t, err)
+	assert.NotNil(t, resp)
+}