misc: Thread-safe symbol cache

Author: enwask

devito/types/basic.py

@@ -16,7 +16,7 @@
 from devito.tools import (Pickable, as_tuple, dtype_to_ctype,
                           frozendict, memoized_meth, sympy_mutex, CustomDtype)
 from devito.types.args import ArgProvider
-from devito.types.caching import Cached, Uncached
+from devito.types.caching import CacheManager, Cached, Uncached
 from devito.types.lazy import Evaluable
 from devito.types.utils import DimensionTuple
 
@@ -559,24 +559,31 @@ def _cache_key(cls, *args, **kwargs):
     def __new__(cls, *args, **kwargs):
         assumptions, kwargs = cls._filter_assumptions(**kwargs)
         key = cls._cache_key(*args, **{**assumptions, **kwargs})
-        obj = cls._cache_get(key)
 
+        # Initial cache lookup (not locked)
+        obj = cls._cache_get(key)
         if obj is not None:
             return obj
 
-        # Not in cache. Create a new Symbol via sympy.Symbol
-        args = list(args)
-        name = kwargs.pop('name', None) or args.pop(0)
-        newobj = cls.__xnew__(cls, name, **assumptions)
+        # Lock against the symbol cache and double-check the cache
+        with CacheManager.lock():
+            obj = cls._cache_get(key)
+            if obj is not None:
+                return obj
 
-        # Initialization
-        newobj._dtype = cls.__dtype_setup__(**kwargs)
-        newobj.__init_finalize__(name, *args, **kwargs)
+            # Not in cache. Create a new Symbol via sympy.Symbol
+            args = list(args)
+            name = kwargs.pop('name', None) or args.pop(0)
+            newobj = cls.__xnew__(cls, name, **assumptions)
 
-        # Store new instance in symbol cache
-        Cached.__init__(newobj, key)
+            # Initialization
+            newobj._dtype = cls.__dtype_setup__(**kwargs)
+            newobj.__init_finalize__(name, *args, **kwargs)
 
-        return newobj
+            # Store new instance in symbol cache
+            Cached.__init__(newobj, key)
+
+            return newobj
 
     __hash__ = Cached.__hash__
 

devito/types/caching.py

@@ -1,4 +1,5 @@
 import gc
+from threading import RLock
 import weakref
 
 import sympy
@@ -10,6 +11,7 @@
 __all__ = ['Cached', 'Uncached', '_SymbolCache', 'CacheManager']
 
 _SymbolCache = {}
+_cache_lock = RLock()
 """The symbol cache."""
 
 
@@ -76,8 +78,10 @@ def _cache_get(cls, key):
             obj = obj_cached()
             if obj is None:
                 # Cleanup _SymbolCache (though practically unnecessary)
-                # does not fail if it's already gone
-                _SymbolCache.pop(key, None)
+                with _cache_lock:
+                    # Ensure another thread hasn't replaced the ref we're evicting
+                    if _SymbolCache.get(key) is obj_cached:
+                        _SymbolCache.pop(key, None)
                 return None
             else:
                 return obj
@@ -196,16 +200,21 @@ def clear(cls, force=True):
                 # We won't call gc.collect() this time
                 cls.ncalls_w_force_false += 1
 
-        for key in cache_copied:
-            obj = _SymbolCache.get(key)
-            if obj is None:
-                # deleted by another thread since we took the copy
-                continue
-            if obj() is None:
-                # (key could be removed in another thread since get() above)
-                _SymbolCache.pop(key, None)
+        for key, obj_cached in cache_copied.items():
+            if obj_cached() is None:
+                with _cache_lock:
+                    # Check if our snapshot of the cached object is still live
+                    if _SymbolCache.get(key) is obj_cached:
+                        _SymbolCache.pop(key, None)
 
         # Maybe trigger garbage collection
         if force:
             del cache_copied
             gc.collect()
+
+    @staticmethod
+    def lock():
+        """
+        Returns the global symbol cache lock for atomic construction.
+        """
+        return _cache_lock