hgext3rd/pullbundle.py
changeset 4130 a1f6b8211016
parent 4129 bc4e62a1cb82
child 4132 afc933d32085
--- a/hgext3rd/pullbundle.py	Sun Sep 23 23:41:08 2018 +0200
+++ b/hgext3rd/pullbundle.py	Sun Sep 23 00:08:02 2018 +0200
@@ -123,7 +123,7 @@
 
     subranges = canonicalsubranges(repo, stablerange, rangeid)
     idx = 0
-    slices =[]
+    slices = []
     nodes.reverse()
     for rangeid in subranges:
         size = rangelength(rangeid)
@@ -158,7 +158,7 @@
             cut += cursor
         if cursor == 1:
             break
-        cursor //=2
+        cursor //= 2
 
     # 2. optimise, bottom part
     if skip != cut:
@@ -216,10 +216,8 @@
     if repo.ui.verbose or source == 'bundle':
         repo.ui.status(_("%d changesets found\n") % len(nodes))
 
-def makeonecgpart(newpart, repo, rangeid, outgoing, version, source,
-                  bundlecaps, filematcher, cgversions):
-    # same as upstream code
-
+def _makenewstream(newpart, repo, outgoing, version, source,
+                   bundlecaps, filematcher, cgversions):
     old = changegroup._changegroupinfo
     try:
         changegroup._changegroupinfo = _changegroupinfo
@@ -229,12 +227,76 @@
     finally:
         changegroup._changegroupinfo = old
 
+    nbchanges = len(outgoing.missing)
+    pversion = None
+    if cgversions:
+        pversion = version
+    return (cgstream, nbchanges, pversion)
+
+def _makepartfromstream(newpart, repo, cgstream, nbchanges, version):
+    # same as upstream code
+
     part = newpart('changegroup', data=cgstream)
-    if cgversions:
+    if version:
         part.addparam('version', version)
 
-    part.addparam('nbchanges', '%d' % len(outgoing.missing),
+    part.addparam('nbchanges', '%d' % nbchanges,
                   mandatory=False)
 
     if 'treemanifest' in repo.requirements:
         part.addparam('treemanifest', '1')
+
+# cache management
+
+def cachedir(repo):
+    return repo.cachevfs.join('pullbundles')
+
+def getcache(repo, bundlename):
+    cdir = cachedir(repo)
+    bundlepath = os.path.join(cdir, bundlename)
+    try:
+        fd = open(bundlepath, 'rb')
+        return util.filechunkiter(fd)
+    except IOError as exc:
+        if exc.errno != errno.ENOENT:
+            raise
+        return None
+
+def cachewriter(repo, bundlename, stream):
+    cdir = cachedir(repo)
+    bundlepath = os.path.join(cdir, bundlename)
+    try:
+        os.makedirs(cdir)
+    except OSError as exc:
+        if exc.errno == errno.EEXIST:
+            pass
+    with util.atomictempfile(bundlepath) as cachefile:
+        for chunk in stream:
+            cachefile.write(chunk)
+            yield chunk
+
+BUNDLEMASK = "%s-%s-%010iskip-%010isize.hg"
+
+def makeonecgpart(newpart, repo, rangeid, outgoing, version, source,
+                  bundlecaps, filematcher, cgversions):
+    bundlename = cachedata = None
+    if rangeid is not None:
+        nbchanges = repo.stablerange.rangelength(repo, rangeid)
+        headnode = nodemod.hex(repo.changelog.node(rangeid[0]))
+        # XXX do we need to use cgversion in there?
+        bundlename = BUNDLEMASK % (version, headnode, rangeid[1], nbchanges)
+        cachedata = getcache(repo, bundlename)
+    if cachedata is None:
+        partdata = _makenewstream(newpart, repo, outgoing, version, source,
+                                  bundlecaps, filematcher, cgversions)
+        if bundlename is not None:
+            cgstream = cachewriter(repo, bundlename, partdata[0])
+            partdata = (cgstream,) + partdata[1:]
+    else:
+        if repo.ui.verbose or source == 'bundle':
+            repo.ui.status(_("%d changesets found in caches\n") % nbchanges)
+        pversion = None
+        if cgversions:
+            pversion = version
+        partdata = (cachedata, nbchanges, pversion)
+    return _makepartfromstream(newpart, repo, *partdata)