github.com/insolar/vanilla@v0.0.0-20201023172447-248fdf805322/ratelimiter/refill_period.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  /****************************************************/
    15  
    16  type PeriodRefiller struct {
    17  	manager     *PeriodManager
    18  	periodCount atomickit.Uint64
    19  }
    20  
    21  func (p *PeriodRefiller) TakeQuota(max int64, state *BucketState, refillFn BucketRefillFunc) int64 {
    22  	currentPeriod := p.manager.currentPeriod.Load()
    23  	for {
    24  		if q := p.TakeQuotaNoWait(max, state, refillFn); q > 0 {
    25  			return q
    26  		}
    27  		p.manager.waitNextPeriod(currentPeriod)
    28  	}
    29  }
    30  
    31  func (p *PeriodRefiller) TakeQuotaNoWait(max int64, state *BucketState, refillFn BucketRefillFunc) int64 {
    32  	return state.TakeQuotaNoWait(max, p.manager.amountScale, refillFn)
    33  }
    34  
    35  func (p *PeriodRefiller) GetRefillCount() uint64 {
    36  	for {
    37  		currentPeriod := p.manager.currentPeriod.Load()
    38  		lastPeriod := p.periodCount.Load()
    39  		if currentPeriod == lastPeriod {
    40  			return 0
    41  		}
    42  		if p.periodCount.CompareAndSwap(lastPeriod, currentPeriod) {
    43  			periods := currentPeriod - lastPeriod
    44  			if periods > math.MaxUint32 {
    45  				// to avoid overflows of following calculations
    46  				return math.MaxUint32
    47  			}
    48  			return periods
    49  		}
    50  	}
    51  }