limiter.go raw

   1  package bridge
   2  
   3  import (
   4  	"sync"
   5  	"time"
   6  )
   7  
   8  // limiter is a token bucket in bytes. It allows a one-second burst and then
   9  // the configured rate. Two independent limiters gate a session's up and down
  10  // streams, so one direction cannot starve the other.
  11  type limiter struct {
  12  	rate int64
  13  
  14  	mu     sync.Mutex
  15  	tokens float64
  16  	last   time.Time
  17  }
  18  
  19  func newLimiter(rate int64) *limiter {
  20  	return &limiter{rate: rate, tokens: float64(rate), last: time.Now()}
  21  }
  22  
  23  func (l *limiter) wait(n int64) {
  24  	if l == nil || l.rate <= 0 || n <= 0 {
  25  		return
  26  	}
  27  	for {
  28  		l.mu.Lock()
  29  		now := time.Now()
  30  		l.tokens += now.Sub(l.last).Seconds() * float64(l.rate)
  31  		if l.tokens > float64(l.rate) {
  32  			l.tokens = float64(l.rate)
  33  		}
  34  		l.last = now
  35  		if l.tokens >= float64(n) {
  36  			l.tokens -= float64(n)
  37  			l.mu.Unlock()
  38  			return
  39  		}
  40  		delay := time.Duration((float64(n) - l.tokens) / float64(l.rate) * float64(time.Second))
  41  		l.mu.Unlock()
  42  		if delay < time.Millisecond {
  43  			delay = time.Millisecond
  44  		}
  45  		time.Sleep(delay)
  46  	}
  47  }
  48