Leaking Bucket Algorithm

import time
from collections import deque

class LeakyBucket:
    def __init__(self, bucket_size, leak_rate):
        """
        Initializes the leaky bucket.
        
        :param bucket_size: Maximum capacity of the bucket (queue size).
        :param leak_rate: Number of requests processed per second (leak rate).
        """
        self.bucket_size = bucket_size
        self.leak_rate = leak_rate
        self.tokens = 0  # Current tokens in the bucket
        self.last_refill_time = time.time()
        self.queue = deque()  # Queue to hold incoming requests
    
    def refill(self):
        """
        Refill the bucket based on the leak rate.
        The bucket leaks at a fixed rate over time.
        """
        current_time = time.time()
        elapsed_time = current_time - self.last_refill_time
        
        # Calculate how many tokens to leak based on elapsed time
        tokens_to_leak = int(elapsed_time * self.leak_rate)
        
        if tokens_to_leak > 0:
            # Leak the tokens
            self.tokens = max(0, self.tokens - tokens_to_leak)
            self.last_refill_time = current_time
    
    def handle_request(self):
        """
        Process a request. If the bucket has space, the request is allowed.
        If the bucket is full, the request is denied.
        """
        self.refill()

        if len(self.queue) < self.bucket_size:
            # Add request to the queue (i.e., in the bucket)
            self.queue.append(time.time())
            return True  # Request allowed
        else:
            # Drop the request as the bucket is full
            return False


class RateLimiter:
    def __init__(self, bucket_size, leak_rate, mode='global'):
        """
        Initializes the rate limiter.
        
        :param bucket_size: Maximum capacity of the bucket (queue size).
        :param leak_rate: Number of requests processed per second.
        :param mode: Mode of rate limiting ('global', 'per_user', 'per_api').
        """
        self.bucket_size = bucket_size
        self.leak_rate = leak_rate
        self.mode = mode
        self.buckets = {}

    def get_bucket_key(self, identifier):
        """
        Generate the key for each rate-limited bucket based on the mode.
        
        :param identifier: Unique identifier (API, user, IP).
        """
        if self.mode == 'global':
            return 'global'
        elif self.mode == 'per_api':
            return identifier  # API endpoint as identifier
        elif self.mode == 'per_user':
            return identifier  # User or IP address as identifier
        else:
            raise ValueError("Unknown mode for rate limiter")

    def handle_request(self, identifier):
        """
        Handle an incoming request. Returns True if allowed, False if denied.
        
        :param identifier: Unique identifier (API, user, IP).
        """
        bucket_key = self.get_bucket_key(identifier)

        if bucket_key not in self.buckets:
            # Create a new bucket if it doesn't exist
            self.buckets[bucket_key] = LeakyBucket(self.bucket_size, self.leak_rate)

        bucket = self.buckets[bucket_key]
        return bucket.handle_request()


# Example usage
if __name__ == "__main__":
    # Create a global rate limiter with bucket size of 5 and leak rate of 1 request per second
    limiter = RateLimiter(bucket_size=5, leak_rate=1, mode='global')

    # Simulate requests from multiple sources
    for i in range(10):
        if limiter.handle_request('global'):
            print(f"Request {i + 1}: Allowed")
        else:
            print(f"Request {i + 1}: Denied")
        time.sleep(0.5)  # Wait for 0.5 seconds between requests

    # Per User example with different identifiers (IP or User ID)
    user_limiter = RateLimiter(bucket_size=3, leak_rate=1, mode='per_user')
    users = ['user1', 'user2', 'user1', 'user3', 'user2']

    for user in users:
        if user_limiter.handle_request(user):
            print(f"User '{user}': Request Allowed")
        else:
            print(f"User '{user}': Request Denied")
        time.sleep(0.5)  # Wait for 0.5 seconds between requests

    # Per API example
    api_limiter = RateLimiter(bucket_size=4, leak_rate=1, mode='per_api')
    apis = ['GET /home', 'POST /login', 'GET /home', 'POST /login', 'GET /home']

    for api in apis:
        if api_limiter.handle_request(api):
            print(f"API '{api}': Request Allowed")
        else:
            print(f"API '{api}': Request Denied")
        time.sleep(0.5)  # Wait for 0.5 seconds between requests