token bucket implemented in memory, in app

2023-08-10 12:28:09 +03:00
commit 224cc4f5ec
5 changed files with 59 additions and 0 deletions
--- a/my_limiter/init.py
+++ b/my_limiter/init.py
--- a/my_limiter/algos.py
+++ b/my_limiter/algos.py
@@ -0,0 +1,33 @@
+import datetime as dt
+
+
+
+TOKEN_BUCKET = {}
+TIME_INTERVAL_SECONDS = 15
+
+class TooManyRequests(Exception):
+    pass
+
+
+def token_bucket(ip: str):
+    """
+    Tokens are put in the bucket at preset rates periodically. 
+    Once the bucket is full, no more tokens are added. 
+    The refiller puts NUM_TOKENS_TO_REFILL tokens into the bucket every minute. 
+    """
+    REFILL_EVERY_SECONDS = TIME_INTERVAL_SECONDS
+    NUM_TOKENS_TO_REFILL = 4
+    MAX_CAPACITY = 4
+
+    entry = TOKEN_BUCKET.get(ip)
+
+    if entry is None:
+        TOKEN_BUCKET[ip] = {'tokens': 4, 'last_refilled': dt.datetime.now().timestamp()}
+    else:
+        if dt.datetime.now().timestamp() >= entry['last_refilled'] + REFILL_EVERY_SECONDS:
+            entry['last_refilled'] = dt.datetime.now().timestamp()
+            entry['tokens'] = min(entry['tokens'] + NUM_TOKENS_TO_REFILL, MAX_CAPACITY)
+    left = TOKEN_BUCKET[ip]['tokens']
+    if left == 0:
+        raise TooManyRequests
+    TOKEN_BUCKET[ip]['tokens'] -= 1
--- a/my_limiter/config.py
+++ b/my_limiter/config.py
@@ -0,0 +1 @@
+
--- a/my_limiter/wsgi.py
+++ b/my_limiter/wsgi.py
@@ -0,0 +1,22 @@
+import flask as f
+
+from . import algos
+
+
+application = f.Flask(__name__)
+
+
+increment_requests_func = algos.token_bucket
+
+
+@application.before_request
+def before_request():
+    ip = f.request.remote_addr
+    try:
+        increment_requests_func(ip)
+    except algos.TooManyRequests:
+        return f.abort(429)
+
+@application.route('/')
+def home():
+    return '<!doctype html><title>Hello</title><h1>Hello</h1>'