github.com/amazechain/amc@v0.1.3/internal/p2p/leaky-bucket/leakybucket.go (about)

     1  /*
     2  Package leakybucket implements a scalable leaky bucket algorithm.
     3  
     4  There are at least two different definitions of the leaky bucket algorithm.
     5  This package implements the leaky bucket as a meter. For more details see:
     6  
     7  https://en.wikipedia.org/wiki/Leaky_bucket#As_a_meter
     8  
     9  This means it is the exact mirror of a token bucket.
    10  
    11  	// New LeakyBucket that leaks at the rate of 0.5/sec and a total capacity of 10.
    12  	b := NewLeakyBucket(0.5, 10)
    13  
    14  	b.Add(5)
    15  	b.Add(5)
    16  	// Bucket is now full!
    17  
    18  	n := b.Add(1)
    19  	// n == 0
    20  
    21  A Collector is a convenient way to keep track of multiple LeakyBucket's.
    22  Buckets are associated with string keys for fast lookup. It can dynamically
    23  add new buckets and automatically remove them as they become empty, freeing
    24  up resources.
    25  
    26  	// New Collector that leaks at 1 MiB/sec, a total capacity of 10 MiB and
    27  	// automatic removal of bucket's when they become empty.
    28  	const megabyte = 1<<20
    29  	c := NewCollector(megabyte, megabyte*10, true)
    30  
    31  	// Attempt to add 100 MiB to a bucket associated with an IP.
    32  	n := c.Add("192.168.0.42", megabyte*100)
    33  
    34  	// 100 MiB is over the capacity, so only 10 MiB is actually added.
    35  	// n equals 10 MiB.
    36  */
    37  package leakybucket
    38  
    39  import (
    40  	"math"
    41  	"time"
    42  )
    43  
    44  // Makes it easy to test time based things.
    45  var now = time.Now
    46  
    47  // LeakyBucket represents a bucket that leaks at a constant rate.
    48  type LeakyBucket struct {
    49  	// The identifying key, used for map lookups.
    50  	key string
    51  
    52  	// How large the bucket is.
    53  	capacity int64
    54  
    55  	// Amount the bucket leaks per time duration.
    56  	rate float64
    57  
    58  	// The priority of the bucket in a min-heap priority queue, where p is the
    59  	// exact time the bucket will have leaked enough to be empty. Buckets that
    60  	// are empty or will be the soonest are at the top of the heap. This allows
    61  	// for quick pruning of empty buckets that scales very well. p is adjusted
    62  	// any time an amount is added to the Queue().
    63  	p time.Time
    64  
    65  	// The time duration through which the leaky bucket is
    66  	// assessed.
    67  	period time.Duration
    68  
    69  	// The index is maintained by the heap.Interface methods.
    70  	index int
    71  }
    72  
    73  // NewLeakyBucket creates a new LeakyBucket with the give rate and capacity.
    74  func NewLeakyBucket(rate float64, capacity int64, period time.Duration) *LeakyBucket {
    75  	return &LeakyBucket{
    76  		rate:     rate,
    77  		capacity: capacity,
    78  		period:   period,
    79  		p:        now(),
    80  	}
    81  }
    82  
    83  // Count returns the bucket's current count.
    84  func (b *LeakyBucket) Count() int64 {
    85  	if !now().Before(b.p) {
    86  		return 0
    87  	}
    88  
    89  	nsRemaining := float64(b.p.Sub(now()))
    90  	nsPerDrip := float64(b.period) / b.rate
    91  	count := int64(math.Ceil(nsRemaining / nsPerDrip))
    92  
    93  	return count
    94  }
    95  
    96  // Rate returns the amount the bucket leaks per second.
    97  func (b *LeakyBucket) Rate() float64 {
    98  	return b.rate
    99  }
   100  
   101  // Capacity returns the bucket's capacity.
   102  func (b *LeakyBucket) Capacity() int64 {
   103  	return b.capacity
   104  }
   105  
   106  // Remaining returns the bucket's remaining capacity.
   107  func (b *LeakyBucket) Remaining() int64 {
   108  	return b.capacity - b.Count()
   109  }
   110  
   111  // ChangeCapacity changes the bucket's capacity.
   112  //
   113  // If the bucket's current count is greater than the new capacity, the count
   114  // will be decreased to match the new capacity.
   115  func (b *LeakyBucket) ChangeCapacity(capacity int64) {
   116  	diff := float64(capacity - b.capacity)
   117  
   118  	if diff < 0 && b.Count() > capacity {
   119  		// We are shrinking the capacity and the new bucket size can't hold all
   120  		// the current contents. Dump the extra and adjust the time till empty.
   121  		nsPerDrip := float64(b.period) / b.rate
   122  		b.p = now().Add(time.Duration(nsPerDrip * float64(capacity)))
   123  	}
   124  	b.capacity = capacity
   125  }
   126  
   127  // TillEmpty returns how much time must pass until the bucket is empty.
   128  func (b *LeakyBucket) TillEmpty() time.Duration {
   129  	return b.p.Sub(now())
   130  }
   131  
   132  // Add 'amount' to the bucket's count, up to it's capacity. Returns how much
   133  // was added to the bucket. If the return is less than 'amount', then the
   134  // bucket's capacity was reached.
   135  func (b *LeakyBucket) Add(amount int64) int64 {
   136  	count := b.Count()
   137  	if count >= b.capacity {
   138  		// The bucket is full.
   139  		return 0
   140  	}
   141  
   142  	if !now().Before(b.p) {
   143  		// The bucket needs to be reset.
   144  		b.p = now()
   145  	}
   146  	remaining := b.capacity - count
   147  	if amount > remaining {
   148  		amount = remaining
   149  	}
   150  	t := time.Duration(float64(b.period) * (float64(amount) / b.rate))
   151  	b.p = b.p.Add(t)
   152  
   153  	return amount
   154  }