github.com/minio/mc@v0.0.0-20240503112107-b471de8d1882/pkg/limiter/limiter.go (about)

     1  // Copyright (c) 2015-2022 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 limiter implements throughput upload and download limits via http.RoundTripper
    19  package limiter
    20  
    21  import (
    22  	"errors"
    23  	"io"
    24  	"net/http"
    25  
    26  	"github.com/juju/ratelimit"
    27  )
    28  
    29  type limiter struct {
    30  	upload    *ratelimit.Bucket
    31  	download  *ratelimit.Bucket
    32  	transport http.RoundTripper // HTTP transport that needs to be intercepted
    33  }
    34  
    35  func (l limiter) limitReader(r io.Reader, b *ratelimit.Bucket) io.Reader {
    36  	if b == nil {
    37  		return r
    38  	}
    39  	return ratelimit.Reader(r, b)
    40  }
    41  
    42  // RoundTrip executes user provided request and response hooks for each HTTP call.
    43  func (l limiter) RoundTrip(req *http.Request) (res *http.Response, err error) {
    44  	if l.transport == nil {
    45  		return nil, errors.New("Invalid Argument")
    46  	}
    47  
    48  	type readCloser struct {
    49  		io.Reader
    50  		io.Closer
    51  	}
    52  
    53  	if req.Body != nil {
    54  		req.Body = &readCloser{
    55  			Reader: l.limitReader(req.Body, l.upload),
    56  			Closer: req.Body,
    57  		}
    58  	}
    59  
    60  	res, err = l.transport.RoundTrip(req)
    61  	if res != nil && res.Body != nil {
    62  		res.Body = &readCloser{
    63  			Reader: l.limitReader(res.Body, l.download),
    64  			Closer: res.Body,
    65  		}
    66  	}
    67  
    68  	return res, err
    69  }
    70  
    71  // New return a ratelimited transport
    72  func New(uploadLimit, downloadLimit int64, transport http.RoundTripper) http.RoundTripper {
    73  	if uploadLimit == 0 && downloadLimit == 0 {
    74  		return transport
    75  	}
    76  
    77  	var (
    78  		uploadBucket   *ratelimit.Bucket
    79  		downloadBucket *ratelimit.Bucket
    80  	)
    81  
    82  	if uploadLimit > 0 {
    83  		uploadBucket = ratelimit.NewBucketWithRate(float64(uploadLimit), uploadLimit)
    84  	}
    85  
    86  	if downloadLimit > 0 {
    87  		downloadBucket = ratelimit.NewBucketWithRate(float64(downloadLimit), downloadLimit)
    88  	}
    89  
    90  	return &limiter{
    91  		upload:    uploadBucket,
    92  		download:  downloadBucket,
    93  		transport: transport,
    94  	}
    95  }