zbparse/flask_app/ConnectionLimiter.py
2024-11-25 09:15:56 +08:00

31 lines
948 B
Python

from threading import Semaphore
from functools import wraps
from flask import jsonify, current_app
class ConnectionLimiter:
_instance = None
def __new__(cls, *args, **kwargs):
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def __init__(self, max_connections=10):
if not hasattr(self, 'semaphore'):
self.semaphore = Semaphore(max_connections)
def limit_connections(self, f):
@wraps(f)
def decorated_function(*args, **kwargs):
if not self.semaphore.acquire(blocking=False):
return jsonify({
'error': 'Server is busy. Maximum number of concurrent connections reached.',
'code': 429
}), 429
try:
return f(*args, **kwargs)
finally:
self.semaphore.release()
return decorated_function