mirror of
https://github.com/langgenius/dify.git
synced 2026-10-04 13:47:55 +08:00
88 lines
3.4 KiB
Python
88 lines
3.4 KiB
Python
"""Cross-process Redis session-advance lock.
|
|
|
|
Acquired in the web process (dispatch) and released in the Celery process
|
|
(the dify_builder advance task). Because acquire and release happen on different
|
|
processes/objects, this cannot use ``redis_client.lock()``'s context
|
|
manager, which ties acquire+release to a single object. Instead this module
|
|
uses an explicit SET-NX-token + Lua compare-del, the same pattern used by
|
|
``services.oauth_device_flow`` for its device/user-code state machine.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from uuid import uuid4
|
|
|
|
from configs import dify_config
|
|
from extensions.ext_redis import redis_client
|
|
from extensions.redis_names import serialize_redis_name
|
|
|
|
_KEY_FMT = "dify_builder:advance:{session_id}"
|
|
|
|
# The dispatch lease covers queue delay. Once a worker starts, its lease must
|
|
# outlive the task's hard limit so it cannot expire while that worker can still
|
|
# mutate the session or draft.
|
|
_HARD_LIMIT_GRACE_SECONDS = 30
|
|
_ACTIVE_LEASE_GRACE_SECONDS = 60
|
|
|
|
# Atomic compare-del: only delete the lock if the caller still holds the
|
|
# token it was given on acquire. Prevents a slow/expired holder from
|
|
# deleting a lock that another process has since acquired.
|
|
_RELEASE_LUA = "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end"
|
|
_ACTIVATE_LUA = """
|
|
if redis.call('get', KEYS[1]) == ARGV[1] then
|
|
return redis.call('pexpire', KEYS[1], ARGV[2])
|
|
else
|
|
return 0
|
|
end
|
|
"""
|
|
|
|
|
|
def _key(session_id: str) -> str:
|
|
return _KEY_FMT.format(session_id=session_id)
|
|
|
|
|
|
def _max_advance_ms() -> int:
|
|
return dify_config.DIFY_BUILDER_MAX_ADVANCE_SECONDS * 1000
|
|
|
|
|
|
def active_advance_time_limit() -> int:
|
|
"""Hard task limit; shorter than the worker's active lock lease."""
|
|
return dify_config.DIFY_BUILDER_MAX_ADVANCE_SECONDS + _HARD_LIMIT_GRACE_SECONDS
|
|
|
|
|
|
def acquire(session_id: str) -> str | None:
|
|
"""Try to acquire the advance lock for ``session_id``.
|
|
|
|
Returns a token to be passed to :func:`release` if the lock was
|
|
acquired, or ``None`` if another process already holds it.
|
|
"""
|
|
token = str(uuid4())
|
|
ok = redis_client.set(_key(session_id), token, nx=True, px=_max_advance_ms())
|
|
return token if ok else None
|
|
|
|
|
|
def activate(session_id: str, token: str) -> bool:
|
|
"""Claim an enqueued command at worker start and extend its lease.
|
|
|
|
A command that sat in the queue past its dispatch lease must not run after
|
|
another command acquired the same session. The compare-and-extend is atomic
|
|
so an expired worker cannot extend a newer worker's token.
|
|
"""
|
|
active_ms = (dify_config.DIFY_BUILDER_MAX_ADVANCE_SECONDS + _ACTIVE_LEASE_GRACE_SECONDS) * 1000
|
|
return bool(redis_client.eval(_ACTIVATE_LUA, 1, serialize_redis_name(_key(session_id)), token, active_ms))
|
|
|
|
|
|
def release(session_id: str, token: str) -> None:
|
|
"""Release the advance lock for ``session_id`` iff ``token`` matches the
|
|
holder currently recorded in Redis (compare-del). A no-op if the lock
|
|
was already released, expired, or is held by a different token.
|
|
"""
|
|
# RedisClientWrapper prefixes .set/.get, but .eval is delegated to the raw client — serialize
|
|
# the key ourselves so it matches acquire()'s written key.
|
|
redis_client.eval(_RELEASE_LUA, 1, serialize_redis_name(_key(session_id)), token)
|
|
|
|
|
|
def exists(session_id: str) -> bool:
|
|
"""Whether the advance lock for ``session_id`` is currently held."""
|
|
return redis_client.get(_key(session_id)) is not None
|