fastapi-stream-lease 0.1.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- fastapi_stream_lease/__init__.py +17 -0
- fastapi_stream_lease/config.py +47 -0
- fastapi_stream_lease/exceptions.py +58 -0
- fastapi_stream_lease/lease.py +95 -0
- fastapi_stream_lease/lua.py +95 -0
- fastapi_stream_lease/manager.py +147 -0
- fastapi_stream_lease/py.typed +1 -0
- fastapi_stream_lease-0.1.0.dist-info/METADATA +213 -0
- fastapi_stream_lease-0.1.0.dist-info/RECORD +11 -0
- fastapi_stream_lease-0.1.0.dist-info/WHEEL +4 -0
- fastapi_stream_lease-0.1.0.dist-info/licenses/LICENSE +21 -0
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from fastapi_stream_lease.config import LeaseConfig
|
|
4
|
+
from fastapi_stream_lease.exceptions import StreamLeaseError, StreamLeaseRejected
|
|
5
|
+
from fastapi_stream_lease.lease import StreamLease
|
|
6
|
+
from fastapi_stream_lease.manager import StreamLeaseManager
|
|
7
|
+
|
|
8
|
+
__version__ = "0.1.0"
|
|
9
|
+
|
|
10
|
+
__all__ = [
|
|
11
|
+
"LeaseConfig",
|
|
12
|
+
"StreamLease",
|
|
13
|
+
"StreamLeaseError",
|
|
14
|
+
"StreamLeaseManager",
|
|
15
|
+
"StreamLeaseRejected",
|
|
16
|
+
"__version__",
|
|
17
|
+
]
|
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@dataclass(frozen=True)
|
|
7
|
+
class LeaseConfig:
|
|
8
|
+
"""Configuration for stream lease management."""
|
|
9
|
+
|
|
10
|
+
lease_seconds: float = 30.0
|
|
11
|
+
"""Duration (in seconds) of each lease before expiring in Redis if not renewed."""
|
|
12
|
+
|
|
13
|
+
max_per_user: int = 3
|
|
14
|
+
"""Maximum concurrent streams allowed for a single user/principal key."""
|
|
15
|
+
|
|
16
|
+
max_global: int = 500
|
|
17
|
+
"""Maximum concurrent streams allowed across the entire cluster."""
|
|
18
|
+
|
|
19
|
+
key_prefix: str = "stream_lease"
|
|
20
|
+
"""Prefix for Redis keys (e.g. stream_lease:user:{id}, stream_lease:global)."""
|
|
21
|
+
|
|
22
|
+
def __post_init__(self) -> None:
|
|
23
|
+
if self.lease_seconds <= 0:
|
|
24
|
+
raise ValueError("lease_seconds must be greater than 0")
|
|
25
|
+
if self.max_per_user < 0:
|
|
26
|
+
raise ValueError("max_per_user cannot be negative")
|
|
27
|
+
if self.max_global < 0:
|
|
28
|
+
raise ValueError("max_global cannot be negative")
|
|
29
|
+
|
|
30
|
+
@property
|
|
31
|
+
def _cluster_prefix(self) -> str:
|
|
32
|
+
"""Ensure prefix uses Redis hash tags {...} for slot affinity in Redis Cluster."""
|
|
33
|
+
if "{" in self.key_prefix and "}" in self.key_prefix:
|
|
34
|
+
return self.key_prefix
|
|
35
|
+
return f"{{{self.key_prefix}}}"
|
|
36
|
+
|
|
37
|
+
def user_key(self, user_id: str | int) -> str:
|
|
38
|
+
return f"{self._cluster_prefix}:user:{user_id}"
|
|
39
|
+
|
|
40
|
+
@property
|
|
41
|
+
def global_key(self) -> str:
|
|
42
|
+
return f"{self._cluster_prefix}:global"
|
|
43
|
+
|
|
44
|
+
@property
|
|
45
|
+
def redis_ttl(self) -> int:
|
|
46
|
+
"""TTL set on Redis keys to ensure dead keys self-clean (twice lease duration)."""
|
|
47
|
+
return max(60, int(self.lease_seconds * 2))
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import Any, Literal
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
class StreamLeaseError(Exception):
|
|
7
|
+
"""Base exception for all stream lease errors."""
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class StreamLeaseRejected(StreamLeaseError):
|
|
11
|
+
"""Raised when a stream lease request is rejected due to concurrency limits."""
|
|
12
|
+
|
|
13
|
+
def __init__(
|
|
14
|
+
self,
|
|
15
|
+
reason: Literal["user_limit", "global_limit"],
|
|
16
|
+
retry_after: int = 5,
|
|
17
|
+
detail: str | None = None,
|
|
18
|
+
) -> None:
|
|
19
|
+
self.reason: Literal["user_limit", "global_limit"] = reason
|
|
20
|
+
self.retry_after: int = retry_after
|
|
21
|
+
self.detail: str = detail or (
|
|
22
|
+
"User stream limit reached" if reason == "user_limit" else "Global stream limit reached"
|
|
23
|
+
)
|
|
24
|
+
super().__init__(f"Stream lease rejected: {self.reason} - {self.detail}")
|
|
25
|
+
|
|
26
|
+
def as_response(self) -> Any:
|
|
27
|
+
"""Convert this rejection into a JSONResponse (HTTP 429)."""
|
|
28
|
+
try:
|
|
29
|
+
from starlette.responses import JSONResponse
|
|
30
|
+
except ImportError as err:
|
|
31
|
+
raise RuntimeError(
|
|
32
|
+
"Starlette or FastAPI must be installed to use as_response()"
|
|
33
|
+
) from err
|
|
34
|
+
|
|
35
|
+
return JSONResponse(
|
|
36
|
+
status_code=429,
|
|
37
|
+
content={
|
|
38
|
+
"code": "stream_connection_limit",
|
|
39
|
+
"reason": self.reason,
|
|
40
|
+
"detail": self.detail,
|
|
41
|
+
},
|
|
42
|
+
headers={"Retry-After": str(self.retry_after)},
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
def as_http_exception(self) -> Any:
|
|
46
|
+
"""Convert this rejection into a FastAPI/Starlette HTTPException (HTTP 429) to be raised."""
|
|
47
|
+
try:
|
|
48
|
+
from starlette.exceptions import HTTPException
|
|
49
|
+
except ImportError as err:
|
|
50
|
+
raise RuntimeError(
|
|
51
|
+
"FastAPI or Starlette must be installed to use as_http_exception()"
|
|
52
|
+
) from err
|
|
53
|
+
|
|
54
|
+
return HTTPException(
|
|
55
|
+
status_code=429,
|
|
56
|
+
detail=self.detail,
|
|
57
|
+
headers={"Retry-After": str(self.retry_after)},
|
|
58
|
+
)
|
|
@@ -0,0 +1,95 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import logging
|
|
5
|
+
import time
|
|
6
|
+
from collections.abc import AsyncIterable, AsyncIterator
|
|
7
|
+
from contextlib import suppress
|
|
8
|
+
from dataclasses import dataclass, field
|
|
9
|
+
from typing import TYPE_CHECKING, Any, TypeVar
|
|
10
|
+
|
|
11
|
+
if TYPE_CHECKING:
|
|
12
|
+
from fastapi_stream_lease.manager import StreamLeaseManager
|
|
13
|
+
|
|
14
|
+
logger = logging.getLogger(__name__)
|
|
15
|
+
|
|
16
|
+
T = TypeVar("T")
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@dataclass
|
|
20
|
+
class StreamLease:
|
|
21
|
+
"""Represents an active, acquired stream lease."""
|
|
22
|
+
|
|
23
|
+
lease_id: str
|
|
24
|
+
user_id: str | int
|
|
25
|
+
user_key: str
|
|
26
|
+
global_key: str
|
|
27
|
+
manager: StreamLeaseManager
|
|
28
|
+
created_at: float = field(default_factory=time.time)
|
|
29
|
+
_is_released: bool = field(default=False, init=False)
|
|
30
|
+
|
|
31
|
+
async def __aenter__(self) -> StreamLease:
|
|
32
|
+
return self
|
|
33
|
+
|
|
34
|
+
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
35
|
+
await self.release()
|
|
36
|
+
|
|
37
|
+
async def renew(self) -> bool:
|
|
38
|
+
"""Manually renew this lease, extending its TTL in Redis."""
|
|
39
|
+
if self._is_released:
|
|
40
|
+
return False
|
|
41
|
+
return await self.manager.renew(self)
|
|
42
|
+
|
|
43
|
+
async def release(self) -> None:
|
|
44
|
+
"""Explicitly release this lease from Redis."""
|
|
45
|
+
if self._is_released:
|
|
46
|
+
return
|
|
47
|
+
self._is_released = True
|
|
48
|
+
await self.manager.release(self)
|
|
49
|
+
|
|
50
|
+
async def wrap(
|
|
51
|
+
self,
|
|
52
|
+
stream: AsyncIterable[T],
|
|
53
|
+
auto_renew: bool = True,
|
|
54
|
+
renew_interval: float | None = None,
|
|
55
|
+
) -> AsyncIterator[T]:
|
|
56
|
+
"""
|
|
57
|
+
Wrap an async stream (e.g. SSE event generator or LLM token stream).
|
|
58
|
+
|
|
59
|
+
Guarantees that:
|
|
60
|
+
1. A background renewal worker keeps the lease alive during long pauses
|
|
61
|
+
(such as slow LLM time-to-first-token or client idle periods).
|
|
62
|
+
2. When the stream terminates, is cancelled, or the client disconnects,
|
|
63
|
+
the lease is guaranteed to be released in the `finally` block.
|
|
64
|
+
"""
|
|
65
|
+
if renew_interval is not None and renew_interval >= self.manager.config.lease_seconds:
|
|
66
|
+
logger.warning(
|
|
67
|
+
"renew_interval (%.1fs) >= lease_seconds (%.1fs) on lease %s",
|
|
68
|
+
renew_interval,
|
|
69
|
+
self.manager.config.lease_seconds,
|
|
70
|
+
self.lease_id,
|
|
71
|
+
)
|
|
72
|
+
|
|
73
|
+
interval = renew_interval or (self.manager.config.lease_seconds / 2.0)
|
|
74
|
+
renew_task: asyncio.Task[None] | None = None
|
|
75
|
+
|
|
76
|
+
async def _auto_renew_worker() -> None:
|
|
77
|
+
while not self._is_released:
|
|
78
|
+
await asyncio.sleep(interval)
|
|
79
|
+
success = await self.renew()
|
|
80
|
+
if not success:
|
|
81
|
+
logger.warning("Stream lease %s was lost during auto-renewal", self.lease_id)
|
|
82
|
+
break
|
|
83
|
+
|
|
84
|
+
if auto_renew:
|
|
85
|
+
renew_task = asyncio.create_task(_auto_renew_worker())
|
|
86
|
+
|
|
87
|
+
try:
|
|
88
|
+
async for chunk in stream:
|
|
89
|
+
yield chunk
|
|
90
|
+
finally:
|
|
91
|
+
if renew_task is not None:
|
|
92
|
+
renew_task.cancel()
|
|
93
|
+
with suppress(asyncio.CancelledError):
|
|
94
|
+
await renew_task
|
|
95
|
+
await self.release()
|
|
@@ -0,0 +1,95 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
# Atomic Lua script to acquire a stream lease:
|
|
4
|
+
# Checks and clears expired leases, validates user and global limits, and adds the lease.
|
|
5
|
+
# Returns:
|
|
6
|
+
# 1 -> Acquired successfully
|
|
7
|
+
# 2 -> User stream limit reached
|
|
8
|
+
# 3 -> Global stream limit reached
|
|
9
|
+
ACQUIRE_SCRIPT = """
|
|
10
|
+
local user_key = KEYS[1]
|
|
11
|
+
local global_key = KEYS[2]
|
|
12
|
+
local now = tonumber(ARGV[1])
|
|
13
|
+
local expires = tonumber(ARGV[2])
|
|
14
|
+
local lease_id = ARGV[3]
|
|
15
|
+
local user_limit = tonumber(ARGV[4])
|
|
16
|
+
local global_limit = tonumber(ARGV[5])
|
|
17
|
+
local ttl = tonumber(ARGV[6])
|
|
18
|
+
|
|
19
|
+
redis.call('ZREMRANGEBYSCORE', user_key, '-inf', now)
|
|
20
|
+
if global_limit > 0 then
|
|
21
|
+
redis.call('ZREMRANGEBYSCORE', global_key, '-inf', now)
|
|
22
|
+
end
|
|
23
|
+
|
|
24
|
+
if user_limit > 0 and redis.call('ZCARD', user_key) >= user_limit then
|
|
25
|
+
return 2
|
|
26
|
+
end
|
|
27
|
+
|
|
28
|
+
if global_limit > 0 and redis.call('ZCARD', global_key) >= global_limit then
|
|
29
|
+
return 3
|
|
30
|
+
end
|
|
31
|
+
|
|
32
|
+
redis.call('ZADD', user_key, expires, lease_id)
|
|
33
|
+
redis.call('EXPIRE', user_key, ttl)
|
|
34
|
+
|
|
35
|
+
if global_limit > 0 then
|
|
36
|
+
redis.call('ZADD', global_key, expires, lease_id)
|
|
37
|
+
redis.call('EXPIRE', global_key, ttl)
|
|
38
|
+
end
|
|
39
|
+
|
|
40
|
+
return 1
|
|
41
|
+
"""
|
|
42
|
+
|
|
43
|
+
# Atomic Lua script to renew an active stream lease:
|
|
44
|
+
# Verifies the lease is still active in sets and extends its expiration.
|
|
45
|
+
# Returns:
|
|
46
|
+
# 1 -> Renewed successfully
|
|
47
|
+
# 0 -> Lease not found or already expired
|
|
48
|
+
RENEW_SCRIPT = """
|
|
49
|
+
local user_key = KEYS[1]
|
|
50
|
+
local global_key = KEYS[2]
|
|
51
|
+
local lease_id = ARGV[1]
|
|
52
|
+
local new_expires = tonumber(ARGV[2])
|
|
53
|
+
local ttl = tonumber(ARGV[3])
|
|
54
|
+
local check_global = tonumber(ARGV[4]) or 1
|
|
55
|
+
|
|
56
|
+
if not redis.call('ZSCORE', user_key, lease_id) then
|
|
57
|
+
return 0
|
|
58
|
+
end
|
|
59
|
+
|
|
60
|
+
if check_global > 0 and not redis.call('ZSCORE', global_key, lease_id) then
|
|
61
|
+
return 0
|
|
62
|
+
end
|
|
63
|
+
|
|
64
|
+
redis.call('ZADD', user_key, new_expires, lease_id)
|
|
65
|
+
redis.call('EXPIRE', user_key, ttl)
|
|
66
|
+
|
|
67
|
+
if check_global > 0 then
|
|
68
|
+
redis.call('ZADD', global_key, new_expires, lease_id)
|
|
69
|
+
redis.call('EXPIRE', global_key, ttl)
|
|
70
|
+
end
|
|
71
|
+
|
|
72
|
+
return 1
|
|
73
|
+
"""
|
|
74
|
+
|
|
75
|
+
# Atomic Lua script to release a stream lease immediately:
|
|
76
|
+
# Removes the lease from both user and global sets.
|
|
77
|
+
RELEASE_SCRIPT = """
|
|
78
|
+
local user_key = KEYS[1]
|
|
79
|
+
local global_key = KEYS[2]
|
|
80
|
+
local lease_id = ARGV[1]
|
|
81
|
+
|
|
82
|
+
redis.call('ZREM', user_key, lease_id)
|
|
83
|
+
redis.call('ZREM', global_key, lease_id)
|
|
84
|
+
|
|
85
|
+
return 1
|
|
86
|
+
"""
|
|
87
|
+
|
|
88
|
+
# Atomic script to clean expired elements and return active stream count:
|
|
89
|
+
COUNT_SCRIPT = """
|
|
90
|
+
local target_key = KEYS[1]
|
|
91
|
+
local now = tonumber(ARGV[1])
|
|
92
|
+
|
|
93
|
+
redis.call('ZREMRANGEBYSCORE', target_key, '-inf', now)
|
|
94
|
+
return redis.call('ZCARD', target_key)
|
|
95
|
+
"""
|
|
@@ -0,0 +1,147 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import logging
|
|
4
|
+
import time
|
|
5
|
+
from collections.abc import AsyncIterator
|
|
6
|
+
from contextlib import asynccontextmanager
|
|
7
|
+
from typing import Any
|
|
8
|
+
from uuid import uuid4
|
|
9
|
+
|
|
10
|
+
from fastapi_stream_lease.config import LeaseConfig
|
|
11
|
+
from fastapi_stream_lease.exceptions import StreamLeaseRejected
|
|
12
|
+
from fastapi_stream_lease.lease import StreamLease
|
|
13
|
+
from fastapi_stream_lease.lua import (
|
|
14
|
+
ACQUIRE_SCRIPT,
|
|
15
|
+
COUNT_SCRIPT,
|
|
16
|
+
RELEASE_SCRIPT,
|
|
17
|
+
RENEW_SCRIPT,
|
|
18
|
+
)
|
|
19
|
+
|
|
20
|
+
logger = logging.getLogger(__name__)
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class StreamLeaseManager:
|
|
24
|
+
"""
|
|
25
|
+
Coordinates distributed stream concurrency leases backed by atomic Redis Lua scripts.
|
|
26
|
+
"""
|
|
27
|
+
|
|
28
|
+
def __init__(self, redis: Any, config: LeaseConfig | None = None) -> None:
|
|
29
|
+
self.redis = redis
|
|
30
|
+
self.config: LeaseConfig = config or LeaseConfig()
|
|
31
|
+
|
|
32
|
+
async def acquire(self, user_id: str | int) -> StreamLease:
|
|
33
|
+
"""
|
|
34
|
+
Acquire a new stream lease for the given user.
|
|
35
|
+
|
|
36
|
+
Raises:
|
|
37
|
+
StreamLeaseRejected: If user or global concurrency limits are exceeded.
|
|
38
|
+
RuntimeError: If Redis evaluation fails unexpectedly.
|
|
39
|
+
"""
|
|
40
|
+
lease_id = uuid4().hex
|
|
41
|
+
user_key = self.config.user_key(user_id)
|
|
42
|
+
global_key = self.config.global_key
|
|
43
|
+
now = time.time()
|
|
44
|
+
expires = now + self.config.lease_seconds
|
|
45
|
+
|
|
46
|
+
result = await self.redis.eval(
|
|
47
|
+
ACQUIRE_SCRIPT,
|
|
48
|
+
2,
|
|
49
|
+
user_key,
|
|
50
|
+
global_key,
|
|
51
|
+
now,
|
|
52
|
+
expires,
|
|
53
|
+
lease_id,
|
|
54
|
+
self.config.max_per_user,
|
|
55
|
+
self.config.max_global,
|
|
56
|
+
self.config.redis_ttl,
|
|
57
|
+
)
|
|
58
|
+
|
|
59
|
+
code = int(result)
|
|
60
|
+
if code == 2:
|
|
61
|
+
raise StreamLeaseRejected(reason="user_limit")
|
|
62
|
+
if code == 3:
|
|
63
|
+
raise StreamLeaseRejected(reason="global_limit")
|
|
64
|
+
if code != 1:
|
|
65
|
+
raise RuntimeError(f"Unexpected stream lease acquisition return code: {code}")
|
|
66
|
+
|
|
67
|
+
return StreamLease(
|
|
68
|
+
lease_id=lease_id,
|
|
69
|
+
user_id=user_id,
|
|
70
|
+
user_key=user_key,
|
|
71
|
+
global_key=global_key,
|
|
72
|
+
manager=self,
|
|
73
|
+
created_at=now,
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
async def renew(self, lease: StreamLease) -> bool:
|
|
77
|
+
"""
|
|
78
|
+
Renew an active lease, extending its TTL in Redis.
|
|
79
|
+
|
|
80
|
+
Returns:
|
|
81
|
+
bool: True if successfully extended, False if the lease expired, was evicted,
|
|
82
|
+
or if a Redis error occurred.
|
|
83
|
+
"""
|
|
84
|
+
new_expires = time.time() + self.config.lease_seconds
|
|
85
|
+
check_global = 1 if self.config.max_global > 0 else 0
|
|
86
|
+
try:
|
|
87
|
+
result = await self.redis.eval(
|
|
88
|
+
RENEW_SCRIPT,
|
|
89
|
+
2,
|
|
90
|
+
lease.user_key,
|
|
91
|
+
lease.global_key,
|
|
92
|
+
lease.lease_id,
|
|
93
|
+
new_expires,
|
|
94
|
+
self.config.redis_ttl,
|
|
95
|
+
check_global,
|
|
96
|
+
)
|
|
97
|
+
return int(result) == 1
|
|
98
|
+
except Exception as exc:
|
|
99
|
+
logger.warning("Failed to renew stream lease %s: %s", lease.lease_id, exc)
|
|
100
|
+
return False
|
|
101
|
+
|
|
102
|
+
async def release(self, lease: StreamLease) -> None:
|
|
103
|
+
"""
|
|
104
|
+
Release an active lease immediately from Redis.
|
|
105
|
+
"""
|
|
106
|
+
try:
|
|
107
|
+
await self.redis.eval(
|
|
108
|
+
RELEASE_SCRIPT,
|
|
109
|
+
2,
|
|
110
|
+
lease.user_key,
|
|
111
|
+
lease.global_key,
|
|
112
|
+
lease.lease_id,
|
|
113
|
+
)
|
|
114
|
+
except Exception as exc:
|
|
115
|
+
logger.warning("Failed to release stream lease %s: %s", lease.lease_id, exc)
|
|
116
|
+
|
|
117
|
+
async def get_active_count(self, user_id: str | int | None = None) -> int:
|
|
118
|
+
"""
|
|
119
|
+
Return the current number of active (non-expired) streams for a user or globally.
|
|
120
|
+
"""
|
|
121
|
+
target_key = (
|
|
122
|
+
self.config.user_key(user_id) if user_id is not None else self.config.global_key
|
|
123
|
+
)
|
|
124
|
+
now = time.time()
|
|
125
|
+
count = await self.redis.eval(COUNT_SCRIPT, 1, target_key, now)
|
|
126
|
+
return int(count)
|
|
127
|
+
|
|
128
|
+
@asynccontextmanager
|
|
129
|
+
async def lease(self, user_id: str | int) -> AsyncIterator[StreamLease]:
|
|
130
|
+
"""
|
|
131
|
+
Context manager for acquiring and safely releasing a stream lease for scoped
|
|
132
|
+
executions (such as WebSockets, background tasks, or pub/sub loops).
|
|
133
|
+
|
|
134
|
+
For HTTP StreamingResponse (SSE / LLM tokens), use `lease = await acquire()`
|
|
135
|
+
and `return StreamingResponse(lease.wrap(...))` instead.
|
|
136
|
+
|
|
137
|
+
Example:
|
|
138
|
+
async with lease_manager.lease(user_id=42) as lease:
|
|
139
|
+
while True:
|
|
140
|
+
msg = await websocket.receive_text()
|
|
141
|
+
...
|
|
142
|
+
"""
|
|
143
|
+
stream_lease = await self.acquire(user_id)
|
|
144
|
+
try:
|
|
145
|
+
yield stream_lease
|
|
146
|
+
finally:
|
|
147
|
+
await stream_lease.release()
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
# Marker file for PEP 561.
|
|
@@ -0,0 +1,213 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: fastapi-stream-lease
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Distributed stream and SSE concurrency lease manager for FastAPI and Starlette backed by atomic Redis Lua scripts.
|
|
5
|
+
Project-URL: Homepage, https://github.com/agustin18/fastapi-stream-lease
|
|
6
|
+
Project-URL: Repository, https://github.com/agustin18/fastapi-stream-lease
|
|
7
|
+
Project-URL: Issues, https://github.com/agustin18/fastapi-stream-lease/issues
|
|
8
|
+
Author-email: Agustin Saiz <agustinsaiz02@gmail.com>
|
|
9
|
+
License: MIT
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Keywords: concurrency,fastapi,llm-streaming,rate-limiting,redis,server-sent-events,sse,starlette,streaming
|
|
12
|
+
Classifier: Development Status :: 4 - Beta
|
|
13
|
+
Classifier: Framework :: FastAPI
|
|
14
|
+
Classifier: Intended Audience :: Developers
|
|
15
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
16
|
+
Classifier: Programming Language :: Python :: 3
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
21
|
+
Classifier: Topic :: Internet :: WWW/HTTP :: HTTP Servers
|
|
22
|
+
Classifier: Typing :: Typed
|
|
23
|
+
Requires-Python: >=3.10
|
|
24
|
+
Requires-Dist: redis>=5.0.0
|
|
25
|
+
Provides-Extra: dev
|
|
26
|
+
Requires-Dist: fakeredis[lua]>=2.20.0; extra == 'dev'
|
|
27
|
+
Requires-Dist: fastapi>=0.100.0; extra == 'dev'
|
|
28
|
+
Requires-Dist: httpx>=0.25.0; extra == 'dev'
|
|
29
|
+
Requires-Dist: mypy>=1.10.0; extra == 'dev'
|
|
30
|
+
Requires-Dist: pytest-asyncio>=0.23.0; extra == 'dev'
|
|
31
|
+
Requires-Dist: pytest-cov>=5.0.0; extra == 'dev'
|
|
32
|
+
Requires-Dist: pytest>=8.0.0; extra == 'dev'
|
|
33
|
+
Requires-Dist: ruff>=0.4.0; extra == 'dev'
|
|
34
|
+
Provides-Extra: fastapi
|
|
35
|
+
Requires-Dist: fastapi>=0.100.0; extra == 'fastapi'
|
|
36
|
+
Description-Content-Type: text/markdown
|
|
37
|
+
|
|
38
|
+
# fastapi-stream-lease
|
|
39
|
+
|
|
40
|
+
[](https://github.com/agustin18/fastapi-stream-lease/actions)
|
|
41
|
+
[](https://pypi.org/project/fastapi-stream-lease/)
|
|
42
|
+
[](https://www.python.org/downloads/)
|
|
43
|
+
[](https://opensource.org/licenses/MIT)
|
|
44
|
+
[](https://github.com/astral-sh/ruff)
|
|
45
|
+
|
|
46
|
+
**Distributed stream and SSE concurrency lease manager for FastAPI and Starlette, backed by atomic Redis Lua scripts.**
|
|
47
|
+
|
|
48
|
+
---
|
|
49
|
+
|
|
50
|
+
## โก The Problem: Why Traditional Rate Limiters Fail for Streams & LLMs
|
|
51
|
+
|
|
52
|
+
Standard rate limiters (such as `fastapi-limiter` or `slowapi`) count **requests per unit of time** (e.g. *5 requests per minute*).
|
|
53
|
+
|
|
54
|
+
While this works for standard REST APIs, it completely breaks down for **long-lived streaming connections** (Server-Sent Events, WebSockets, or streaming LLM tokens from OpenAI / Claude / Ollama):
|
|
55
|
+
|
|
56
|
+
1. **Duration Blindness:** A user can make a single request that stays open for 15 minutes, consuming a server socket the entire time. A rate limiter considers this "1 request" and allows the user to open 50 more tabs.
|
|
57
|
+
2. **Zombie Connection Leaks:** When mobile users switch networks or close tabs abruptly without clean TCP closure, worker connections remain blocked until timeout, causing connection pool exhaustion and denial of service.
|
|
58
|
+
3. **Multi-Worker Desynchronization:** In-memory concurrency limiters (like `asyncio.Semaphore`) fail across multi-process deployments (Gunicorn / Docker containers) because workers cannot share state.
|
|
59
|
+
|
|
60
|
+
```
|
|
61
|
+
Traditional Rate Limiter: fastapi-stream-lease:
|
|
62
|
+
โโโโโโโโโโโโโโโโโโโโโโโโโ โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
|
|
63
|
+
โ Request 1 -> ALLOWED โ โ Stream 1 (Active) -> LEASE ACQUIRED (1/2) โ
|
|
64
|
+
โ Request 2 -> ALLOWED โ โ Stream 2 (Active) -> LEASE ACQUIRED (2/2) โ
|
|
65
|
+
โ (Both streams active โ โ Stream 3 (Attempt) -> REJECTED: HTTP 429 โ
|
|
66
|
+
โ for 10 minutes, โ โ Retry-After: 5 โ
|
|
67
|
+
โ server sockets exhausted!) โ Stream 1 disconnects -> LEASE RELEASED โ
|
|
68
|
+
โโโโโโโโโโโโโโโโโโโโโโโโโ โ Stream 3 re-attempt -> LEASE ACQUIRED (2/2) โ
|
|
69
|
+
โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
`fastapi-stream-lease` solves this with **Sliding Distributed Leases** inside **atomic Redis Lua scripts**:
|
|
73
|
+
- Enforces strict concurrency limits **per user** (`max_per_user`) and **globally** (`max_global`).
|
|
74
|
+
- Leases automatically self-expire if the client or worker dies without clean closure (zero zombies).
|
|
75
|
+
- A background renewal task keeps long-running streams alive even during slow Time-To-First-Token (TTFT) pauses.
|
|
76
|
+
- Released immediately when the stream finishes or client disconnects.
|
|
77
|
+
|
|
78
|
+
---
|
|
79
|
+
|
|
80
|
+
## ๐ Installation
|
|
81
|
+
|
|
82
|
+
```bash
|
|
83
|
+
pip install fastapi-stream-lease
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
Or using `uv`:
|
|
87
|
+
|
|
88
|
+
```bash
|
|
89
|
+
uv add fastapi-stream-lease
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
*(Requires Redis 5.0+ and Python 3.10+)*
|
|
93
|
+
|
|
94
|
+
---
|
|
95
|
+
|
|
96
|
+
## ๐ก Quickstart
|
|
97
|
+
|
|
98
|
+
Protect an SSE or LLM streaming endpoint in just a few lines:
|
|
99
|
+
|
|
100
|
+
```python
|
|
101
|
+
from fastapi import FastAPI, Depends, Request
|
|
102
|
+
from fastapi.responses import StreamingResponse
|
|
103
|
+
import redis.asyncio as redis
|
|
104
|
+
|
|
105
|
+
from fastapi_stream_lease import (
|
|
106
|
+
StreamLeaseManager,
|
|
107
|
+
LeaseConfig,
|
|
108
|
+
StreamLeaseRejected,
|
|
109
|
+
)
|
|
110
|
+
|
|
111
|
+
app = FastAPI()
|
|
112
|
+
redis_client = redis.from_url("redis://localhost:6379")
|
|
113
|
+
|
|
114
|
+
# Configure lease boundaries:
|
|
115
|
+
# Each user can hold at most 2 concurrent streams; cluster max is 500.
|
|
116
|
+
lease_manager = StreamLeaseManager(
|
|
117
|
+
redis=redis_client,
|
|
118
|
+
config=LeaseConfig(
|
|
119
|
+
max_per_user=2,
|
|
120
|
+
max_global=500,
|
|
121
|
+
lease_seconds=30.0,
|
|
122
|
+
),
|
|
123
|
+
)
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
# Convert lease rejections into clean HTTP 429 Too Many Requests responses:
|
|
127
|
+
@app.exception_handler(StreamLeaseRejected)
|
|
128
|
+
async def lease_rejected_handler(request: Request, exc: StreamLeaseRejected):
|
|
129
|
+
return exc.as_response()
|
|
130
|
+
|
|
131
|
+
|
|
132
|
+
@app.get("/api/chat/stream")
|
|
133
|
+
async def chat_stream(user_id: str = "user_123"):
|
|
134
|
+
# 1. Acquire lease (raises StreamLeaseRejected if limit reached)
|
|
135
|
+
lease = await lease_manager.acquire(user_id)
|
|
136
|
+
|
|
137
|
+
async def token_generator():
|
|
138
|
+
# Example: streaming tokens from an LLM
|
|
139
|
+
for word in ["Hello", "world", "this", "is", "streamed!"]:
|
|
140
|
+
yield f"data: {word}\n\n"
|
|
141
|
+
|
|
142
|
+
# 2. Wrap generator: guarantees background auto-renewal and release on disconnect
|
|
143
|
+
return StreamingResponse(
|
|
144
|
+
lease.wrap(token_generator()),
|
|
145
|
+
media_type="text/event-stream",
|
|
146
|
+
)
|
|
147
|
+
```
|
|
148
|
+
|
|
149
|
+
### Protecting WebSockets
|
|
150
|
+
|
|
151
|
+
For WebSockets and scoped async routines, use the `lease_manager.lease(...)` context manager:
|
|
152
|
+
|
|
153
|
+
```python
|
|
154
|
+
@app.websocket("/ws/chat/{user_id}")
|
|
155
|
+
async def websocket_chat(websocket: WebSocket, user_id: str):
|
|
156
|
+
await websocket.accept()
|
|
157
|
+
# Acquires lease on enter, automatically releases when socket closes or disconnects
|
|
158
|
+
async with lease_manager.lease(user_id):
|
|
159
|
+
while True:
|
|
160
|
+
msg = await websocket.receive_text()
|
|
161
|
+
await websocket.send_text(f"Echo: {msg}")
|
|
162
|
+
```
|
|
163
|
+
|
|
164
|
+
---
|
|
165
|
+
|
|
166
|
+
## ๐ ๏ธ How It Works (Algorithmic Math)
|
|
167
|
+
|
|
168
|
+
All concurrency validations, expirations, and insertions run inside **atomic Lua scripts** on Redis:
|
|
169
|
+
|
|
170
|
+
1. **Sorted Sets (`ZSET`):** Active streams are stored in Redis `ZSET`s where the value is a unique `lease_id` and the score is the epoch expiration timestamp (`now + lease_seconds`).
|
|
171
|
+
2. **Atomic Eviction:** Before checking capacity, `ZREMRANGEBYSCORE` purges all expired entries in $O(\log N + M)$.
|
|
172
|
+
3. **Capacity Check:** `ZCARD` verifies current stream count in $O(1)$ against `max_per_user` and `max_global`. Setting either to `0` disables that limit.
|
|
173
|
+
4. **Redis Cluster Slot Affinity:** Keys automatically use `{prefix}` hash tags (e.g. `{stream_lease}:user:123` and `{stream_lease}:global`), guaranteeing zero `CROSSSLOT` errors across distributed Redis clusters.
|
|
174
|
+
5. **Acquisition:** If capacity permits, `ZADD` registers the lease in $O(\log N)$ and updates the key TTL.
|
|
175
|
+
6. **Auto-Renewal:** While the stream is active, `lease.wrap()` spawns a lightweight background worker that calls `ZADD` to advance the expiration score every `lease_seconds / 2`.
|
|
176
|
+
7. **Guaranteed Release:** When the stream completes or the client disconnects, `ZREM` removes the lease immediately in the `finally:` block.
|
|
177
|
+
|
|
178
|
+
---
|
|
179
|
+
|
|
180
|
+
## โ๏ธ Configuration Options
|
|
181
|
+
|
|
182
|
+
Customize `LeaseConfig`:
|
|
183
|
+
|
|
184
|
+
```python
|
|
185
|
+
from fastapi_stream_lease import LeaseConfig
|
|
186
|
+
|
|
187
|
+
config = LeaseConfig(
|
|
188
|
+
lease_seconds=30.0, # Lease expiration window (seconds)
|
|
189
|
+
max_per_user=3, # Max active streams per user (set 0 to disable)
|
|
190
|
+
max_global=1000, # Max active streams cluster-wide (set 0 to disable)
|
|
191
|
+
key_prefix="my_app:sse", # Custom Redis key prefix (hash-tag safe)
|
|
192
|
+
)
|
|
193
|
+
```
|
|
194
|
+
|
|
195
|
+
---
|
|
196
|
+
|
|
197
|
+
## ๐งช Testing & Observability
|
|
198
|
+
|
|
199
|
+
You can inspect the live count of active streams at any time:
|
|
200
|
+
|
|
201
|
+
```python
|
|
202
|
+
# Active streams for a specific user:
|
|
203
|
+
active_user_streams = await lease_manager.get_active_count("user_123")
|
|
204
|
+
|
|
205
|
+
# Active streams across the entire cluster:
|
|
206
|
+
active_global_streams = await lease_manager.get_active_count()
|
|
207
|
+
```
|
|
208
|
+
|
|
209
|
+
---
|
|
210
|
+
|
|
211
|
+
## ๐ License
|
|
212
|
+
|
|
213
|
+
This project is licensed under the [MIT License](LICENSE).
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
fastapi_stream_lease/__init__.py,sha256=cKkJzlD-no7fUIaVbcyNmkdvxh1GPSyHuDsje270BPU,453
|
|
2
|
+
fastapi_stream_lease/config.py,sha256=VUAiWiOaTcliZpvLvHUHsKTuZnXRR5lqzZZg15NDiAg,1624
|
|
3
|
+
fastapi_stream_lease/exceptions.py,sha256=uNBDyckMf1irbglLCZh2f6LlPtxCS5UCSYfEz-9iEgo,1989
|
|
4
|
+
fastapi_stream_lease/lease.py,sha256=bObEIpm19mn0Xw_-lqPYK_GbOV0eYNCTPSoE1hzPOxk,3108
|
|
5
|
+
fastapi_stream_lease/lua.py,sha256=kWlIeZ2rC2KC8n24_nih_wzVSJxrIIE3K9nIGbzTPWg,2482
|
|
6
|
+
fastapi_stream_lease/manager.py,sha256=33KVnyXZpIW83wGMfIw0NAkVgpkkoNm7VfYx60Ak4KQ,4777
|
|
7
|
+
fastapi_stream_lease/py.typed,sha256=bWew9mHgMy8LqMu7RuqQXFXLBxh2CRx0dUbSx-3wE48,27
|
|
8
|
+
fastapi_stream_lease-0.1.0.dist-info/METADATA,sha256=dc6ZTTg_eJtu9NYUQo2bP0f3a9drwLzyCvh1Wu3m2Rs,9446
|
|
9
|
+
fastapi_stream_lease-0.1.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
|
|
10
|
+
fastapi_stream_lease-0.1.0.dist-info/licenses/LICENSE,sha256=lIUyM6GLlCSB712vCX1qjIh556nlLYlwChOhrhazGb8,1069
|
|
11
|
+
fastapi_stream_lease-0.1.0.dist-info/RECORD,,
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Agustin Saiz
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|