Add atomically_replaced_file for safer output

[bup.git] / cmd / midx-cmd.py

diff --git a/cmd/midx-cmd.py b/cmd/midx-cmd.py
index 38a5fdfca500639894111181edeb3a1007aea38a..a18707d8b3eb6a0e9dcb37a1cfdc053de0844990 100755 (executable)
--- a/cmd/midx-cmd.py
+++ b/cmd/midx-cmd.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 import sys, math, struct, glob, resource
 import tempfile
-from bup import options, git, midx, _helpers
+from bup import options, git, midx, _helpers, xstat
 from bup.helpers import *
 
 PAGE_SIZE=4096
@@ -81,61 +81,71 @@ def _do_midx(outdir, outfilename, infilenames, prefixstr):
     inp = []
     total = 0
     allfilenames = []
-    for name in infilenames:
-        ix = git.open_idx(name)
-        inp.append((
-            ix.map,
-            len(ix),
-            ix.sha_ofs,
-            isinstance(ix, midx.PackMidx) and ix.which_ofs or 0,
-            len(allfilenames),
-        ))
-        for n in ix.idxnames:
-            allfilenames.append(os.path.basename(n))
-        total += len(ix)
-    inp.sort(lambda x,y: cmp(str(y[0][y[2]:y[2]+20]),str(x[0][x[2]:x[2]+20])))
-
-    if not _first: _first = outdir
-    dirprefix = (_first != outdir) and git.repo_rel(outdir)+': ' or ''
-    debug1('midx: %s%screating from %d files (%d objects).\n'
-           % (dirprefix, prefixstr, len(infilenames), total))
-    if (opt.auto and (total < 1024 and len(infilenames) < 3)) \
-       or ((opt.auto or opt.force) and len(infilenames) < 2) \
-       or (opt.force and not total):
-        debug1('midx: nothing to do.\n')
-        return
-
-    pages = int(total/SHA_PER_PAGE) or 1
-    bits = int(math.ceil(math.log(pages, 2)))
-    entries = 2**bits
-    debug1('midx: table size: %d (%d bits)\n' % (entries*4, bits))
-
-    unlink(outfilename)
-    f = open(outfilename + '.tmp', 'w+b')
-    f.write('MIDX')
-    f.write(struct.pack('!II', midx.MIDX_VERSION, bits))
-    assert(f.tell() == 12)
-
-    f.truncate(12 + 4*entries + 20*total + 4*total)
-
-    fmap = mmap_readwrite(f, close=False)
-
-    count = merge_into(fmap, bits, total, inp)
-    del fmap
-
-    f.seek(0, git.SEEK_END)
-    f.write('\0'.join(allfilenames))
-    f.close()
-    os.rename(outfilename + '.tmp', outfilename)
-
-    # this is just for testing
+    midxs = []
+    try:
+        for name in infilenames:
+            ix = git.open_idx(name)
+            midxs.append(ix)
+            inp.append((
+                ix.map,
+                len(ix),
+                ix.sha_ofs,
+                isinstance(ix, midx.PackMidx) and ix.which_ofs or 0,
+                len(allfilenames),
+            ))
+            for n in ix.idxnames:
+                allfilenames.append(os.path.basename(n))
+            total += len(ix)
+        inp.sort(lambda x,y: cmp(str(y[0][y[2]:y[2]+20]),str(x[0][x[2]:x[2]+20])))
+
+        if not _first: _first = outdir
+        dirprefix = (_first != outdir) and git.repo_rel(outdir)+': ' or ''
+        debug1('midx: %s%screating from %d files (%d objects).\n'
+               % (dirprefix, prefixstr, len(infilenames), total))
+        if (opt.auto and (total < 1024 and len(infilenames) < 3)) \
+           or ((opt.auto or opt.force) and len(infilenames) < 2) \
+           or (opt.force and not total):
+            debug1('midx: nothing to do.\n')
+            return
+
+        pages = int(total/SHA_PER_PAGE) or 1
+        bits = int(math.ceil(math.log(pages, 2)))
+        entries = 2**bits
+        debug1('midx: table size: %d (%d bits)\n' % (entries*4, bits))
+
+        unlink(outfilename)
+        with atomically_replaced_file(outfilename, 'wb') as f:
+            f.write('MIDX')
+            f.write(struct.pack('!II', midx.MIDX_VERSION, bits))
+            assert(f.tell() == 12)
+
+            f.truncate(12 + 4*entries + 20*total + 4*total)
+            f.flush()
+            fdatasync(f.fileno())
+
+            fmap = mmap_readwrite(f, close=False)
+
+            count = merge_into(fmap, bits, total, inp)
+            del fmap # Assume this calls msync() now.
+            f.seek(0, os.SEEK_END)
+            f.write('\0'.join(allfilenames))
+    finally:
+        for ix in midxs:
+            if isinstance(ix, midx.PackMidx):
+                ix.close()
+        midxs = None
+        inp = None
+
+
+    # This is just for testing (if you enable this, don't clear inp above)
     if 0:
         p = midx.PackMidx(outfilename)
         assert(len(p.idxnames) == len(infilenames))
         print p.idxnames
         assert(len(p) == total)
         for pe, e in p, git.idxmerge(inp, final_progress=False):
-            assert(i == pi.next())
+            pin = pi.next()
+            assert(i == pin)
             assert(p.exists(i))
 
     return total, outfilename
@@ -162,7 +172,7 @@ def do_midx_dir(path):
                     
         # sort the biggest+newest midxes first, so that we can eliminate
         # smaller (or older) redundant ones that come later in the list
-        midxs.sort(key=lambda ix: (-sizes[ix], -os.stat(ix).st_mtime))
+        midxs.sort(key=lambda ix: (-sizes[ix], -xstat.stat(ix).st_mtime))
         
         for mname in midxs:
             any = 0