minstrel/internal/backup/throttle.go

package backup

import (
	"context"
	"io"
	"time"
)

// Throttle returns a reader that reads r no faster than perSecond bytes a
// second, so an upload leaves room on the connection for everything else.
// perSecond of 0 or less returns r unchanged. A wait ends early when ctx does,
// and the read then returns ctx's error.
func Throttle(ctx context.Context, r io.Reader, perSecond int64) io.Reader {
	if perSecond <= 0 {
		return r
	}
	return &throttled{ctx: ctx, r: r, perSecond: perSecond}
}

type throttled struct {
	ctx       context.Context
	r         io.Reader
	perSecond int64
	start     time.Time
	sent      int64
}

func (t *throttled) Read(p []byte) (int, error) {
	if t.start.IsZero() {
		t.start = time.Now()
	}
	// Reads are cut to a tenth of a second's worth, so the pace is even
	// rather than a burst followed by a long pause.
	if chunk := max(t.perSecond/10, 1); int64(len(p)) > chunk {
		p = p[:chunk]
	}
	n, err := t.r.Read(p)
	t.sent += int64(n)
	at := t.start.Add(time.Duration(float64(t.sent) / float64(t.perSecond) * float64(time.Second)))
	if wait := time.Until(at); wait > 0 {
		timer := time.NewTimer(wait)
		defer timer.Stop()
		select {
		case <-t.ctx.Done():
			return n, t.ctx.Err()
		case <-timer.C:
		}
	}
	return n, err
}