Harden pickle temp-file cleanup and lock fallback

Author: shaypal5

src/cachier/cores/pickle.py

@@ -107,10 +107,18 @@ def cache_fpath(self) -> str:
 
     @property
     def _shared_lock_fpath(self) -> str:
-        lock_dir = os.path.join(tempfile.gettempdir(), "cachier-locks")
-        os.makedirs(lock_dir, exist_ok=True)
         cache_hash = hashlib.sha256(self.cache_fpath.encode("utf-8")).hexdigest()
-        return os.path.join(lock_dir, f"{cache_hash}{self._SHARED_LOCK_SUFFIX}")
+        candidate_dirs = (
+            os.path.join(tempfile.gettempdir(), "cachier-locks"),
+            os.path.join(os.path.dirname(self.cache_fpath), ".cachier-locks"),
+        )
+        for lock_dir in candidate_dirs:
+            try:
+                os.makedirs(lock_dir, exist_ok=True)
+                return os.path.join(lock_dir, f"{cache_hash}{self._SHARED_LOCK_SUFFIX}")
+            except OSError:
+                continue
+        return os.path.join(os.path.dirname(self.cache_fpath), f".{cache_hash}{self._SHARED_LOCK_SUFFIX}")
 
     @staticmethod
     def _convert_legacy_cache_entry(
@@ -203,20 +211,22 @@ def _save_cache(
                     pickle.dump(cache, cast(IO[bytes], cache_file), protocol=4)
             else:
                 with portalocker.Lock(self._shared_lock_fpath, mode="a+b"):
-                    with tempfile.NamedTemporaryFile(
-                        mode="wb",
-                        dir=parent_dir,
-                        delete=False,
-                    ) as temp_file:
-                        temp_path = temp_file.name
-                        pickle.dump(cache, cast(IO[bytes], temp_file), protocol=4)
-                        temp_file.flush()
-                        os.fsync(temp_file.fileno())
+                    temp_path = ""
                     try:
+                        with tempfile.NamedTemporaryFile(
+                            mode="wb",
+                            dir=parent_dir,
+                            delete=False,
+                        ) as temp_file:
+                            temp_path = temp_file.name
+                            pickle.dump(cache, cast(IO[bytes], temp_file), protocol=4)
+                            temp_file.flush()
+                            os.fsync(temp_file.fileno())
                         os.replace(temp_path, fpath)
                     finally:
-                        with suppress(FileNotFoundError):
-                            os.remove(temp_path)
+                        if temp_path:
+                            with suppress(FileNotFoundError):
+                                os.remove(temp_path)
             # the same as check for separate_file, but changed for typing
             if isinstance(cache, dict):
                 self._cache_dict = cache

tests/pickle_tests/test_pickle_core.py

@@ -709,6 +709,130 @@ def mock_func():
         core._save_cache({"key": "value"}, separate_file_key="test_key")
 
 
+@pytest.mark.pickle
+def test_save_cache_removes_temp_file_when_fsync_fails(tmp_path):
+    """Test _save_cache removes the temp file when fsync fails."""
+    core = _PickleCore(
+        hash_func=None,
+        cache_dir=tmp_path,
+        pickle_reload=False,
+        wait_for_calc_timeout=10,
+        separate_files=False,
+    )
+
+    def mock_func():
+        pass
+
+    core.set_func(mock_func)
+
+    with patch("cachier.cores.pickle.os.fsync", side_effect=OSError("fsync failed")), pytest.raises(
+        OSError, match="fsync failed"
+    ):
+        core._save_cache(
+            {"key": CacheEntry(value="value", time=datetime.now(), stale=False, _processing=False)}
+        )
+
+    assert list(tmp_path.iterdir()) == []
+
+
+@pytest.mark.pickle
+def test_save_cache_propagates_tempfile_creation_failure_without_cleanup_error(tmp_path):
+    """Test _save_cache handles temp-file creation failures before temp_path exists."""
+    core = _PickleCore(
+        hash_func=None,
+        cache_dir=tmp_path,
+        pickle_reload=False,
+        wait_for_calc_timeout=10,
+        separate_files=False,
+    )
+
+    def mock_func():
+        pass
+
+    core.set_func(mock_func)
+
+    with patch("cachier.cores.pickle.tempfile.NamedTemporaryFile", side_effect=OSError("tempfile failed")), patch(
+        "cachier.cores.pickle.os.replace"
+    ) as mock_replace, patch("cachier.cores.pickle.os.remove") as mock_remove, pytest.raises(
+        OSError, match="tempfile failed"
+    ):
+        core._save_cache(
+            {"key": CacheEntry(value="value", time=datetime.now(), stale=False, _processing=False)}
+        )
+
+    mock_replace.assert_not_called()
+    mock_remove.assert_not_called()
+    assert list(tmp_path.iterdir()) == []
+
+
+@pytest.mark.pickle
+def test_shared_lock_fpath_falls_back_to_cache_dir_when_temp_dir_unwritable(tmp_path):
+    """Test _shared_lock_fpath falls back when the system temp dir is not writable."""
+    core = _PickleCore(
+        hash_func=None,
+        cache_dir=tmp_path,
+        pickle_reload=False,
+        wait_for_calc_timeout=10,
+        separate_files=False,
+    )
+
+    def mock_func():
+        pass
+
+    core.set_func(mock_func)
+
+    temp_lock_dir = os.path.join("/non-writable-temp", "cachier-locks")
+    fallback_lock_dir = os.path.join(core.cache_dir, ".cachier-locks")
+
+    def mock_makedirs(path, exist_ok=False):
+        if path in (core.cache_dir, fallback_lock_dir):
+            return None
+        if path == temp_lock_dir:
+            raise PermissionError("temp dir not writable")
+        raise AssertionError(f"Unexpected os.makedirs path: {path}")
+
+    with patch("cachier.cores.pickle.tempfile.gettempdir", return_value="/non-writable-temp"), patch(
+        "cachier.cores.pickle.os.makedirs", side_effect=mock_makedirs
+    ):
+        assert core._shared_lock_fpath == os.path.join(
+            fallback_lock_dir,
+            f"{hashlib.sha256(core.cache_fpath.encode('utf-8')).hexdigest()}{core._SHARED_LOCK_SUFFIX}",
+        )
+
+
+@pytest.mark.pickle
+def test_shared_lock_fpath_uses_cache_dir_file_when_lock_dirs_unwritable(tmp_path):
+    """Test _shared_lock_fpath falls back to a lockfile in the cache dir."""
+    core = _PickleCore(
+        hash_func=None,
+        cache_dir=tmp_path,
+        pickle_reload=False,
+        wait_for_calc_timeout=10,
+        separate_files=False,
+    )
+
+    def mock_func():
+        pass
+
+    core.set_func(mock_func)
+
+    temp_lock_dir = os.path.join("/non-writable-temp", "cachier-locks")
+    fallback_lock_dir = os.path.join(core.cache_dir, ".cachier-locks")
+    cache_hash = hashlib.sha256(core.cache_fpath.encode("utf-8")).hexdigest()
+
+    def mock_makedirs(path, exist_ok=False):
+        if path == core.cache_dir:
+            return None
+        if path in (temp_lock_dir, fallback_lock_dir):
+            raise PermissionError("lock dir not writable")
+        raise AssertionError(f"Unexpected os.makedirs path: {path}")
+
+    with patch("cachier.cores.pickle.tempfile.gettempdir", return_value="/non-writable-temp"), patch(
+        "cachier.cores.pickle.os.makedirs", side_effect=mock_makedirs
+    ):
+        assert core._shared_lock_fpath == os.path.join(core.cache_dir, f".{cache_hash}{core._SHARED_LOCK_SUFFIX}")
+
+
 @pytest.mark.pickle
 def test_set_entry_should_not_store(tmp_path):
     """Test set_entry when value should not be stored."""