feat: add per-request upload and download bandwidth limiting

client.go

@@ -1898,6 +1898,22 @@ func (c *Client) roundTrip(r *Request) (resp *Response, err error) {
 	if r.unReplayableBody != nil {
 		getBody = nil
 	}
+	if r.uploadLimit > 0 {
+		limit := r.uploadLimit
+		if reqBody != nil {
+			reqBody = newRateLimitedReadCloser(reqBody, limit)
+		}
+		if getBody != nil {
+			inner := getBody
+			getBody = func() (io.ReadCloser, error) {
+				rc, err := inner()
+				if err != nil || rc == nil {
+					return rc, err
+				}
+				return newRateLimitedReadCloser(rc, limit), nil
+			}
+		}
+	}
 	req := &http.Request{
 		Method:        r.Method,
 		Header:        r.Headers.Clone(),
@@ -1914,8 +1930,15 @@ func (c *Client) roundTrip(r *Request) (resp *Response, err error) {
 	for _, cookie := range r.Cookies {
 		req.AddCookie(cookie)
 	}
+	var bodyWraps []wrapResponseBodyFunc
+	if r.downloadLimit > 0 {
+		limit := r.downloadLimit
+		bodyWraps = append(bodyWraps, func(rc io.ReadCloser) io.ReadCloser {
+			return newRateLimitedReadCloser(rc, limit)
+		})
+	}
 	if r.isSaveResponse && r.downloadCallback != nil {
-		var wrap wrapResponseBodyFunc = func(rc io.ReadCloser) io.ReadCloser {
+		bodyWraps = append(bodyWraps, func(rc io.ReadCloser) io.ReadCloser {
 			return &callbackReader{
 				ReadCloser: rc,
 				callback: func(read int64) {
@@ -1927,6 +1950,14 @@ func (c *Client) roundTrip(r *Request) (resp *Response, err error) {
 				lastTime: time.Now(),
 				interval: r.downloadCallbackInterval,
 			}
+		})
+	}
+	if len(bodyWraps) > 0 {
+		var wrap wrapResponseBodyFunc = func(rc io.ReadCloser) io.ReadCloser {
+			for _, w := range bodyWraps {
+				rc = w(rc)
+			}
+			return rc
 		}
 		if ctx == nil {
 			ctx = context.Background()

ratelimit.go

@@ -0,0 +1,44 @@
+package req
+
+import (
+	"io"
+	"time"
+)
+
+// rateLimitedReadCloser wraps an io.ReadCloser and paces reads so that the
+// average throughput does not exceed limit bytes per second. It is used for
+// both upload (the transport reads the request body through it) and download
+// (the response body is read through it) bandwidth limiting.
+//
+// After each read it sleeps for however long that many bytes should have taken
+// at the configured rate, minus the time the read itself already took. Data is
+// therefore never delayed longer than necessary and the average rate converges
+// to the limit, at the cost of allowing a single read (at most one buffer) to
+// burst through before the pause.
+type rateLimitedReadCloser struct {
+	rc    io.ReadCloser
+	limit float64 // bytes per second, always > 0
+}
+
+func newRateLimitedReadCloser(rc io.ReadCloser, bytesPerSecond int64) *rateLimitedReadCloser {
+	return &rateLimitedReadCloser{
+		rc:    rc,
+		limit: float64(bytesPerSecond),
+	}
+}
+
+func (l *rateLimitedReadCloser) Read(p []byte) (int, error) {
+	start := time.Now()
+	n, err := l.rc.Read(p)
+	if n > 0 {
+		expected := time.Duration(float64(n) / l.limit * float64(time.Second))
+		if d := expected - time.Since(start); d > 0 {
+			time.Sleep(d)
+		}
+	}
+	return n, err
+}
+
+func (l *rateLimitedReadCloser) Close() error {
+	return l.rc.Close()
+}

ratelimit_test.go

@@ -0,0 +1,102 @@
+package req
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/imroc/req/v3/internal/tests"
+)
+
+func TestRateLimitedReadCloser(t *testing.T) {
+	const size = 20000
+	const limit = 40000 // bytes per second, so ~0.5s for the payload
+	src := bytes.Repeat([]byte("x"), size)
+
+	r := newRateLimitedReadCloser(io.NopCloser(bytes.NewReader(src)), limit)
+	start := time.Now()
+	got, err := io.ReadAll(r)
+	elapsed := time.Since(start)
+	tests.AssertNoError(t, err)
+	tests.AssertEqual(t, size, len(got))
+	if !bytes.Equal(src, got) {
+		t.Fatal("data read through the limiter does not match the source")
+	}
+
+	min := time.Duration(float64(size)/limit*float64(time.Second)) / 2
+	if elapsed < min {
+		t.Errorf("expected the read to be throttled to at least %s, but it took %s", min, elapsed)
+	}
+}
+
+func TestSetDownloadLimit(t *testing.T) {
+	const size = 40 * 1024
+	const limit = 80 * 1024 // ~0.5s for the payload
+	payload := bytes.Repeat([]byte("a"), size)
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write(payload)
+	}))
+	defer ts.Close()
+
+	// Without a limit the download should be near instant.
+	start := time.Now()
+	resp, err := C().R().Get(ts.URL)
+	fast := time.Since(start)
+	tests.AssertNoError(t, err)
+	tests.AssertEqual(t, size, len(resp.Bytes()))
+	if fast > 200*time.Millisecond {
+		t.Fatalf("unlimited download was unexpectedly slow (%s), test environment too slow to be meaningful", fast)
+	}
+
+	start = time.Now()
+	resp, err = C().R().SetDownloadLimit(limit).Get(ts.URL)
+	slow := time.Since(start)
+	tests.AssertNoError(t, err)
+	tests.AssertEqual(t, size, len(resp.Bytes()))
+	if !bytes.Equal(payload, resp.Bytes()) {
+		t.Fatal("limited download body does not match")
+	}
+	min := time.Duration(float64(size)/limit*float64(time.Second)) / 2
+	if slow < min {
+		t.Errorf("expected the download to be throttled to at least %s, but it took %s", min, slow)
+	}
+}
+
+func TestSetUploadLimit(t *testing.T) {
+	const size = 40 * 1024
+	const limit = 80 * 1024 // ~0.5s for the payload
+	var received int64
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		n, _ := io.Copy(io.Discard, r.Body)
+		received = n
+	}))
+	defer ts.Close()
+
+	body := strings.Repeat("b", size)
+	start := time.Now()
+	resp, err := C().R().SetUploadLimit(limit).SetBodyString(body).Post(ts.URL)
+	elapsed := time.Since(start)
+	tests.AssertNoError(t, err)
+	tests.AssertEqual(t, http.StatusOK, resp.StatusCode)
+	tests.AssertEqual(t, int64(size), received)
+
+	min := time.Duration(float64(size)/limit*float64(time.Second)) / 2
+	if elapsed < min {
+		t.Errorf("expected the upload to be throttled to at least %s, but it took %s", min, elapsed)
+	}
+}
+
+func TestSetLimitDisable(t *testing.T) {
+	r := C().R().SetDownloadLimit(1024).SetUploadLimit(1024)
+	tests.AssertEqual(t, int64(1024), r.downloadLimit)
+	tests.AssertEqual(t, int64(1024), r.uploadLimit)
+
+	// A zero or negative value clears the limit.
+	r.SetDownloadLimit(0).SetUploadLimit(-1)
+	tests.AssertEqual(t, int64(0), r.downloadLimit)
+	tests.AssertEqual(t, int64(0), r.uploadLimit)
+}

request.go

@@ -55,6 +55,8 @@ type Request struct {
 	uploadCallbackInterval   time.Duration
 	downloadCallback         DownloadCallback
 	downloadCallbackInterval time.Duration
+	uploadLimit              int64
+	downloadLimit            int64
 	// maxResponseSize, when non-nil, overrides Client.maxResponseSize for this
 	// request. A pointed-to value of 0 means no limit for this request.
 	maxResponseSize    *int64
@@ -405,6 +407,33 @@ func (r *Request) SetDownloadCallbackWithInterval(callback DownloadCallback, min
 	return r
 }
 
+// SetUploadLimit limits the request body upload speed to at most
+// bytesPerSecond, which is useful to avoid saturating the network. The limit
+// is applied to the raw bytes sent over the wire and is reapplied when a
+// request body is replayed on retry or redirect. A value of 0 or less removes
+// the limit.
+func (r *Request) SetUploadLimit(bytesPerSecond int64) *Request {
+	if bytesPerSecond <= 0 {
+		r.uploadLimit = 0
+		return r
+	}
+	r.uploadLimit = bytesPerSecond
+	return r
+}
+
+// SetDownloadLimit limits the response body download speed to at most
+// bytesPerSecond. The limit is applied to the raw bytes received from the
+// server, before any content decoding, so it reflects actual network usage. A
+// value of 0 or less removes the limit.
+func (r *Request) SetDownloadLimit(bytesPerSecond int64) *Request {
+	if bytesPerSecond <= 0 {
+		r.downloadLimit = 0
+		return r
+	}
+	r.downloadLimit = bytesPerSecond
+	return r
+}
+
 // SetResult set the result that response Body will be unmarshalled to if
 // no error occurs and Response.ResultState() returns SuccessState, by default
 // it requires HTTP status `code >= 200 && code <= 299`, you can also use

request_wrapper.go

@@ -534,6 +534,18 @@ func SetDownloadCallbackWithInterval(callback DownloadCallback, minInterval time
 	return defaultClient.R().SetDownloadCallbackWithInterval(callback, minInterval)
 }
 
+// SetUploadLimit is a global wrapper methods which delegated
+// to the default client, create a request and SetUploadLimit for request.
+func SetUploadLimit(bytesPerSecond int64) *Request {
+	return defaultClient.R().SetUploadLimit(bytesPerSecond)
+}
+
+// SetDownloadLimit is a global wrapper methods which delegated
+// to the default client, create a request and SetDownloadLimit for request.
+func SetDownloadLimit(bytesPerSecond int64) *Request {
+	return defaultClient.R().SetDownloadLimit(bytesPerSecond)
+}
+
 // EnableCloseConnection is a global wrapper methods which delegated
 // to the default client, create a request and EnableCloseConnection for request.
 func EnableCloseConnection() *Request {