limiter.go raw

   1  package bce
   2  
   3  import (
   4  	"context"
   5  	"time"
   6  
   7  	"github.com/baidubce/bce-sdk-go/rate"
   8  )
   9  
  10  const (
  11  	RateLimiterSlotRx int = iota
  12  	RateLimiterSlotTx
  13  	RateLimiterSlots
  14  )
  15  
  16  type RateLimiter struct {
  17  	Bandwidth int64 // Byte/S
  18  	Limiter   *rate.Limiter
  19  }
  20  
  21  type RateLimiters [RateLimiterSlots]*RateLimiter
  22  
  23  func newRateLimiter(bandwidth int64) *RateLimiter {
  24  	return &RateLimiter{
  25  		Bandwidth: bandwidth,
  26  		Limiter:   newTokenBucket(bandwidth),
  27  	}
  28  }
  29  
  30  func newTokenBucket(bandwidth int64) *rate.Limiter {
  31  	const defaultMaxBurstSize = 4 * 1024 * 1024
  32  	maxBurstSize := (bandwidth * defaultMaxBurstSize) / (256 * 1024 * 1024)
  33  	if maxBurstSize < defaultMaxBurstSize {
  34  		maxBurstSize = defaultMaxBurstSize
  35  	}
  36  	tb := rate.NewLimiter(rate.Limit(bandwidth), int(maxBurstSize))
  37  	tb.AllowN(time.Now(), int(maxBurstSize))
  38  	return tb
  39  }
  40  
  41  func (rl *RateLimiter) LimitBandwidth(n int) {
  42  	rl.Limiter.WaitN(context.Background(), n)
  43  }
  44