minio/internal/bucket/bandwidth/reader.go
2022-01-02 09:15:06 -08:00

99 lines
2.8 KiB
Go

// Copyright (c) 2015-2021 MinIO, Inc.
//
// This file is part of MinIO Object Storage stack
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
package bandwidth
import (
"context"
"io"
"math"
)
// MonitoredReader represents a throttled reader subject to bandwidth monitoring
type MonitoredReader struct {
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
HeaderSize int
}
// Read implements a throttled read
func (r *MonitoredReader) Read(buf []byte) (n int, err error) {
if r.throttle == nil {
return r.r.Read(buf)
}
if r.lastErr != nil {
err = r.lastErr
return
}
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
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 {
return
}
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
}