github.com/minio/minio@v0.0.0-20240328213742-3f72439b8a27/internal/bucket/bandwidth/reader.go (about)

     1  // Copyright (c) 2015-2021 MinIO, Inc.
     2  //
     3  // This file is part of MinIO Object Storage stack
     4  //
     5  // This program is free software: you can redistribute it and/or modify
     6  // it under the terms of the GNU Affero General Public License as published by
     7  // the Free Software Foundation, either version 3 of the License, or
     8  // (at your option) any later version.
     9  //
    10  // This program is distributed in the hope that it will be useful
    11  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    12  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13  // GNU Affero General Public License for more details.
    14  //
    15  // You should have received a copy of the GNU Affero General Public License
    16  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17  
    18  package bandwidth
    19  
    20  import (
    21  	"context"
    22  	"io"
    23  	"math"
    24  )
    25  
    26  // MonitoredReader represents a throttled reader subject to bandwidth monitoring
    27  type MonitoredReader struct {
    28  	r        io.Reader
    29  	throttle *bucketThrottle
    30  	ctx      context.Context // request context
    31  	lastErr  error           // last error reported, if this non-nil all reads will fail.
    32  	m        *Monitor
    33  	opts     *MonitorReaderOptions
    34  }
    35  
    36  // BucketOptions represents the bucket and optionally its replication target pair.
    37  type BucketOptions struct {
    38  	Name           string
    39  	ReplicationARN string // This is optional, and not mandatory.
    40  }
    41  
    42  // MonitorReaderOptions provides configurable options for monitor reader implementation.
    43  type MonitorReaderOptions struct {
    44  	BucketOptions
    45  	HeaderSize int
    46  }
    47  
    48  // Read implements a throttled read
    49  func (r *MonitoredReader) Read(buf []byte) (n int, err error) {
    50  	if r.throttle == nil {
    51  		return r.r.Read(buf)
    52  	}
    53  	if r.lastErr != nil {
    54  		err = r.lastErr
    55  		return
    56  	}
    57  	b := r.throttle.Burst()  // maximum available tokens
    58  	need := len(buf)         // number of bytes requested by caller
    59  	hdr := r.opts.HeaderSize // remaining header bytes
    60  	var tokens int           // number of tokens to request
    61  
    62  	if hdr > 0 { // available tokens go towards header first
    63  		if hdr < b { // all of header can be accommodated
    64  			r.opts.HeaderSize = 0
    65  			need = int(math.Min(float64(b-hdr), float64(need))) // use remaining tokens towards payload
    66  			tokens = need + hdr
    67  
    68  		} else { // part of header can be accommodated
    69  			r.opts.HeaderSize -= b - 1
    70  			need = 1 // to ensure we read at least one byte for every Read
    71  			tokens = b
    72  		}
    73  	} else { // all tokens go towards payload
    74  		need = int(math.Min(float64(b), float64(need)))
    75  		tokens = need
    76  	}
    77  
    78  	err = r.throttle.WaitN(r.ctx, tokens)
    79  	if err != nil {
    80  		return
    81  	}
    82  
    83  	n, err = r.r.Read(buf[:need])
    84  	if err != nil {
    85  		r.lastErr = err
    86  		return
    87  	}
    88  	r.m.updateMeasurement(r.opts.BucketOptions, uint64(tokens))
    89  	return
    90  }
    91  
    92  // NewMonitoredReader returns reference to a monitored reader that throttles reads to configured bandwidth for the
    93  // bucket.
    94  func NewMonitoredReader(ctx context.Context, m *Monitor, r io.Reader, opts *MonitorReaderOptions) *MonitoredReader {
    95  	reader := MonitoredReader{
    96  		r:        r,
    97  		throttle: m.throttle(opts.BucketOptions),
    98  		m:        m,
    99  		opts:     opts,
   100  		ctx:      ctx,
   101  	}
   102  	reader.m.init(opts.BucketOptions)
   103  	return &reader
   104  }