125 lines
3.7 KiB
Python
125 lines
3.7 KiB
Python
# Copyright (c) 2024, Frappe Technologies Pvt. Ltd. and Contributors
|
|
# License: MIT. See LICENSE
|
|
|
|
"""
|
|
Concurrency limiter for expensive whitelisted methods.
|
|
|
|
Provides a @frappe.concurrent_limit() decorator that limits the number of
|
|
simultaneous in-flight executions of a function across all gunicorn workers
|
|
using a Redis-backed semaphore (LIST + BLPOP).
|
|
|
|
Usage::
|
|
|
|
@frappe.whitelist(allow_guest=True)
|
|
@frappe.concurrent_limit(limit=3)
|
|
def download_pdf(...):
|
|
...
|
|
|
|
"""
|
|
|
|
from collections.abc import Callable
|
|
from functools import wraps
|
|
|
|
import frappe
|
|
from frappe.exceptions import ServiceUnavailableError
|
|
from frappe.utils import cint
|
|
from frappe.utils.caching import redis_cache
|
|
from frappe.utils.redis_semaphore import RedisSemaphore
|
|
|
|
# Default wait timeout (seconds) before returning 503 to the caller.
|
|
_DEFAULT_WAIT_TIMEOUT = 10
|
|
|
|
|
|
@redis_cache(shared=True)
|
|
def _default_limit() -> int:
|
|
"""Derive a sensible default concurrency limit from gunicorn's max concurrency."""
|
|
return max(1, gunicorn_max_concurrency() // 2)
|
|
|
|
|
|
def gunicorn_max_concurrency() -> int:
|
|
"""Detect max concurrent requests from the running gunicorn master's cmdline."""
|
|
import os
|
|
|
|
fallback = 4
|
|
|
|
try:
|
|
ppid = os.getppid()
|
|
with open(f"/proc/{ppid}/cmdline", "rb") as f:
|
|
args = f.read().rstrip(b"\0").decode().split("\0")
|
|
|
|
if not any("gunicorn" in a for a in args):
|
|
return fallback
|
|
|
|
workers = _extract_cli_int(args, "-w", "--workers") or fallback
|
|
threads = _extract_cli_int(args, "--threads") or 1
|
|
return workers * threads
|
|
except OSError:
|
|
return fallback
|
|
|
|
|
|
def _extract_cli_int(args: list[str], *flags: str) -> int | None:
|
|
"""Return the integer value for a CLI flag from a split argument list.
|
|
|
|
Handles both ``--flag value`` and ``--flag=value`` forms.
|
|
"""
|
|
for i, arg in enumerate(args):
|
|
for flag in flags:
|
|
if arg == flag and i + 1 < len(args):
|
|
return int(args[i + 1])
|
|
if arg.startswith(f"{flag}="):
|
|
return int(arg.split("=", 1)[1])
|
|
return None
|
|
|
|
|
|
def concurrent_limit(limit: int | None = None, wait_timeout: int = _DEFAULT_WAIT_TIMEOUT):
|
|
"""Decorator that limits simultaneous in-flight executions of the wrapped function.
|
|
|
|
:param limit: Maximum number of concurrent executions. Defaults to half of ``workers x threads``
|
|
as detected from the gunicorn master process.
|
|
:param wait_timeout: Seconds to wait for a free slot before returning 503.
|
|
Defaults to 10 s.
|
|
|
|
The limiter is skipped entirely for background jobs, CLI commands, and
|
|
tests that call functions directly (i.e. outside of an HTTP request).
|
|
"""
|
|
|
|
def decorator(fn: Callable) -> Callable:
|
|
@wraps(fn)
|
|
def wrapper(*args, **kwargs):
|
|
# Skip concurrency limiting outside of HTTP requests (background jobs,
|
|
# CLI commands, tests that call functions directly, etc.).
|
|
if getattr(frappe.local, "request", None) is None:
|
|
return fn(*args, **kwargs)
|
|
|
|
_limit = cint(limit) if limit is not None else _default_limit()
|
|
key = f"concurrency:{fn.__module__}.{fn.__qualname__}"
|
|
|
|
sem = RedisSemaphore(key, _limit, wait_timeout, shared=True)
|
|
token = sem.acquire()
|
|
if not token:
|
|
retry_after = max(1, int(wait_timeout))
|
|
if (headers := getattr(frappe.local, "response_headers", None)) is not None:
|
|
headers.set("Retry-After", str(retry_after))
|
|
exc = ServiceUnavailableError(frappe._("Server is busy. Please try again in a few seconds."))
|
|
exc.retry_after = retry_after
|
|
raise exc
|
|
|
|
try:
|
|
return fn(*args, **kwargs)
|
|
finally:
|
|
sem.release(token)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
@frappe.whitelist()
|
|
def get_stats() -> dict:
|
|
frappe.only_for("System Manager")
|
|
cached_limit = _default_limit()
|
|
gunicorn_limit = gunicorn_max_concurrency()
|
|
return {
|
|
"cached_limit": cached_limit,
|
|
"gunicorn_limit": gunicorn_limit,
|
|
}
|