Ensure errors are propagated to the main thread

Author: rhpvorderman

src/isal/igzip_threaded.py

@@ -150,6 +150,8 @@ def __init__(self,
                  level: int = isal_zlib.ISAL_DEFAULT_COMPRESSION,
                  threads: int = 1,
                  queue_size: int = 2):
+        self.lock = threading.Lock()
+        self.exception = None
         self.raw = fp
         self.level = level
         self.previous_block = b""
@@ -189,14 +191,17 @@ def start(self):
         for worker in self.compression_workers:
             worker.start()
 
-    def stop_immediately(self):
+    def stop(self):
         """Stop, but do not care for remaining work"""
         self.running = False
-        self.output_worker.join()
         for worker in self.compression_workers:
             worker.join()
+        self.output_worker.join()
 
     def write(self, b) -> int:
+        with self.lock:
+            if self.exception:
+                raise self.exception
         if self._closed:
             raise IOError("Can not write closed file")
         index = self.index
@@ -221,7 +226,9 @@ def flush(self):
 
     def close(self) -> None:
         self.flush()
-        self.stop_immediately()
+        self.stop()
+        if self.exception:
+            raise self.exception
         # Write an empty deflate block with a lost block marker.
         self.raw.write(isal_zlib.compress(b"", wbits=-15))
         trailer = struct.pack("<II", self._crc, self._size & 0xFFFFFFFF)
@@ -237,16 +244,31 @@ def closed(self) -> bool:
     def _compress(self, index: int):
         in_queue = self.input_queues[index]
         out_queue = self.output_queues[index]
-        while self.running:
+        while True:
             try:
                 data, zdict = in_queue.get(timeout=0.05)
             except queue.Empty:
+                if not self.running:
+                    return
                 continue
-            compressor = isal_zlib.compressobj(
-                self.level, wbits=-15, zdict=zdict)
-            compressed = compressor.compress(data) + compressor.flush(
-                isal_zlib.Z_SYNC_FLUSH)
-            crc = isal_zlib.crc32(data)
+            try:
+                compressor = isal_zlib.compressobj(
+                    self.level, wbits=-15, zdict=zdict)
+                compressed = compressor.compress(data) + compressor.flush(
+                    isal_zlib.Z_SYNC_FLUSH)
+                crc = isal_zlib.crc32(data)
+            except Exception as e:
+                with self.lock:
+                    self.exception = e
+                    # Abort everything and empty the queue
+                    in_queue.task_done()
+                    self.running = False
+                    while True:
+                        try:
+                            _ = in_queue.get(timeout=0.05)
+                            in_queue.task_done()
+                        except queue.Empty:
+                            return
             data_length = len(data)
             out_queue.put((compressed, crc, data_length))
             in_queue.task_done()
@@ -257,20 +279,22 @@ def _write(self):
         fp = self.raw
         total_crc = 0
         size = 0
-        while self.running:
+        while True:
             out_index = index % self.threads
             output_queue = output_queues[out_index]
             try:
                 compressed, crc, data_length = output_queue.get(timeout=0.05)
             except queue.Empty:
+                if not self.running:
+                    self._crc = total_crc
+                    self._size = size
+                    return
                 continue
             total_crc = isal_zlib.crc32_combine(total_crc, crc, data_length)
             size += data_length
             fp.write(compressed)
             output_queue.task_done()
             index += 1
-        self._crc = total_crc
-        self._size = size
 
     def writable(self) -> bool:
         return True