steindamm 0.7.0__py3-none-any.whl

steindamm/__init__.py

@@ -0,0 +1,28 @@
+"""
+Various token bucket and semaphore implementations using a Redis or local backend.
+
+Use SyncTokenBucket or AsyncTokenBucket to automatically select between Redis-based
+and local in-memory implementations based on whether a Redis connection is provided.
+
+For explicit control over the implementation, import and use
+SyncRedisTokenBucket, AsyncRedisTokenBucket, SyncLocalTokenBucket, or AsyncLocalTokenBucket directly.
+"""
+
+# TODO: Add local semaphore implementation and update docs accordingly
+from redis_limiters.exceptions import MaxSleepExceededError
+from redis_limiters.semaphore import AsyncSemaphore, SyncSemaphore
+from redis_limiters.token_bucket.local_token_bucket import AsyncLocalTokenBucket, SyncLocalTokenBucket
+from redis_limiters.token_bucket.redis_token_bucket import AsyncRedisTokenBucket, SyncRedisTokenBucket
+from redis_limiters.token_bucket.token_bucket import AsyncTokenBucket, SyncTokenBucket
+
+__all__ = (
+    "AsyncLocalTokenBucket",
+    "AsyncRedisTokenBucket",
+    "AsyncSemaphore",
+    "AsyncTokenBucket",
+    "MaxSleepExceededError",
+    "SyncLocalTokenBucket",
+    "SyncRedisTokenBucket",
+    "SyncSemaphore",
+    "SyncTokenBucket",
+)

steindamm/base.py

@@ -0,0 +1,49 @@
+"""Base classes for Redis Lua script handling."""
+
+from pathlib import Path
+from typing import Any, ClassVar
+
+from pydantic import BaseModel, ConfigDict
+from redis import Redis as SyncRedis
+from redis.asyncio import Redis as AsyncRedis
+from redis.asyncio.cluster import RedisCluster as AsyncRedisCluster
+from redis.cluster import RedisCluster as SyncRedisCluster
+from redis.commands.core import AsyncScript, Script
+
+
+class SyncLuaScriptBase(BaseModel):
+    """Base class for synchronous Redis Lua script handling."""
+
+    model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
+
+    connection: SyncRedis | SyncRedisCluster
+    script_name: ClassVar[str]
+    script: Script = None  # type: ignore[assignment]
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the Lua script base class and load the script."""
+        super().__init__(**kwargs)
+
+        # https://github.com/redis/redis-py/issues/3712
+        # Load script on initialization
+        with open(Path(__file__).parent / self.script_name) as f:
+            self.script = self.connection.register_script(f.read())  # type: ignore
+
+
+class AsyncLuaScriptBase(BaseModel):
+    """Base class for asynchronous Redis Lua script handling."""
+
+    model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
+
+    connection: AsyncRedis | AsyncRedisCluster
+    script_name: ClassVar[str]
+    script: AsyncScript = None  # type: ignore[assignment]
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the Lua script base class and load the script."""
+        super().__init__(**kwargs)
+
+        # https://github.com/redis/redis-py/issues/3712
+        # Load script on initialization
+        with open(Path(__file__).parent / self.script_name) as f:
+            self.script = self.connection.register_script(f.read())  # type: ignore

steindamm/exceptions.py

@@ -0,0 +1,7 @@
+"""Exceptions for redis_limiters package."""
+
+
+class MaxSleepExceededError(Exception):
+    """Raised when we've slept for longer than the `max_sleep` specified limit."""
+
+    pass

steindamm/semaphore.lua

@@ -0,0 +1,45 @@
+--- Script called from the Semaphore implementation.
+---
+--- Lua scripts are run atomically by default, and since redis
+--- is single threaded, there are no race conditions to worry about.
+---
+--- The script checks if a list exists for the Semaphore, and
+--- creates one of length `capacity` if it doesn't.
+---
+--- keys:
+--- * key: The key to use for the list
+--- * exists: The key to use for the string we use to check if the lists exists
+---
+--- args:
+--- * capacity: The capacity of the semaphore (i.e., the length of the list)
+---
+--- returns:
+--- * 1 if created, else 0 (but the return value isn't used; only useful for debugging)
+
+redis.replicate_commands()
+
+-- Init config variables
+local key = tostring(KEYS[1])
+local exists = tostring(KEYS[2])
+local capacity = tonumber(ARGV[1])
+
+-- Check if list exists
+-- Note, we cannot use `EXISTS` or `LLEN` directly on the `key` below,
+-- as a list in Redis will "stop existing" if it's empty (empty state will occur
+-- whenever the Semaphore is fully utilized). Instead, we use a separate
+-- key to check whether a list has been created for our `key` or not.
+local does_not_exist = redis.call('SETNX', string.format(exists, key), 1)
+
+-- Create the list if none exists
+if does_not_exist == 1 then
+    -- Add '1' as an argument equal to the capacity of the semaphore
+    -- If capacity is 5 here, we generate `{RPUSH, 1, 1, 1, 1, 1}`.
+    local args = { 'RPUSH', key }
+    for _ = 1, capacity do
+        table.insert(args, 1)
+    end
+    redis.call(unpack(args))
+    return true
+end
+
+return false

steindamm/semaphore.py

@@ -0,0 +1,127 @@
+"""Semaphore limiter implementation."""
+
+from datetime import datetime
+from logging import getLogger
+from types import TracebackType
+from typing import Annotated, ClassVar
+
+from pydantic import BaseModel, Field
+from redis.asyncio.client import Pipeline
+from redis.asyncio.cluster import ClusterPipeline
+
+from redis_limiters import MaxSleepExceededError
+from redis_limiters.base import AsyncLuaScriptBase, SyncLuaScriptBase
+
+logger = getLogger(__name__)
+
+PositiveInt = Annotated[int, Field(gt=0)]
+NonNegativeFloat = Annotated[float, Field(ge=0)]
+
+
+# TODO: Implement local semaphore as done with token bucket.
+class SemaphoreBase(BaseModel):  # noqa: D101 TODO: Fix after local semaphore is added
+    name: str
+    capacity: PositiveInt = 5
+    expiry: PositiveInt = 60
+    max_sleep: NonNegativeFloat = 30
+
+    @property
+    def key(self) -> str:
+        """Key to use for the Semaphore list."""
+        return f"{{limiter}}:semaphore:{self.name}"
+
+    @property
+    def exists(self) -> str:
+        """Key to use when checking if the Semaphore list has been created or not."""
+        return f"{{limiter}}:semaphore:{self.name}-exists"
+
+    def __str__(self) -> str:
+        return f"Semaphore instance for queue {self.key}"
+
+
+class SyncSemaphore(SemaphoreBase, SyncLuaScriptBase):  # noqa: D101 TODO: Fix after local semaphore is added
+    script_name: ClassVar[str] = "semaphore.lua"
+
+    def __enter__(self) -> None:
+        """Call the semaphore Lua script to create a semaphore, then call BLPOP to acquire it."""
+        # Retrieve timestamp for when to wake up from Redis
+        # To understand what exists does, check the Lua script
+        if self.script(
+            keys=[self.key, self.exists],
+            args=[self.capacity],
+        ):
+            logger.info("Created new semaphore `%s` with capacity %s", self.name, self.capacity)
+        else:
+            logger.debug("Skipped creating semaphore, since one exists")
+
+        start = datetime.now()
+
+        self.connection.blpop([self.key], self.max_sleep)
+        pipeline = self.connection.pipeline()
+        pipeline.expire(self.key, self.expiry)
+        pipeline.expire(self.exists, self.expiry)
+        pipeline.execute()
+
+        # Raise an exception if we exceeded `max_sleep`
+        if 0.0 < self.max_sleep < (datetime.now() - start).total_seconds():
+            raise MaxSleepExceededError("Max sleep exceeded waiting for Semaphore")
+
+        logger.debug("Acquired semaphore %s", self.name)
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        pipeline = self.connection.pipeline()
+        pipeline.lpush(self.key, 1)
+        pipeline.expire(self.key, self.expiry)
+        pipeline.expire(self.exists, self.expiry)
+        pipeline.execute()
+
+        logger.debug("Released semaphore %s", self.name)
+
+
+class AsyncSemaphore(SemaphoreBase, AsyncLuaScriptBase):  # noqa: D101 TODO: Fix after local semaphore is added
+    script_name: ClassVar[str] = "semaphore.lua"
+
+    async def __aenter__(self) -> None:
+        """Call the semaphore Lua script to create a semaphore, then call BLPOP to acquire it."""
+        # Retrieve timestamp for when to wake up from Redis
+
+        if await self.script(
+            keys=[self.key, self.exists],
+            args=[self.capacity],
+        ):
+            logger.info("Created new semaphore `%s` with capacity %s", self.name, self.capacity)
+        else:
+            logger.debug("Skipped creating semaphore, since one exists")
+
+        start = datetime.now()
+
+        await self.connection.blpop([self.key], self.max_sleep)  # type: ignore[union-attr]
+        pipeline: Pipeline | ClusterPipeline = self.connection.pipeline()
+        pipeline.expire(self.key, self.expiry)  # type: ignore[union-attr]
+        pipeline.expire(self.exists, self.expiry)  # type: ignore[union-attr]
+        await pipeline.execute()
+
+        # Raise an exception if we waited too long
+        if 0.0 < self.max_sleep < (datetime.now() - start).total_seconds():
+            raise MaxSleepExceededError(f"Max sleep ({self.max_sleep}s) exceeded waiting for Semaphore")
+
+        logger.debug("Acquired semaphore %s", self.name)
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        pipeline: Pipeline[str] | ClusterPipeline[str] = self.connection.pipeline()
+        pipeline.lpush(self.key, 1)  # type: ignore[union-attr]
+        pipeline.expire(self.key, self.expiry)  # type: ignore[union-attr]
+        pipeline.expire(self.exists, self.expiry)  # type: ignore[union-attr]
+        await pipeline.execute()
+
+        logger.debug("Released semaphore %s", self.name)

steindamm/token_bucket/local_token_bucket.py

@@ -0,0 +1,118 @@
+"""Synchronous and Asynchronous local token bucket implementations."""
+
+import asyncio
+import time
+from threading import Lock
+from types import TracebackType
+from typing import ClassVar
+
+from redis_limiters.token_bucket.token_bucket_base import TokenBucketBase
+
+
+class SyncLocalTokenBucket(TokenBucketBase):
+    """
+    Synchronous local token bucket.
+
+    Args:
+        name: Unique identifier for this token bucket.
+        capacity: Maximum number of tokens the bucket can hold.
+        refill_frequency: Time in seconds between token refills.
+        initial_tokens: Starting number of tokens. Defaults to capacity if not specified.
+        refill_amount: Number of tokens added per refill.
+        max_sleep: Maximum seconds to sleep when rate limited. 0 means no limit.
+        expiry: Key expiry time in seconds - currently not implemented for local buckets.
+        tokens_to_consume: Number of tokens to consume per operation.
+
+    Example:
+        .. code-block:: python
+
+           bucket = SyncLocalTokenBucket(name="api", capacity=10)
+            with bucket:
+                make_api_call()
+
+    """
+
+    # Class-level storage for bucket state (shared across instances)
+    # TODO: Currently there's no cleanup of old buckets.
+    # Consider adding periodic cleanup based on expiry.
+    _buckets: ClassVar[dict[str, dict]] = {}
+    _locks: ClassVar[dict[str, Lock]] = {}
+    _main_lock: ClassVar[Lock] = Lock()
+
+    def _get_lock(self) -> Lock:
+        # This is not safe in free threaded python
+        # Not acquiring main lock to improve performance in CPython with GIL
+        if self.key not in self._locks:
+            with self._main_lock:
+                if self.key not in self._locks:
+                    self._locks[self.key] = Lock()
+        return self._locks[self.key]
+
+    def __enter__(self) -> None:
+        """Acquire token(s) from the token bucket and sleep until they are available."""
+        # Execute token bucket logic with thread safety
+        with self._get_lock():
+            timestamp = self.execute_local_token_bucket_logic(self._buckets)
+
+        # Parse timestamp and sleep
+        sleep_time = self.parse_timestamp(timestamp)
+        time.sleep(sleep_time)
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        return
+
+
+class AsyncLocalTokenBucket(TokenBucketBase):
+    """
+    Asynchronous local token bucket.
+
+    Args:
+        name: Unique identifier for this token bucket.
+        capacity: Maximum number of tokens the bucket can hold.
+        refill_frequency: Time in seconds between token refills.
+        initial_tokens: Starting number of tokens. Defaults to capacity if not specified.
+        refill_amount: Number of tokens added per refill.
+        max_sleep: Maximum seconds to sleep when rate limited. 0 means no limit.
+        expiry: Key expiry time in seconds - currently not implemented for local buckets.
+        tokens_to_consume: Number of tokens to consume per operation.
+
+    Example:
+        .. code-block:: python
+
+            bucket = AsyncLocalTokenBucket(name="api", capacity=10)
+            async with bucket:
+                await make_api_call()
+
+    Note: If you need to use this class from multiple threads (multiple event loops),
+    consider using SyncLocalTokenBucket instead, which provides proper thread safety.
+
+    """
+
+    # Class-level storage for bucket state (shared across instances)
+    # TODO: Currently there's no cleanup of old buckets.
+    # Consider adding periodic cleanup based on expiry.
+    _buckets: ClassVar[dict[str, dict]] = {}
+
+    async def __aenter__(self) -> None:
+        """Acquire token(s) from the token bucket and sleep until they are available."""
+        # Execute token bucket logic
+        # No lock needed: asyncio is single-threaded and execute_local_token_bucket_logic
+        # has no await points, making it atomic from asyncio's perspective
+        timestamp = self.execute_local_token_bucket_logic(self._buckets)
+
+        # Parse timestamp and sleep
+        sleep_time = self.parse_timestamp(timestamp)
+        await asyncio.sleep(sleep_time)
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        return

steindamm/token_bucket/redis_token_bucket.py

@@ -0,0 +1,137 @@
+"""Synchronous and Asynchronous Redis-backed (Standalone or Cluster) token bucket implementations."""
+
+import asyncio
+import time
+from types import TracebackType
+from typing import ClassVar, cast
+
+from redis_limiters.base import AsyncLuaScriptBase, SyncLuaScriptBase
+from redis_limiters.token_bucket.token_bucket_base import TokenBucketBase, get_current_time_ms
+
+
+class SyncRedisTokenBucket(TokenBucketBase, SyncLuaScriptBase):
+    """
+    Synchronous Redis-backed (Standalone or Cluster) token bucket.
+
+    Args:
+        name: Unique identifier for this token bucket.
+        connection: Redis connection (SyncRedis or SyncRedisCluster).
+        capacity: Maximum number of tokens the bucket can hold.
+        refill_frequency: Time in seconds between token refills.
+        initial_tokens: Starting number of tokens. Defaults to capacity if not specified.
+        refill_amount: Number of tokens added per refill.
+        max_sleep: Maximum seconds to sleep when rate limited. 0 means no limit.
+        expiry: Key expiry time in seconds.
+        tokens_to_consume: Number of tokens to consume per operation.
+
+    Example:
+        .. code-block:: python
+
+            from redis import Redis  # or from redis.cluster import RedisCluster
+            redis_conn = Redis(host='localhost', port=6379)
+            bucket = SyncRedisTokenBucket(connection=redis_conn, name="api", capacity=10)
+            with bucket:
+                make_api_call()
+
+    """
+
+    script_name: ClassVar[str] = "token_bucket/token_bucket.lua"
+
+    def __enter__(self) -> float:
+        """Acquire token(s) from the token bucket and sleep until they are available."""
+        # Retrieve timestamp for when to wake up from Redis Lua script
+        milliseconds = get_current_time_ms()
+        timestamp: int = cast(
+            int,
+            self.script(
+                keys=[self.key],
+                args=[
+                    self.capacity,
+                    self.refill_amount,
+                    self.initial_tokens or self.capacity,
+                    self.refill_frequency,
+                    milliseconds,
+                    self.expiry,
+                    self.tokens_to_consume,
+                ],
+            ),
+        )
+
+        # Estimate sleep time
+        sleep_time = self.parse_timestamp(timestamp)
+
+        # Sleep before returning
+        time.sleep(sleep_time)
+
+        return sleep_time
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        return
+
+
+class AsyncRedisTokenBucket(TokenBucketBase, AsyncLuaScriptBase):
+    """
+    Asynchronous Redis-backed (Standalone or Cluster) token bucket.
+
+    Args:
+        name: Unique identifier for this token bucket.
+        connection: Redis connection (AsyncRedis or AsyncRedisCluster).
+        capacity: Maximum number of tokens the bucket can hold.
+        refill_frequency: Time in seconds between token refills.
+        initial_tokens: Starting number of tokens. Defaults to capacity if not specified.
+        refill_amount: Number of tokens added per refill.
+        max_sleep: Maximum seconds to sleep when rate limited. 0 means no limit.
+        expiry: Key expiry time in seconds.
+        tokens_to_consume: Number of tokens to consume per operation.
+
+    Example:
+        .. code-block:: python
+
+            from redis.asyncio import Redis  # or from redis.asyncio.cluster import RedisCluster
+            redis_conn = Redis(host='localhost', port=6379)
+            bucket = AsyncRedisTokenBucket(connection=redis_conn, name="api", capacity=10)
+            async with bucket:
+                await make_api_call()
+
+    """
+
+    script_name: ClassVar[str] = "token_bucket/token_bucket.lua"
+
+    async def __aenter__(self) -> None:
+        """Acquire token(s) from the token bucket and sleep until they are available."""
+        # Retrieve timestamp for when to wake up from Redis Lua script
+        milliseconds = get_current_time_ms()
+        timestamp: int = cast(
+            int,
+            await self.script(
+                keys=[self.key],
+                args=[
+                    self.capacity,
+                    self.refill_amount,
+                    self.initial_tokens or self.capacity,
+                    self.refill_frequency,
+                    milliseconds,
+                    self.expiry,
+                    self.tokens_to_consume,
+                ],
+            ),
+        )
+
+        # Estimate sleep time
+        sleep_time = self.parse_timestamp(timestamp)
+
+        # Sleep before returning
+        await asyncio.sleep(sleep_time)
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        return

steindamm/token_bucket/token_bucket.lua

@@ -0,0 +1,89 @@
+--- Lua scripts are run atomically by default, and since redis
+--- is single threaded, there are no race conditions to worry about.
+---
+--- This script does three things, in order:
+--- 1. Retrieves token bucket state, which means the last slot assigned,
+---    and how many tokens are left to be assigned for that slot
+--- 2. Works out whether we need to move to the next slot(s), or consume
+---    tokens from the current one.
+--- 3. Saves the token bucket state and returns the slot. The state is a
+---    combination of the last slot assigned (timestamp) and the number of tokens left.
+---
+--- The token bucket implementation is forward looking, so we're really just handing
+--- out the next time there would be tokens in the bucket, and letting the client
+--- decide wait until that time.
+---
+--- returns:
+--- * The assigned slot, as a millisecond timestamp
+
+redis.replicate_commands()
+
+-- Arguments
+local capacity = tonumber(ARGV[1])
+local refill_amount = tonumber(ARGV[2])
+local initial_tokens = tonumber(ARGV[3])
+local time_between_slots = tonumber(ARGV[4]) * 1000 -- Convert to milliseconds
+local milliseconds = tonumber(ARGV[5])
+local expiry = tonumber(ARGV[6])
+local tokens_to_consume = tonumber(ARGV[7]) -- Number of tokens to consume
+
+-- Validate that tokens_to_consume doesn't exceed capacity
+if tokens_to_consume > capacity then
+    return redis.error_reply("Requested tokens exceed bucket capacity")
+end
+
+-- Validate that tokens_to_consume is positive
+if tokens_to_consume <= 0 then
+    return redis.error_reply("Must consume at least 1 token")
+end
+
+-- Keys
+local data_key = KEYS[1]
+
+-- Get current time in milliseconds
+local now = milliseconds
+
+-- Default bucket values (used if no bucket exists yet)
+local tokens = math.min(initial_tokens, capacity)
+local slot = now
+
+-- Retrieve stored state, if any
+local data = redis.call('GET', data_key)
+if data then
+    local last_slot, stored_tokens = data:match('(%S+) (%S+)')
+    slot = tonumber(last_slot)
+    tokens = tonumber(stored_tokens)
+
+    -- Calculate the number of slots that have passed since the last update
+    local slots_passed = math.floor((now - slot) / time_between_slots)
+    if slots_passed > 0 then
+        -- Refill the tokens based on the number of slots passed, capped by capacity
+        tokens = math.min(tokens + slots_passed * refill_amount, capacity)
+        -- Update the slot to this run
+        -- The previously added +20 ms execution penalty was removed as it was not needed
+        -- and all it did was add additional latency to all requests and in our case,
+        -- timing is handled gracefully with the condition used (wake_up_time < now)
+        slot = now
+    end
+end
+
+-- If not enough tokens are available, move to the next slot(s) and refill accordingly
+if tokens < tokens_to_consume then
+    -- Calculate how many additional tokens we need
+    local needed_tokens = tokens_to_consume - tokens
+    -- Calculate how many slots we need to move forward to get enough tokens
+    local needed_slots = math.ceil(needed_tokens / refill_amount)
+    slot = slot + needed_slots * time_between_slots
+    tokens = tokens + needed_slots * refill_amount
+    -- Clamp tokens to capacity
+    tokens = math.min(tokens, capacity)
+end
+
+-- Consume tokens
+tokens = tokens - tokens_to_consume
+
+-- Save updated state and set expiry
+redis.call('SETEX', data_key, expiry, string.format('%d %d', slot, tokens))
+
+-- Return the slot when the next token(s) will be available
+return slot