golimiter/golimiter.go

// golimiter.go
package golimiter

import (
	"sync"
	"time"
)

// TokenBucket represents a rate limiting mechanism based on the token bucket algorithm.
type TokenBucket struct {
	capacity      int           // Maximum number of tokens the bucket can hold
	tokens        int           // Current number of tokens in the bucket
	fillRate      int           // Rate at which tokens are added to the bucket (tokens per second)
	lastTimestamp time.Time     // Timestamp of the last token retrieval or addition
	expirationTime time.Time   // Time when the bucket will expire and be removed from the manager
	mu            sync.Mutex   // Mutex for synchronization
}

// NewTokenBucket creates a new TokenBucket instance with the specified parameters.
func NewTokenBucket(capacity, fillRate int, bucketLifetime time.Duration) *TokenBucket {
    return &TokenBucket{
        capacity:      capacity,
        tokens:        capacity,
        fillRate:      fillRate,
        lastTimestamp: time.Now(),
        expirationTime: time.Now().Add(bucketLifetime),
    }
}

// TakeToken tries to retrieve a token from the bucket. Returns true if successful, false otherwise.

// AddTokensToBucket adds tokens to the bucket based on the elapsed time and fill rate.
func (tb *TokenBucket) AddTokensToBucket() {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	currentTime := time.Now()
	elapsed := currentTime.Sub(tb.lastTimestamp).Milliseconds()

	fillAmount := int(elapsed) * tb.fillRate / 1000
	tb.tokens = tb.tokens + fillAmount
	if tb.tokens > tb.capacity {
		tb.tokens = tb.capacity
	}
	tb.lastTimestamp = time.Now()
}

// TakeToken tries to retrieve a token from the bucket. Returns true if successful, false otherwise.
func (tb *TokenBucket) TakeToken() bool {

	tb.AddTokensToBucket()
	tb.mu.Lock()
	defer tb.mu.Unlock()

	if tb.tokens > 0 {
		tb.tokens--
		return true
	}

	return false
}