github.com/blend/go-sdk@v1.20220411.3/ratelimiter/queue.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package ratelimiter
     9  
    10  import (
    11  	"time"
    12  
    13  	"github.com/blend/go-sdk/collections"
    14  )
    15  
    16  // NewQueue returns a new queue based rate limiter.
    17  func NewQueue(numberOfActions int, quantum time.Duration) *Queue {
    18  	return &Queue{
    19  		NumberOfActions: numberOfActions,
    20  		Quantum:         quantum,
    21  		Limits:          map[string]collections.Queue{},
    22  		Now:             func() time.Time { return time.Now().UTC() },
    23  	}
    24  }
    25  
    26  // Queue is a simple implementation of a rate checker.
    27  type Queue struct {
    28  	NumberOfActions int
    29  	Quantum         time.Duration
    30  	Limits          map[string]collections.Queue
    31  	Now             func() time.Time
    32  }
    33  
    34  // Check returns true if it has been called NumberOfActions times or more in Quantum or smaller duration.
    35  func (q *Queue) Check(id string) bool {
    36  	queue, hasQueue := q.Limits[id]
    37  	if !hasQueue {
    38  		queue = collections.NewRingBufferWithCapacity(q.NumberOfActions)
    39  		q.Limits[id] = queue
    40  	}
    41  
    42  	now := q.Now()
    43  	queue.Enqueue(now)
    44  	if queue.Len() < q.NumberOfActions {
    45  		return false
    46  	}
    47  
    48  	oldest := queue.Dequeue().(time.Time)
    49  	return now.Sub(oldest) < q.Quantum
    50  }