Make lazy importing thread safe

When there are many threads that are lazily importing modules at the same time, we have seen occurrences of them failing to find methods in a module. This change should make sure a module can only be imported once at the same time.

PiperOrigin-RevId: 702432284

CHANGELOG.md

@@ -8,6 +8,9 @@ Changelog follow https://keepachangelog.com/ format.
 
 ## [Unreleased]
 
+*   `epy`:
+    *   `epy.lazy_api_imports`: Make lazy importing thread safe.
+
 ## [1.11.0] - 2024-11-27
 
 *   `enp`:

etils/epy/lazy_imports_utils.py

@@ -24,11 +24,13 @@
 from __future__ import annotations
 
 import builtins
+import collections
 import contextlib
 import dataclasses
 import functools
 import importlib
 import sys
+import threading
 import types
 from typing import Any, Callable, ContextManager, Iterator
 
@@ -39,6 +41,10 @@
 _ErrorCallback = Callable[[Exception], None]
 _SuccessCallback = Callable[[str], None]
 
+# Store a lock per module to avoid problems with multiple threads trying to
+# import the same module at the same time.
+_LOCK_PER_MODULE = collections.defaultdict(threading.Lock)
+
 
 @dataclasses.dataclass(kw_only=True)
 class LazyModule:
@@ -67,7 +73,10 @@ def _module(self) -> types.ModuleType:
     # Try to import the module, eventually replaying the adhoc scope
     with self._maybe_adhoc():
       try:
-        module = importlib.import_module(self.module_name)
+        # When multiple threads try to import the same module, make sure only
+        # one of them is importing it at the same time.
+        with _LOCK_PER_MODULE[self.module_name]:
+          module = importlib.import_module(self.module_name)
       except ImportError as e:
         if self.error_callback is not None:
           if isinstance(self.error_callback, str):