commit 224cc4f5ec9a08bbdb55bdccab01374b0c45aa3c Author: Zev Averbach Date: Thu Aug 10 12:28:09 2023 +0300 token bucket implemented in memory, in app 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