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