fix: Thread-safety improvements (#2186)

* Thread-safety - add locks to engine/memf

* Use reentrant locks

* [skip ci] Types

Author: kukushking

awswrangler/_distributed.py

@@ -3,6 +3,7 @@
 # pylint: disable=import-outside-toplevel
 
 import importlib.util
+import threading
 from collections import defaultdict
 from enum import Enum, unique
 from functools import wraps
@@ -37,6 +38,7 @@ class Engine:
     _engine: Optional[EngineEnum] = None
     _initialized_engine: Optional[EngineEnum] = None
     _registry: Dict[EngineLiteral, Dict[str, Callable[..., Any]]] = defaultdict(dict)
+    _lock: threading.RLock = threading.RLock()
 
     @classmethod
     def get_installed(cls) -> EngineEnum:
@@ -64,27 +66,31 @@ def get(cls) -> EngineEnum:
         str
             The distribution engine configured.
         """
-        return cls._engine if cls._engine else cls.get_installed()
+        with cls._lock:
+            return cls._engine if cls._engine else cls.get_installed()
 
     @classmethod
     def set(cls, name: EngineLiteral) -> None:
         """Set the distribution engine."""
-        cls._engine = EngineEnum._member_map_[  # type: ignore[assignment]  # pylint: disable=protected-access,no-member
-            name.upper()
-        ]
+        with cls._lock:
+            cls._engine = EngineEnum._member_map_[  # type: ignore[assignment]  # pylint: disable=protected-access,no-member
+                name.upper()
+            ]
 
     @classmethod
     def dispatch_func(cls, source_func: FunctionType, value: Optional[EngineLiteral] = None) -> FunctionType:
         """Dispatch a func based on value or the distribution engine and the source function."""
         try:
-            return cls._registry[value or cls.get().value][source_func.__name__]  # type: ignore[return-value]
+            with cls._lock:
+                return cls._registry[value or cls.get().value][source_func.__name__]  # type: ignore[return-value]
         except KeyError:
             return getattr(source_func, "_source_func", source_func)
 
     @classmethod
     def register_func(cls, source_func: Callable[..., Any], destination_func: Callable[..., Any]) -> Callable[..., Any]:
         """Register a func based on the distribution engine and source function."""
-        cls._registry[cls.get().value][source_func.__name__] = destination_func
+        with cls._lock:
+            cls._registry[cls.get().value][source_func.__name__] = destination_func
         return destination_func
 
     @classmethod
@@ -102,38 +108,42 @@ def wrapper(*args: Any, **kw: Dict[str, Any]) -> Any:
     @classmethod
     def register(cls, name: Optional[EngineLiteral] = None) -> None:
         """Register the distribution engine dispatch methods."""
-        engine_name = cast(EngineLiteral, name or cls.get_installed().value)
-        cls.set(engine_name)
-        cls._registry.clear()
+        with cls._lock:
+            engine_name = cast(EngineLiteral, name or cls.get_installed().value)
+            cls.set(engine_name)
+            cls._registry.clear()
 
-        if engine_name == EngineEnum.RAY.value:
-            from awswrangler.distributed.ray._register import register_ray
+            if engine_name == EngineEnum.RAY.value:
+                from awswrangler.distributed.ray._register import register_ray
 
-            register_ray()
+                register_ray()
 
     @classmethod
     def initialize(cls, name: Optional[EngineLiteral] = None) -> None:
         """Initialize the distribution engine."""
-        engine_name = cast(EngineLiteral, name or cls.get_installed().value)
-        if engine_name == EngineEnum.RAY.value:
-            from awswrangler.distributed.ray import initialize_ray
+        with cls._lock:
+            engine_name = cast(EngineLiteral, name or cls.get_installed().value)
+            if engine_name == EngineEnum.RAY.value:
+                from awswrangler.distributed.ray import initialize_ray
 
-            initialize_ray()
-        cls.register(engine_name)
-        cls._initialized_engine = cls.get()
+                initialize_ray()
+            cls.register(engine_name)
+            cls._initialized_engine = cls.get()
 
     @classmethod
     def is_initialized(cls, name: Optional[EngineLiteral] = None) -> bool:
         """Check if the distribution engine is initialized."""
-        engine_name = cast(EngineLiteral, name or cls.get_installed().value)
+        with cls._lock:
+            engine_name = cast(EngineLiteral, name or cls.get_installed().value)
 
-        return False if not cls._initialized_engine else cls._initialized_engine.value == engine_name
+            return False if not cls._initialized_engine else cls._initialized_engine.value == engine_name
 
 
 class MemoryFormat:
     """Memory format configuration class."""
 
     _enum: Optional[MemoryFormatEnum] = None
+    _lock: threading.RLock = threading.RLock()
 
     @classmethod
     def get_installed(cls) -> MemoryFormatEnum:
@@ -161,14 +171,16 @@ def get(cls) -> MemoryFormatEnum:
         Enum
             The memory format configured.
         """
-        return cls._enum if cls._enum else cls.get_installed()
+        with cls._lock:
+            return cls._enum if cls._enum else cls.get_installed()
 
     @classmethod
     def set(cls, name: EngineLiteral) -> None:
         """Set the memory format."""
-        cls._enum = MemoryFormatEnum._member_map_[name.upper()]  # type: ignore[assignment]  # pylint: disable=protected-access,no-member
+        with cls._lock:
+            cls._enum = MemoryFormatEnum._member_map_[name.upper()]  # type: ignore[assignment]  # pylint: disable=protected-access,no-member
 
-        _reload()
+            _reload()
 
 
 def _reload() -> None: