github.com/fzfile/BaiduPCS-Go@v0.0.0-20200606205115-4408961cf336/requester/rio/speeds/ratelimit.go (about)

     1  package speeds
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	"time"
     7  )
     8  
     9  type (
    10  	RateLimit struct {
    11  		MaxRate int64
    12  
    13  		count           int64
    14  		interval        time.Duration
    15  		ticker          *time.Ticker
    16  		muChan          chan struct{}
    17  		closeChan       chan struct{}
    18  		backServiceOnce sync.Once
    19  	}
    20  
    21  	// AddCountFunc func() (count int64)
    22  )
    23  
    24  func NewRateLimit(maxRate int64) *RateLimit {
    25  	return &RateLimit{
    26  		MaxRate: maxRate,
    27  	}
    28  }
    29  
    30  func (rl *RateLimit) SetInterval(i time.Duration) {
    31  	if i <= 0 {
    32  		i = 1 * time.Second
    33  	}
    34  	rl.interval = i
    35  	if rl.ticker != nil {
    36  		rl.ticker.Stop()
    37  		rl.ticker = time.NewTicker(i)
    38  	}
    39  }
    40  
    41  func (rl *RateLimit) Stop() {
    42  	if rl.ticker != nil {
    43  		rl.ticker.Stop()
    44  	}
    45  	if rl.closeChan != nil {
    46  		close(rl.closeChan)
    47  	}
    48  	return
    49  }
    50  
    51  func (rl *RateLimit) resetChan() {
    52  	if rl.muChan != nil {
    53  		close(rl.muChan)
    54  	}
    55  	rl.muChan = make(chan struct{})
    56  }
    57  
    58  func (rl *RateLimit) backService() {
    59  	if rl.interval <= 0 {
    60  		rl.interval = 1 * time.Second
    61  	}
    62  	rl.ticker = time.NewTicker(rl.interval)
    63  	rl.closeChan = make(chan struct{})
    64  	rl.resetChan()
    65  	go func() {
    66  		for {
    67  			select {
    68  			case <-rl.ticker.C:
    69  				rl.resetChan()
    70  				atomic.StoreInt64(&rl.count, 0)
    71  			case <-rl.closeChan:
    72  				return
    73  			}
    74  		}
    75  	}()
    76  }
    77  
    78  func (rl *RateLimit) Add(count int64) {
    79  	rl.backServiceOnce.Do(rl.backService)
    80  	for {
    81  		if atomic.LoadInt64(&rl.count) >= rl.MaxRate { // 超出最大限额
    82  			// 阻塞
    83  			<-rl.muChan
    84  			continue
    85  		}
    86  		atomic.AddInt64(&rl.count, count)
    87  		break
    88  	}
    89  }