Enhance ConcurrentDictionary with get_locked and key_lock methods for improved thread safety; update README and tests accordingly. Bump version to 2.0.0.

Author: alelom

README.md

@@ -72,16 +72,41 @@ print(list(bag))  # [1, 2, 3, 4]
 
 ### ConcurrentDictionary
 
-A thread-safe dictionary. For atomic compound updates, use `update_atomic`.
+A thread-safe dictionary. It has a few notable methods:
+
+- assign_atomic()
+- get_locked()
+- update_atomic()
+
+#### ConcurrentDictionary's `assign_atomic()`
+
+Assigns a dictionary value under a key in a thread-safe way.
+While `dict["somekey"] = value` is allowed, it's best to use `assign_atomic()` for clarity of intent. Using normal assignment will work but raise a UserWarning.
+
+
+#### ConcurrentDictionary's `get_locked()`
+
+When working with `ConcurrentDictionary`, you should use the `get_locked` method to safely read or update the value for a specific key in a multi-threaded environment. This ensures that only one thread can access or modify the value for a given key at a time, preventing race conditions.
 
 ```python
 from concurrent_collections import ConcurrentDictionary
 
-d = ConcurrentDictionary({'x': 1})
-d['y'] = 2  # Simple assignment is thread-safe
-# For atomic updates:
-d.update_atomic('x', lambda v: v + 1)
-print(d['x'])  # 2
+d = ConcurrentDictionary({'x': "some value" })
+
+# Safely read and update the value for 'x'
+with d.get_locked('x') as value:
+    # value is locked for this thread
+    d['x'] = "new value"
+```
+
+#### ConcurrentDictionary's `update_atomic()`
+
+Performs a thread-safe, in-place update to an existing value under a key.
+
+```python
+
+d = ConcurrentDictionary({'x': 1 })
+d.update_atomic("x", lambda v: v + 1) # d now contains 2 under the 'x' key.
 ```
 
 ### ConcurrentQueue

concurrent_collections/concurrent_dict.py

@@ -1,7 +1,6 @@
 import threading
-from typing import Any, Callable, Dict, Iterator, List, Optional, TypeVar, Generic, Tuple
+from typing import Any, Callable, Dict, Iterator, List, Optional, TypeVar, Generic, Tuple, ContextManager
 import warnings
-import sys
 
 T = TypeVar('T')
 K = TypeVar('K')
@@ -21,6 +20,47 @@ class ConcurrentDictionary(Generic[K, V]):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         self._lock = threading.RLock()
         self._dict: Dict[K, V] = dict(*args, **kwargs)  # type: ignore
+        self._key_locks: Dict[K, threading.RLock] = {}
+
+    def _get_key_lock(self, key: K) -> threading.RLock:
+        with self._lock:
+            if key not in self._key_locks:
+                self._key_locks[key] = threading.RLock()
+            return self._key_locks[key]
+
+    class _KeyLockContext:
+        def __init__(self, outer : "ConcurrentDictionary[K,V]", key: K):
+            self._outer = outer
+            self._key = key
+            self._lock = outer._get_key_lock(key)
+
+        def __enter__(self) -> V:
+            self._lock.acquire()
+            return self._outer._dict[self._key]
+
+        def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
+            self._lock.release()
+
+    def get_locked(self, key: K) -> ContextManager[V]:
+        """
+        Context manager: lock the key, yield its value, unlock on exit.
+
+        Usage:
+            with d.get_locked('x') as value:
+                # safely read/update value for 'x'
+        """
+        return self._KeyLockContext(self, key)
+
+    def key_lock(self, key: K):
+        """
+        Context manager: lock the key, yield nothing, unlock on exit.
+
+        Usage:
+            with d.key_lock('x'):
+                # safely update d['x'] or perform multiple operations
+        """
+        lock = self._get_key_lock(key)
+        return lock
 
     def __getitem__(self, key: K) -> V:
         with self._lock:

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "concurrent_collections"
-version = "1.4.0"
+version = "2.0.0"
 description = "A brief description of concurrent-collections"
 authors = [{ name = "Alessio Lombardi", email = "work@alelom.com" }]
 license = { file = "MIT" }

tests/concurrent_dict_test.py

@@ -95,7 +95,87 @@ def worker():
 
     # No thread safety errors should occur
     assert not errors, f"Thread safety errors occurred: {errors}"
-      
+    
+def test_get_locked_context_manager_returns_value_and_locks():
+    d: ConcurrentDictionary[str, int] = ConcurrentDictionary({'a': 1})
+    with d.get_locked('a') as value:
+        assert value == 1
+        # Update inside lock
+        d['a'] = value + 1
+        assert d['a'] == 2
+
+def test_get_locked_raises_keyerror_for_missing_key():
+    d: ConcurrentDictionary[str, int] = ConcurrentDictionary()
+    try:
+        with d.get_locked('missing'):
+            assert False, "Should raise KeyError for missing key"
+    except KeyError:
+        pass
+
+def test_key_lock_context_manager_locks_and_unlocks():
+    d: ConcurrentDictionary[str, int] = ConcurrentDictionary({'b': 10})
+    lock = d.key_lock('b')
+    # Should be a context manager (RLock)
+    assert hasattr(lock, '__enter__') and hasattr(lock, '__exit__')
+    with lock:
+        d['b'] = 20
+        assert d['b'] == 20
+
+def test_get_locked_allows_nested_access():
+    d: ConcurrentDictionary[str, int] = ConcurrentDictionary({'c': 5})
+    with d.get_locked('c') as value:
+        assert value == 5
+        # Nested get_locked should not deadlock
+        with d.get_locked('c') as value2:
+            assert value2 == 5
+
+def test_get_locked_thread_safety():
+    d: ConcurrentDictionary[str, int] = ConcurrentDictionary({'x': 0})
+    errors : List[Exception] = []
+
+    def worker():
+        try:
+            for _ in range(1000):
+                with d.get_locked('x') as v:
+                    d['x'] = v + 1
+        except Exception as e:
+            errors.append(e)
+
+    threads = [threading.Thread(target=worker) for _ in range(4)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    assert not errors, f"Thread safety errors occurred: {errors}"
+    assert d['x'] == 4000
+
+def test_modify_without_get_locked_causes_race_condition():
+    """
+    Demonstrates that modifying a value at a key without using get_locked
+    can result in a race condition and incorrect result.
+    """
+    d: ConcurrentDictionary[str, int] = ConcurrentDictionary({'x': 0})
+
+    def worker():
+        # Not using get_locked or any locking
+        for _ in range(1000):
+            v = d['x']  # Read without lock
+            # Simulate context switch
+            time.sleep(0.000001)
+            d['x'] = v + 1  # Write without lock
+
+    threads = [threading.Thread(target=worker) for _ in range(4)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    # The correct value should be 4000, but due to race conditions, it will likely be less
+    assert d['x'] != 4000, (
+        "Modifying without get_locked should result in incorrect value due to race conditions, "
+        f"but got {d['x']}"
+    )
 
 if __name__ == "__main__":
     pytest.main([__file__])