lockable_dict.py

import threading
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, wait
from typing import (Any, Dict, Generic, Hashable, ItemsView, Iterator,
                    Optional, TypeVar, Union, overload)

KT = TypeVar("KT", bound=Hashable)
VT = TypeVar("VT")
VT_co = TypeVar("VT_co")


class LockableDict(Generic[KT, VT]):
    def __init__(self, data: Optional[Dict[KT, VT]] = None) -> None:
        self._data: Dict[KT, VT] = data or dict()
        self._keylevel_lock: Dict[int, threading.Lock] = defaultdict(threading.Lock)
        self._toplevel_lock = threading.Lock()

    def acquire(self, kid: Optional[int] = None) -> bool:
        lock = self._get_lock(kid)
        return lock.acquire()

    def release(self, kid: Optional[int] = None) -> None:
        lock = self._get_lock(kid)
        lock.release()

    def __enter__(self) -> "LockableDict":
        self.acquire()
        return self

    def __exit__(self, *args: Any, **kwargs: Any) -> None:
        self.release()

    def __getitem__(self, key: KT) -> VT:
        return self._data[key]

    def __setitem__(self, key: KT, value: VT) -> None:
        self._data[key] = value

    def __iter__(self) -> Iterator[KT]:
        yield from self._data

    def items(self) -> ItemsView[KT, VT]:
        return self._data.items()

    @overload
    def get(self, key: KT) -> Optional[VT]:
        ...

    @overload
    def get(self, key: KT, with_lock: bool) -> Optional[VT]:
        ...

    @overload
    def get(self, key: KT, default: VT_co) -> Union[VT, VT_co]:
        ...

    @overload
    def get(self, key: KT, default: VT_co, with_lock: bool) -> Union[VT, VT_co]:
        ...

    def get(self, key, default=None, with_lock=False):
        if with_lock:
            with self.lock(key):
                return self._data.get(key, default)
        else:
            return self._data.get(key, default)

    def _get_lock(self, kid: Optional[int] = None) -> threading.Lock:
        if kid is None:
            lock = self._toplevel_lock
        else:
            with self._toplevel_lock:
                lock = self._keylevel_lock[kid]
        return lock

    def lock(self, key: KT) -> threading.Lock:
        return self._get_lock(hash(key))


if __name__ == "__main__":
    import random
    import time

    data = LockableDict[str, int]({"a": 0, "b": 0})

    def task(key, index):
        time.sleep(random.uniform(0, 0.1))

        with data.lock(key):
            print(f"task_{key}_{index} - lock {key}")
            count = data[key]
            time.sleep(random.uniform(0, 0.5))
            count += 1
            data[key] = count
            print(f"task_{key}_{index} - release {key}")
        print(f"task_{key}_{index} - count: {count}")

    with ThreadPoolExecutor(max_workers=5) as executor:
        futures = [executor.submit(task, "a", i) for i in range(3)]
        futures += [executor.submit(task, "b", i) for i in range(3)]
        wait(futures)