]> arthur.barton.de Git - bup.git/blob - lib/bup/git.py
git.PackWriter: avoid pack corruption if interrupted by a signal.
[bup.git] / lib / bup / git.py
1 import os, errno, zlib, time, sha, subprocess, struct, stat, re, tempfile
2 import heapq
3 from bup.helpers import *
4
5 verbose = 0
6 ignore_midx = 0
7 home_repodir = os.path.expanduser('~/.bup')
8 repodir = None
9
10 _typemap =  { 'blob':3, 'tree':2, 'commit':1, 'tag':4 }
11 _typermap = { 3:'blob', 2:'tree', 1:'commit', 4:'tag' }
12
13
14 class GitError(Exception):
15     pass
16
17
18 def repo(sub = ''):
19     global repodir
20     if not repodir:
21         raise GitError('You should call check_repo_or_die()')
22     gd = os.path.join(repodir, '.git')
23     if os.path.exists(gd):
24         repodir = gd
25     return os.path.join(repodir, sub)
26
27
28 def _encode_packobj(type, content):
29     szout = ''
30     sz = len(content)
31     szbits = (sz & 0x0f) | (_typemap[type]<<4)
32     sz >>= 4
33     while 1:
34         if sz: szbits |= 0x80
35         szout += chr(szbits)
36         if not sz:
37             break
38         szbits = sz & 0x7f
39         sz >>= 7
40     z = zlib.compressobj(1)
41     yield szout
42     yield z.compress(content)
43     yield z.flush()
44
45
46 def _encode_looseobj(type, content):
47     z = zlib.compressobj(1)
48     yield z.compress('%s %d\0' % (type, len(content)))
49     yield z.compress(content)
50     yield z.flush()
51
52
53 def _decode_looseobj(buf):
54     assert(buf);
55     s = zlib.decompress(buf)
56     i = s.find('\0')
57     assert(i > 0)
58     l = s[:i].split(' ')
59     type = l[0]
60     sz = int(l[1])
61     content = s[i+1:]
62     assert(type in _typemap)
63     assert(sz == len(content))
64     return (type, content)
65
66
67 def _decode_packobj(buf):
68     assert(buf)
69     c = ord(buf[0])
70     type = _typermap[(c & 0x70) >> 4]
71     sz = c & 0x0f
72     shift = 4
73     i = 0
74     while c & 0x80:
75         i += 1
76         c = ord(buf[i])
77         sz |= (c & 0x7f) << shift
78         shift += 7
79         if not (c & 0x80):
80             break
81     return (type, zlib.decompress(buf[i+1:]))
82
83
84 class PackIdx:
85     def __init__(self, filename):
86         self.name = filename
87         self.map = mmap_read(open(filename))
88         assert(str(self.map[0:8]) == '\377tOc\0\0\0\2')
89         self.fanout = list(struct.unpack('!256I',
90                                          str(buffer(self.map, 8, 256*4))))
91         self.fanout.append(0)  # entry "-1"
92         nsha = self.fanout[255]
93         self.ofstable = buffer(self.map,
94                                8 + 256*4 + nsha*20 + nsha*4,
95                                nsha*4)
96         self.ofs64table = buffer(self.map,
97                                  8 + 256*4 + nsha*20 + nsha*4 + nsha*4)
98
99     def _ofs_from_idx(self, idx):
100         ofs = struct.unpack('!I', str(buffer(self.ofstable, idx*4, 4)))[0]
101         if ofs & 0x80000000:
102             idx64 = ofs & 0x7fffffff
103             ofs = struct.unpack('!I',
104                                 str(buffer(self.ofs64table, idx64*8, 8)))[0]
105         return ofs
106
107     def _idx_from_hash(self, hash):
108         assert(len(hash) == 20)
109         b1 = ord(hash[0])
110         start = self.fanout[b1-1] # range -1..254
111         end = self.fanout[b1] # range 0..255
112         buf = buffer(self.map, 8 + 256*4, end*20)
113         want = str(hash)
114         while start < end:
115             mid = start + (end-start)/2
116             v = str(buf[mid*20:(mid+1)*20])
117             if v < want:
118                 start = mid+1
119             elif v > want:
120                 end = mid
121             else: # got it!
122                 return mid
123         return None
124         
125     def find_offset(self, hash):
126         idx = self._idx_from_hash(hash)
127         if idx != None:
128             return self._ofs_from_idx(idx)
129         return None
130
131     def exists(self, hash):
132         return hash and (self._idx_from_hash(hash) != None) and True or None
133
134     def __iter__(self):
135         for i in xrange(self.fanout[255]):
136             yield buffer(self.map, 8 + 256*4 + 20*i, 20)
137
138     def __len__(self):
139         return int(self.fanout[255])
140
141
142 def extract_bits(buf, bits):
143     mask = (1<<bits) - 1
144     v = struct.unpack('!I', buf[0:4])[0]
145     v = (v >> (32-bits)) & mask
146     return v
147
148
149 class PackMidx:
150     def __init__(self, filename):
151         self.name = filename
152         assert(filename.endswith('.midx'))
153         self.map = mmap_read(open(filename))
154         if str(self.map[0:8]) == 'MIDX\0\0\0\1':
155             log('Warning: ignoring old-style midx %r\n' % filename)
156             self.bits = 0
157             self.entries = 1
158             self.fanout = buffer('\0\0\0\0')
159             self.shalist = buffer('\0'*20)
160             self.idxnames = []
161         else:
162             assert(str(self.map[0:8]) == 'MIDX\0\0\0\2')
163             self.bits = struct.unpack('!I', self.map[8:12])[0]
164             self.entries = 2**self.bits
165             self.fanout = buffer(self.map, 12, self.entries*4)
166             shaofs = 12 + self.entries*4
167             nsha = self._fanget(self.entries-1)
168             self.shalist = buffer(self.map, shaofs, nsha*20)
169             self.idxnames = str(self.map[shaofs + 20*nsha:]).split('\0')
170
171     def _fanget(self, i):
172         start = i*4
173         s = self.fanout[start:start+4]
174         return struct.unpack('!I', s)[0]
175     
176     def exists(self, hash):
177         want = str(hash)
178         el = extract_bits(want, self.bits)
179         if el:
180             start = self._fanget(el-1)
181         else:
182             start = 0
183         end = self._fanget(el)
184         while start < end:
185             mid = start + (end-start)/2
186             v = str(self.shalist[mid*20:(mid+1)*20])
187             if v < want:
188                 start = mid+1
189             elif v > want:
190                 end = mid
191             else: # got it!
192                 return True
193         return None
194     
195     def __iter__(self):
196         for i in xrange(self._fanget(self.entries-1)):
197             yield buffer(self.shalist, i*20, 20)
198     
199     def __len__(self):
200         return int(self._fanget(self.entries-1))
201
202
203 _mpi_count = 0
204 class PackIdxList:
205     def __init__(self, dir):
206         global _mpi_count
207         assert(_mpi_count == 0) # these things suck tons of VM; don't waste it
208         _mpi_count += 1
209         self.dir = dir
210         self.also = {}
211         self.packs = []
212         self.refresh()
213
214     def __del__(self):
215         global _mpi_count
216         _mpi_count -= 1
217         assert(_mpi_count == 0)
218
219     def __iter__(self):
220         return iter(idxmerge(self.packs))
221
222     def exists(self, hash):
223         if hash in self.also:
224             return True
225         for i in range(len(self.packs)):
226             p = self.packs[i]
227             if p.exists(hash):
228                 # reorder so most recently used packs are searched first
229                 self.packs = [p] + self.packs[:i] + self.packs[i+1:]
230                 return p.name
231         return None
232
233     def refresh(self, skip_midx = False, forget_packs = False):
234         if forget_packs:
235             self.packs = []
236         skip_midx = skip_midx or ignore_midx
237         d = dict((p.name, 1) for p in self.packs)
238         if os.path.exists(self.dir):
239             if not skip_midx:
240                 midxl = []
241                 for f in os.listdir(self.dir):
242                     full = os.path.join(self.dir, f)
243                     if f.endswith('.midx') and not d.get(full):
244                         mx = PackMidx(full)
245                         (mxd, mxf) = os.path.split(mx.name)
246                         broken = 0
247                         for n in mx.idxnames:
248                             if not os.path.exists(os.path.join(mxd, n)):
249                                 log(('warning: index %s missing\n' +
250                                     '  used by %s\n') % (n, mxf))
251                                 broken += 1
252                         if not broken:
253                             midxl.append(mx)
254                 midxl.sort(lambda x,y: -cmp(len(x),len(y)))
255                 for ix in midxl:
256                     any = 0
257                     for sub in ix.idxnames:
258                         if not d.get(os.path.join(self.dir, sub)):
259                             self.packs.append(ix)
260                             d[ix.name] = 1
261                             for name in ix.idxnames:
262                                 d[os.path.join(self.dir, name)] = 1
263                             any += 1
264                             break
265                     if not any:
266                         log('midx: removing redundant: %s\n' 
267                             % os.path.basename(ix.name))
268                         unlink(ix.name)
269             for f in os.listdir(self.dir):
270                 full = os.path.join(self.dir, f)
271                 if f.endswith('.idx') and not d.get(full):
272                     self.packs.append(PackIdx(full))
273                     d[full] = 1
274         log('PackIdxList: using %d index%s.\n' 
275             % (len(self.packs), len(self.packs)!=1 and 'es' or ''))
276
277     def add(self, hash):
278         self.also[hash] = 1
279
280     def zap_also(self):
281         self.also = {}
282
283
284 def calc_hash(type, content):
285     header = '%s %d\0' % (type, len(content))
286     sum = sha.sha(header)
287     sum.update(content)
288     return sum.digest()
289
290
291 def _shalist_sort_key(ent):
292     (mode, name, id) = ent
293     if stat.S_ISDIR(int(mode, 8)):
294         return name + '/'
295     else:
296         return name
297
298
299 def idxmerge(idxlist):
300     total = sum(len(i) for i in idxlist)
301     iters = (iter(i) for i in idxlist)
302     heap = [(next(it), it) for it in iters]
303     heapq.heapify(heap)
304     count = 0
305     last = None
306     while heap:
307         if (count % 10024) == 0:
308             progress('Reading indexes: %.2f%% (%d/%d)\r'
309                      % (count*100.0/total, count, total))
310         (e, it) = heap[0]
311         if e != last:
312             yield e
313             last = e
314         count += 1
315         e = next(it)
316         if e:
317             heapq.heapreplace(heap, (e, it))
318         else:
319             heapq.heappop(heap)
320     log('Reading indexes: %.2f%% (%d/%d), done.\n' % (100, total, total))
321
322     
323 class PackWriter:
324     def __init__(self, objcache_maker=None):
325         self.count = 0
326         self.outbytes = 0
327         self.filename = None
328         self.file = None
329         self.objcache_maker = objcache_maker
330         self.objcache = None
331
332     def __del__(self):
333         self.close()
334
335     def _make_objcache(self):
336         if not self.objcache:
337             if self.objcache_maker:
338                 self.objcache = self.objcache_maker()
339             else:
340                 self.objcache = PackIdxList(repo('objects/pack'))
341
342     def _open(self):
343         if not self.file:
344             self._make_objcache()
345             (fd,name) = tempfile.mkstemp(suffix='.pack', dir=repo('objects'))
346             self.file = os.fdopen(fd, 'w+b')
347             assert(name.endswith('.pack'))
348             self.filename = name[:-5]
349             self.file.write('PACK\0\0\0\2\0\0\0\0')
350
351     def _raw_write(self, datalist):
352         self._open()
353         f = self.file
354         # in case we get interrupted (eg. KeyboardInterrupt), it's best if
355         # the file never has a *partial* blob.  So let's make sure it's
356         # all-or-nothing.  (The blob shouldn't be very big anyway, thanks
357         # to our hashsplit algorithm.)  f.write() does its own buffering,
358         # but that's okay because we'll flush it in _end().
359         oneblob = ''.join(datalist)
360         f.write(oneblob)
361         self.outbytes += len(oneblob)
362         self.count += 1
363
364     def _write(self, bin, type, content):
365         if verbose:
366             log('>')
367         self._raw_write(_encode_packobj(type, content))
368         return bin
369
370     def breakpoint(self):
371         id = self._end()
372         self.outbytes = self.count = 0
373         return id
374
375     def write(self, type, content):
376         return self._write(calc_hash(type, content), type, content)
377
378     def exists(self, id):
379         if not self.objcache:
380             self._make_objcache()
381         return self.objcache.exists(id)
382
383     def maybe_write(self, type, content):
384         bin = calc_hash(type, content)
385         if not self.exists(bin):
386             self._write(bin, type, content)
387             self.objcache.add(bin)
388         return bin
389
390     def new_blob(self, blob):
391         return self.maybe_write('blob', blob)
392
393     def new_tree(self, shalist):
394         shalist = sorted(shalist, key = _shalist_sort_key)
395         l = []
396         for (mode,name,bin) in shalist:
397             assert(mode)
398             assert(mode != '0')
399             assert(mode[0] != '0')
400             assert(name)
401             assert(len(bin) == 20)
402             l.append('%s %s\0%s' % (mode,name,bin))
403         return self.maybe_write('tree', ''.join(l))
404
405     def _new_commit(self, tree, parent, author, adate, committer, cdate, msg):
406         l = []
407         if tree: l.append('tree %s' % tree.encode('hex'))
408         if parent: l.append('parent %s' % parent.encode('hex'))
409         if author: l.append('author %s %s' % (author, _git_date(adate)))
410         if committer: l.append('committer %s %s' % (committer, _git_date(cdate)))
411         l.append('')
412         l.append(msg)
413         return self.maybe_write('commit', '\n'.join(l))
414
415     def new_commit(self, parent, tree, msg):
416         now = time.time()
417         userline = '%s <%s@%s>' % (userfullname(), username(), hostname())
418         commit = self._new_commit(tree, parent,
419                                   userline, now, userline, now,
420                                   msg)
421         return commit
422
423     def abort(self):
424         f = self.file
425         if f:
426             self.file = None
427             f.close()
428             os.unlink(self.filename + '.pack')
429
430     def _end(self):
431         f = self.file
432         if not f: return None
433         self.file = None
434         self.objcache = None
435
436         # update object count
437         f.seek(8)
438         cp = struct.pack('!i', self.count)
439         assert(len(cp) == 4)
440         f.write(cp)
441
442         # calculate the pack sha1sum
443         f.seek(0)
444         sum = sha.sha()
445         while 1:
446             b = f.read(65536)
447             sum.update(b)
448             if not b: break
449         f.write(sum.digest())
450         
451         f.close()
452
453         p = subprocess.Popen(['git', 'index-pack', '-v',
454                               '--index-version=2',
455                               self.filename + '.pack'],
456                              preexec_fn = _gitenv,
457                              stdout = subprocess.PIPE)
458         out = p.stdout.read().strip()
459         _git_wait('git index-pack', p)
460         if not out:
461             raise GitError('git index-pack produced no output')
462         nameprefix = repo('objects/pack/%s' % out)
463         if os.path.exists(self.filename + '.map'):
464             os.unlink(self.filename + '.map')
465         os.rename(self.filename + '.pack', nameprefix + '.pack')
466         os.rename(self.filename + '.idx', nameprefix + '.idx')
467         return nameprefix
468
469     def close(self):
470         return self._end()
471
472
473 def _git_date(date):
474     return time.strftime('%s %z', time.localtime(date))
475
476
477 def _gitenv():
478     os.environ['GIT_DIR'] = os.path.abspath(repo())
479
480
481 def list_refs(refname = None):
482     argv = ['git', 'show-ref', '--']
483     if refname:
484         argv += [refname]
485     p = subprocess.Popen(argv, preexec_fn = _gitenv, stdout = subprocess.PIPE)
486     out = p.stdout.read().strip()
487     rv = p.wait()  # not fatal
488     if rv:
489         assert(not out)
490     if out:
491         for d in out.split('\n'):
492             (sha, name) = d.split(' ', 1)
493             yield (name, sha.decode('hex'))
494
495
496 def read_ref(refname):
497     l = list(list_refs(refname))
498     if l:
499         assert(len(l) == 1)
500         return l[0][1]
501     else:
502         return None
503
504
505 def rev_list(ref, count=None):
506     assert(not ref.startswith('-'))
507     opts = []
508     if count:
509         opts += ['-n', str(atoi(count))]
510     argv = ['git', 'rev-list', '--pretty=format:%ct'] + opts + [ref, '--']
511     p = subprocess.Popen(argv, preexec_fn = _gitenv, stdout = subprocess.PIPE)
512     commit = None
513     for row in p.stdout:
514         s = row.strip()
515         if s.startswith('commit '):
516             commit = s[7:].decode('hex')
517         else:
518             date = int(s)
519             yield (date, commit)
520     rv = p.wait()  # not fatal
521     if rv:
522         raise GitError, 'git rev-list returned error %d' % rv
523
524
525 def rev_get_date(ref):
526     for (date, commit) in rev_list(ref, count=1):
527         return date
528     raise GitError, 'no such commit %r' % ref
529
530
531 def update_ref(refname, newval, oldval):
532     if not oldval:
533         oldval = ''
534     assert(refname.startswith('refs/heads/'))
535     p = subprocess.Popen(['git', 'update-ref', refname,
536                           newval.encode('hex'), oldval.encode('hex')],
537                          preexec_fn = _gitenv)
538     _git_wait('git update-ref', p)
539
540
541 def guess_repo(path=None):
542     global repodir
543     if path:
544         repodir = path
545     if not repodir:
546         repodir = os.environ.get('BUP_DIR')
547         if not repodir:
548             repodir = os.path.expanduser('~/.bup')
549
550
551 def init_repo(path=None):
552     guess_repo(path)
553     d = repo()
554     if os.path.exists(d) and not os.path.isdir(os.path.join(d, '.')):
555         raise GitError('"%d" exists but is not a directory\n' % d)
556     p = subprocess.Popen(['git', '--bare', 'init'], stdout=sys.stderr,
557                          preexec_fn = _gitenv)
558     _git_wait('git init', p)
559     p = subprocess.Popen(['git', 'config', 'pack.indexVersion', '2'],
560                          stdout=sys.stderr, preexec_fn = _gitenv)
561     _git_wait('git config', p)
562
563
564 def check_repo_or_die(path=None):
565     guess_repo(path)
566     if not os.path.isdir(repo('objects/pack/.')):
567         if repodir == home_repodir:
568             init_repo()
569         else:
570             log('error: %r is not a bup/git repository\n' % repo())
571             sys.exit(15)
572
573
574 def _treeparse(buf):
575     ofs = 0
576     while ofs < len(buf):
577         z = buf[ofs:].find('\0')
578         assert(z > 0)
579         spl = buf[ofs:ofs+z].split(' ', 1)
580         assert(len(spl) == 2)
581         sha = buf[ofs+z+1:ofs+z+1+20]
582         ofs += z+1+20
583         yield (spl[0], spl[1], sha)
584
585
586 _ver = None
587 def ver():
588     global _ver
589     if not _ver:
590         p = subprocess.Popen(['git', '--version'],
591                              stdout=subprocess.PIPE)
592         gvs = p.stdout.read()
593         _git_wait('git --version', p)
594         m = re.match(r'git version (\S+.\S+)', gvs)
595         if not m:
596             raise GitError('git --version weird output: %r' % gvs)
597         _ver = tuple(m.group(1).split('.'))
598     needed = ('1','5', '3', '1')
599     if _ver < needed:
600         raise GitError('git version %s or higher is required; you have %s'
601                        % ('.'.join(needed), '.'.join(_ver)))
602     return _ver
603
604
605 def _git_wait(cmd, p):
606     rv = p.wait()
607     if rv != 0:
608         raise GitError('%s returned %d' % (cmd, rv))
609
610
611 def _git_capture(argv):
612     p = subprocess.Popen(argv, stdout=subprocess.PIPE, preexec_fn = _gitenv)
613     r = p.stdout.read()
614     _git_wait(repr(argv), p)
615     return r
616
617
618 _ver_warned = 0
619 class CatPipe:
620     def __init__(self):
621         global _ver_warned
622         wanted = ('1','5','6')
623         if ver() < wanted:
624             if not _ver_warned:
625                 log('warning: git version < %s; bup will be slow.\n'
626                     % '.'.join(wanted))
627                 _ver_warned = 1
628             self.get = self._slow_get
629         else:
630             self.p = subprocess.Popen(['git', 'cat-file', '--batch'],
631                                       stdin=subprocess.PIPE, 
632                                       stdout=subprocess.PIPE,
633                                       preexec_fn = _gitenv)
634             self.get = self._fast_get
635             self.inprogress = None
636
637     def _fast_get(self, id):
638         if self.inprogress:
639             log('_fast_get: opening %r while %r is open' 
640                 % (id, self.inprogress))
641         assert(not self.inprogress)
642         assert(id.find('\n') < 0)
643         assert(id.find('\r') < 0)
644         assert(id[0] != '-')
645         self.inprogress = id
646         self.p.stdin.write('%s\n' % id)
647         hdr = self.p.stdout.readline()
648         if hdr.endswith(' missing\n'):
649             raise KeyError('blob %r is missing' % id)
650         spl = hdr.split(' ')
651         if len(spl) != 3 or len(spl[0]) != 40:
652             raise GitError('expected blob, got %r' % spl)
653         (hex, type, size) = spl
654
655         def ondone():
656             assert(self.p.stdout.readline() == '\n')
657             self.inprogress = None
658
659         it = AutoFlushIter(chunkyreader(self.p.stdout, int(spl[2])),
660                            ondone = ondone)
661         yield type
662         for blob in it:
663             yield blob
664         del it
665
666     def _slow_get(self, id):
667         assert(id.find('\n') < 0)
668         assert(id.find('\r') < 0)
669         assert(id[0] != '-')
670         type = _git_capture(['git', 'cat-file', '-t', id]).strip()
671         yield type
672
673         p = subprocess.Popen(['git', 'cat-file', type, id],
674                              stdout=subprocess.PIPE,
675                              preexec_fn = _gitenv)
676         for blob in chunkyreader(p.stdout):
677             yield blob
678         _git_wait('git cat-file', p)
679
680     def _join(self, it):
681         type = it.next()
682         if type == 'blob':
683             for blob in it:
684                 yield blob
685         elif type == 'tree':
686             treefile = ''.join(it)
687             for (mode, name, sha) in _treeparse(treefile):
688                 for blob in self.join(sha.encode('hex')):
689                     yield blob
690         elif type == 'commit':
691             treeline = ''.join(it).split('\n')[0]
692             assert(treeline.startswith('tree '))
693             for blob in self.join(treeline[5:]):
694                 yield blob
695         else:
696             raise GitError('invalid object type %r: expected blob/tree/commit'
697                            % type)
698
699     def join(self, id):
700         try:
701             for d in self._join(self.get(id)):
702                 yield d
703         except StopIteration:
704             log('booger!\n')
705         
706
707 def cat(id):
708     c = CatPipe()
709     for d in c.join(id):
710         yield d