github.com/insolar/vanilla@v0.0.0-20201023172447-248fdf805322/ratelimiter/bucket_state.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  	"math"
    10  
    11  	"github.com/insolar/vanilla/atomickit"
    12  )
    13  
    14  type BucketState struct {
    15  	filledAmount   atomickit.Uint32
    16  	residualAmount atomickit.Uint32 // scaled value!
    17  	bucketConfig
    18  }
    19  type bucketConfig = BucketConfig
    20  
    21  type BucketRefillFunc func(need uint32) uint32
    22  
    23  func (p *BucketState) TakeQuotaNoWait(max int64, scale uint32, refillFn BucketRefillFunc) int64 {
    24  	if max <= 0 {
    25  		return 0
    26  	}
    27  	for {
    28  		switch n := p.residualAmount.Load(); {
    29  		case int64(n) >= max:
    30  			if p.residualAmount.CompareAndSwap(n, n-uint32(max)) {
    31  				return max
    32  			}
    33  		case n > 0:
    34  			if p.residualAmount.CompareAndSwap(n, 0) {
    35  				return int64(n) + p._takeQuotaNoWait(max-int64(n), scale, refillFn)
    36  			}
    37  		default:
    38  			return p._takeQuotaNoWait(max, scale, refillFn)
    39  		}
    40  	}
    41  }
    42  
    43  func (p *BucketState) _takeQuotaNoWait(max int64, scale uint32, refillFn BucketRefillFunc) int64 {
    44  	maxUnscaled := (uint64(max) + uint64(scale) - 1) / uint64(scale)
    45  	if maxUnscaled > math.MaxUint32 {
    46  		maxUnscaled = math.MaxUint32
    47  	}
    48  
    49  	allocated := int64(p.takeFilled(uint32(maxUnscaled), refillFn)) * int64(scale)
    50  	if allocated <= max {
    51  		return allocated
    52  	}
    53  	p.addResidual(uint64(allocated - max))
    54  	return max
    55  }
    56  
    57  func (p *BucketState) addResidual(residual uint64) {
    58  	for {
    59  		n := p.residualAmount.Load()
    60  		x := residual + uint64(n)
    61  		if x > math.MaxUint32 {
    62  			x = math.MaxUint32
    63  		}
    64  		if p.residualAmount.CompareAndSwap(n, uint32(x)) {
    65  			return
    66  		}
    67  	}
    68  }
    69  
    70  func (p *BucketState) TakeQuotaNoScale(max uint32, refillFn BucketRefillFunc) uint32 {
    71  	if max == 0 {
    72  		return 0
    73  	}
    74  	return p.takeFilled(max, refillFn)
    75  }
    76  
    77  func (p *BucketState) takeFilled(max uint32, refillFn BucketRefillFunc) uint32 {
    78  	allocate := p.takeQuantizedFilled(quantizeCeiling(max, p.Quantum), refillFn)
    79  	if allocate <= max {
    80  		return allocate
    81  	}
    82  	p.residualAmount.Add(allocate - max)
    83  	return max
    84  }
    85  
    86  func (p *BucketState) takeQuantizedFilled(maxAligned uint32, refillFn BucketRefillFunc) uint32 {
    87  	for allocated := uint32(0); ; {
    88  		n := p.filledAmount.Load()
    89  
    90  		switch needed := maxAligned - allocated; {
    91  		case n < needed:
    92  			switch allocated += refillFn(needed - n); {
    93  			case allocated < maxAligned:
    94  				//
    95  			case allocated == maxAligned:
    96  				return allocated
    97  			default:
    98  				p.filledAmount.Add(allocated - maxAligned)
    99  				return allocated
   100  			}
   101  		case p.filledAmount.CompareAndSwap(n, n-needed):
   102  			return maxAligned
   103  		default:
   104  			continue
   105  		}
   106  
   107  		switch q := quantizeFloor(allocated+n, p.Quantum); {
   108  		case q > allocated:
   109  			if !p.filledAmount.CompareAndSwap(n, n+allocated-q) {
   110  				continue
   111  			}
   112  			return q
   113  		case q == 0:
   114  			if allocated > 0 {
   115  				p.filledAmount.Add(allocated)
   116  			}
   117  			return 0
   118  		case q == allocated:
   119  			return q
   120  		default:
   121  			p.filledAmount.Add(allocated - q)
   122  			return q
   123  		}
   124  	}
   125  }
   126  
   127  func (p *BucketState) PeriodsToRefill(periods uint64) uint32 {
   128  	if delta := uint64(p.RefillAmount) * periods; delta < uint64(p.MaxAmount) {
   129  		return uint32(delta)
   130  	}
   131  	return p.MaxAmount
   132  }
   133  
   134  func (p *BucketState) ForceRefill(x uint32) {
   135  	if x == 0 {
   136  		return
   137  	}
   138  	for {
   139  		n := p.filledAmount.Load()
   140  		v := n + x
   141  		if v > p.MaxAmount {
   142  			v = p.MaxAmount
   143  		}
   144  		if p.filledAmount.CompareAndSwap(n, v) {
   145  			return
   146  		}
   147  	}
   148  }
   149  
   150  func (p *BucketState) config() bucketConfig {
   151  	return bucketConfig{
   152  		RefillAmount: p.RefillAmount,
   153  		Quantum:      p.Quantum,
   154  		MaxAmount:    p.MaxAmount,
   155  	}
   156  }