lock once for metadata and body in separate body cache

cachecontrol/cache.py

@@ -57,14 +57,37 @@ class SeparateBodyBaseCache(BaseCache):
     In this variant, the body is not stored mixed in with the metadata, but is
     passed in (as a bytes-like object) in a separate call to ``set_body()``.
 
-    That is, the expected interaction pattern is::
+    Cache implementations can override ``set_with_body()`` and
+    ``get_with_body()`` when the two operations need to be synchronized. The
+    default implementations retain the separate calls for compatibility.
+
+    The low-level interaction pattern is::
 
         cache.set(key, serialized_metadata)
         cache.set_body(key)
 
     Similarly, the body should be loaded separately via ``get_body()``.
     """
 
+    def get_with_body(self, key: str) -> tuple[bytes | None, IO[bytes] | None]:
+        """Return an entry's metadata and body together."""
+        metadata = self.get(key)
+        if metadata is None:
+            return None, None
+        return metadata, self.get_body(key)
+
+    def set_with_body(
+        self,
+        key: str,
+        metadata: bytes,
+        body: bytes | None,
+        expires: int | datetime | None = None,
+    ) -> None:
+        """Store metadata and replace the body when one is provided."""
+        self.set(key, metadata, expires)
+        if body is not None:
+            self.set_body(key, body)
+
     def set_body(self, key: str, body: bytes) -> None:
         raise NotImplementedError()
 

cachecontrol/caches/file_cache.py

@@ -86,14 +86,17 @@ def _write(self, path: str, data: bytes) -> None:
         os.makedirs(dirname, self.dirmode, exist_ok=True)
 
         with self.lock_class(path + ".lock"):
-            # Write our actual file
-            (fd, name) = tempfile.mkstemp(dir=dirname)
-            try:
-                os.write(fd, data)
-            finally:
-                os.close(fd)
-            os.chmod(name, self.filemode)
-            os.replace(name, path)
+            self._write_unlocked(path, data)
+
+    def _write_unlocked(self, path: str, data: bytes) -> None:
+        """Atomically replace a file while its lock is held."""
+        (fd, name) = tempfile.mkstemp(dir=os.path.dirname(path))
+        try:
+            os.write(fd, data)
+        finally:
+            os.close(fd)
+        os.chmod(name, self.filemode)
+        os.replace(name, path)
 
     def _delete(self, key: str, suffix: str) -> None:
         name = self._fn(key) + suffix
@@ -120,6 +123,25 @@ class SeparateBodyFileCache(_FileCacheMixin, SeparateBodyBaseCache):
     peak memory usage.
     """
 
+    def get_with_body(self, key: str) -> tuple[bytes | None, IO[bytes] | None]:
+        name = self._fn(key)
+        with self.lock_class(name + ".lock"):
+            return super().get_with_body(key)
+
+    def set_with_body(
+        self,
+        key: str,
+        metadata: bytes,
+        body: bytes | None,
+        expires: int | datetime | None = None,
+    ) -> None:
+        name = self._fn(key)
+        os.makedirs(os.path.dirname(name), self.dirmode, exist_ok=True)
+        with self.lock_class(name + ".lock"):
+            self._write_unlocked(name, metadata)
+            if body is not None:
+                self._write_unlocked(name + ".body", body)
+
     def get_body(self, key: str) -> IO[bytes] | None:
         name = self._fn(key) + ".body"
         try:
@@ -128,12 +150,18 @@ def get_body(self, key: str) -> IO[bytes] | None:
             return None
 
     def set_body(self, key: str, body: bytes) -> None:
-        name = self._fn(key) + ".body"
-        self._write(name, body)
+        name = self._fn(key)
+        os.makedirs(os.path.dirname(name), self.dirmode, exist_ok=True)
+        with self.lock_class(name + ".lock"):
+            self._write_unlocked(name + ".body", body)
 
     def delete(self, key: str) -> None:
-        self._delete(key, "")
-        self._delete(key, ".body")
+        if self.forever:
+            return
+        name = self._fn(key)
+        with self.lock_class(name + ".lock"):
+            self._delete(key, "")
+            self._delete(key, ".body")
 
 
 def url_to_file_path(url: str, filecache: FileCache) -> str:

cachecontrol/controller.py

@@ -151,15 +151,14 @@ def _load_from_cache(self, request: PreparedRequest) -> HTTPResponse | None:
 
         cache_url = request.url
         assert cache_url is not None
-        cache_data = self.cache.get(cache_url)
-        if cache_data is None:
-            logger.debug("No cache entry available")
-            return None
-
         if isinstance(self.cache, SeparateBodyBaseCache):
-            body_file = self.cache.get_body(cache_url)
+            cache_data, body_file = self.cache.get_with_body(cache_url)
         else:
+            cache_data = self.cache.get(cache_url)
             body_file = None
+        if cache_data is None:
+            logger.debug("No cache entry available")
+            return None
 
         result = self.serializer.loads(request, cache_data, body_file)
         if result is None:
@@ -305,17 +304,12 @@ def _cache_set(
         Store the data in the cache.
         """
         if isinstance(self.cache, SeparateBodyBaseCache):
-            # We pass in the body separately; just put a placeholder empty
-            # string in the metadata.
-            self.cache.set(
+            self.cache.set_with_body(
                 cache_url,
                 self.serializer.dumps(request, response, b""),
+                body,
                 expires=expires_time,
             )
-            # body is None can happen when, for example, we're only updating
-            # headers, as is the case in update_cached_response().
-            if body is not None:
-                self.cache.set_body(cache_url, body)
         else:
             self.cache.set(
                 cache_url,

tests/test_cache_control.py

@@ -8,7 +8,7 @@
 
 import os
 import time
-from unittest.mock import ANY, Mock
+from unittest.mock import ANY, Mock, patch
 
 import pytest
 
@@ -127,6 +127,28 @@ def test_no_cache_with_vary_star(self, cc):
 
         assert not cc.cache.set.called
 
+    def test_separate_body_cache_uses_paired_operations(self, tmp_path):
+        cache = SeparateBodyFileCache(os.fsdecode(tmp_path))
+        serializer = Mock()
+        serializer.dumps.return_value = b"metadata"
+        controller = CacheController(cache, serializer=serializer)
+        request = self.req()
+        response = self.resp()
+
+        with (
+            patch.object(cache, "set_with_body") as set_with_body,
+            patch.object(
+                cache, "get_with_body", return_value=(b"metadata", None)
+            ) as get_with_body,
+        ):
+            controller._cache_set(self.url, request, response, b"body", 60)
+            controller._load_from_cache(request)
+
+        set_with_body.assert_called_once_with(
+            self.url, b"metadata", b"body", expires=60
+        )
+        get_with_body.assert_called_once_with(self.url)
+
     def test_update_cached_response_no_local_cache(self):
         """
         If the local cache doesn't have the given URL, just reuse the response

tests/test_storage_filecache.py

@@ -8,8 +8,8 @@
 
 import os
 import string
-
 from random import randint, sample
+from unittest.mock import MagicMock, call, patch
 
 import pytest
 import requests
@@ -148,6 +148,48 @@ class TestSeparateBodyFileCache(FileCacheTestsMixin):
 
     FileCacheClass = SeparateBodyFileCache
 
+    def test_entry_write_uses_one_lock(self, sess):
+        operations = MagicMock()
+        cache = self.cache
+        key = self.url
+        name = cache._fn(key)
+
+        with (
+            patch.object(cache, "lock_class", operations.lock),
+            patch.object(cache, "_write_unlocked", operations.write),
+        ):
+            cache.set_with_body(key, b"metadata", b"body")
+
+        assert operations.mock_calls == [
+            call.lock(name + ".lock"),
+            call.lock().__enter__(),
+            call.write(name, b"metadata"),
+            call.write(name + ".body", b"body"),
+            call.lock().__exit__(None, None, None),
+        ]
+
+    def test_entry_read_uses_one_lock(self, sess):
+        operations = MagicMock()
+        cache = self.cache
+        key = self.url
+        name = cache._fn(key)
+        operations.get.return_value = b"metadata"
+        operations.get_body.return_value = None
+        with (
+            patch.object(cache, "lock_class", operations.lock),
+            patch.object(cache, "get", operations.get),
+            patch.object(cache, "get_body", operations.get_body),
+        ):
+            cache.get_with_body(key)
+
+        assert operations.mock_calls == [
+            call.lock(name + ".lock"),
+            call.lock().__enter__(),
+            call.get(key),
+            call.get_body(key),
+            call.lock().__exit__(None, None, None),
+        ]
+
     def test_body_actually_stored_separately(self, sess):
         """
         Body is stored and can be retrieved from the SeparateBodyFileCache, with assurances