Rate limited io.Reader

rate-limited-reader.go

package main

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"time"

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

const burstLimit = 1000 * 1000 * 1000

type Reader struct {
	r io.Reader
	l *rate.Limiter
}

func NewReader(r io.Reader) *Reader {
	return &Reader{
		r: r,
		l: rate.NewLimiter(rate.Inf, burstLimit),
	}
}

// SetRateLimit sets the rate limit as in bytes-per-second.
func (r *Reader) SetRateLimit(n float64) {
	r.l.SetLimit(rate.Limit(n))
	r.l.AllowN(time.Now(), burstLimit)
}

func (r *Reader) Read(p []byte) (int, error) {
	n, err := r.r.Read(p)
	if err != nil {
		return n, err
	}

	if err := r.l.WaitN(context.Background(), n); err != nil {
		return n, err
	}

	return n, nil
}

func main() {
	const limit = 256 * 1024 //bytes
	const size = 1024 * 1024 // bytes
	r := NewReader(bytes.NewReader(make([]byte, size)))
	r.SetRateLimit(limit)

	now := time.Now()
	fmt.Println("Gonna copy 1024 kB by throttling with 256 kB/s")
	n, err := io.Copy(ioutil.Discard, r)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("Copied %v bytes in %.02f seconds. Avg: %.02f kBps\n", n, time.Since(now).Seconds(), float64(n/1024)/time.Since(now).Seconds())
}