Merge pull request #212 from pycompression/blocking_threads

Fix issue where threads would block python exit when an error is thrown and the file cannot be closed.

Author: rhpvorderman

CHANGELOG.rst

@@ -7,6 +7,11 @@ Changelog
 .. This document is user facing. Please word the changes in such a way
 .. that users understand how the changes affect the new version.
 
+version 1.7.1-dev
+-----------------
++ Prevent threaded opening from blocking python exit when an error is thrown
+  in the calling thread.
+
 version 1.7.0
 -----------------
 + Include a patched ISA-L version 2.31. The applied patches make compilation

src/isal/igzip_threaded.py

@@ -101,6 +101,7 @@ def __init__(self, filename, queue_size=2, block_size=1024 * 1024):
         self.worker = threading.Thread(target=self._decompress)
         self._closed = False
         self.running = True
+        self._calling_thread = threading.current_thread()
         self.worker.start()
 
     def _check_closed(self, msg=None):
@@ -110,15 +111,15 @@ def _check_closed(self, msg=None):
     def _decompress(self):
         block_size = self.block_size
         block_queue = self.queue
-        while self.running:
+        while self.running and self._calling_thread.is_alive():
             try:
                 data = self.fileobj.read(block_size)
             except Exception as e:
                 self.exception = e
                 return
             if not data:
                 return
-            while self.running:
+            while self.running and self._calling_thread.is_alive():
                 try:
                     block_queue.put(data, timeout=0.05)
                     break
@@ -215,6 +216,7 @@ def __init__(self,
         if "b" not in mode:
             mode += "b"
         self.lock = threading.Lock()
+        self._calling_thread = threading.current_thread()
         self.exception: Optional[Exception] = None
         self.level = level
         self.previous_block = b""
@@ -348,7 +350,7 @@ def _compress(self, index: int):
             try:
                 data, zdict = in_queue.get(timeout=0.05)
             except queue.Empty:
-                if not self.running:
+                if not (self.running and self._calling_thread.is_alive()):
                     return
                 continue
             try:
@@ -373,7 +375,7 @@ def _write(self):
             try:
                 compressed, crc, data_length = output_queue.get(timeout=0.05)
             except queue.Empty:
-                if not self.running:
+                if not (self.running and self._calling_thread.is_alive()):
                     self._crc = total_crc
                     self._size = size
                     return
@@ -396,7 +398,7 @@ def _compress_and_write(self):
             try:
                 data, zdict = in_queue.get(timeout=0.05)
             except queue.Empty:
-                if not self.running:
+                if not (self.running and self._calling_thread.is_alive()):
                     self._crc = total_crc
                     self._size = size
                     return

tests/test_igzip_threaded.py

@@ -9,6 +9,8 @@
 import io
 import itertools
 import os
+import subprocess
+import sys
 import tempfile
 from pathlib import Path
 
@@ -218,3 +220,22 @@ def test_threaded_writer_does_not_close_stream():
     assert not test_stream.closed
     test_stream.seek(0)
     assert gzip.decompress(test_stream.read()) == b"thisisatest"
+
+
+@pytest.mark.timeout(5)
+@pytest.mark.parametrize(
+    ["mode", "threads"], itertools.product(["rb", "wb"], [1, 2]))
+def test_threaded_program_can_exit_on_error(tmp_path, mode, threads):
+    program = tmp_path / "no_context_manager.py"
+    test_file = tmp_path / "output.gz"
+    # Write 40 mb input data to saturate read buffer. Because of the repetitive
+    # nature the resulting gzip file is very small (~40 KiB).
+    test_file.write_bytes(gzip.compress(b"test" * (10 * 1024 * 1024)))
+    with open(program, "wt") as f:
+        f.write("from isal import igzip_threaded\n")
+        f.write(
+            f"f = igzip_threaded.open('{test_file}', "
+            f"mode='{mode}', threads={threads})\n"
+        )
+        f.write("raise Exception('Error')\n")
+    subprocess.run([sys.executable, str(program)])