github.com/insolar/vanilla@v0.0.0-20201023172447-248fdf805322/ratelimiter/refill_quota.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  	"github.com/insolar/vanilla/throw"
    11  )
    12  
    13  type RefillQuota struct {
    14  	refillQuota atomickit.Uint32 // used by child bucket to avoid exhaustion of parent
    15  }
    16  
    17  func (p *RefillQuota) Refill(needed uint32, periods uint64, state *BucketState, parentRefillFn func(uint32) uint32) uint32 {
    18  
    19  	q := p.getParentQuota(needed)
    20  	if q < needed && periods > 0 {
    21  		allowed := state.PeriodsToRefill(periods)
    22  		if lacks := needed - q; lacks >= allowed {
    23  			q += allowed
    24  		} else {
    25  			q = needed
    26  			p.addToRefillQuota(allowed-lacks, state.MaxAmount)
    27  		}
    28  	}
    29  
    30  	allocated := parentRefillFn(q)
    31  	if q > allocated {
    32  		panic(throw.Impossible())
    33  	}
    34  	return allocated
    35  
    36  }
    37  
    38  func (p *RefillQuota) getParentQuota(needed uint32) uint32 {
    39  	for {
    40  		switch q := p.refillQuota.Load(); {
    41  		case q > needed:
    42  			if !p.refillQuota.CompareAndSwap(q, q-needed) {
    43  				continue
    44  			}
    45  			return needed
    46  		case q > 0:
    47  			if !p.refillQuota.CompareAndSwap(q, 0) {
    48  				continue
    49  			}
    50  			return q
    51  		}
    52  		return 0
    53  	}
    54  }
    55  
    56  func (p *RefillQuota) addToRefillQuota(x, max uint32) {
    57  	for {
    58  		v := p.refillQuota.Load()
    59  		v2 := v + x
    60  		if v2 > max || v2 < v /* overflow */ {
    61  			v2 = max
    62  		}
    63  		if p.refillQuota.CompareAndSwap(v, v2) {
    64  			return
    65  		}
    66  	}
    67  }