src/backend/base/langflow/utils/concurrency.py

import re
import threading
from contextlib import contextmanager
from pathlib import Path

from filelock import FileLock
from platformdirs import user_cache_dir


class KeyedMemoryLockManager:
    """A manager for acquiring and releasing memory locks based on a key."""

    def __init__(self) -> None:
        self.locks: dict[str, threading.Lock] = {}
        self.global_lock = threading.Lock()

    def _get_lock(self, key: str):
        with self.global_lock:
            if key not in self.locks:
                self.locks[key] = threading.Lock()
            return self.locks[key]

    @contextmanager
    def lock(self, key: str):
        lock = self._get_lock(key)
        lock.acquire()
        try:
            yield
        finally:
            lock.release()


class KeyedWorkerLockManager:
    """A manager for acquiring locks between workers based on a key."""

    def __init__(self) -> None:
        self.locks_dir = Path(user_cache_dir("langflow"), ensure_exists=True) / "worker_locks"

    def _validate_key(self, key: str) -> bool:
        """Validate that the string only contains alphanumeric characters and underscores.

        Parameters:
        s (str): The string to validate.

        Returns:
        bool: True if the string is valid, False otherwise.
        """
        pattern = re.compile(r"^\w+$")
        return bool(pattern.match(key))

    @contextmanager
    def lock(self, key: str):
        if not self._validate_key(key):
            msg = f"Invalid key: {key}"
            raise ValueError(msg)

        lock = FileLock(self.locks_dir / key)
        with lock:
            yield