Redis distributed lock with asyncio and Lua script

import asyncio
import uuid
from contextlib import asynccontextmanager
from redis.asyncio import Redis

# Lua script for atomic lock release (check-and-delete)
RELEASE_LOCK_SCRIPT = """
if redis.call('GET', KEYS[1]) == ARGV[1] then
    return redis.call('DEL', KEYS[1])
else
    return 0
end
"""

class DistributedLock:
    def __init__(self, redis: Redis, name: str, ttl: int = 30):
        self.redis = redis
        self.key = f'lock:{name}'
        self.ttl = ttl
        self.token = str(uuid.uuid4())  # Unique token prevents releasing others' locks

    async def acquire(self, timeout: float = 10.0) -> bool:
        deadline = asyncio.get_event_loop().time() + timeout
        while asyncio.get_event_loop().time() < deadline:
            # SET NX EX: set if not exists, with TTL
            acquired = await self.redis.set(
                self.key, self.token,
                nx=True,  # Only set if key doesn't exist
                ex=self.ttl,  # Auto-expire after TTL seconds
            )
            if acquired:
                return True
            await asyncio.sleep(0.1)  # Poll interval
        return False

    async def release(self) -> bool:
        result = await self.redis.eval(RELEASE_LOCK_SCRIPT, 1, self.key, self.token)
        return bool(result)

    async def extend(self, additional_ttl: int) -> bool:
        # Extend TTL atomically — only if we still hold the lock
        script = """
        if redis.call('GET', KEYS[1]) == ARGV[1] then
            return redis.call('EXPIRE', KEYS[1], ARGV[2])
        else
            return 0
        end
        """
        result = await self.redis.eval(script, 1, self.key, self.token, additional_ttl)
        return bool(result)

@asynccontextmanager
async def distributed_lock(redis: Redis, name: str, ttl: int = 30, timeout: float = 10.0):
    lock = DistributedLock(redis, name, ttl)
    acquired = await lock.acquire(timeout=timeout)
    if not acquired:
        raise TimeoutError(f'Could not acquire lock: {name}')
    try:
        yield lock
    finally:
        await lock.release()

# Usage
async def process_trace(trace_id: str, redis: Redis) -> None:
    async with distributed_lock(redis, f'trace:{trace_id}', ttl=60) as lock:
        # Only one worker runs this block per trace_id
        trace = await fetch_trace(trace_id)
        embedding = await generate_embedding(trace.text)
        # Extend lock if processing takes longer
        if embedding_took_long:
            await lock.extend(30)
        await save_embedding(trace_id, embedding)