From 224cc4f5ec9a08bbdb55bdccab01374b0c45aa3c Mon Sep 17 00:00:00 2001 From: Zev Averbach Date: Thu, 10 Aug 2023 12:28:09 +0300 Subject: [PATCH] token bucket implemented in memory, in app --- .gitignore | 3 +++ my_limiter/__init__.py | 0 my_limiter/algos.py | 33 +++++++++++++++++++++++++++++++++ my_limiter/config.py | 1 + my_limiter/wsgi.py | 22 ++++++++++++++++++++++ 5 files changed, 59 insertions(+) create mode 100644 .gitignore create mode 100644 my_limiter/__init__.py create mode 100644 my_limiter/algos.py create mode 100644 my_limiter/config.py create mode 100644 my_limiter/wsgi.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..33f7cf2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +env/ +*.pyc diff --git a/my_limiter/__init__.py b/my_limiter/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/my_limiter/algos.py b/my_limiter/algos.py new file mode 100644 index 0000000..f7b0642 --- /dev/null +++ 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 \ No newline at end of file diff --git a/my_limiter/config.py b/my_limiter/config.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/my_limiter/config.py @@ -0,0 +1 @@ + diff --git a/my_limiter/wsgi.py b/my_limiter/wsgi.py new file mode 100644 index 0000000..d139e82 --- /dev/null +++ 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 'Hello

Hello

' \ No newline at end of file