github.com/aliyun/aliyun-oss-go-sdk@v3.0.2+incompatible/oss/limit_reader_1_7.go (about)

     1  //go:build go1.7
     2  // +build go1.7
     3  
     4  package oss
     5  
     6  import (
     7  	"fmt"
     8  	"io"
     9  	"math"
    10  	"time"
    11  
    12  	"golang.org/x/time/rate"
    13  )
    14  
    15  const (
    16  	perTokenBandwidthSize int = 1024
    17  )
    18  
    19  // OssLimiter wrapper rate.Limiter
    20  type OssLimiter struct {
    21  	limiter *rate.Limiter
    22  }
    23  
    24  // GetOssLimiter create OssLimiter
    25  // uploadSpeed KB/s
    26  func GetOssLimiter(uploadSpeed int) (ossLimiter *OssLimiter, err error) {
    27  	limiter := rate.NewLimiter(rate.Limit(uploadSpeed), uploadSpeed)
    28  
    29  	// first consume the initial full token,the limiter will behave more accurately
    30  	limiter.AllowN(time.Now(), uploadSpeed)
    31  
    32  	return &OssLimiter{
    33  		limiter: limiter,
    34  	}, nil
    35  }
    36  
    37  // LimitSpeedReader for limit bandwidth upload
    38  type LimitSpeedReader struct {
    39  	io.ReadCloser
    40  	reader     io.Reader
    41  	ossLimiter *OssLimiter
    42  }
    43  
    44  // Read
    45  func (r *LimitSpeedReader) Read(p []byte) (n int, err error) {
    46  	n = 0
    47  	err = nil
    48  	start := 0
    49  	burst := r.ossLimiter.limiter.Burst()
    50  	var end int
    51  	var tmpN int
    52  	var tc int
    53  	for start < len(p) {
    54  		if start+burst*perTokenBandwidthSize < len(p) {
    55  			end = start + burst*perTokenBandwidthSize
    56  		} else {
    57  			end = len(p)
    58  		}
    59  
    60  		tmpN, err = r.reader.Read(p[start:end])
    61  		if tmpN > 0 {
    62  			n += tmpN
    63  			start = n
    64  		}
    65  
    66  		if err != nil {
    67  			return
    68  		}
    69  
    70  		tc = int(math.Ceil(float64(tmpN) / float64(perTokenBandwidthSize)))
    71  		now := time.Now()
    72  		re := r.ossLimiter.limiter.ReserveN(now, tc)
    73  		if !re.OK() {
    74  			err = fmt.Errorf("LimitSpeedReader.Read() failure,ReserveN error,start:%d,end:%d,burst:%d,perTokenBandwidthSize:%d",
    75  				start, end, burst, perTokenBandwidthSize)
    76  			return
    77  		}
    78  		timeDelay := re.Delay()
    79  		time.Sleep(timeDelay)
    80  	}
    81  	return
    82  }
    83  
    84  // Close ...
    85  func (r *LimitSpeedReader) Close() error {
    86  	rc, ok := r.reader.(io.ReadCloser)
    87  	if ok {
    88  		return rc.Close()
    89  	}
    90  	return nil
    91  }