cmd/bloom-cmd.py

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