sheriff-limiter 0.1.0__tar.gz
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.
- sheriff_limiter-0.1.0/.gitignore +141 -0
- sheriff_limiter-0.1.0/PKG-INFO +153 -0
- sheriff_limiter-0.1.0/README.md +130 -0
- sheriff_limiter-0.1.0/pyproject.toml +63 -0
- sheriff_limiter-0.1.0/src/sheriff/__init__.py +10 -0
- sheriff_limiter-0.1.0/src/sheriff/core.py +217 -0
- sheriff_limiter-0.1.0/src/sheriff/exceptions.py +19 -0
- sheriff_limiter-0.1.0/tests/__init__.py +1 -0
- sheriff_limiter-0.1.0/tests/conftest.py +9 -0
- sheriff_limiter-0.1.0/tests/test_limiter.py +193 -0
|
@@ -0,0 +1,141 @@
|
|
|
1
|
+
# Byte-compiled / optimized / DLL files
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*$py.class
|
|
5
|
+
|
|
6
|
+
# C extensions
|
|
7
|
+
*.so
|
|
8
|
+
|
|
9
|
+
# Distribution / packaging
|
|
10
|
+
bin/
|
|
11
|
+
build/
|
|
12
|
+
develop-eggs/
|
|
13
|
+
dist/
|
|
14
|
+
downloads/
|
|
15
|
+
eggs/
|
|
16
|
+
.eggs/
|
|
17
|
+
lib/
|
|
18
|
+
lib64/
|
|
19
|
+
parts/
|
|
20
|
+
sdist/
|
|
21
|
+
var/
|
|
22
|
+
wheels/
|
|
23
|
+
share/python-wheels/
|
|
24
|
+
*.egg-info/
|
|
25
|
+
.installed.cfg
|
|
26
|
+
*.egg
|
|
27
|
+
MANIFEST
|
|
28
|
+
|
|
29
|
+
# PyInstaller
|
|
30
|
+
# Usually these files are written by a python script from a template
|
|
31
|
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
|
32
|
+
*.manifest
|
|
33
|
+
*.spec
|
|
34
|
+
|
|
35
|
+
# Installer logs
|
|
36
|
+
pip-log.txt
|
|
37
|
+
pip-delete-this-directory.txt
|
|
38
|
+
|
|
39
|
+
# Unit test / coverage reports
|
|
40
|
+
htmlcov/
|
|
41
|
+
.tox/
|
|
42
|
+
.nox/
|
|
43
|
+
.coverage
|
|
44
|
+
.coverage.*
|
|
45
|
+
.cache
|
|
46
|
+
nosetests.xml
|
|
47
|
+
coverage.xml
|
|
48
|
+
*.cover
|
|
49
|
+
*.py,cover
|
|
50
|
+
.htmlcov/
|
|
51
|
+
.pytest_cache/
|
|
52
|
+
.ruff_cache/
|
|
53
|
+
.mypy_cache/
|
|
54
|
+
|
|
55
|
+
# Translations
|
|
56
|
+
*.mo
|
|
57
|
+
*.pot
|
|
58
|
+
|
|
59
|
+
# Django stuff:
|
|
60
|
+
*.log
|
|
61
|
+
local_settings.py
|
|
62
|
+
db.sqlite3
|
|
63
|
+
db.sqlite3-journal
|
|
64
|
+
|
|
65
|
+
# Sphinx documentation
|
|
66
|
+
doc/_build/
|
|
67
|
+
|
|
68
|
+
# PyBuilder
|
|
69
|
+
.pybuilder/
|
|
70
|
+
target/
|
|
71
|
+
|
|
72
|
+
# Jupyter Notebook
|
|
73
|
+
.ipynb_checkpoints
|
|
74
|
+
|
|
75
|
+
# IPython
|
|
76
|
+
profile_default/
|
|
77
|
+
ipython_config.py
|
|
78
|
+
|
|
79
|
+
# pyenv
|
|
80
|
+
# For a library or app, you might want to share your .python-version.
|
|
81
|
+
# Comment/uncomment to share.
|
|
82
|
+
#.python-version
|
|
83
|
+
|
|
84
|
+
# pipenv
|
|
85
|
+
# According to pypa/pipenv#1402, Pipfile.lock prevents deterministic
|
|
86
|
+
# builds in applications but is recommended for libraries.
|
|
87
|
+
# So, for a library you should keep it.
|
|
88
|
+
#Pipfile.lock
|
|
89
|
+
|
|
90
|
+
# poetry
|
|
91
|
+
# Similarly, poetry.lock should be committed for applications, but is optional/not recommended for libraries.
|
|
92
|
+
#poetry.lock
|
|
93
|
+
|
|
94
|
+
# pdm
|
|
95
|
+
# Similar to Pipfile.lock, pdm.lock is usage-specific.
|
|
96
|
+
#pdm.lock
|
|
97
|
+
|
|
98
|
+
# PEP 582; project local packages directory (used by pdm)
|
|
99
|
+
__pypackages__/
|
|
100
|
+
|
|
101
|
+
# Celery stuff
|
|
102
|
+
celerybeat-schedule
|
|
103
|
+
celerybeat.pid
|
|
104
|
+
|
|
105
|
+
# SageMath parsed files
|
|
106
|
+
*.sage.py
|
|
107
|
+
|
|
108
|
+
# Environments
|
|
109
|
+
.env
|
|
110
|
+
.venv
|
|
111
|
+
env/
|
|
112
|
+
venv/
|
|
113
|
+
ENV/
|
|
114
|
+
env.bak/
|
|
115
|
+
venv.bak/
|
|
116
|
+
|
|
117
|
+
# Spyder project settings
|
|
118
|
+
.spyderproject
|
|
119
|
+
.spyproject
|
|
120
|
+
|
|
121
|
+
# Rope project settings
|
|
122
|
+
.ropeproject
|
|
123
|
+
|
|
124
|
+
# mkdocs documentation
|
|
125
|
+
/site
|
|
126
|
+
|
|
127
|
+
# mypy
|
|
128
|
+
.mypy_cache/
|
|
129
|
+
.nosetests
|
|
130
|
+
/nosetests.xml
|
|
131
|
+
/.pytest_cache
|
|
132
|
+
.coverage
|
|
133
|
+
|
|
134
|
+
# Cython debug symbols
|
|
135
|
+
cython_debug/
|
|
136
|
+
|
|
137
|
+
# IDE files
|
|
138
|
+
.idea/
|
|
139
|
+
.vscode/
|
|
140
|
+
*.swp
|
|
141
|
+
*.swo
|
|
@@ -0,0 +1,153 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: sheriff-limiter
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: An elegant, thread-safe, in-memory rate limiter for Python
|
|
5
|
+
Author-email: Vahsi Bati <info@vahsibati.com.tr>
|
|
6
|
+
License: MIT
|
|
7
|
+
Classifier: Development Status :: 3 - Alpha
|
|
8
|
+
Classifier: Intended Audience :: Developers
|
|
9
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.8
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Requires-Python: >=3.8
|
|
17
|
+
Provides-Extra: dev
|
|
18
|
+
Requires-Dist: mypy>=1.0; extra == 'dev'
|
|
19
|
+
Requires-Dist: pytest-cov>=4.0; extra == 'dev'
|
|
20
|
+
Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
21
|
+
Requires-Dist: ruff>=0.1.0; extra == 'dev'
|
|
22
|
+
Description-Content-Type: text/markdown
|
|
23
|
+
|
|
24
|
+
# Sheriff 🤠
|
|
25
|
+
|
|
26
|
+
An elegant, thread-safe, in-memory rate limiter for Python.
|
|
27
|
+
|
|
28
|
+
`sheriff` implements the **Token Bucket** algorithm, ensuring complete thread-safety with fine-grained locking and zero-leak memory management. It is designed to be lightweight, dependency-free, and extremely easy to integrate into any application or web framework (like FastAPI).
|
|
29
|
+
|
|
30
|
+
---
|
|
31
|
+
|
|
32
|
+
## Features
|
|
33
|
+
|
|
34
|
+
- 🔒 **Thread-Safe**: Uses fine-grained concurrent locks to ensure rate-limiting consistency across multiple threads.
|
|
35
|
+
- 🪣 **Token Bucket Algorithm**: Standard token bucket rate limiting with lazy, high-precision token replenishment.
|
|
36
|
+
- 🧹 **Self-Cleaning (Lazy Cleanup)**: Prunes stale/fully-replenished buckets from memory automatically to prevent memory leaks.
|
|
37
|
+
- ⚡ **Zero Dependencies**: Pure Python, built using standard library tools.
|
|
38
|
+
- 🚀 **FastAPI / Web Ready**: Fits perfectly into FastAPI's dependency injection (`Depends`) system.
|
|
39
|
+
|
|
40
|
+
---
|
|
41
|
+
|
|
42
|
+
## Installation
|
|
43
|
+
|
|
44
|
+
Install using `pip`:
|
|
45
|
+
|
|
46
|
+
```bash
|
|
47
|
+
pip install sheriff-limiter
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
---
|
|
51
|
+
|
|
52
|
+
## Quick Start
|
|
53
|
+
|
|
54
|
+
### Basic Usage
|
|
55
|
+
|
|
56
|
+
Use `is_allowed` for a simple boolean check:
|
|
57
|
+
|
|
58
|
+
```python
|
|
59
|
+
from sheriff import RateLimiter
|
|
60
|
+
|
|
61
|
+
# Default: 10 requests capacity, replenishes 1 token per second
|
|
62
|
+
limiter = RateLimiter()
|
|
63
|
+
|
|
64
|
+
# Check if allowed
|
|
65
|
+
if limiter.is_allowed("user_ip_address"):
|
|
66
|
+
print("Request allowed!")
|
|
67
|
+
else:
|
|
68
|
+
print("Rate limit exceeded.")
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
### Configuration Options
|
|
72
|
+
|
|
73
|
+
Initialize the limiter with custom parameters:
|
|
74
|
+
|
|
75
|
+
```python
|
|
76
|
+
from sheriff import RateLimiter
|
|
77
|
+
|
|
78
|
+
# Configured for max 100 requests per minute
|
|
79
|
+
limiter = RateLimiter(max_requests=100, period=60.0)
|
|
80
|
+
|
|
81
|
+
# Or set capacity and refill rate directly
|
|
82
|
+
# Capacity of 5 tokens, refilling 0.5 tokens/sec
|
|
83
|
+
limiter = RateLimiter(capacity=5.0, refill_rate=0.5)
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
---
|
|
87
|
+
|
|
88
|
+
## Advanced Features
|
|
89
|
+
|
|
90
|
+
### 1. Raising Exceptions on Exceeding Limits
|
|
91
|
+
|
|
92
|
+
You can use `.check()` which raises a `RateLimitExceeded` exception. The exception contains a `retry_after` parameter telling you how long to wait in seconds.
|
|
93
|
+
|
|
94
|
+
```python
|
|
95
|
+
from sheriff import RateLimiter, RateLimitExceeded
|
|
96
|
+
|
|
97
|
+
limiter = RateLimiter(max_requests=5, period=10.0)
|
|
98
|
+
|
|
99
|
+
try:
|
|
100
|
+
# Consume 1 token
|
|
101
|
+
limiter.check("client_1")
|
|
102
|
+
except RateLimitExceeded as e:
|
|
103
|
+
print(f"Rate limit exceeded! Retry after {e.retry_after:.2f} seconds.")
|
|
104
|
+
```
|
|
105
|
+
|
|
106
|
+
### 2. Manual Resets
|
|
107
|
+
|
|
108
|
+
Clear specific keys or reset all rate limits entirely:
|
|
109
|
+
|
|
110
|
+
```python
|
|
111
|
+
# Reset a single client
|
|
112
|
+
limiter.reset("client_1")
|
|
113
|
+
|
|
114
|
+
# Reset all clients and clear the memory cache
|
|
115
|
+
limiter.reset_all()
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
---
|
|
119
|
+
|
|
120
|
+
## FastAPI Integration
|
|
121
|
+
|
|
122
|
+
`sheriff` is perfect for FastAPI dependencies. Here is how you can use it to rate-limit endpoints by IP address:
|
|
123
|
+
|
|
124
|
+
```python
|
|
125
|
+
from fastapi import FastAPI, Depends, Request, HTTPException, status
|
|
126
|
+
from sheriff import RateLimiter, RateLimitExceeded
|
|
127
|
+
|
|
128
|
+
app = FastAPI()
|
|
129
|
+
|
|
130
|
+
# 100 requests per minute limit
|
|
131
|
+
limiter = RateLimiter(max_requests=100, period=60.0)
|
|
132
|
+
|
|
133
|
+
def rate_limit(request: Request):
|
|
134
|
+
client_ip = request.client.host if request.client else "unknown"
|
|
135
|
+
try:
|
|
136
|
+
limiter.check(client_ip)
|
|
137
|
+
except RateLimitExceeded as e:
|
|
138
|
+
raise HTTPException(
|
|
139
|
+
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
140
|
+
detail="Too many requests. Please slow down.",
|
|
141
|
+
headers={"Retry-After": str(int(e.retry_after or 0))}
|
|
142
|
+
)
|
|
143
|
+
|
|
144
|
+
@app.get("/items", dependencies=[Depends(rate_limit)])
|
|
145
|
+
async def read_items():
|
|
146
|
+
return {"status": "ok"}
|
|
147
|
+
```
|
|
148
|
+
|
|
149
|
+
---
|
|
150
|
+
|
|
151
|
+
## License
|
|
152
|
+
|
|
153
|
+
MIT License. See [LICENSE](LICENSE) for details.
|
|
@@ -0,0 +1,130 @@
|
|
|
1
|
+
# Sheriff 🤠
|
|
2
|
+
|
|
3
|
+
An elegant, thread-safe, in-memory rate limiter for Python.
|
|
4
|
+
|
|
5
|
+
`sheriff` implements the **Token Bucket** algorithm, ensuring complete thread-safety with fine-grained locking and zero-leak memory management. It is designed to be lightweight, dependency-free, and extremely easy to integrate into any application or web framework (like FastAPI).
|
|
6
|
+
|
|
7
|
+
---
|
|
8
|
+
|
|
9
|
+
## Features
|
|
10
|
+
|
|
11
|
+
- 🔒 **Thread-Safe**: Uses fine-grained concurrent locks to ensure rate-limiting consistency across multiple threads.
|
|
12
|
+
- 🪣 **Token Bucket Algorithm**: Standard token bucket rate limiting with lazy, high-precision token replenishment.
|
|
13
|
+
- 🧹 **Self-Cleaning (Lazy Cleanup)**: Prunes stale/fully-replenished buckets from memory automatically to prevent memory leaks.
|
|
14
|
+
- ⚡ **Zero Dependencies**: Pure Python, built using standard library tools.
|
|
15
|
+
- 🚀 **FastAPI / Web Ready**: Fits perfectly into FastAPI's dependency injection (`Depends`) system.
|
|
16
|
+
|
|
17
|
+
---
|
|
18
|
+
|
|
19
|
+
## Installation
|
|
20
|
+
|
|
21
|
+
Install using `pip`:
|
|
22
|
+
|
|
23
|
+
```bash
|
|
24
|
+
pip install sheriff-limiter
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
---
|
|
28
|
+
|
|
29
|
+
## Quick Start
|
|
30
|
+
|
|
31
|
+
### Basic Usage
|
|
32
|
+
|
|
33
|
+
Use `is_allowed` for a simple boolean check:
|
|
34
|
+
|
|
35
|
+
```python
|
|
36
|
+
from sheriff import RateLimiter
|
|
37
|
+
|
|
38
|
+
# Default: 10 requests capacity, replenishes 1 token per second
|
|
39
|
+
limiter = RateLimiter()
|
|
40
|
+
|
|
41
|
+
# Check if allowed
|
|
42
|
+
if limiter.is_allowed("user_ip_address"):
|
|
43
|
+
print("Request allowed!")
|
|
44
|
+
else:
|
|
45
|
+
print("Rate limit exceeded.")
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
### Configuration Options
|
|
49
|
+
|
|
50
|
+
Initialize the limiter with custom parameters:
|
|
51
|
+
|
|
52
|
+
```python
|
|
53
|
+
from sheriff import RateLimiter
|
|
54
|
+
|
|
55
|
+
# Configured for max 100 requests per minute
|
|
56
|
+
limiter = RateLimiter(max_requests=100, period=60.0)
|
|
57
|
+
|
|
58
|
+
# Or set capacity and refill rate directly
|
|
59
|
+
# Capacity of 5 tokens, refilling 0.5 tokens/sec
|
|
60
|
+
limiter = RateLimiter(capacity=5.0, refill_rate=0.5)
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
---
|
|
64
|
+
|
|
65
|
+
## Advanced Features
|
|
66
|
+
|
|
67
|
+
### 1. Raising Exceptions on Exceeding Limits
|
|
68
|
+
|
|
69
|
+
You can use `.check()` which raises a `RateLimitExceeded` exception. The exception contains a `retry_after` parameter telling you how long to wait in seconds.
|
|
70
|
+
|
|
71
|
+
```python
|
|
72
|
+
from sheriff import RateLimiter, RateLimitExceeded
|
|
73
|
+
|
|
74
|
+
limiter = RateLimiter(max_requests=5, period=10.0)
|
|
75
|
+
|
|
76
|
+
try:
|
|
77
|
+
# Consume 1 token
|
|
78
|
+
limiter.check("client_1")
|
|
79
|
+
except RateLimitExceeded as e:
|
|
80
|
+
print(f"Rate limit exceeded! Retry after {e.retry_after:.2f} seconds.")
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
### 2. Manual Resets
|
|
84
|
+
|
|
85
|
+
Clear specific keys or reset all rate limits entirely:
|
|
86
|
+
|
|
87
|
+
```python
|
|
88
|
+
# Reset a single client
|
|
89
|
+
limiter.reset("client_1")
|
|
90
|
+
|
|
91
|
+
# Reset all clients and clear the memory cache
|
|
92
|
+
limiter.reset_all()
|
|
93
|
+
```
|
|
94
|
+
|
|
95
|
+
---
|
|
96
|
+
|
|
97
|
+
## FastAPI Integration
|
|
98
|
+
|
|
99
|
+
`sheriff` is perfect for FastAPI dependencies. Here is how you can use it to rate-limit endpoints by IP address:
|
|
100
|
+
|
|
101
|
+
```python
|
|
102
|
+
from fastapi import FastAPI, Depends, Request, HTTPException, status
|
|
103
|
+
from sheriff import RateLimiter, RateLimitExceeded
|
|
104
|
+
|
|
105
|
+
app = FastAPI()
|
|
106
|
+
|
|
107
|
+
# 100 requests per minute limit
|
|
108
|
+
limiter = RateLimiter(max_requests=100, period=60.0)
|
|
109
|
+
|
|
110
|
+
def rate_limit(request: Request):
|
|
111
|
+
client_ip = request.client.host if request.client else "unknown"
|
|
112
|
+
try:
|
|
113
|
+
limiter.check(client_ip)
|
|
114
|
+
except RateLimitExceeded as e:
|
|
115
|
+
raise HTTPException(
|
|
116
|
+
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
117
|
+
detail="Too many requests. Please slow down.",
|
|
118
|
+
headers={"Retry-After": str(int(e.retry_after or 0))}
|
|
119
|
+
)
|
|
120
|
+
|
|
121
|
+
@app.get("/items", dependencies=[Depends(rate_limit)])
|
|
122
|
+
async def read_items():
|
|
123
|
+
return {"status": "ok"}
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
---
|
|
127
|
+
|
|
128
|
+
## License
|
|
129
|
+
|
|
130
|
+
MIT License. See [LICENSE](LICENSE) for details.
|
|
@@ -0,0 +1,63 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "sheriff-limiter"
|
|
7
|
+
dynamic = ["version"]
|
|
8
|
+
description = "An elegant, thread-safe, in-memory rate limiter for Python"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.8"
|
|
11
|
+
license = {text = "MIT"}
|
|
12
|
+
authors = [
|
|
13
|
+
{name = "Vahsi Bati", email = "info@vahsibati.com.tr"}
|
|
14
|
+
]
|
|
15
|
+
classifiers = [
|
|
16
|
+
"Development Status :: 3 - Alpha",
|
|
17
|
+
"Intended Audience :: Developers",
|
|
18
|
+
"License :: OSI Approved :: MIT License",
|
|
19
|
+
"Programming Language :: Python :: 3",
|
|
20
|
+
"Programming Language :: Python :: 3.8",
|
|
21
|
+
"Programming Language :: Python :: 3.9",
|
|
22
|
+
"Programming Language :: Python :: 3.10",
|
|
23
|
+
"Programming Language :: Python :: 3.11",
|
|
24
|
+
"Programming Language :: Python :: 3.12",
|
|
25
|
+
]
|
|
26
|
+
dependencies = []
|
|
27
|
+
|
|
28
|
+
[project.optional-dependencies]
|
|
29
|
+
dev = [
|
|
30
|
+
"pytest>=7.0",
|
|
31
|
+
"pytest-cov>=4.0",
|
|
32
|
+
"ruff>=0.1.0",
|
|
33
|
+
"mypy>=1.0",
|
|
34
|
+
]
|
|
35
|
+
|
|
36
|
+
[tool.hatch.version]
|
|
37
|
+
path = "src/sheriff/__init__.py"
|
|
38
|
+
|
|
39
|
+
[tool.hatch.build.targets.wheel]
|
|
40
|
+
packages = ["src/sheriff"]
|
|
41
|
+
|
|
42
|
+
[tool.ruff]
|
|
43
|
+
line-length = 88
|
|
44
|
+
target-version = "py38"
|
|
45
|
+
|
|
46
|
+
[tool.ruff.lint]
|
|
47
|
+
select = [
|
|
48
|
+
"E", # pycodestyle errors
|
|
49
|
+
"W", # pycodestyle warnings
|
|
50
|
+
"F", # pyflakes
|
|
51
|
+
"I", # isort
|
|
52
|
+
"B", # flake8-bugbear
|
|
53
|
+
"C4", # flake8-comprehensions
|
|
54
|
+
"UP", # pyupgrade
|
|
55
|
+
]
|
|
56
|
+
ignore = []
|
|
57
|
+
|
|
58
|
+
[tool.pytest.ini_options]
|
|
59
|
+
minversion = "7.0"
|
|
60
|
+
addopts = "-ra -q --tb=short"
|
|
61
|
+
testpaths = [
|
|
62
|
+
"tests",
|
|
63
|
+
]
|
|
@@ -0,0 +1,217 @@
|
|
|
1
|
+
import time
|
|
2
|
+
from threading import Lock
|
|
3
|
+
from typing import Dict, Optional, Tuple
|
|
4
|
+
|
|
5
|
+
from sheriff.exceptions import RateLimitExceeded
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class TokenBucket:
|
|
9
|
+
"""Represents a single Token Bucket for rate limiting a specific key."""
|
|
10
|
+
|
|
11
|
+
def __init__(self, capacity: float, refill_rate: float):
|
|
12
|
+
self.capacity = capacity
|
|
13
|
+
self.refill_rate = refill_rate
|
|
14
|
+
self.tokens = capacity
|
|
15
|
+
self.last_updated = time.monotonic()
|
|
16
|
+
self.lock = Lock()
|
|
17
|
+
|
|
18
|
+
def consume(self, tokens: float = 1.0) -> Tuple[bool, float]:
|
|
19
|
+
"""Consume tokens from the bucket.
|
|
20
|
+
|
|
21
|
+
Returns:
|
|
22
|
+
Tuple[bool, float]: (allowed, retry_after)
|
|
23
|
+
where allowed is True if consumed, False otherwise.
|
|
24
|
+
retry_after is the number of seconds to wait before there's enough tokens.
|
|
25
|
+
"""
|
|
26
|
+
with self.lock:
|
|
27
|
+
now = time.monotonic()
|
|
28
|
+
elapsed = now - self.last_updated
|
|
29
|
+
if elapsed > 0:
|
|
30
|
+
self.tokens = min(
|
|
31
|
+
self.capacity, self.tokens + elapsed * self.refill_rate
|
|
32
|
+
)
|
|
33
|
+
self.last_updated = now
|
|
34
|
+
|
|
35
|
+
if self.tokens >= tokens:
|
|
36
|
+
self.tokens -= tokens
|
|
37
|
+
return True, 0.0
|
|
38
|
+
|
|
39
|
+
needed = tokens - self.tokens
|
|
40
|
+
retry_after = needed / self.refill_rate
|
|
41
|
+
return False, retry_after
|
|
42
|
+
|
|
43
|
+
def get_tokens(self) -> float:
|
|
44
|
+
"""Returns the current number of tokens in the bucket after replenishment."""
|
|
45
|
+
with self.lock:
|
|
46
|
+
now = time.monotonic()
|
|
47
|
+
elapsed = now - self.last_updated
|
|
48
|
+
if elapsed > 0:
|
|
49
|
+
self.tokens = min(
|
|
50
|
+
self.capacity, self.tokens + elapsed * self.refill_rate
|
|
51
|
+
)
|
|
52
|
+
self.last_updated = now
|
|
53
|
+
return self.tokens
|
|
54
|
+
|
|
55
|
+
def reset(self) -> None:
|
|
56
|
+
"""Resets the bucket to its full capacity."""
|
|
57
|
+
with self.lock:
|
|
58
|
+
self.tokens = self.capacity
|
|
59
|
+
self.last_updated = time.monotonic()
|
|
60
|
+
|
|
61
|
+
def is_full(self, now: float) -> bool:
|
|
62
|
+
"""Check if the bucket is fully replenished.
|
|
63
|
+
Must be called under RateLimiter container lock or self.lock.
|
|
64
|
+
"""
|
|
65
|
+
with self.lock:
|
|
66
|
+
elapsed = now - self.last_updated
|
|
67
|
+
current_tokens = min(
|
|
68
|
+
self.capacity, self.tokens + elapsed * self.refill_rate
|
|
69
|
+
)
|
|
70
|
+
return current_tokens >= self.capacity
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
class RateLimiter:
|
|
74
|
+
"""Core class representing the Sheriff thread-safe, in-memory rate limiter."""
|
|
75
|
+
|
|
76
|
+
def __init__(
|
|
77
|
+
self,
|
|
78
|
+
capacity: float = 10.0,
|
|
79
|
+
refill_rate: float = 1.0,
|
|
80
|
+
max_requests: Optional[int] = None,
|
|
81
|
+
period: Optional[float] = None,
|
|
82
|
+
cleanup_interval: float = 60.0,
|
|
83
|
+
):
|
|
84
|
+
"""Initializes the RateLimiter.
|
|
85
|
+
|
|
86
|
+
Args:
|
|
87
|
+
capacity: Maximum number of tokens a bucket can hold. Defaults to 10.0.
|
|
88
|
+
refill_rate: Number of tokens added to the bucket per second.
|
|
89
|
+
Defaults to 1.0.
|
|
90
|
+
max_requests: Optional parameter to initialize capacity using requests.
|
|
91
|
+
period: Optional parameter to specify the period in seconds for
|
|
92
|
+
max_requests.
|
|
93
|
+
cleanup_interval: Time in seconds between periodic cleanup sweeps of
|
|
94
|
+
fully replenished buckets. Defaults to 60.0.
|
|
95
|
+
"""
|
|
96
|
+
if max_requests is not None:
|
|
97
|
+
capacity = float(max_requests)
|
|
98
|
+
if period is not None:
|
|
99
|
+
refill_rate = capacity / period
|
|
100
|
+
else:
|
|
101
|
+
refill_rate = capacity
|
|
102
|
+
|
|
103
|
+
if capacity <= 0:
|
|
104
|
+
raise ValueError("Capacity must be greater than zero.")
|
|
105
|
+
if refill_rate <= 0:
|
|
106
|
+
raise ValueError("Refill rate must be greater than zero.")
|
|
107
|
+
if cleanup_interval <= 0:
|
|
108
|
+
raise ValueError("Cleanup interval must be greater than zero.")
|
|
109
|
+
|
|
110
|
+
self.capacity = capacity
|
|
111
|
+
self.refill_rate = refill_rate
|
|
112
|
+
self.cleanup_interval = cleanup_interval
|
|
113
|
+
|
|
114
|
+
self.buckets: Dict[str, TokenBucket] = {}
|
|
115
|
+
self.lock = Lock()
|
|
116
|
+
self.last_cleanup = time.monotonic()
|
|
117
|
+
|
|
118
|
+
def _get_bucket(self, key: str) -> TokenBucket:
|
|
119
|
+
"""Thread-safe retrieval or creation of a TokenBucket for a given key.
|
|
120
|
+
Also triggers lazy cleanup if the cleanup interval has elapsed.
|
|
121
|
+
"""
|
|
122
|
+
with self.lock:
|
|
123
|
+
self._maybe_cleanup()
|
|
124
|
+
if key not in self.buckets:
|
|
125
|
+
self.buckets[key] = TokenBucket(self.capacity, self.refill_rate)
|
|
126
|
+
return self.buckets[key]
|
|
127
|
+
|
|
128
|
+
def _maybe_cleanup(self) -> None:
|
|
129
|
+
"""Prunes fully replenished buckets from memory.
|
|
130
|
+
Must be called with self.lock held.
|
|
131
|
+
"""
|
|
132
|
+
now = time.monotonic()
|
|
133
|
+
if now - self.last_cleanup >= self.cleanup_interval:
|
|
134
|
+
keys_to_delete = []
|
|
135
|
+
for key, bucket in self.buckets.items():
|
|
136
|
+
if bucket.is_full(now):
|
|
137
|
+
keys_to_delete.append(key)
|
|
138
|
+
for key in keys_to_delete:
|
|
139
|
+
del self.buckets[key]
|
|
140
|
+
self.last_cleanup = now
|
|
141
|
+
|
|
142
|
+
def is_allowed(self, key: str, tokens: float = 1.0) -> bool:
|
|
143
|
+
"""Check if the request is allowed under the rate limit.
|
|
144
|
+
|
|
145
|
+
Args:
|
|
146
|
+
key: Unique identifier for the client or bucket.
|
|
147
|
+
tokens: Number of tokens to consume. Defaults to 1.0.
|
|
148
|
+
|
|
149
|
+
Returns:
|
|
150
|
+
bool: True if the request is allowed, False otherwise.
|
|
151
|
+
"""
|
|
152
|
+
bucket = self._get_bucket(key)
|
|
153
|
+
allowed, _ = bucket.consume(tokens)
|
|
154
|
+
return allowed
|
|
155
|
+
|
|
156
|
+
def check(self, key: str, tokens: float = 1.0) -> None:
|
|
157
|
+
"""Check if the request is allowed under the rate limit.
|
|
158
|
+
Raises RateLimitExceeded if not.
|
|
159
|
+
|
|
160
|
+
Args:
|
|
161
|
+
key: Unique identifier for the client or bucket.
|
|
162
|
+
tokens: Number of tokens to consume. Defaults to 1.0.
|
|
163
|
+
|
|
164
|
+
Raises:
|
|
165
|
+
RateLimitExceeded: If the key is rate-limited.
|
|
166
|
+
"""
|
|
167
|
+
bucket = self._get_bucket(key)
|
|
168
|
+
allowed, retry_after = bucket.consume(tokens)
|
|
169
|
+
if not allowed:
|
|
170
|
+
raise RateLimitExceeded(
|
|
171
|
+
message=f"Rate limit exceeded for key: {key}",
|
|
172
|
+
retry_after=retry_after,
|
|
173
|
+
)
|
|
174
|
+
|
|
175
|
+
def consume(self, key: str, tokens: float = 1.0) -> Tuple[bool, float]:
|
|
176
|
+
"""Consume tokens from the bucket for the given key.
|
|
177
|
+
|
|
178
|
+
Args:
|
|
179
|
+
key: Unique identifier for the client or bucket.
|
|
180
|
+
tokens: Number of tokens to consume. Defaults to 1.0.
|
|
181
|
+
|
|
182
|
+
Returns:
|
|
183
|
+
Tuple[bool, float]: (allowed, retry_after)
|
|
184
|
+
where allowed is True if consumed, False otherwise.
|
|
185
|
+
retry_after is the number of seconds to wait before there are
|
|
186
|
+
enough tokens.
|
|
187
|
+
"""
|
|
188
|
+
bucket = self._get_bucket(key)
|
|
189
|
+
return bucket.consume(tokens)
|
|
190
|
+
|
|
191
|
+
def get_tokens(self, key: str) -> float:
|
|
192
|
+
"""Returns the current number of tokens available in the bucket for the key.
|
|
193
|
+
|
|
194
|
+
Args:
|
|
195
|
+
key: Unique identifier for the client or bucket.
|
|
196
|
+
|
|
197
|
+
Returns:
|
|
198
|
+
float: Current number of tokens.
|
|
199
|
+
"""
|
|
200
|
+
bucket = self._get_bucket(key)
|
|
201
|
+
return bucket.get_tokens()
|
|
202
|
+
|
|
203
|
+
def reset(self, key: str) -> None:
|
|
204
|
+
"""Reset the rate limit bucket for the given key.
|
|
205
|
+
|
|
206
|
+
Args:
|
|
207
|
+
key: Unique identifier for the client or bucket.
|
|
208
|
+
"""
|
|
209
|
+
with self.lock:
|
|
210
|
+
bucket = self.buckets.get(key)
|
|
211
|
+
if bucket is not None:
|
|
212
|
+
bucket.reset()
|
|
213
|
+
|
|
214
|
+
def reset_all(self) -> None:
|
|
215
|
+
"""Reset all rate limit buckets, clearing the internal cache."""
|
|
216
|
+
with self.lock:
|
|
217
|
+
self.buckets.clear()
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
from typing import Optional
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
class SheriffError(Exception):
|
|
5
|
+
"""Base exception for all Sheriff rate limiter errors."""
|
|
6
|
+
|
|
7
|
+
pass
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class RateLimitExceeded(SheriffError):
|
|
11
|
+
"""Exception raised when a rate limit is exceeded."""
|
|
12
|
+
|
|
13
|
+
def __init__(
|
|
14
|
+
self,
|
|
15
|
+
message: str = "Rate limit exceeded",
|
|
16
|
+
retry_after: Optional[float] = None,
|
|
17
|
+
):
|
|
18
|
+
super().__init__(message)
|
|
19
|
+
self.retry_after = retry_after
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
# Test suite for sheriff rate limiter
|
|
@@ -0,0 +1,193 @@
|
|
|
1
|
+
import threading
|
|
2
|
+
import time
|
|
3
|
+
from unittest.mock import patch
|
|
4
|
+
|
|
5
|
+
import pytest
|
|
6
|
+
|
|
7
|
+
from sheriff.core import RateLimiter
|
|
8
|
+
from sheriff.exceptions import RateLimitExceeded
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def test_limiter_creation_defaults(limiter):
|
|
12
|
+
assert isinstance(limiter, RateLimiter)
|
|
13
|
+
assert limiter.capacity == 10.0
|
|
14
|
+
assert limiter.refill_rate == 1.0
|
|
15
|
+
assert limiter.cleanup_interval == 60.0
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def test_limiter_creation_custom():
|
|
19
|
+
limiter = RateLimiter(capacity=20.0, refill_rate=2.5, cleanup_interval=30.0)
|
|
20
|
+
assert limiter.capacity == 20.0
|
|
21
|
+
assert limiter.refill_rate == 2.5
|
|
22
|
+
assert limiter.cleanup_interval == 30.0
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def test_limiter_creation_max_requests():
|
|
26
|
+
limiter = RateLimiter(max_requests=100, period=60.0)
|
|
27
|
+
assert limiter.capacity == 100.0
|
|
28
|
+
assert limiter.refill_rate == 100.0 / 60.0
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def test_limiter_creation_max_requests_no_period():
|
|
32
|
+
limiter = RateLimiter(max_requests=50)
|
|
33
|
+
assert limiter.capacity == 50.0
|
|
34
|
+
assert limiter.refill_rate == 50.0
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def test_invalid_parameters():
|
|
38
|
+
with pytest.raises(ValueError, match="Capacity must be greater than zero"):
|
|
39
|
+
RateLimiter(capacity=0)
|
|
40
|
+
with pytest.raises(ValueError, match="Capacity must be greater than zero"):
|
|
41
|
+
RateLimiter(capacity=-10)
|
|
42
|
+
with pytest.raises(ValueError, match="Refill rate must be greater than zero"):
|
|
43
|
+
RateLimiter(refill_rate=0)
|
|
44
|
+
with pytest.raises(ValueError, match="Refill rate must be greater than zero"):
|
|
45
|
+
RateLimiter(refill_rate=-1)
|
|
46
|
+
with pytest.raises(ValueError, match="Cleanup interval must be greater than zero"):
|
|
47
|
+
RateLimiter(cleanup_interval=0)
|
|
48
|
+
with pytest.raises(ValueError, match="Cleanup interval must be greater than zero"):
|
|
49
|
+
RateLimiter(cleanup_interval=-5)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def test_is_allowed_basic():
|
|
53
|
+
limiter = RateLimiter(capacity=3.0, refill_rate=1.0)
|
|
54
|
+
|
|
55
|
+
assert limiter.is_allowed("user-1", tokens=1.0) is True
|
|
56
|
+
assert limiter.is_allowed("user-1", tokens=2.0) is True
|
|
57
|
+
# Now empty
|
|
58
|
+
assert limiter.is_allowed("user-1", tokens=1.0) is False
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def test_token_replenishment():
|
|
62
|
+
limiter = RateLimiter(capacity=5.0, refill_rate=2.0)
|
|
63
|
+
|
|
64
|
+
start_time = 100.0
|
|
65
|
+
with patch("time.monotonic", return_value=start_time):
|
|
66
|
+
assert limiter.is_allowed("user-2", tokens=5.0) is True
|
|
67
|
+
assert limiter.is_allowed("user-2", tokens=1.0) is False
|
|
68
|
+
|
|
69
|
+
# After 1.5 seconds, we should replenish 1.5 * 2 = 3.0 tokens
|
|
70
|
+
with patch("time.monotonic", return_value=start_time + 1.5):
|
|
71
|
+
assert limiter.is_allowed("user-2", tokens=3.0) is True
|
|
72
|
+
assert limiter.is_allowed("user-2", tokens=1.0) is False
|
|
73
|
+
|
|
74
|
+
# After another 2.5 seconds, we should replenish up to capacity (max 5)
|
|
75
|
+
with patch("time.monotonic", return_value=start_time + 4.0):
|
|
76
|
+
assert limiter.is_allowed("user-2", tokens=5.0) is True
|
|
77
|
+
assert limiter.is_allowed("user-2", tokens=1.0) is False
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
def test_check_raises_exception():
|
|
81
|
+
limiter = RateLimiter(capacity=2.0, refill_rate=1.0)
|
|
82
|
+
|
|
83
|
+
start_time = 100.0
|
|
84
|
+
with patch("time.monotonic", return_value=start_time):
|
|
85
|
+
limiter.check("user-3", tokens=2.0)
|
|
86
|
+
|
|
87
|
+
with pytest.raises(RateLimitExceeded) as exc_info:
|
|
88
|
+
limiter.check("user-3", tokens=1.0)
|
|
89
|
+
|
|
90
|
+
assert exc_info.value.retry_after == 1.0
|
|
91
|
+
|
|
92
|
+
# If we need 2 tokens, it will require 2 seconds
|
|
93
|
+
with pytest.raises(RateLimitExceeded) as exc_info2:
|
|
94
|
+
limiter.check("user-3", tokens=2.0)
|
|
95
|
+
|
|
96
|
+
assert exc_info2.value.retry_after == 2.0
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def test_consume_returns_tuple():
|
|
100
|
+
limiter = RateLimiter(capacity=2.0, refill_rate=0.5)
|
|
101
|
+
|
|
102
|
+
start_time = 100.0
|
|
103
|
+
with patch("time.monotonic", return_value=start_time):
|
|
104
|
+
allowed, retry_after = limiter.consume("user-4", tokens=1.5)
|
|
105
|
+
assert allowed is True
|
|
106
|
+
assert retry_after == 0.0
|
|
107
|
+
|
|
108
|
+
allowed, retry_after = limiter.consume("user-4", tokens=1.0)
|
|
109
|
+
assert allowed is False
|
|
110
|
+
# Needs 1.0 - 0.5 = 0.5 tokens. At refill_rate 0.5, needs 1.0 second.
|
|
111
|
+
assert retry_after == 1.0
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
def test_get_tokens():
|
|
115
|
+
limiter = RateLimiter(capacity=10.0, refill_rate=2.0)
|
|
116
|
+
|
|
117
|
+
start_time = 100.0
|
|
118
|
+
with patch("time.monotonic", return_value=start_time):
|
|
119
|
+
assert limiter.get_tokens("user-5") == 10.0
|
|
120
|
+
assert limiter.is_allowed("user-5", tokens=4.0) is True
|
|
121
|
+
assert limiter.get_tokens("user-5") == 6.0
|
|
122
|
+
|
|
123
|
+
with patch("time.monotonic", return_value=start_time + 1.5):
|
|
124
|
+
# 6.0 + 1.5 * 2 = 9.0
|
|
125
|
+
assert limiter.get_tokens("user-5") == 9.0
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
def test_reset_and_reset_all():
|
|
129
|
+
limiter = RateLimiter(capacity=5.0, refill_rate=1.0)
|
|
130
|
+
|
|
131
|
+
limiter.is_allowed("user-a", tokens=5.0)
|
|
132
|
+
limiter.is_allowed("user-b", tokens=5.0)
|
|
133
|
+
|
|
134
|
+
assert limiter.is_allowed("user-a", tokens=1.0) is False
|
|
135
|
+
assert limiter.is_allowed("user-b", tokens=1.0) is False
|
|
136
|
+
|
|
137
|
+
limiter.reset("user-a")
|
|
138
|
+
assert limiter.is_allowed("user-a", tokens=5.0) is True
|
|
139
|
+
assert limiter.is_allowed("user-b", tokens=1.0) is False
|
|
140
|
+
|
|
141
|
+
limiter.is_allowed("user-a", tokens=5.0)
|
|
142
|
+
limiter.reset_all()
|
|
143
|
+
assert limiter.is_allowed("user-a", tokens=5.0) is True
|
|
144
|
+
assert limiter.is_allowed("user-b", tokens=5.0) is True
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def test_lazy_cleanup():
|
|
148
|
+
limiter = RateLimiter(capacity=5.0, refill_rate=1.0, cleanup_interval=0.1)
|
|
149
|
+
|
|
150
|
+
# Access a key to create a bucket
|
|
151
|
+
assert limiter.is_allowed("key1", tokens=1.0) is True
|
|
152
|
+
assert "key1" in limiter.buckets
|
|
153
|
+
|
|
154
|
+
# Wait, but not long enough to fully replenish
|
|
155
|
+
time.sleep(0.15)
|
|
156
|
+
|
|
157
|
+
# Access key2 to trigger cleanup check
|
|
158
|
+
assert limiter.is_allowed("key2", tokens=1.0) is True
|
|
159
|
+
# key1 is not fully replenished, so it shouldn't be deleted
|
|
160
|
+
assert "key1" in limiter.buckets
|
|
161
|
+
|
|
162
|
+
# Wait long enough to fully replenish key1 (needs 1.0 second to recover 1.0 token)
|
|
163
|
+
time.sleep(1.0)
|
|
164
|
+
|
|
165
|
+
# Access key2 to trigger cleanup
|
|
166
|
+
assert limiter.is_allowed("key2", tokens=1.0) is True
|
|
167
|
+
# key1 is now fully replenished, so it should be deleted
|
|
168
|
+
assert "key1" not in limiter.buckets
|
|
169
|
+
|
|
170
|
+
|
|
171
|
+
def test_concurrent_consumption():
|
|
172
|
+
# Thread safety check: Ensure no double-consumption
|
|
173
|
+
limiter = RateLimiter(capacity=50.0, refill_rate=0.0001)
|
|
174
|
+
|
|
175
|
+
successes = [0]
|
|
176
|
+
lock = threading.Lock()
|
|
177
|
+
|
|
178
|
+
def worker():
|
|
179
|
+
for _ in range(10):
|
|
180
|
+
if limiter.is_allowed("concurrent-key", tokens=1.0):
|
|
181
|
+
with lock:
|
|
182
|
+
successes[0] += 1
|
|
183
|
+
|
|
184
|
+
threads = [threading.Thread(target=worker) for _ in range(10)]
|
|
185
|
+
for t in threads:
|
|
186
|
+
t.start()
|
|
187
|
+
for t in threads:
|
|
188
|
+
t.join()
|
|
189
|
+
|
|
190
|
+
# With 10 threads doing 10 attempts each, total 100 attempts,
|
|
191
|
+
# but capacity is 50 and refill rate is negligible.
|
|
192
|
+
# Therefore, exactly 50 should succeed.
|
|
193
|
+
assert successes[0] == 50
|