github.com/keltia/go-ipfs@v0.3.8-0.20150909044612-210793031c63/p2p/net/mock/ratelimiter.go (about)

     1  package mocknet
     2  
     3  import (
     4  	"time"
     5  )
     6  
     7  //  A ratelimiter is used by a link to determine how long to wait before sending
     8  //  data given a bandwidth cap.
     9  type ratelimiter struct {
    10  	bandwidth    float64       // bytes per nanosecond
    11  	allowance    float64       // in bytes
    12  	maxAllowance float64       // in bytes
    13  	lastUpdate   time.Time     // when allowance was updated last
    14  	count        int           // number of times rate limiting was applied
    15  	duration     time.Duration // total delay introduced due to rate limiting
    16  }
    17  
    18  //  Creates a new ratelimiter with bandwidth (in bytes/sec)
    19  func NewRatelimiter(bandwidth float64) *ratelimiter {
    20  	//  convert bandwidth to bytes per nanosecond
    21  	b := bandwidth / float64(time.Second)
    22  	return &ratelimiter{
    23  		bandwidth:    b,
    24  		allowance:    0,
    25  		maxAllowance: bandwidth,
    26  		lastUpdate:   time.Now(),
    27  	}
    28  }
    29  
    30  //  Changes bandwidth of a ratelimiter and resets its allowance
    31  func (r *ratelimiter) UpdateBandwidth(bandwidth float64) {
    32  	//  Convert bandwidth from bytes/second to bytes/nanosecond
    33  	b := bandwidth / float64(time.Second)
    34  	r.bandwidth = b
    35  	//  Reset allowance
    36  	r.allowance = 0
    37  	r.maxAllowance = bandwidth
    38  	r.lastUpdate = time.Now()
    39  }
    40  
    41  //  Returns how long to wait before sending data with length 'dataSize' bytes
    42  func (r *ratelimiter) Limit(dataSize int) time.Duration {
    43  	//  update time
    44  	var duration time.Duration = time.Duration(0)
    45  	if r.bandwidth == 0 {
    46  		return duration
    47  	}
    48  	current := time.Now()
    49  	elapsedTime := current.Sub(r.lastUpdate)
    50  	r.lastUpdate = current
    51  
    52  	allowance := r.allowance + float64(elapsedTime)*r.bandwidth
    53  	//  allowance can't exceed bandwidth
    54  	if allowance > r.maxAllowance {
    55  		allowance = r.maxAllowance
    56  	}
    57  
    58  	allowance -= float64(dataSize)
    59  	if allowance < 0 {
    60  		//  sleep until allowance is back to 0
    61  		duration = time.Duration(-allowance / r.bandwidth)
    62  		//  rate limiting was applied, record stats
    63  		r.count++
    64  		r.duration += duration
    65  	}
    66  
    67  	r.allowance = allowance
    68  	return duration
    69  }