github.com/BGrewell/tail@v1.0.1-0.20210309152823-689d25348e0e/ratelimiter/leakybucket.go (about)

     1  // Package ratelimiter implements the Leaky Bucket ratelimiting algorithm with memcached and in-memory backends.
     2  package ratelimiter
     3  
     4  import (
     5  	"time"
     6  )
     7  
     8  type LeakyBucket struct {
     9  	Size         uint16
    10  	Fill         float64
    11  	LeakInterval time.Duration // time.Duration for 1 unit of size to leak
    12  	Lastupdate   time.Time
    13  	Now          func() time.Time
    14  }
    15  
    16  func NewLeakyBucket(size uint16, leakInterval time.Duration) *LeakyBucket {
    17  	bucket := LeakyBucket{
    18  		Size:         size,
    19  		Fill:         0,
    20  		LeakInterval: leakInterval,
    21  		Now:          time.Now,
    22  		Lastupdate:   time.Now(),
    23  	}
    24  
    25  	return &bucket
    26  }
    27  
    28  func (b *LeakyBucket) updateFill() {
    29  	now := b.Now()
    30  	if b.Fill > 0 {
    31  		elapsed := now.Sub(b.Lastupdate)
    32  
    33  		b.Fill -= float64(elapsed) / float64(b.LeakInterval)
    34  		if b.Fill < 0 {
    35  			b.Fill = 0
    36  		}
    37  	}
    38  	b.Lastupdate = now
    39  }
    40  
    41  func (b *LeakyBucket) Pour(amount uint16) bool {
    42  	b.updateFill()
    43  
    44  	var newfill float64 = b.Fill + float64(amount)
    45  
    46  	if newfill > float64(b.Size) {
    47  		return false
    48  	}
    49  
    50  	b.Fill = newfill
    51  
    52  	return true
    53  }
    54  
    55  // The time at which this bucket will be completely drained
    56  func (b *LeakyBucket) DrainedAt() time.Time {
    57  	return b.Lastupdate.Add(time.Duration(b.Fill * float64(b.LeakInterval)))
    58  }
    59  
    60  // The duration until this bucket is completely drained
    61  func (b *LeakyBucket) TimeToDrain() time.Duration {
    62  	return b.DrainedAt().Sub(b.Now())
    63  }
    64  
    65  func (b *LeakyBucket) TimeSinceLastUpdate() time.Duration {
    66  	return b.Now().Sub(b.Lastupdate)
    67  }
    68  
    69  type LeakyBucketSer struct {
    70  	Size         uint16
    71  	Fill         float64
    72  	LeakInterval time.Duration // time.Duration for 1 unit of size to leak
    73  	Lastupdate   time.Time
    74  }
    75  
    76  func (b *LeakyBucket) Serialise() *LeakyBucketSer {
    77  	bucket := LeakyBucketSer{
    78  		Size:         b.Size,
    79  		Fill:         b.Fill,
    80  		LeakInterval: b.LeakInterval,
    81  		Lastupdate:   b.Lastupdate,
    82  	}
    83  
    84  	return &bucket
    85  }
    86  
    87  func (b *LeakyBucketSer) DeSerialise() *LeakyBucket {
    88  	bucket := LeakyBucket{
    89  		Size:         b.Size,
    90  		Fill:         b.Fill,
    91  		LeakInterval: b.LeakInterval,
    92  		Lastupdate:   b.Lastupdate,
    93  		Now:          time.Now,
    94  	}
    95  
    96  	return &bucket
    97  }