go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/rate/checker.go (about)

     1  /*
     2  
     3  Copyright (c) 2024 - Present. Will Charczuk. All rights reserved.
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.
     5  
     6  */
     7  
     8  package rate
     9  
    10  import (
    11  	"time"
    12  
    13  	"go.charczuk.com/sdk/collections"
    14  	"go.charczuk.com/sdk/errutil"
    15  )
    16  
    17  const (
    18  	ErrRateCheckFailure errutil.Class = "rate check failure"
    19  )
    20  
    21  // NewRateChecker returns a new rate checker.
    22  func NewRateChecker(numActions int, quantum time.Duration) *RateChecker {
    23  	return &RateChecker{
    24  		numActions:  numActions,
    25  		quantum:     quantum,
    26  		events:      new(collections.Queue[time.Time]),
    27  		nowProvider: time.Now,
    28  	}
    29  }
    30  
    31  // RateChecker is a simple queue based rate measurer
    32  // that will return an error if `Check` is called too often.
    33  type RateChecker struct {
    34  	numActions  int
    35  	quantum     time.Duration
    36  	events      *collections.Queue[time.Time]
    37  	nowProvider func() time.Time
    38  }
    39  
    40  // Check returns an error if it is called more than
    41  // numActions times in quantum duration.
    42  func (rc *RateChecker) Check() error {
    43  	now := rc.nowProvider()
    44  	rc.events.Push(now)
    45  	if rc.events.Len() < rc.numActions {
    46  		return nil
    47  	}
    48  
    49  	oldest, _ := rc.events.Pop()
    50  	if now.Sub(oldest) < rc.quantum {
    51  		return ErrRateCheckFailure
    52  	}
    53  	return nil
    54  }