Use rate.Limiter for bandwidth monitoring (#12506)

Bonus: fixes a hang when bandwidth caps are enabled for
synchronous replication

internal/bucket/bandwidth/reader.go

@@ -20,61 +20,78 @@ package bandwidth
 import (
 	"context"
 	"io"
+	"math"
 )
 
-// MonitoredReader monitors the bandwidth
+// MonitoredReader represents a throttled reader subject to bandwidth monitoring
 type MonitoredReader struct {
-	opts              *MonitorReaderOptions
-	bucketMeasurement *bucketMeasurement // bucket measurement object
-	reader            io.Reader          // Reader to wrap
-	throttle          *throttle          // throttle the rate at which replication occur
-	monitor           *Monitor           // Monitor reference
-	lastErr           error              // last error reported, if this non-nil all reads will fail.
+	r        io.Reader
+	throttle *throttle
+	ctx      context.Context // request context
+	lastErr  error           // last error reported, if this non-nil all reads will fail.
+	m        *Monitor
+	opts     *MonitorReaderOptions
 }
 
 // MonitorReaderOptions provides configurable options for monitor reader implementation.
 type MonitorReaderOptions struct {
-	Bucket               string
-	Object               string
-	HeaderSize           int
-	BandwidthBytesPerSec int64
-	ClusterBandwidth     int64
+	Bucket     string
+	HeaderSize int
 }
 
-// NewMonitoredReader returns a io.Reader that reports bandwidth details.
-func NewMonitoredReader(ctx context.Context, monitor *Monitor, reader io.Reader, opts *MonitorReaderOptions) *MonitoredReader {
-	return &MonitoredReader{
-		opts:              opts,
-		bucketMeasurement: monitor.track(opts.Bucket, opts.Object),
-		reader:            reader,
-		throttle:          monitor.throttleBandwidth(ctx, opts.Bucket, opts.BandwidthBytesPerSec, opts.ClusterBandwidth),
-		monitor:           monitor,
-	}
-}
-
-// Read wraps the read reader
-func (m *MonitoredReader) Read(p []byte) (n int, err error) {
-	if m.lastErr != nil {
-		err = m.lastErr
+// Read implements a throttled read
+func (r *MonitoredReader) Read(buf []byte) (n int, err error) {
+	if r.lastErr != nil {
+		err = r.lastErr
 		return
 	}
 
-	p = p[:m.throttle.GetLimitForBytes(int64(len(p)))]
+	b := r.throttle.Burst()  // maximum available tokens
+	need := len(buf)         // number of bytes requested by caller
+	hdr := r.opts.HeaderSize // remaining header bytes
+	var tokens int           // number of tokens to request
 
-	n, err = m.reader.Read(p)
+	if hdr > 0 { // available tokens go towards header first
+		if hdr < b { // all of header can be accommodated
+			r.opts.HeaderSize = 0
+			need = int(math.Min(float64(b-hdr), float64(need))) // use remaining tokens towards payload
+			tokens = need + hdr
+
+		} else { // part of header can be accommodated
+			r.opts.HeaderSize -= b - 1
+			need = 1 // to ensure we read at least one byte for every Read
+			tokens = b
+		}
+
+	} else { // all tokens go towards payload
+		need = int(math.Min(float64(b), float64(need)))
+		tokens = need
+	}
+
+	err = r.throttle.WaitN(r.ctx, tokens)
 	if err != nil {
-		m.lastErr = err
+		return
 	}
 
-	update := n + m.opts.HeaderSize
-	unused := len(p) - update
-
-	m.bucketMeasurement.incrementBytes(uint64(update))
-	m.opts.HeaderSize = 0 // Set to 0 post first read
-
-	if unused > 0 {
-		m.throttle.ReleaseUnusedBandwidth(int64(unused))
+	n, err = r.r.Read(buf[:need])
+	if err != nil {
+		r.lastErr = err
+		return
 	}
-
+	r.m.updateMeasurement(r.opts.Bucket, uint64(tokens))
 	return
 }
+
+// NewMonitoredReader returns reference to a monitored reader that throttles reads to configured bandwidth for the
+// bucket.
+func NewMonitoredReader(ctx context.Context, m *Monitor, r io.Reader, opts *MonitorReaderOptions) *MonitoredReader {
+	reader := MonitoredReader{
+		r:        r,
+		throttle: m.throttle(opts.Bucket),
+		m:        m,
+		opts:     opts,
+		ctx:      ctx,
+	}
+	reader.m.track(opts.Bucket)
+	return &reader
+}