life-echo/api/app/core/redis_lock.py

"""Small Redis lock helpers for background tasks."""

import threading
import uuid
from dataclasses import dataclass

import redis

from app.core.config import settings

_redis_lock_client: redis.Redis | None = None
_redis_lock_init_lock = threading.Lock()


def _get_redis_lock_client() -> redis.Redis:
    """进程内复用单个 Redis 客户端（decode_responses=False，与锁 token 字节一致）。"""
    global _redis_lock_client
    if _redis_lock_client is None:
        with _redis_lock_init_lock:
            if _redis_lock_client is None:
                _redis_lock_client = redis.from_url(
                    settings.redis_url, decode_responses=False
                )
    return _redis_lock_client


@dataclass(frozen=True)
class RedisLockHandle:
    key: str
    token: bytes


def acquire_redis_lock(key: str, *, ttl_seconds: int) -> RedisLockHandle | None:
    """Acquire a single-owner Redis lock or return None when unavailable."""
    client = _get_redis_lock_client()
    token = uuid.uuid4().hex.encode("utf-8")
    if not client.set(key, token, nx=True, ex=ttl_seconds):
        return None
    return RedisLockHandle(key=key, token=token)


def release_redis_lock(handle: RedisLockHandle | None) -> None:
    """Release the lock only if we still own it."""
    if handle is None:
        return
    _get_redis_lock_client().eval(
        """
        if redis.call("GET", KEYS[1]) == ARGV[1] then
            return redis.call("DEL", KEYS[1])
        end
        return 0
        """,
        1,
        handle.key,
        handle.token,
    )