138 lines
4.1 KiB
Python
138 lines
4.1 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 atomic counter (semaphore).
|
|
|
|
Usage::
|
|
|
|
@frappe.whitelist(allow_guest=True)
|
|
@frappe.concurrent_limit(limit=3)
|
|
def download_pdf(...):
|
|
...
|
|
|
|
"""
|
|
|
|
import time
|
|
from collections.abc import Callable
|
|
from functools import wraps
|
|
|
|
import frappe
|
|
|
|
# Safety TTL (seconds) for the Redis key — prevents leaked semaphore slots if a
|
|
# worker crashes mid-request. Should be larger than any realistic execution time.
|
|
_SLOT_TTL = 120
|
|
|
|
# Default wait timeout (seconds) before returning 503 to the caller.
|
|
_DEFAULT_WAIT_TIMEOUT = 10
|
|
|
|
# Polling interval (seconds) while waiting for a slot to open.
|
|
_POLL_INTERVAL = 0.25
|
|
|
|
|
|
def _default_limit() -> int:
|
|
"""Derive a sensible default concurrency limit from the number of gunicorn workers."""
|
|
import multiprocessing
|
|
|
|
workers = frappe.conf.get("gunicorn_workers") or (multiprocessing.cpu_count() * 2 + 1)
|
|
return max(1, int(workers) // 2)
|
|
|
|
|
|
def concurrent_limit(limit: int | None = None, wait_timeout: int | None = None):
|
|
"""Decorator that limits simultaneous in-flight executions of the wrapped function.
|
|
|
|
:param limit: Maximum number of concurrent executions. Defaults to
|
|
``gunicorn_workers // 2`` (or the value in ``concurrency_limits`` site config).
|
|
:param wait_timeout: Seconds to wait for a free slot before returning 503.
|
|
Defaults to 10 s. Suppressed for background jobs.
|
|
"""
|
|
|
|
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 not getattr(frappe.local, "request", None):
|
|
return fn(*args, **kwargs)
|
|
|
|
effective_limit = int(limit) if limit is not None else _default_limit()
|
|
effective_wait = (
|
|
wait_timeout
|
|
if wait_timeout is not None
|
|
else frappe.conf.get("concurrency_wait_timeout", _DEFAULT_WAIT_TIMEOUT)
|
|
)
|
|
|
|
cache_key = frappe.cache.make_key(f"concurrency:{fn.__module__}.{fn.__qualname__}")
|
|
|
|
acquired = _acquire(cache_key, effective_limit, effective_wait)
|
|
if not acquired:
|
|
from frappe.exceptions import ServiceUnavailableError
|
|
|
|
retry_after = max(1, int(effective_wait))
|
|
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:
|
|
_release(cache_key)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
def _acquire(cache_key: str, limit: int, wait_timeout: float) -> bool:
|
|
"""Increment the counter and return True if we got a slot within *wait_timeout* seconds.
|
|
|
|
The counter is incremented first; if the new value exceeds *limit* the
|
|
increment is undone and we wait before retrying. This avoids a separate
|
|
check-then-act race condition — INCRBY is atomic.
|
|
"""
|
|
deadline = time.monotonic() + wait_timeout
|
|
|
|
while True:
|
|
try:
|
|
current = frappe.cache.incrby(cache_key, 1)
|
|
except Exception:
|
|
# Redis unavailable — fail open to avoid breaking the endpoint entirely.
|
|
frappe.log_error("Concurrency limiter: Redis unavailable, skipping limit")
|
|
return True
|
|
|
|
# Refresh TTL on every successful increment so that a slow request
|
|
# doesn't let the slot expire before it finishes.
|
|
try:
|
|
frappe.cache.expire(cache_key, _SLOT_TTL)
|
|
except Exception:
|
|
pass
|
|
|
|
if current <= limit:
|
|
return True
|
|
|
|
# Over the limit — give back the slot and wait.
|
|
try:
|
|
frappe.cache.incrby(cache_key, -1)
|
|
except Exception:
|
|
pass
|
|
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
return False
|
|
|
|
time.sleep(min(_POLL_INTERVAL, remaining))
|
|
|
|
|
|
def _release(cache_key: str) -> None:
|
|
"""Decrement the counter, clamping at 0 to guard against double-release."""
|
|
try:
|
|
new_val = frappe.cache.incrby(cache_key, -1)
|
|
if new_val < 0:
|
|
# Shouldn't happen, but clamp to prevent permanently negative counters.
|
|
frappe.cache.incrby(cache_key, -new_val)
|
|
except Exception:
|
|
pass
|