github.com/insolar/vanilla@v0.0.0-20201023172447-248fdf805322/ratelimiter/bucket.go (about)

     1  // Copyright 2020 Insolar Network Ltd.
     2  // All rights reserved.
     3  // This material is licensed under the Insolar License version 1.0,
     4  // available at https://github.com/insolar/assured-ledger/blob/master/LICENSE.md.
     5  
     6  package ratelimiter
     7  
     8  import (
     9  	"github.com/insolar/vanilla/atomickit"
    10  )
    11  
    12  type RateQuota interface {
    13  	TakeQuota(max int64) int64
    14  	TakeQuotaNoWait(max int64) int64
    15  }
    16  
    17  func NewBucket(bm *PeriodManager, cfg BucketConfig) *Bucket {
    18  	cfg.Check()
    19  
    20  	return &Bucket{state: BucketState{
    21  		filledAmount: atomickit.NewUint32(cfg.RefillAmount),
    22  		bucketConfig: cfg,
    23  	}, refiller: PeriodRefiller{
    24  		manager:     bm,
    25  		periodCount: atomickit.NewUint64(bm.currentPeriod.Load()),
    26  	}}
    27  }
    28  
    29  func NewChildBucket(parent ParentBucket, cfg BucketConfig) *ChildBucket {
    30  	cfg.Check()
    31  
    32  	bm := parent.(periodBucket).GetManager()
    33  	b := &ChildBucket{state: BucketState{
    34  		bucketConfig: cfg,
    35  	}, refiller: PeriodRefiller{
    36  		manager:     bm,
    37  		periodCount: atomickit.NewUint64(bm.currentPeriod.Load()),
    38  	}, parent: parent}
    39  	return b
    40  }
    41  
    42  func NewChildBucketRatio(parent ParentBucket, refillRatio, maxRatio float32) *ChildBucket {
    43  	return NewChildBucket(parent, parent.CopyConfig().Scale(refillRatio, maxRatio))
    44  }
    45  
    46  type ParentBucket interface {
    47  	CopyConfig() BucketConfig
    48  	TakeQuotaNoScale(uint32) uint32
    49  }
    50  
    51  type periodBucket interface {
    52  	GetManager() *PeriodManager
    53  }
    54  
    55  /******************************************************/
    56  
    57  var _ ParentBucket = &Bucket{}
    58  var _ periodBucket = &Bucket{}
    59  
    60  type Bucket struct {
    61  	state    BucketState
    62  	refiller PeriodRefiller
    63  }
    64  
    65  func (p *Bucket) TakeQuota(max int64) int64 {
    66  	return p.refiller.TakeQuota(max, &p.state, p.refill)
    67  }
    68  
    69  func (p *Bucket) TakeQuotaNoWait(max int64) int64 {
    70  	return p.refiller.TakeQuotaNoWait(max, &p.state, p.refill)
    71  }
    72  
    73  func (p *Bucket) TakeQuotaNoScale(max uint32) uint32 {
    74  	return p.state.TakeQuotaNoScale(max, p.refill)
    75  }
    76  
    77  func (p *Bucket) CopyConfig() BucketConfig {
    78  	return p.state.config()
    79  }
    80  
    81  func (p *Bucket) GetManager() *PeriodManager {
    82  	return p.refiller.manager
    83  }
    84  
    85  func (p *Bucket) refill(uint32) uint32 {
    86  	x := p.refiller.GetRefillCount()
    87  	return p.state.PeriodsToRefill(x)
    88  }
    89  
    90  /******************************************************/
    91  
    92  var _ ParentBucket = &ChildBucket{}
    93  var _ periodBucket = &ChildBucket{}
    94  
    95  type ChildBucket struct {
    96  	state    BucketState
    97  	refiller PeriodRefiller
    98  	quota    RefillQuota
    99  	parent   ParentBucket
   100  }
   101  
   102  func (p *ChildBucket) TakeQuota(max int64) int64 {
   103  	return p.refiller.TakeQuota(max, &p.state, p.refill)
   104  }
   105  
   106  func (p *ChildBucket) TakeQuotaNoWait(max int64) int64 {
   107  	return p.refiller.TakeQuotaNoWait(max, &p.state, p.refill)
   108  }
   109  
   110  func (p *ChildBucket) TakeQuotaNoScale(max uint32) uint32 {
   111  	return p.state.TakeQuotaNoScale(max, p.refill)
   112  }
   113  
   114  func (p *ChildBucket) CopyConfig() BucketConfig {
   115  	return p.state.config()
   116  }
   117  
   118  func (p *ChildBucket) GetManager() *PeriodManager {
   119  	return p.refiller.manager
   120  }
   121  
   122  func (p *ChildBucket) refill(needed uint32) uint32 {
   123  	x := p.refiller.GetRefillCount()
   124  	return p.quota.Refill(needed, x, &p.state, p.parent.TakeQuotaNoScale)
   125  }