github.com/influxdata/telegraf@v1.30.3/internal/limiter/limiter.go (about)

     1  package limiter
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  )
     7  
     8  // NewRateLimiter returns a rate limiter that will emit from the C
     9  // channel only 'n' times every 'rate' seconds.
    10  func NewRateLimiter(n int, rate time.Duration) *rateLimiter {
    11  	r := &rateLimiter{
    12  		C:        make(chan bool),
    13  		rate:     rate,
    14  		n:        n,
    15  		shutdown: make(chan bool),
    16  	}
    17  	r.wg.Add(1)
    18  	go r.limiter()
    19  	return r
    20  }
    21  
    22  type rateLimiter struct {
    23  	C    chan bool
    24  	rate time.Duration
    25  	n    int
    26  
    27  	shutdown chan bool
    28  	wg       sync.WaitGroup
    29  }
    30  
    31  func (r *rateLimiter) Stop() {
    32  	close(r.shutdown)
    33  	r.wg.Wait()
    34  	close(r.C)
    35  }
    36  
    37  func (r *rateLimiter) limiter() {
    38  	defer r.wg.Done()
    39  	ticker := time.NewTicker(r.rate)
    40  	defer ticker.Stop()
    41  	counter := 0
    42  	for {
    43  		select {
    44  		case <-r.shutdown:
    45  			return
    46  		case <-ticker.C:
    47  			counter = 0
    48  		default:
    49  			if counter < r.n {
    50  				select {
    51  				case r.C <- true:
    52  					counter++
    53  				case <-r.shutdown:
    54  					return
    55  				}
    56  			}
    57  		}
    58  	}
    59  }