A simple rate limit for go http client

ratelimit.go

package ratelimit

import (
	"context"
	"net/http"
	"time"

	"golang.org/x/time/rate"
)

// New adds the rate suppport to http client.
// it limits by the number of request per second.
func New(num int) *HttpClient {
	return &HttpClient{
		client: http.DefaultClient,
		rl:     rate.NewLimiter(rate.Every(time.Second), num),
	}
}

// HttpClient adds the rate limit to http client.
type HttpClient struct {
	client *http.Client
	rl     *rate.Limiter
}

func (c *HttpClient) Do(ctx context.Context, req *http.Request) (*http.Response, error) {
	if err := c.rl.Wait(ctx); err != nil { // This is a blocking call
		return nil, err
	}
	return c.client.Do(req)
}

ratelimit_test.go

package ratelimit

import (
	"context"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"
)

func TestNew(t *testing.T) {
	var cnt int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&cnt, 1)
		w.Write([]byte("OK"))
	}))
	ctx, cancel := context.WithCancel(context.Background())
	wantRate := 100
	dur := time.Second
	go func() {
		client := New(wantRate)
		for {
			select {
			case <-time.After(dur):
				cancel()
				break
			default:
				req, _ := http.NewRequest("GET", srv.URL, nil)
				client.Do(ctx, req)
			}
		}
	}()
	<-time.After(dur)
	srv.Close()

	if cnt != int32(wantRate) {
		t.Errorf("rate: want %d - got %d", wantRate, cnt)
	}
}