Create chunks, strata in cache via temporary files

This avoids problems with files with the right names but partial
content, if morph is killed in the middle of writing the file.

9 files changed, 235 insertions, 22 deletions

diff --git a/morphlib/__init__.py b/morphlib/__init__.py
index 11b6cc9a..7c2f66bb 100644
--- a/morphlib/__init__.py
+++ b/morphlib/__init__.py
@@ -31,8 +31,10 @@ import execute
 import git
 import morphology
 import morphologyloader
+import savefile
 import sourcemanager
 import stopwatch
 import tempdir
 import tester
 import util
+
diff --git a/morphlib/bins.py b/morphlib/bins.py
index b4b5396e..0d704e49 100644
--- a/morphlib/bins.py
+++ b/morphlib/bins.py
@@ -27,8 +27,7 @@ import re
 import tarfile
 
 
-def create_chunk(rootdir, chunk_filename, regexps, ex,
-                 dump_memory_profile=None):
+def create_chunk(rootdir, f, regexps, ex, dump_memory_profile=None):
     '''Create a chunk from the contents of a directory.
     
     Only files and directories that match at least one of the regular
@@ -36,6 +35,8 @@ def create_chunk(rootdir, chunk_filename, regexps, ex,
     anchored to the beginning of the string, but not the end. The 
     filenames are relative to rootdir.
     
+    ``f`` is an open file handle, to which the tar file is written.
+    
     '''
 
     dump_memory_profile = dump_memory_profile or (lambda msg: None )
@@ -57,7 +58,7 @@ def create_chunk(rootdir, chunk_filename, regexps, ex,
             yield filename
 
     logging.debug('Creating chunk file %s from %s with regexps %s' % 
-                    (chunk_filename, rootdir, regexps))
+                    (f.name, rootdir, regexps))
     dump_memory_profile('at beginning of create_chunk')
 
     compiled = [re.compile(x) for x in regexps]
@@ -79,7 +80,7 @@ def create_chunk(rootdir, chunk_filename, regexps, ex,
     dump_memory_profile('after walking')
 
     include = sorted(include) # get dirs before contents
-    tar = tarfile.open(name=chunk_filename, mode='w:gz')
+    tar = tarfile.open(fileobj=f, mode='w:gz')
     for filename in include:
         tar.add(filename, arcname=mkrel(filename), recursive=False)
     tar.close()
@@ -94,11 +95,12 @@ def create_chunk(rootdir, chunk_filename, regexps, ex,
     dump_memory_profile('after removing in create_chunks')
 
 
-def create_stratum(rootdir, stratum_filename, ex):
+def create_stratum(rootdir, f, ex):
     '''Create a stratum from the contents of a directory.'''
-    logging.debug('Creating stratum file %s from %s' % 
-                    (stratum_filename, rootdir))
-    ex.runv(['tar', '-C', rootdir, '-caf', stratum_filename, '.'])
+    logging.debug('Creating stratum file %s from %s' % (f.name, rootdir))
+    tar = tarfile.open(fileobj=f, mode='w:gz')
+    tar.add(rootdir, arcname='.')
+    tar.close()
 
 
 def unpack_binary(filename, dirname, ex):
diff --git a/morphlib/bins_tests.py b/morphlib/bins_tests.py
index 86b55746..90c8cc0a 100644
--- a/morphlib/bins_tests.py
+++ b/morphlib/bins_tests.py
@@ -66,9 +66,11 @@ class ChunkTests(BinsTest):
         self.tempdir = tempfile.mkdtemp()
         self.instdir = os.path.join(self.tempdir, 'inst')
         self.chunk_file = os.path.join(self.tempdir, 'chunk')
+        self.chunk_f = open(self.chunk_file, 'wb')
         self.unpacked = os.path.join(self.tempdir, 'unpacked')
         
     def tearDown(self):
+        self.chunk_f.close()
         shutil.rmtree(self.tempdir)
 
     def populate_instdir(self):
@@ -92,8 +94,8 @@ class ChunkTests(BinsTest):
         
     def test_empties_everything(self):
         self.populate_instdir()
-        morphlib.bins.create_chunk(self.instdir, self.chunk_file, ['.'],
-                                   self.ex)
+        morphlib.bins.create_chunk(self.instdir, self.chunk_f, ['.'], self.ex)
+        self.chunk_f.flush()
         empty = os.path.join(self.tempdir, 'empty')
         os.mkdir(empty)
         self.assertEqual([x for x,y in self.recursive_lstat(self.instdir)],
@@ -102,16 +104,17 @@ class ChunkTests(BinsTest):
     def test_creates_and_unpacks_chunk_exactly(self):
         self.populate_instdir()
         orig_files = self.recursive_lstat(self.instdir)
-        morphlib.bins.create_chunk(self.instdir, self.chunk_file, ['.'],
-                                   self.ex)
+        morphlib.bins.create_chunk(self.instdir, self.chunk_f, ['.'], self.ex)
+        self.chunk_f.flush()
         os.mkdir(self.unpacked)
         morphlib.bins.unpack_binary(self.chunk_file, self.unpacked, self.ex)
         self.assertEqual(orig_files, self.recursive_lstat(self.unpacked))
 
     def test_uses_only_matching_names(self):
         self.populate_instdir()
-        morphlib.bins.create_chunk(self.instdir, self.chunk_file, ['bin'],
+        morphlib.bins.create_chunk(self.instdir, self.chunk_f, ['bin'],
                                    self.ex)
+        self.chunk_f.flush()
         os.mkdir(self.unpacked)
         morphlib.bins.unpack_binary(self.chunk_file, self.unpacked, self.ex)
         self.assertEqual([x for x,y in self.recursive_lstat(self.unpacked)],
@@ -119,6 +122,7 @@ class ChunkTests(BinsTest):
         self.assertEqual([x for x,y in self.recursive_lstat(self.instdir)],
                          ['.', 'lib', 'lib/libfoo.so'])
 
+
 class StratumTests(BinsTest):
 
     def setUp(self):
@@ -126,9 +130,11 @@ class StratumTests(BinsTest):
         self.tempdir = tempfile.mkdtemp()
         self.instdir = os.path.join(self.tempdir, 'inst')
         self.stratum_file = os.path.join(self.tempdir, 'stratum')
+        self.stratum_f = open(self.stratum_file, 'wb')
         self.unpacked = os.path.join(self.tempdir, 'unpacked')
         
     def tearDown(self):
+        self.stratum_f.close()
         shutil.rmtree(self.tempdir)
 
     def populate_instdir(self):
@@ -137,7 +143,8 @@ class StratumTests(BinsTest):
 
     def test_creates_and_unpacks_stratum_exactly(self):
         self.populate_instdir()
-        morphlib.bins.create_stratum(self.instdir, self.stratum_file, self.ex)
+        morphlib.bins.create_stratum(self.instdir, self.stratum_f, self.ex)
+        self.stratum_f.flush()
         os.mkdir(self.unpacked)
         morphlib.bins.unpack_binary(self.stratum_file, self.unpacked, self.ex)
         self.assertEqual(self.recursive_lstat(self.instdir),
diff --git a/morphlib/builder.py b/morphlib/builder.py
index f08ae259..3a46cf54 100644
--- a/morphlib/builder.py
+++ b/morphlib/builder.py
@@ -66,6 +66,8 @@ class BlobBuilder(object):
         self.staging = None
         self.settings = None
         self.real_msg = None
+        self.cachedir = None
+        self.cache_basename = None
         self.cache_prefix = None
         self.tempdir = None
         self.logfile = None
@@ -168,8 +170,7 @@ class BlobBuilder(object):
 
     def write_cache_metadata(self, meta):
         self.msg('Writing metadata to the cache')
-        filename = '%s.meta' % self.cache_prefix
-        with open(filename, 'w') as f:
+        with self.cachedir.open(self.cache_basename + '.meta') as f:
             json.dump(meta, f, indent=4)
             f.write('\n')
 
@@ -432,9 +433,12 @@ class ChunkBuilder(BlobBuilder):
                 patterns = self.blob.chunks[chunk_name]
                 patterns += [r'baserock/%s\.' % chunk_name]
                 filename = self.filename(chunk_name)
+                basename = os.path.basename(filename)
                 self.msg('Creating binary for %s' % chunk_name)
-                morphlib.bins.create_chunk(self.destdir, filename, patterns,
-                                           self.ex, self.dump_memory_profile)
+                with self.cachedir.open(filename) as f:
+                    morphlib.bins.create_chunk(self.destdir, f, patterns,
+                                               self.ex, 
+                                               self.dump_memory_profile)
                 chunks.append((chunk_name, filename))
 
         files = os.listdir(self.destdir)
@@ -461,7 +465,9 @@ class StratumBuilder(BlobBuilder):
             self.prepare_binary_metadata(self.blob.morph.name)
             self.msg('Creating binary for %s' % self.blob.morph.name)
             filename = self.filename(self.blob.morph.name)
-            morphlib.bins.create_stratum(self.destdir, filename, ex)
+            basename = os.path.basename(filename)
+            with self.cachedir.open(basename) as f:
+                morphlib.bins.create_stratum(self.destdir, f, ex)
         return { self.blob.morph.name: filename }
 
 
@@ -712,7 +718,9 @@ class Builder(object):
         builder.destdir = os.path.join(s, '%s.inst' % blob.morph.name)
         builder.settings = self.settings
         builder.real_msg = self.msg
+        builder.cachedir = self.cachedir
         builder.cache_prefix = self.cachedir.name(cache_id)
+        builder.cache_basename = os.path.basename(builder.cache_prefix)
         builder.tempdir = self.tempdir
         builder.dump_memory_profile = self.dump_memory_profile
         
diff --git a/morphlib/cachedir.py b/morphlib/cachedir.py
index cc528a43..d0763c4c 100644
--- a/morphlib/cachedir.py
+++ b/morphlib/cachedir.py
@@ -17,6 +17,8 @@
 import hashlib
 import os
 
+import morphlib
+
 
 class CacheDir(object):
 
@@ -62,3 +64,18 @@ class CacheDir(object):
 
         return os.path.join(self.dirname, key + suffix)
 
+    def open(self, relative_name):
+        '''Open a file for writing in the cache.
+        
+        The file will be written with a temporary name, and renamed to
+        the final name when the file is closed. Additionally, if the
+        caller decides, mid-writing, that they don't want to write the
+        file after all (e.g., a log file), then the ``abort`` method
+        in the returned file handle can be called to remove the
+        temporary file.
+        
+        '''
+        
+        path = os.path.join(self.dirname, relative_name)
+        return morphlib.savefile.SaveFile(path, 'w')
+
diff --git a/morphlib/cachedir_tests.py b/morphlib/cachedir_tests.py
index 172b7d86..f6cccc97 100644
--- a/morphlib/cachedir_tests.py
+++ b/morphlib/cachedir_tests.py
@@ -14,18 +14,27 @@
 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 
 
-import unittest
 import os
+import shutil
+import tempfile
+import unittest
 
 import morphlib
 
 
 class CacheDirTests(unittest.TestCase):
 
+    def cat(self, relative_name):
+        with open(os.path.join(self.cachedir.dirname, relative_name)) as f:
+            return f.read()
+
     def setUp(self):
-        self.dirname = 'cache/dir'
+        self.dirname = tempfile.mkdtemp()
         self.cachedir = morphlib.cachedir.CacheDir(self.dirname)
 
+    def tearDown(self):
+        shutil.rmtree(self.dirname)
+
     def test_sets_dirname_attribute(self):
         self.assertEqual(self.cachedir.dirname, os.path.abspath(self.dirname))
 
@@ -88,3 +97,16 @@ class CacheDirTests(unittest.TestCase):
         pathname = self.cachedir.name(dict_key)
         self.assert_(pathname.startswith(self.cachedir.dirname + '/'))
 
+    def test_allows_file_to_be_written(self):
+        f = self.cachedir.open('foo')
+        f.write('bar')
+        f.close()
+        self.assertEqual(self.cat('foo'), 'bar')
+
+    def test_allows_file_to_be_aborted(self):
+        f = self.cachedir.open('foo')
+        f.write('bar')
+        f.abort()
+        pathname = os.path.join(self.cachedir.dirname, 'foo')
+        self.assertFalse(os.path.exists(pathname))
+
diff --git a/morphlib/savefile.py b/morphlib/savefile.py
new file mode 100644
index 00000000..902a9923
--- /dev/null
+++ b/morphlib/savefile.py
@@ -0,0 +1,66 @@
+# Copyright (C) 2012  Codethink Limited
+# 
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; version 2 of the License.
+# 
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+# 
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+
+import os
+import tempfile
+
+class SaveFile(file):
+
+    '''Save files with a temporary name and rename when they're ready.
+    
+    This class acts exactly like the normal ``file`` class, except that
+    it is meant only for saving data to files. The data is written to
+    a temporary file, which gets renamed to the target name when the
+    open file is closed. This avoids readers of the file from getting
+    an incomplete file.
+    
+    Example:
+    
+        f = SaveFile('foo', 'w')
+        f.write(stuff)
+        f.close()
+        
+    The file will be called something like ``tmpCAFEBEEF`` until ``close``
+    is called, at which point it gets renamed to ``foo``.
+    
+    If the writer decides the file is not worth saving, they can call the
+    ``abort`` method, which deletes the temporary file.
+    
+    '''
+
+    def __init__(self, filename, *args, **kwargs):
+        self._savefile_filename = filename
+        dirname = os.path.dirname(filename)
+        fd, self._savefile_tempname = tempfile.mkstemp(dir=dirname)
+        os.close(fd)
+        file.__init__(self, self._savefile_tempname, *args, **kwargs)
+
+    def abort(self):
+        '''Abort file saving.
+        
+        The temporary file will be removed, and the universe is almost
+        exactly as if the file save had never started.
+        
+        '''
+
+        os.remove(self._savefile_tempname)
+        return file.close(self)
+
+    def close(self):
+        ret = file.close(self)
+        os.rename(self._savefile_tempname, self._savefile_filename)
+        return ret
+
diff --git a/morphlib/savefile_tests.py b/morphlib/savefile_tests.py
new file mode 100644
index 00000000..442a0779
--- /dev/null
+++ b/morphlib/savefile_tests.py
@@ -0,0 +1,90 @@
+# Copyright (C) 2012  Codethink Limited
+# 
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; version 2 of the License.
+# 
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+# 
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+
+import os
+import shutil
+import tempfile
+import unittest
+
+import savefile
+
+
+class SaveFileTests(unittest.TestCase):
+
+    def cat(self, filename):
+        with open(filename) as f:
+            return f.read()
+
+    def mkfile(self, filename, contents):
+        with open(filename, 'w') as f:
+            f.write(contents)
+
+    def setUp(self):
+        self.tempdir = tempfile.mkdtemp()
+        self.basename = 'filename'
+        self.filename = os.path.join(self.tempdir, self.basename)
+        
+    def tearDown(self):
+        shutil.rmtree(self.tempdir)
+    
+    def test_there_are_no_files_initially(self):
+        self.assertEqual(os.listdir(self.tempdir), [])
+
+    def test_saves_new_file(self):
+        f = savefile.SaveFile(self.filename, 'w')
+        f.write('foo')
+        f.close()
+        self.assertEqual(os.listdir(self.tempdir), [self.basename])
+        self.assertEqual(self.cat(self.filename), 'foo')
+
+    def test_overwrites_existing_file(self):
+        self.mkfile(self.filename, 'yo!')
+        f = savefile.SaveFile(self.filename, 'w')
+        f.write('foo')
+        f.close()
+        self.assertEqual(os.listdir(self.tempdir), [self.basename])
+        self.assertEqual(self.cat(self.filename), 'foo')
+
+    def test_leaves_no_file_after_aborted_new_file(self):
+        f = savefile.SaveFile(self.filename, 'w')
+        f.write('foo')
+        f.abort()
+        self.assertEqual(os.listdir(self.tempdir), [])
+
+    def test_leaves_original_file_after_aborted_overwrite(self):
+        self.mkfile(self.filename, 'yo!')
+        f = savefile.SaveFile(self.filename, 'w')
+        f.write('foo')
+        f.abort()
+        self.assertEqual(os.listdir(self.tempdir), [self.basename])
+        self.assertEqual(self.cat(self.filename), 'yo!')
+
+    def test_saves_normally_with_with(self):
+        with savefile.SaveFile(self.filename, 'w') as f:
+            f.write('foo')
+        self.assertEqual(os.listdir(self.tempdir), [self.basename])
+        self.assertEqual(self.cat(self.filename), 'foo')
+
+    def test_saves_normally_with_exception_within_with(self):
+        try:
+            with savefile.SaveFile(self.filename, 'w') as f:
+                f.write('foo')
+                raise Exception()
+        except Exception:
+            pass
+        self.assertEqual(os.listdir(self.tempdir), [self.basename])
+        self.assertEqual(self.cat(self.filename), 'foo')
+
diff --git a/morphlib/sourcemanager_tests.py b/morphlib/sourcemanager_tests.py
index 6a41986a..eb18dcc3 100644
--- a/morphlib/sourcemanager_tests.py
+++ b/morphlib/sourcemanager_tests.py
@@ -55,7 +55,6 @@ class SourceManagerTests(unittest.TestCase):
         shutil.rmtree(self.temprepodir)
 
     def test_uses_provided_cache_dir(self):
-        return
         app = DummyApp()
         
         tempdir = '/bla/bla/bla'