126 lines
3.6 KiB
Python
126 lines
3.6 KiB
Python
# Copyright (c) 2024, Frappe Technologies Pvt. Ltd. and Contributors
|
|
# License: MIT. See LICENSE
|
|
|
|
"""Distributed counting semaphore backed by a Redis LIST."""
|
|
|
|
import frappe
|
|
|
|
|
|
class RedisSemaphore:
|
|
"""A distributed counting semaphore backed by a Redis LIST.
|
|
|
|
Allows up to *limit* concurrent holders across all processes sharing the
|
|
same Redis instance. The token pool is lazily initialized via an atomic
|
|
Lua script and self-heals after crashes thanks to a TTL on the capacity
|
|
key.
|
|
|
|
Usage as a context manager::
|
|
|
|
sem = RedisSemaphore("my-resource", limit=5, wait_timeout=10)
|
|
with sem:
|
|
... # at most 5 concurrent holders
|
|
|
|
Or acquire/release manually::
|
|
|
|
token = sem.acquire()
|
|
if token is None:
|
|
raise Exception("Too busy")
|
|
try:
|
|
...
|
|
finally:
|
|
sem.release(token)
|
|
"""
|
|
|
|
# Safety TTL (seconds) for the capacity key — allows the pool to self-heal
|
|
# after a worker crash that leaked a token.
|
|
CAPACITY_TTL = 3600 # 1 hour
|
|
|
|
# Lua script that atomically initializes the token pool.
|
|
# KEYS[1] = capacity key, KEYS[2] = token list key
|
|
# ARGV[1] = limit, ARGV[2] = TTL
|
|
_INIT_SCRIPT = """\
|
|
if redis.call('SET', KEYS[1], ARGV[1], 'NX', 'EX', ARGV[2]) then
|
|
redis.call('DEL', KEYS[2])
|
|
local n = tonumber(ARGV[1])
|
|
for i = 1, n do
|
|
redis.call('RPUSH', KEYS[2], tostring(i))
|
|
end
|
|
end
|
|
"""
|
|
|
|
def __init__(self, key: str, limit: int, wait_timeout: float = 0):
|
|
"""
|
|
:param key: A unique Redis key name for this semaphore (will be
|
|
prefixed by the cache layer).
|
|
:param limit: Maximum number of concurrent holders.
|
|
:param wait_timeout: Seconds to block waiting for a free slot.
|
|
0 means non-blocking (immediate return if unavailable).
|
|
"""
|
|
self.key = key
|
|
self.limit = limit
|
|
self.wait_timeout = wait_timeout
|
|
self._token: str | None = None
|
|
|
|
def acquire(self) -> str | None:
|
|
"""Try to acquire a token from the pool.
|
|
|
|
Returns a token string on success, ``None`` if no slot was
|
|
available within *wait_timeout*, or ``"fallback"`` if Redis is
|
|
unreachable (fail-open).
|
|
"""
|
|
try:
|
|
self._ensure_tokens()
|
|
|
|
if self.wait_timeout <= 0:
|
|
result = frappe.cache.lpop(self.key)
|
|
return self._decode(result) if result is not None else None
|
|
|
|
if result := frappe.cache.blpop(self.key, timeout=int(self.wait_timeout)):
|
|
return self._decode(result[1])
|
|
return None
|
|
|
|
except Exception:
|
|
frappe.log_error(f"RedisSemaphore({self.key}): Redis unavailable, skipping limit")
|
|
return "fallback"
|
|
|
|
def release(self, token: str) -> None:
|
|
"""Return *token* to the pool."""
|
|
if token == "fallback":
|
|
return
|
|
try:
|
|
frappe.cache.lpush(self.key, token)
|
|
except Exception:
|
|
frappe.log_error(f"RedisSemaphore({self.key}): Failed to release token {token}")
|
|
|
|
# -- context-manager protocol ------------------------------------------
|
|
|
|
def __enter__(self):
|
|
self._token = self.acquire()
|
|
return self._token
|
|
|
|
def __exit__(self, *exc_info):
|
|
if self._token is not None:
|
|
self.release(self._token)
|
|
self._token = None
|
|
|
|
# -- internals ---------------------------------------------------------
|
|
|
|
def _ensure_tokens(self) -> None:
|
|
"""Lazily initialize the token pool via an atomic Lua script."""
|
|
try:
|
|
prefixed_cap_key = frappe.cache.make_key(f"{self.key}:capacity")
|
|
prefixed_key = frappe.cache.make_key(self.key)
|
|
frappe.cache.eval(
|
|
self._INIT_SCRIPT,
|
|
2,
|
|
prefixed_cap_key,
|
|
prefixed_key,
|
|
str(self.limit),
|
|
str(self.CAPACITY_TTL),
|
|
)
|
|
except Exception:
|
|
frappe.log_error(f"RedisSemaphore({self.key}): Failed to initialize tokens")
|
|
|
|
@staticmethod
|
|
def _decode(result):
|
|
return result.decode() if isinstance(result, bytes) else result
|