Close file handles correctly on exceptions when unpacking strata

This prevents us from leaving file handles open when code throws an
exception. When the file handle is on a loopback mount, this is a real
problem because Morph then cannot unmount the image in its exception
handler. In most cases 'with' is the best option.

2 files changed, 25 insertions, 23 deletions

diff --git a/morphlib/bins.py b/morphlib/bins.py
index 71483172..95566f39 100644
--- a/morphlib/bins.py
+++ b/morphlib/bins.py
@@ -187,11 +187,12 @@ def unpack_binary_from_file(f, dirname):  # pragma: no cover
     tf.makedev = monkey_patcher(tf.makedev)
     tf.makelink = monkey_patcher(tf.makelink)
 
-    tf.extractall(path=dirname)
-    tf.close
+    try:
+        tf.extractall(path=dirname)
+    finally:
+        tf.close()
 
 
 def unpack_binary(filename, dirname):
-    f = open(filename, "rb")
-    unpack_binary_from_file(f, dirname)
-    f.close()
+    with open(filename, "rb") as f:
+        unpack_binary_from_file(f, dirname)
diff --git a/morphlib/builder2.py b/morphlib/builder2.py
index d829474c..d8ed829a 100644
--- a/morphlib/builder2.py
+++ b/morphlib/builder2.py
@@ -488,6 +488,24 @@ class SystemKindBuilder(BuilderBase):  # pragma: no cover
 
     '''
 
+    def unpack_one_stratum(self, stratum_artifact, target):
+        '''Unpack a single stratum into a target directory'''
+
+        cache = self.local_artifact_cache
+        with cache.get(stratum_artifact) as stratum_file:
+            artifact_list = json.load(stratum_file)
+            for chunk in (ArtifactCacheReference(a) for a in artifact_list):
+                self.app.status(msg='Unpacking chunk %(basename)s',
+                                basename=chunk.basename(), chatty=True)
+                with cache.get(chunk) as chunk_file:
+                    morphlib.bins.unpack_binary_from_file(chunk_file, target)
+
+        target_metadata = os.path.join(
+                target, 'baserock', '%s.meta' % stratum_artifact.name)
+        with cache.get_artifact_metadata(stratum_artifact, 'meta') as meta_src:
+            with morphlib.savefile.SaveFile(target_metadata, 'w') as meta_dst:
+                shutil.copyfileobj(meta_src, meta_dst)
+
     def unpack_strata(self, path):
         '''Unpack strata into a directory.'''
 
@@ -523,22 +541,7 @@ class SystemKindBuilder(BuilderBase):  # pragma: no cover
 
             # unpack it from the local artifact cache
             for stratum_artifact in self.artifact.dependencies:
-                f = self.local_artifact_cache.get(stratum_artifact)
-                for chunk in (ArtifactCacheReference(a) for a in json.load(f)):
-                    self.app.status(msg='Unpacking chunk %(basename)s',
-                                    basename=chunk.basename(), chatty=True)
-                    chunk_handle = self.local_artifact_cache.get(chunk)
-                    morphlib.bins.unpack_binary_from_file(chunk_handle, path)
-                    chunk_handle.close()
-                f.close()
-                meta = self.local_artifact_cache.get_artifact_metadata(
-                    stratum_artifact, 'meta')
-                dst = morphlib.savefile.SaveFile(
-                    os.path.join(path, 'baserock',
-                                 '%s.meta' % stratum_artifact.name), 'w')
-                shutil.copyfileobj(meta, dst)
-                dst.close()
-                meta.close()
+                self.unpack_one_stratum(stratum_artifact, path)
 
             ldconfig(self.app.runcmd, path)
 
@@ -819,5 +822,3 @@ class DiskImageBuilder(SystemKindBuilder):  # pragma: no cover
                         filename=image_name, chatty=True)
         with self.build_watch('undo-device-mapper'):
             morphlib.fsutils.undo_device_mapping(self.app.runcmd, image_name)
-
-