github.com/insolar/vanilla@v0.0.0-20201023172447-248fdf805322/ratelimiter/throttle.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  	"context"
    10  	"time"
    11  
    12  	"github.com/insolar/vanilla/atomickit"
    13  	"github.com/insolar/vanilla/throw"
    14  )
    15  
    16  func NewController(trafficUnit uint32, refillPeriod time.Duration) *Controller {
    17  	switch {
    18  	case trafficUnit == 0:
    19  		panic(throw.IllegalValue())
    20  	case refillPeriod <= 0:
    21  		panic(throw.IllegalValue())
    22  	}
    23  	c := &Controller{}
    24  	c.root.amountScale = trafficUnit
    25  	c.period = refillPeriod
    26  	c.root.init()
    27  	return c
    28  }
    29  
    30  type Controller struct {
    31  	root     PeriodManager
    32  	period   time.Duration
    33  	throttle atomickit.Uint32
    34  	started  bool
    35  }
    36  
    37  func (p *Controller) Init() {
    38  	if p.started {
    39  		panic(throw.IllegalState())
    40  	}
    41  	p.root.init()
    42  }
    43  
    44  func (p *Controller) SetIncrement(level uint32) {
    45  	p.throttle.Store(level)
    46  }
    47  
    48  func (p *Controller) GetIncrement() uint32 {
    49  	return p.throttle.Load()
    50  }
    51  
    52  func (p *Controller) Start(ctx context.Context) {
    53  	if p.started {
    54  		panic(throw.IllegalState())
    55  	}
    56  	if ctx == nil {
    57  		panic(throw.IllegalValue())
    58  	}
    59  	p.started = true
    60  	go p.run(ctx)
    61  }
    62  
    63  func (p *Controller) Root() *PeriodManager {
    64  	return &p.root
    65  }
    66  
    67  func (p *Controller) run(ctx context.Context) {
    68  	ticker := time.NewTicker(p.period)
    69  	defer ticker.Stop()
    70  
    71  	for {
    72  		select {
    73  		case <-ticker.C:
    74  			step := p.throttle.Load()
    75  			if step > 0 {
    76  				p.root.nextPeriod(uint(step))
    77  			}
    78  		case <-ctx.Done():
    79  			return
    80  		}
    81  	}
    82  }