changegroup: decompress GZ algorithm in larger chunks for better performance

mercurial/changegroup.py

@@ -1,256 +1,256
 # changegroup.py - Mercurial changegroup manipulation functions
 #
 #  Copyright 2006 Matt Mackall <[e-mail]>
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 from i18n import _
 from node import nullrev
 import mdiff, util
 import struct, os, bz2, zlib, tempfile
 
 _BUNDLE10_DELTA_HEADER = "20s20s20s20s"
 
 def readexactly(stream, n):
     '''read n bytes from stream.read and abort if less was available'''
     s = stream.read(n)
     if len(s) < n:
         raise util.Abort(_("stream ended unexpectedly"
                            " (got %d bytes, expected %d)")
                           % (len(s), n))
     return s
 
 def getchunk(stream):
     """return the next chunk from stream as a string"""
     d = readexactly(stream, 4)
     l = struct.unpack(">l", d)[0]
     if l <= 4:
         if l:
             raise util.Abort(_("invalid chunk length %d") % l)
         return ""
     return readexactly(stream, l - 4)
 
 def chunkheader(length):
     """return a changegroup chunk header (string)"""
     return struct.pack(">l", length + 4)
 
 def closechunk():
     """return a changegroup chunk header (string) for a zero-length chunk"""
     return struct.pack(">l", 0)
 
 class nocompress(object):
     def compress(self, x):
         return x
     def flush(self):
         return ""
 
 bundletypes = {
     "": ("", nocompress), # only when using unbundle on ssh and old http servers
                           # since the unification ssh accepts a header but there
                           # is no capability signaling it.
     "HG10UN": ("HG10UN", nocompress),
     "HG10BZ": ("HG10", lambda: bz2.BZ2Compressor()),
     "HG10GZ": ("HG10GZ", lambda: zlib.compressobj()),
 }
 
 # hgweb uses this list to communicate its preferred type
 bundlepriority = ['HG10GZ', 'HG10BZ', 'HG10UN']
 
 def writebundle(cg, filename, bundletype):
     """Write a bundle file and return its filename.
 
     Existing files will not be overwritten.
     If no filename is specified, a temporary file is created.
     bz2 compression can be turned off.
     The bundle file will be deleted in case of errors.
     """
 
     fh = None
     cleanup = None
     try:
         if filename:
             fh = open(filename, "wb")
         else:
             fd, filename = tempfile.mkstemp(prefix="hg-bundle-", suffix=".hg")
             fh = os.fdopen(fd, "wb")
         cleanup = filename
 
         header, compressor = bundletypes[bundletype]
         fh.write(header)
         z = compressor()
 
         # parse the changegroup data, otherwise we will block
         # in case of sshrepo because we don't know the end of the stream
 
         # an empty chunkgroup is the end of the changegroup
         # a changegroup has at least 2 chunkgroups (changelog and manifest).
         # after that, an empty chunkgroup is the end of the changegroup
         empty = False
         count = 0
         while not empty or count <= 2:
             empty = True
             count += 1
             while True:
                 chunk = getchunk(cg)
                 if not chunk:
                     break
                 empty = False
                 fh.write(z.compress(chunkheader(len(chunk))))
                 pos = 0
                 while pos < len(chunk):
                     next = pos + 2**20
                     fh.write(z.compress(chunk[pos:next]))
                     pos = next
             fh.write(z.compress(closechunk()))
         fh.write(z.flush())
         cleanup = None
         return filename
     finally:
         if fh is not None:
             fh.close()
         if cleanup is not None:
             os.unlink(cleanup)
 
 def decompressor(fh, alg):
     if alg == 'UN':
         return fh
     elif alg == 'GZ':
         def generator(f):
             zd = zlib.decompressobj()
-            for chunk in f:
+            for chunk in util.filechunkiter(f):
                 yield zd.decompress(chunk)
     elif alg == 'BZ':
         def generator(f):
             zd = bz2.BZ2Decompressor()
             zd.decompress("BZ")
             for chunk in util.filechunkiter(f, 4096):
                 yield zd.decompress(chunk)
     else:
         raise util.Abort("unknown bundle compression '%s'" % alg)
     return util.chunkbuffer(generator(fh))
 
 class unbundle10(object):
     deltaheader = _BUNDLE10_DELTA_HEADER
     deltaheadersize = struct.calcsize(deltaheader)
     def __init__(self, fh, alg):
         self._stream = decompressor(fh, alg)
         self._type = alg
         self.callback = None
     def compressed(self):
         return self._type != 'UN'
     def read(self, l):
         return self._stream.read(l)
     def seek(self, pos):
         return self._stream.seek(pos)
     def tell(self):
         return self._stream.tell()
     def close(self):
         return self._stream.close()
 
     def chunklength(self):
         d = readexactly(self._stream, 4)
         l = struct.unpack(">l", d)[0]
         if l <= 4:
             if l:
                 raise util.Abort(_("invalid chunk length %d") % l)
             return 0
         if self.callback:
             self.callback()
         return l - 4
 
     def changelogheader(self):
         """v10 does not have a changelog header chunk"""
         return {}
 
     def manifestheader(self):
         """v10 does not have a manifest header chunk"""
         return {}
 
     def filelogheader(self):
         """return the header of the filelogs chunk, v10 only has the filename"""
         l = self.chunklength()
         if not l:
             return {}
         fname = readexactly(self._stream, l)
         return dict(filename=fname)
 
     def _deltaheader(self, headertuple, prevnode):
         node, p1, p2, cs = headertuple
         if prevnode is None:
             deltabase = p1
         else:
             deltabase = prevnode
         return node, p1, p2, deltabase, cs
 
     def deltachunk(self, prevnode):
         l = self.chunklength()
         if not l:
             return {}
         headerdata = readexactly(self._stream, self.deltaheadersize)
         header = struct.unpack(self.deltaheader, headerdata)
         delta = readexactly(self._stream, l - self.deltaheadersize)
         node, p1, p2, deltabase, cs = self._deltaheader(header, prevnode)
         return dict(node=node, p1=p1, p2=p2, cs=cs,
                     deltabase=deltabase, delta=delta)
 
 class headerlessfixup(object):
     def __init__(self, fh, h):
         self._h = h
         self._fh = fh
     def read(self, n):
         if self._h:
             d, self._h = self._h[:n], self._h[n:]
             if len(d) < n:
                 d += readexactly(self._fh, n - len(d))
             return d
         return readexactly(self._fh, n)
 
 def readbundle(fh, fname):
     header = readexactly(fh, 6)
 
     if not fname:
         fname = "stream"
         if not header.startswith('HG') and header.startswith('\0'):
             fh = headerlessfixup(fh, header)
             header = "HG10UN"
 
     magic, version, alg = header[0:2], header[2:4], header[4:6]
 
     if magic != 'HG':
         raise util.Abort(_('%s: not a Mercurial bundle') % fname)
     if version != '10':
         raise util.Abort(_('%s: unknown bundle version %s') % (fname, version))
     return unbundle10(fh, alg)
 
 class bundle10(object):
     deltaheader = _BUNDLE10_DELTA_HEADER
     def __init__(self, lookup):
         self._lookup = lookup
     def close(self):
         return closechunk()
     def fileheader(self, fname):
         return chunkheader(len(fname)) + fname
     def revchunk(self, revlog, rev, prev):
         node = revlog.node(rev)
         p1, p2 = revlog.parentrevs(rev)
         base = prev
 
         prefix = ''
         if base == nullrev:
             delta = revlog.revision(node)
             prefix = mdiff.trivialdiffheader(len(delta))
         else:
             delta = revlog.revdiff(base, rev)
         linknode = self._lookup(revlog, node)
         p1n, p2n = revlog.parents(node)
         basenode = revlog.node(base)
         meta = self.builddeltaheader(node, p1n, p2n, basenode, linknode)
         meta += prefix
         l = len(meta) + len(delta)
         yield chunkheader(l)
         yield meta
         yield delta
     def builddeltaheader(self, node, p1n, p2n, basenode, linknode):
         # do nothing with basenode, it is implicitly the previous one in HG10
         return struct.pack(self.deltaheader, node, p1n, p2n, linknode)