Merge pull request #241 from neutrinoceros/bug/boundary_registry_thread_safety

BUG: fix thread safety for `BoundaryRegistry.register`

Author: neutrinoceros

src/gpgi/_boundaries.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
+from threading import Lock
 from typing import TYPE_CHECKING, Any, Literal, cast
 
 if TYPE_CHECKING:
@@ -26,6 +27,7 @@
 class BoundaryRegistry:
     def __init__(self) -> None:
         self._registry: dict[str, BoundaryRecipeT] = {}
+        self._lock = Lock()
         for key, recipe in _base_registry.items():
             self.register(key, recipe, skip_validation=True)
 
@@ -99,20 +101,21 @@ def register(
           multiple times either under the same key or another, unused key, is
           always safe so it does not raise.
         """
-        if key in self._registry:
-            if recipe is self._registry[key]:
-                return
-            elif not allow_unsafe_override:
-                raise ValueError(
-                    f"Another function is already registered with {key=!r}. "
-                    "If you meant to override the existing function, "
-                    "consider setting allow_unsafe_override=True"
-                )
-
-        if not skip_validation:
-            self._validate_recipe(recipe)
-
-        self._registry[key] = recipe
+        with self._lock:
+            if key in self._registry:
+                if recipe is self._registry[key]:
+                    return
+                elif not allow_unsafe_override:
+                    raise ValueError(
+                        f"Another function is already registered with {key=!r}. "
+                        "If you meant to override the existing function, "
+                        "consider setting allow_unsafe_override=True"
+                    )
+
+            if not skip_validation:
+                self._validate_recipe(recipe)
+
+            self._registry[key] = recipe
 
     def __getitem__(self, key: str) -> BoundaryRecipeT:
         return self._registry[key]

tests/test_concurrent.py

@@ -95,8 +95,14 @@ def closure():
 
 class TestBoundaryRegistry:
     def check(self, results):
-        # Check results: verify that all threads get the same obj
-        assert len(set(results)) == 1
+        # only one thread can succeed registration, all others should raise.
+        expected_msg = (
+            "Another function is already registered with key='test'. "
+            "If you meant to override the existing function, "
+            "consider setting allow_unsafe_override=True"
+        )
+        assert len(results) == N_THREADS - 1
+        assert results.count(expected_msg) == N_THREADS - 1
 
     def test_concurrent_threading(self):
         # Defines a thread barrier that will be spawned before parallel execution
@@ -106,28 +112,32 @@ def test_concurrent_threading(self):
         # This object will be shared by all the threads.
         registry = BoundaryRegistry()
 
-        def test_recipe(
-            same_side_active_layer,
-            same_side_ghost_layer,
-            opposite_side_active_layer,
-            opposite_side_ghost_layer,
-            weight_same_side_active_layer,
-            weight_same_side_ghost_layer,
-            weight_opposite_side_active_layer,
-            weight_opposite_side_ghost_layer,
-            side,
-            metadata,
-        ): ...
-
         results = []
 
         def closure():
+            def test_recipe(
+                same_side_active_layer,
+                same_side_ghost_layer,
+                opposite_side_active_layer,
+                opposite_side_ghost_layer,
+                weight_same_side_active_layer,
+                weight_same_side_ghost_layer,
+                weight_opposite_side_active_layer,
+                weight_opposite_side_ghost_layer,
+                side,
+                metadata,
+            ): ...
+
             assert "test" not in registry
 
             # Ensure that all threads reach this point before concurrent execution.
             barrier.wait()
-            registry.register("test", test_recipe)
-            results.append(registry["test"])
+            try:
+                registry.register("test", test_recipe)
+            except ValueError as exc:
+                msg, *_ = exc.args
+                results.append(msg)
+            assert "test" in registry
 
         # Spawn n threads that call _setup_host_cell_index concurrently.
         workers = []
@@ -150,31 +160,30 @@ def test_concurrent_pool(self):
         # This object will be shared by all the threads.
         registry = BoundaryRegistry()
 
-        def test_recipe(
-            same_side_active_layer,
-            same_side_ghost_layer,
-            opposite_side_active_layer,
-            opposite_side_ghost_layer,
-            weight_same_side_active_layer,
-            weight_same_side_ghost_layer,
-            weight_opposite_side_active_layer,
-            weight_opposite_side_ghost_layer,
-            side,
-            metadata,
-        ): ...
-
-        results = []
-
         def closure():
+            def test_recipe(
+                same_side_active_layer,
+                same_side_ghost_layer,
+                opposite_side_active_layer,
+                opposite_side_ghost_layer,
+                weight_same_side_active_layer,
+                weight_same_side_ghost_layer,
+                weight_opposite_side_active_layer,
+                weight_opposite_side_ghost_layer,
+                side,
+                metadata,
+            ): ...
+
             assert "test" not in registry
 
             # Ensure that all threads reach this point before concurrent execution.
             barrier.wait()
             registry.register("test", test_recipe)
-            results.append(registry["test"])
 
         with ThreadPoolExecutor(max_workers=N_THREADS) as executor:
             futures = [executor.submit(closure) for _ in range(N_THREADS)]
 
-        results = [f.result() for f in futures]
+        assert "test" in registry
+        exceptions = [f.exception() for f in futures]
+        results = [exc.args[0] for exc in exceptions if exc is not None]
         self.check(results)