Repository: numba/numba

Issue:
Cache file collision 
Trigger:

* with files named like:

```
foo/__init__.py
foo/foo.py
```

* both files define the same function (i.e. `bar`) on the same line.
* both files have the same modification timestamps and file sizes.

Details:

The `foo/__init__.py` and `foo/foo.py` will both save to the same cache file at `foo.<funcname>-<lineno><pyver>m.nbi`.  Even though there's a "version" check, it only considers the timestamp and file size, which could be equivalent for two completely different file content.  As a result, usage of the cached function may load the wrong implementation.
Cache file collision 
Trigger:

* with files named like:

```
foo/__init__.py
foo/foo.py
```

* both files define the same function (i.e. `bar`) on the same line.
* both files have the same modification timestamps and file sizes.

Details:

The `foo/__init__.py` and `foo/foo.py` will both save to the same cache file at `foo.<funcname>-<lineno><pyver>m.nbi`.  Even though there's a "version" check, it only considers the timestamp and file size, which could be equivalent for two completely different file content.  As a result, usage of the cached function may load the wrong implementation.


Discussion:

Resolving patch:
diff --git a/numba/caching.py b/numba/caching.py
--- a/numba/caching.py
+++ b/numba/caching.py
@@ -329,9 +329,10 @@ def __init__(self, py_func):
             raise RuntimeError("cannot cache function %r: no locator available "
                                "for file %r" % (qualname, source_path))
         self._locator = locator
-        # Keep the last dotted component, since the package name is already
-        # encoded in the directory.
-        modname = py_func.__module__.split('.')[-1]
+        # Use filename base name as module name to avoid conflict between
+        # foo/__init__.py and foo/foo.py
+        filename = inspect.getfile(py_func)
+        modname = os.path.splitext(os.path.basename(filename))[0]
         fullname = "%s.%s" % (modname, qualname)
         abiflags = getattr(sys, 'abiflags', '')
         self._filename_base = self.get_filename_base(fullname, abiflags)


Test patch:
diff --git a/numba/tests/npyufunc/test_caching.py b/numba/tests/npyufunc/test_caching.py
--- a/numba/tests/npyufunc/test_caching.py
+++ b/numba/tests/npyufunc/test_caching.py
@@ -3,7 +3,6 @@
 import sys
 import os.path
 import re
-from contextlib import contextmanager
 import subprocess
 
 import numpy as np
@@ -11,7 +10,7 @@
 from numba import unittest_support as unittest
 from numba import config
 
-from ..support import captured_stdout
+from ..support import capture_cache_log
 from ..test_dispatcher import BaseCacheTest
 
 
@@ -30,13 +29,6 @@ class UfuncCacheTest(BaseCacheTest):
     regex_data_loaded = re.compile(r'\[cache\] data loaded from')
     regex_index_loaded = re.compile(r'\[cache\] index loaded from')
 
-    @contextmanager
-    def capture_cache_log(self):
-        with captured_stdout() as out:
-            old, config.DEBUG_CACHE = config.DEBUG_CACHE, True
-            yield out
-            config.DEBUG_CACHE = old
-
     def check_cache_saved(self, cachelog, count):
         """
         Check number of cache-save were issued
@@ -63,13 +55,13 @@ def check_ufunc_cache(self, usecase_name, n_overloads, **kwargs):
         mod = self.import_module()
         usecase = getattr(mod, usecase_name)
         # New cache entry saved
-        with self.capture_cache_log() as out:
+        with capture_cache_log() as out:
             new_ufunc = usecase(**kwargs)
         cachelog = out.getvalue()
         self.check_cache_saved(cachelog, count=n_overloads)
 
         # Use cached version
-        with self.capture_cache_log() as out:
+        with capture_cache_log() as out:
             cached_ufunc = usecase(**kwargs)
         cachelog = out.getvalue()
         self.check_cache_loaded(cachelog, count=n_overloads)
@@ -114,16 +106,16 @@ def check_dufunc_usecase(self, usecase_name):
         mod = self.import_module()
         usecase = getattr(mod, usecase_name)
         # Create dufunc
-        with self.capture_cache_log() as out:
+        with capture_cache_log() as out:
             ufunc = usecase()
         self.check_cache_saved(out.getvalue(), count=0)
         # Compile & cache
-        with self.capture_cache_log() as out:
+        with capture_cache_log() as out:
             ufunc(np.arange(10))
         self.check_cache_saved(out.getvalue(), count=1)
         self.check_cache_loaded(out.getvalue(), count=0)
         # Use cached
-        with self.capture_cache_log() as out:
+        with capture_cache_log() as out:
             ufunc = usecase()
             ufunc(np.arange(10))
         self.check_cache_loaded(out.getvalue(), count=1)
@@ -147,7 +139,7 @@ class TestGUfuncCacheTest(UfuncCacheTest):
     def test_filename_prefix(self):
         mod = self.import_module()
         usecase = getattr(mod, "direct_gufunc_cache_usecase")
-        with self.capture_cache_log() as out:
+        with capture_cache_log() as out:
             usecase()
         cachelog = out.getvalue()
         # find number filename with "guf-" prefix
diff --git a/numba/tests/support.py b/numba/tests/support.py
--- a/numba/tests/support.py
+++ b/numba/tests/support.py
@@ -424,6 +424,29 @@ def override_config(name, value):
         setattr(config, name, old_value)
 
 
+@contextlib.contextmanager
+def override_env_config(name, value):
+    """
+    Return a context manager that temporarily sets an Numba config environment
+    *name* to *value*.
+    """
+    old = os.environ.get(name)
+    os.environ[name] = value
+    config.reload_config()
+
+    try:
+        yield
+    finally:
+        if old is None:
+            # If it wasn't set originally, delete the environ var
+            del os.environ[name]
+        else:
+            # Otherwise, restore to the old value
+            os.environ[name] = old
+        # Always reload config
+        config.reload_config()
+
+
 def compile_function(name, code, globs):
     """
     Given a *code* string, compile it with globals *globs* and return
@@ -560,6 +583,13 @@ def captured_stderr():
     return captured_output("stderr")
 
 
+@contextlib.contextmanager
+def capture_cache_log():
+    with captured_stdout() as out:
+        with override_config('DEBUG_CACHE', True):
+            yield out
+
+
 class MemoryLeak(object):
 
     __enable_leak_check = True
diff --git a/numba/tests/test_debug.py b/numba/tests/test_debug.py
--- a/numba/tests/test_debug.py
+++ b/numba/tests/test_debug.py
@@ -7,7 +7,8 @@
 
 import numpy as np
 
-from .support import TestCase, override_config, captured_stdout, forbid_codegen
+from .support import (TestCase, override_config, override_env_config,
+                      captured_stdout, forbid_codegen)
 from numba import unittest_support as unittest
 from numba import jit, jitclass, types
 from numba.compiler import compile_isolated, Flags
@@ -193,23 +194,20 @@ class TestEnvironmentOverride(FunctionDebugTestBase):
     """
     Test that environment variables are reloaded by Numba when modified.
     """
-    
+
     # mutates env with os.environ so must be run serially
     _numba_parallel_test_ = False
 
     def test_debug(self):
         out = self.compile_simple_nopython()
         self.assertFalse(out)
-        os.environ['NUMBA_DEBUG'] = '1'
-        try:
+        with override_env_config('NUMBA_DEBUG', '1'):
             out = self.compile_simple_nopython()
             # Note that all variables dependent on NUMBA_DEBUG are
             # updated too.
             self.check_debug_output(out, ['ir', 'typeinfer',
                                           'llvm', 'func_opt_llvm',
                                           'optimized_llvm', 'assembly'])
-        finally:
-            del os.environ['NUMBA_DEBUG']
         out = self.compile_simple_nopython()
         self.assertFalse(out)
 
@@ -234,17 +232,13 @@ def check_parfors_warning(self, warn_list):
 
     @skip_unsupported
     def test_warns(self):
-        try:
+        with override_env_config('NUMBA_WARNINGS', '1'):
             arr_ty = types.Array(types.float64, 2, "C")
-            os.environ['NUMBA_WARNINGS'] = '1'
             with warnings.catch_warnings(record=True) as w:
                 warnings.simplefilter("always", NumbaWarning)
                 cres = compile_isolated(unsupported_parfor, (arr_ty, arr_ty),
                                         flags=force_parallel_flags)
-
             self.check_parfors_warning(w)
-        finally:
-            del os.environ['NUMBA_WARNINGS']
 
 
 if __name__ == '__main__':
diff --git a/numba/tests/test_dispatcher.py b/numba/tests/test_dispatcher.py
--- a/numba/tests/test_dispatcher.py
+++ b/numba/tests/test_dispatcher.py
@@ -16,11 +16,13 @@
 from numba import config
 from numba import _dispatcher
 from numba.errors import NumbaWarning
-from .support import TestCase, tag, temp_directory, import_dynamic
+from .support import (TestCase, tag, temp_directory, import_dynamic,
+                      override_env_config, capture_cache_log)
 from numba.targets import codegen
 
 import llvmlite.binding as ll
 
+
 def dummy(x):
     return x
 
@@ -1047,11 +1049,9 @@ def test_user_set_cpu_name(self):
 
         mtimes = self.get_cache_mtimes()
         # Change CPU name to generic
-        try:
-            os.environ['NUMBA_CPU_NAME'] = 'generic'
+        with override_env_config('NUMBA_CPU_NAME', 'generic'):
             self.run_in_separate_process()
-        finally:
-            del os.environ['NUMBA_CPU_NAME']
+
         self.check_later_mtimes(mtimes)
         self.assertGreater(len(self.cache_contents()), cache_size)
         # Check cache index
@@ -1082,11 +1082,8 @@ def test_user_set_cpu_features(self):
         system_features = codegen.get_host_cpu_features()
 
         self.assertNotEqual(system_features, my_cpu_features)
-        try:
-            os.environ['NUMBA_CPU_FEATURES'] = my_cpu_features
+        with override_env_config('NUMBA_CPU_FEATURES', my_cpu_features):
             self.run_in_separate_process()
-        finally:
-            del os.environ['NUMBA_CPU_FEATURES']
         self.check_later_mtimes(mtimes)
         self.assertGreater(len(self.cache_contents()), cache_size)
         # Check cache index
@@ -1136,5 +1133,135 @@ def test_multiprocessing(self):
         self.assertEqual(res, n * (n - 1) // 2)
 
 
+class TestCacheFileCollision(unittest.TestCase):
+    _numba_parallel_test_ = False
+
+    here = os.path.dirname(__file__)
+    usecases_file = os.path.join(here, "cache_usecases.py")
+    modname = "caching_file_loc_fodder"
+    source_text_1 = """
+from numba import njit
+@njit(cache=True)
+def bar():
+    return 123
+"""
+    source_text_2 = """
+from numba import njit
+@njit(cache=True)
+def bar():
+    return 321
+"""
+
+    def setUp(self):
+        self.tempdir = temp_directory('test_cache_file_loc')
+        sys.path.insert(0, self.tempdir)
+        self.modname = 'module_name_that_is_unlikely'
+        self.assertNotIn(self.modname, sys.modules)
+        self.modname_bar1 = self.modname
+        self.modname_bar2 = '.'.join([self.modname, 'foo'])
+        foomod = os.path.join(self.tempdir, self.modname)
+        os.mkdir(foomod)
+        with open(os.path.join(foomod, '__init__.py'), 'w') as fout:
+            print(self.source_text_1, file=fout)
+        with open(os.path.join(foomod, 'foo.py'), 'w') as fout:
+            print(self.source_text_2, file=fout)
+
+    def tearDown(self):
+        sys.modules.pop(self.modname_bar1, None)
+        sys.modules.pop(self.modname_bar2, None)
+        sys.path.remove(self.tempdir)
+
+    def import_bar1(self):
+        return import_dynamic(self.modname_bar1).bar
+
+    def import_bar2(self):
+        return import_dynamic(self.modname_bar2).bar
+
+    def test_file_location(self):
+        bar1 = self.import_bar1()
+        bar2 = self.import_bar2()
+        # Check that the cache file is named correctly
+        idxname1 = bar1._cache._cache_file._index_name
+        idxname2 = bar2._cache._cache_file._index_name
+        self.assertNotEqual(idxname1, idxname2)
+        self.assertTrue(idxname1.startswith("__init__.bar-3.py"))
+        self.assertTrue(idxname2.startswith("foo.bar-3.py"))
+
+    @unittest.skipUnless(hasattr(multiprocessing, 'get_context'),
+                         'Test requires multiprocessing.get_context')
+    def test_no_collision(self):
+        bar1 = self.import_bar1()
+        bar2 = self.import_bar2()
+        with capture_cache_log() as buf:
+            res1 = bar1()
+        cachelog = buf.getvalue()
+        # bar1 should save new index and data
+        self.assertEqual(cachelog.count('index saved'), 1)
+        self.assertEqual(cachelog.count('data saved'), 1)
+        self.assertEqual(cachelog.count('index loaded'), 0)
+        self.assertEqual(cachelog.count('data loaded'), 0)
+        with capture_cache_log() as buf:
+            res2 = bar2()
+        cachelog = buf.getvalue()
+        # bar2 should save new index and data
+        self.assertEqual(cachelog.count('index saved'), 1)
+        self.assertEqual(cachelog.count('data saved'), 1)
+        self.assertEqual(cachelog.count('index loaded'), 0)
+        self.assertEqual(cachelog.count('data loaded'), 0)
+        self.assertNotEqual(res1, res2)
+
+        try:
+            # Make sure we can spawn new process without inheriting
+            # the parent context.
+            mp = multiprocessing.get_context('spawn')
+        except ValueError:
+            print("missing spawn context")
+
+        q = mp.Queue()
+        # Start new process that calls `cache_file_collision_tester`
+        proc = mp.Process(target=cache_file_collision_tester,
+                          args=(q, self.tempdir,
+                                self.modname_bar1,
+                                self.modname_bar2))
+        proc.start()
+        # Get results from the process
+        log1 = q.get()
+        got1 = q.get()
+        log2 = q.get()
+        got2 = q.get()
+        proc.join()
+
+        # The remote execution result of bar1() and bar2() should match
+        # the one executed locally.
+        self.assertEqual(got1, res1)
+        self.assertEqual(got2, res2)
+
+        # The remote should have loaded bar1 from cache
+        self.assertEqual(log1.count('index saved'), 0)
+        self.assertEqual(log1.count('data saved'), 0)
+        self.assertEqual(log1.count('index loaded'), 1)
+        self.assertEqual(log1.count('data loaded'), 1)
+
+        # The remote should have loaded bar2 from cache
+        self.assertEqual(log2.count('index saved'), 0)
+        self.assertEqual(log2.count('data saved'), 0)
+        self.assertEqual(log2.count('index loaded'), 1)
+        self.assertEqual(log2.count('data loaded'), 1)
+
+
+def cache_file_collision_tester(q, tempdir, modname_bar1, modname_bar2):
+    sys.path.insert(0, tempdir)
+    bar1 = import_dynamic(modname_bar1).bar
+    bar2 = import_dynamic(modname_bar2).bar
+    with capture_cache_log() as buf:
+        r1 = bar1()
+    q.put(buf.getvalue())
+    q.put(r1)
+    with capture_cache_log() as buf:
+        r2 = bar2()
+    q.put(buf.getvalue())
+    q.put(r2)
+
+
 if __name__ == '__main__':
     unittest.main()