Close callback buffers even when caching fails

cachecontrol/filewrapper.py

@@ -71,31 +71,33 @@ def __is_fp_closed(self) -> bool:
 
     def _close(self) -> None:
         result: Buffer
-        if self.__callback:
-            if self.__buf.tell() == 0:
-                # Empty file:
-                result = b""
-            else:
-                # Return the data without actually loading it into memory,
-                # relying on Python's buffer API and mmap(). mmap() just gives
-                # a view directly into the filesystem's memory cache, so it
-                # doesn't result in duplicate memory use.
-                self.__buf.seek(0, 0)
-                result = memoryview(
-                    mmap.mmap(self.__buf.fileno(), 0, access=mmap.ACCESS_READ)
-                )
-            self.__callback(result)
-
-        # We assign this to None here, because otherwise we can get into
-        # really tricky problems where the CPython interpreter dead locks
-        # because the callback is holding a reference to something which
-        # has a __del__ method. Setting this to None breaks the cycle
-        # and allows the garbage collector to do it's thing normally.
-        self.__callback = None
-
-        # Closing the temporary file releases memory and frees disk space.
-        # Important when caching big files.
-        self.__buf.close()
+        try:
+            if self.__callback:
+                if self.__buf.tell() == 0:
+                    # Empty file:
+                    result = b""
+                else:
+                    # Return the data without actually loading it into memory,
+                    # relying on Python's buffer API and mmap(). mmap() just gives
+                    # a view directly into the filesystem's memory cache, so it
+                    # doesn't result in duplicate memory use.
+                    self.__buf.seek(0, 0)
+                    result = memoryview(
+                        mmap.mmap(self.__buf.fileno(), 0, access=mmap.ACCESS_READ)
+                    )
+                self.__callback(result)
+
+        finally:
+            # We assign this to None here, because otherwise we can get into
+            # really tricky problems where the CPython interpreter dead locks
+            # because the callback is holding a reference to something which
+            # has a __del__ method. Setting this to None breaks the cycle
+            # and allows the garbage collector to do it's thing normally.
+            self.__callback = None
+
+            # Close the temporary file even if preparing or invoking the
+            # callback failed.
+            self.__buf.close()
 
     def read(self, amt: int | None = None) -> bytes:
         data: bytes = self.__fp.read(amt)

tests/test_filewrapper.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: 2015 Eric Larson
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import io
+from http.client import HTTPResponse
+from tempfile import NamedTemporaryFile
+from unittest.mock import Mock
+
+import pytest
+
+from cachecontrol.filewrapper import CallbackFileWrapper
+
+
+@pytest.mark.parametrize("body", [b"", b"response body"])
+@pytest.mark.parametrize("error_type", [None, OSError, KeyboardInterrupt])
+def test_close_buffer(monkeypatch, tmp_path, body, error_type):
+    buffers = []
+
+    def temporary_file(*args, **kwargs):
+        buffer = NamedTemporaryFile(*args, dir=tmp_path, **kwargs)
+        buffers.append(buffer)
+        return buffer
+
+    monkeypatch.setattr("cachecontrol.filewrapper.NamedTemporaryFile", temporary_file)
+    wire = b"HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n" % len(body) + body
+    response = HTTPResponse(Mock(makefile=Mock(return_value=io.BytesIO(wire))))
+    response.begin()
+    error = error_type("cache write failed") if error_type else None
+    callback = Mock(side_effect=error)
+    wrapper = CallbackFileWrapper(response, callback)
+
+    if error_type is None:
+        assert wrapper.read() == body
+    else:
+        with pytest.raises(error_type) as raised:
+            wrapper.read()
+        assert raised.value is error
+    assert buffers[0].closed
+    assert not list(tmp_path.iterdir())
+    assert bytes(callback.call_args.args[0]) == body
+    assert wrapper.read() == b""
+    callback.assert_called_once()