lib/bup/cmd/bloom.py

   1
   2 from __future__ import absolute_import
   3
   4 import os, glob, sys
   5
   6 from bup import options, git, bloom
   7 from bup.compat import argv_bytes, hexstr
   8 from bup.helpers import (add_error, debug1, log, progress, qprogress,
   9                          saved_errors)
  10 from bup.io import path_msg
  11
  12
  13 optspec = """
  14 bup bloom [options...]
  15 --
  16 ruin       ruin the specified bloom file (clearing the bitfield)
  17 f,force    ignore existing bloom file and regenerate it from scratch
  18 o,output=  output bloom filename (default: auto)
  19 d,dir=     input directory to look for idx files (default: auto)
  20 k,hashes=  number of hash functions to use (4 or 5) (default: auto)
  21 c,check=   check the given .idx file against the bloom filter
  22 """
  23
  24
  25 def ruin_bloom(bloomfilename):
  26     rbloomfilename = git.repo_rel(bloomfilename)
  27     if not os.path.exists(bloomfilename):
  28         log(path_msg(bloomfilename) + '\n')
  29         add_error('bloom: %s not found to ruin\n' % path_msg(rbloomfilename))
  30         return
  31     with bloom.ShaBloom(bloomfilename, readwrite=True, expected=1) as b:
  32         b.map[16 : 16 + 2**b.bits] = b'\0' * 2**b.bits
  33
  34
  35 def check_bloom(path, bloomfilename, idx):
  36     rbloomfilename = git.repo_rel(bloomfilename)
  37     ridx = git.repo_rel(idx)
  38     if not os.path.exists(bloomfilename):
  39         log('bloom: %s: does not exist.\n' % path_msg(rbloomfilename))
  40         return
  41     with bloom.ShaBloom(bloomfilename) as b:
  42         if not b.valid():
  43             add_error('bloom: %r is invalid.\n' % path_msg(rbloomfilename))
  44             return
  45         base = os.path.basename(idx)
  46         if base not in b.idxnames:
  47             log('bloom: %s does not contain the idx.\n' % path_msg(rbloomfilename))
  48             return
  49         if base == idx:
  50             idx = os.path.join(path, idx)
  51         log('bloom: bloom file: %s\n' % path_msg(rbloomfilename))
  52         log('bloom:   checking %s\n' % path_msg(ridx))
  53         with git.open_idx(idx) as oids:
  54             for oid in oids:
  55                 if not b.exists(oid):
  56                     add_error('bloom: ERROR: object %s missing' % hexstr(oid))
  57
  58
  59 _first = None
  60 def do_bloom(path, outfilename, k, force):
  61     global _first
  62     assert k in (None, 4, 5)
  63     b = None
  64     try:
  65         if os.path.exists(outfilename) and not force:
  66             b = bloom.ShaBloom(outfilename)
  67             if not b.valid():
  68                 debug1("bloom: Existing invalid bloom found, regenerating.\n")
  69                 b.close()
  70                 b = None
  71
  72         add = []
  73         rest = []
  74         add_count = 0
  75         rest_count = 0
  76         for i, name in enumerate(glob.glob(b'%s/*.idx' % path)):
  77             progress('bloom: counting: %d\r' % i)
  78             with git.open_idx(name) as ix:
  79                 ixbase = os.path.basename(name)
  80                 if b is not None and (ixbase in b.idxnames):
  81                     rest.append(name)
  82                     rest_count += len(ix)
  83                 else:
  84                     add.append(name)
  85                     add_count += len(ix)
  86
  87         if not add:
  88             debug1("bloom: nothing to do.\n")
  89             return
  90
  91         if b is not None:
  92             if len(b) != rest_count:
  93                 debug1("bloom: size %d != idx total %d, regenerating\n"
  94                        % (len(b), rest_count))
  95                 b, b_tmp = None, b
  96                 b_tmp.close()
  97             elif k is not None and k != b.k:
  98                 debug1("bloom: new k %d != existing k %d, regenerating\n"
  99                        % (k, b.k))
 100                 b, b_tmp = None, b
 101                 b_tmp.close()
 102             elif (b.bits < bloom.MAX_BLOOM_BITS[b.k] and
 103                   b.pfalse_positive(add_count) > bloom.MAX_PFALSE_POSITIVE):
 104                 debug1("bloom: regenerating: adding %d entries gives "
 105                        "%.2f%% false positives.\n"
 106                        % (add_count, b.pfalse_positive(add_count)))
 107                 b, b_tmp = None, b
 108                 b_tmp.close()
 109             else:
 110                 b, b_tmp = None, b
 111                 b_tmp.close()
 112                 b = bloom.ShaBloom(outfilename, readwrite=True,
 113                                    expected=add_count)
 114         if b is None: # Need all idxs to build from scratch
 115             add += rest
 116             add_count += rest_count
 117         del rest
 118         del rest_count
 119
 120         msg = b is None and 'creating from' or 'adding'
 121         if not _first: _first = path
 122         dirprefix = (_first != path) and git.repo_rel(path) + b': ' or b''
 123         progress('bloom: %s%s %d file%s (%d object%s).\r'
 124             % (path_msg(dirprefix), msg,
 125                len(add), len(add)!=1 and 's' or '',
 126                add_count, add_count!=1 and 's' or ''))
 127
 128         tfname = None
 129         if b is None:
 130             tfname = os.path.join(path, b'bup.tmp.bloom')
 131             b = bloom.create(tfname, expected=add_count, k=k)
 132         count = 0
 133         icount = 0
 134         for name in add:
 135             with git.open_idx(name) as ix:
 136                 qprogress('bloom: writing %.2f%% (%d/%d objects)\r'
 137                           % (icount*100.0/add_count, icount, add_count))
 138                 b.add_idx(ix)
 139                 count += 1
 140                 icount += len(ix)
 141
 142     finally:  # This won't handle pending exceptions correctly in py2
 143         # Currently, there's an open file object for tfname inside b.
 144         # Make sure it's closed before rename.
 145         if b is not None: b.close()
 146
 147     if tfname:
 148         os.rename(tfname, outfilename)
 149
 150
 151 def main(argv):
 152     o = options.Options(optspec)
 153     opt, flags, extra = o.parse_bytes(argv[1:])
 154
 155     if extra:
 156         o.fatal('no positional parameters expected')
 157
 158     if not opt.check and opt.k and opt.k not in (4,5):
 159         o.fatal('only k values of 4 and 5 are supported')
 160
 161     if opt.check:
 162         opt.check = argv_bytes(opt.check)
 163
 164     git.check_repo_or_die()
 165
 166     output = argv_bytes(opt.output) if opt.output else None
 167     path = opt.dir and argv_bytes(opt.dir) or git.repo(b'objects/pack')
 168     debug1('bloom: scanning %s\n' % path_msg(path))
 169     outfilename = output or os.path.join(path, b'bup.bloom')
 170     if opt.check:
 171         check_bloom(path, outfilename, opt.check)
 172     elif opt.ruin:
 173         ruin_bloom(outfilename)
 174     else:
 175         do_bloom(path, outfilename, opt.k, opt.force)
 176
 177     if saved_errors:
 178         log('WARNING: %d errors encountered during bloom.\n' % len(saved_errors))
 179         sys.exit(1)
 180     elif opt.check:
 181         log('All tests passed.\n')