Clean up multiprocessing pools on failure

src/bagit/__init__.py

@@ -896,12 +896,12 @@
             if processes == 1:
                 hash_results = [_calc_hashes(i) for i in args]
             else:
-                pool = multiprocessing.Pool(
-                    processes if processes else None, initializer=worker_init
+                hash_results = _multiprocessing_pool_map(
+                    _calc_hashes,
+                    args,
+                    processes if processes else None,
+                    initializer=worker_init,
                 )
-                hash_results = pool.map(_calc_hashes, args)
-                pool.close()
-                pool.join()
 
         # Any unhandled exceptions are probably fatal
         except:
@@ -1037,6 +1037,25 @@
     signal.signal(signal.SIGINT, signal.SIG_IGN)
 
 
+def _multiprocessing_pool_map(func, iterable, processes, initializer=None):
+    """Run ``Pool.map()`` and always clean up the pool.
+
+    This ensures worker processes are closed or terminated, then joined, under
+    all conditions.
+    """
+    pool = multiprocessing.Pool(processes=processes, initializer=initializer)
+    try:
+        results = pool.map(func, iterable)
+    except BaseException:
+        pool.terminate()
+        raise
+    else:
+        pool.close()
+        return results
+    finally:
+        pool.join()
+
+
 # The Unicode normalization form used here doesn't matter – all we care about
 # is consistency since the input value will be preserved:
 
@@ -1245,10 +1264,9 @@
     manifest_line_generator = partial(generate_manifest_lines, algorithms=algorithms)
 
     if processes > 1:
-        pool = multiprocessing.Pool(processes=processes)
-        checksums = pool.map(manifest_line_generator, _walk(data_dir))
-        pool.close()
-        pool.join()
+        checksums = _multiprocessing_pool_map(
+            manifest_line_generator, _walk(data_dir), processes=processes
+        )
     else:
         checksums = [manifest_line_generator(i) for i in _walk(data_dir)]
 
@@ -1583,7 +1601,7 @@
                 else:
                     LOGGER.info(_("%s is valid"), bag_dir)
             except BagError as e:
                 LOGGER.error(
                     _("%(bag)s is invalid: %(error)s"), {"bag": bag_dir, "error": e}
                 )
                 rc = 1
@@ -1598,7 +1616,7 @@
                     checksums=args.checksums,
                 )
             except Exception as exc:
                 LOGGER.error(
                     _("Failed to create bag in %(bag_directory)s: %(error)s"),
                     {"bag_directory": bag_dir, "error": exc},
                     exc_info=True,

test.py

@@ -13,10 +13,9 @@
 import tempfile
 import unicodedata
 import unittest
+from io import StringIO
 from os.path import join as j
-
 from unittest import mock
-from io import StringIO
 
 import bagit
 
@@ -458,6 +457,23 @@ def validate(self, bag, *args, **kwargs):
             bag, *args, processes=2, **kwargs
         )
 
+    @mock.patch("bagit.multiprocessing.Pool")
+    def test_validate_multiprocessing_terminates_and_joins_pool_on_failure(self, pool):
+        pool.return_value.map.side_effect = RuntimeError("boom")
+        bag = bagit.make_bag(self.tmpdir)
+
+        with self.assertRaises(RuntimeError):
+            self.validate(bag)
+
+        self.assertEqual(
+            pool.return_value.mock_calls,
+            [
+                mock.call.map(mock.ANY, mock.ANY),
+                mock.call.terminate(),
+                mock.call.join(),
+            ],
+        )
+
     @mock.patch("bagit.multiprocessing.Pool")
     def test_validate_pool_error(self, pool):
         # Simulate the Pool constructor raising a RuntimeError.
@@ -745,6 +761,21 @@ def test_make_bag_multiprocessing(self):
         bagit.make_bag(self.tmpdir, processes=2)
         self.assertTrue(os.path.isdir(j(self.tmpdir, "data")))
 
+    @mock.patch("bagit.multiprocessing.Pool")
+    def test_make_bag_multiprocessing_terminates_and_joins_pool_on_failure(self, pool):
+        pool.return_value.map.side_effect = RuntimeError("boom")
+        with self.assertRaises(RuntimeError):
+            bagit.make_bag(self.tmpdir, processes=2)
+
+        self.assertEqual(
+            pool.return_value.mock_calls,
+            [
+                mock.call.map(mock.ANY, mock.ANY),
+                mock.call.terminate(),
+                mock.call.join(),
+            ],
+        )
+
     def test_multiple_meta_values(self):
         baginfo = {"Multival-Meta": [7, 4, 8, 6, 8]}
         bag = bagit.make_bag(self.tmpdir, baginfo)