golimiter/bucket_manager.go

// bucket_manager.go
package golimiter

import (
	"sync"
	"time"
	"fmt"
)

// TokenBucketManager manages multiple token buckets.
type TokenBucketManager struct {
	tokenBuckets map[string]*TokenBucket // Map of IP addresses to their corresponding token buckets
	mu           sync.Mutex   // Mutex for synchronization
}

// NewTokenBucketManager creates a new TokenBucketManager instance.
func NewTokenBucketManager() *TokenBucketManager {
	return &TokenBucketManager{
		tokenBuckets: make(map[string]*TokenBucket),
	}
}

// GetTokenBucket checks whether a TokenBucket exists for the specified IP address.
// If a bucket exists and is still valid, it returns the corresponding TokenBucket; otherwise, it returns nil.
func (tm *TokenBucketManager) GetTokenBucket(ip string) *TokenBucket {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	tb, exists := tm.tokenBuckets[ip]
	if exists {
		if tb.expirationTime.Before(time.Now()) {
			delete(tm.tokenBuckets, ip) // Remove expired bucket from the map
			return nil
		}

		return tb
	}
	return nil
}

// CreateTokenBucket creates a new TokenBucket for the given IP address and stores it in the manager.
// If a bucket already exists, it returns an error.
func (tm *TokenBucketManager) CreateTokenBucket(ip string, capacity, fillRate int, bucketLifetime time.Duration) (*TokenBucket, error) {
    tm.mu.Lock()
    defer tm.mu.Unlock()

    if _, exists := tm.tokenBuckets[ip]; exists {
        return nil, fmt.Errorf("token bucket for IP %s already exists", ip)
    }

    tb := NewTokenBucket(capacity, fillRate, bucketLifetime)
    tm.tokenBuckets[ip] = tb
    return tb, nil
}

// GetRemainingTokens returns the remaining token count for the specified IP address.
// Returns -1 if the bucket does not exist.
func (tm *TokenBucketManager) GetRemainingTokens(ip string) int {
	tb := tm.GetTokenBucket(ip)
	if tb == nil {
		return -1
	}

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

	tb.AddTokensToBucket()
	return tb.tokens
}

// cleanupExpiredTokenBuckets removes expired token buckets from the manager.
func (tm *TokenBucketManager) cleanupExpiredTokenBuckets() {
    tm.mu.Lock()
    defer tm.mu.Unlock()

    currentTime := time.Now()
    for ip, tb := range tm.tokenBuckets {
        if tb.expirationTime.Before(currentTime) {
            delete(tm.tokenBuckets, ip)
        }
    }
}

// StartCleanupTask starts a periodic task to clean up expired token buckets at a specified interval.
func (tm *TokenBucketManager) StartCleanupTask(cleanupInterval time.Duration) {
    ticker := time.NewTicker(cleanupInterval)
    defer ticker.Stop()

    for {
        select {
        case <-ticker.C:
            tm.cleanupExpiredTokenBuckets()
        }
    }
}