evolve: add a pull path using bundle2
authorPierre-Yves David <pierre-yves.david@fb.com>
Tue, 20 May 2014 14:39:13 -0700
changeset 960 0c1142059a82
parent 959 d39a5c8c82ad
child 961 8de88b323fb6
evolve: add a pull path using bundle2 Instead of using a dedicated wireprotocol commands, we use bundle2 to transmit an obs marker parts. This aims at both testing bundle2 more and to limit the amount of special code we needs to put in simple for server to fit our needs. The massive test changes comes from the fact we can use this fast path for both remote and local push.
hgext/evolve.py
--- a/hgext/evolve.py	Tue Jun 03 02:13:02 2014 -0700
+++ b/hgext/evolve.py	Tue May 20 14:39:13 2014 -0700
@@ -26,6 +26,7 @@
 import random
 from StringIO import StringIO
 import struct
+import urllib
 
 import mercurial
 from mercurial import util
@@ -2276,7 +2277,7 @@
             markers = []
         if not markers:
             repo.ui.status("OBSEXC: no marker to push\n")
-        elif remote.capable('_evoext_b2x_obsmarkers_0_pushonly'):
+        elif remote.capable('_evoext_b2x_obsmarkers_0'):
             obsdata = StringIO()
             _encodemarkersstream(obsdata, markers)
             obsdata.seek(0)
@@ -2398,12 +2399,39 @@
         common = findcommonobsmarkers(repo.ui, repo, remote, revs)
     return {'heads': pullop.pulledsubset, 'common': common}
 
+@eh.uisetup
+def addgetbundleargs(self):
+    if gboptsmap is not None:
+        gboptsmap['evo_obsmarker'] = 'plain'
+        gboptsmap['evo_obscommon'] = 'plain'
+        gboptsmap['evo_obsheads'] = 'plain'
+    else:
+        gboptslist.append('evo_obsheads')
+        gboptslist.append('evo_obscommon')
+        gboptslist.append('evo_obsmarker')
+
+
+
+@eh.wrapfunction(exchange, '_getbundleextrapart')
+def _getbundleextrapart(orig, bundler, repo, source, **kwargs):
+    if int(kwargs.pop('evo_obsmarker', False)):
+        common = kwargs.pop('evo_obscommon')
+        common = wireproto.decodelist(common)
+        heads = kwargs.pop('evo_obsheads')
+        heads = wireproto.decodelist(heads)
+        obsdata = _getobsmarkersstream(repo, common=common, heads=heads)
+        if len(obsdata.getvalue()) > 5:
+            obspart = bundle2.bundlepart('EVOLVE:B2X:OBSMARKERV1', data=obsdata)
+            bundler.addpart(obspart)
+    orig(bundler, repo, source)
 
 @eh.wrapfunction(exchange, '_pullobsolete')
 def _pullobsolete(orig, pullop):
     if not obsolete._enabled:
         return None
-    if not pullop.remote.capable('_evoext_pullobsmarkers_0'):
+    b2xpull = pullop.remote.capable('_evoext_b2x_obsmarkers_0')
+    wirepull = pullop.remote.capable('_evoext_pullobsmarkers_0')
+    if not (b2xpull or wirepull):
         return orig(pullop)
     if 'obsolete' not in pullop.remote.listkeys('namespaces'):
         return None # remote opted out of obsolescence marker exchange
@@ -2411,21 +2439,50 @@
     ui = pullop.repo.ui
     ui.status("OBSEXC: pull obsolescence markers\n")
     boundaries = _buildpullobsmerkersboundaries(pullop)
-    obsdata = pullop.remote.evoext_pullobsmarkers_0(**boundaries)
-    obsdata = obsdata.read()
-    if len(obsdata) > 5:
-        ui.status("OBSEXC: merging obsolescence markers (%i bytes)\n"
-                       % len(obsdata))
-        tr = pullop.gettransaction()
-        old = len(pullop.repo.obsstore._all)
-        pullop.repo.obsstore.mergemarkers(tr, obsdata)
-        new = len(pullop.repo.obsstore._all) - old
-        ui.status("OBSEXC: %i markers added\n" % new)
-        if new:
-            pullop.repo.invalidatevolatilesets()
-    else:
-        ui.status("OBSEXC: no unknown remote markers\n")
-    ui.status("OBSEXC: DONE\n")
+    new = 0
+
+    if b2xpull:
+        kwargs = {'bundlecaps': set(['HG2X'])}
+        capsblob = bundle2.encodecaps(pullop.repo.bundle2caps)
+        kwargs['bundlecaps'].add('bundle2=' + urllib.quote(capsblob))
+        kwargs['heads'] = [nullid]
+        kwargs['common'] = [nullid]
+        kwargs['evo_obsmarker'] = '1'
+        kwargs['evo_obscommon'] = wireproto.encodelist(boundaries['common'])
+        kwargs['evo_obsheads'] = wireproto.encodelist(boundaries['heads'])
+        bundle = pullop.remote.getbundle('pull', **kwargs)
+        try:
+            op = bundle2.processbundle(pullop.repo, bundle, pullop.gettransaction)
+        except bundle2.UnknownPartError, exc:
+            raise util.Abort('missing support for %s' % exc)
+        bytes = new = 0
+        for entry in op.records['evo_obsmarkers']:
+            bytes += entry.get('bytes', 0)
+            new += entry.get('new', 0)
+        if 5 < bytes:
+            ui.status("OBSEXC: merging obsolescence markers (%i bytes)\n"
+                      % bytes)
+            ui.status("OBSEXC: %i markers added\n" % new)
+            tr = op.gettransaction()
+        else:
+            ui.status("OBSEXC: no unknown remote markers\n")
+        ui.status("OBSEXC: DONE\n")
+    elif wirepull:
+        obsdata = pullop.remote.evoext_pullobsmarkers_0(**boundaries)
+        obsdata = obsdata.read()
+        if len(obsdata) > 5:
+            ui.status("OBSEXC: merging obsolescence markers (%i bytes)\n"
+                           % len(obsdata))
+            tr = pullop.gettransaction()
+            old = len(pullop.repo.obsstore._all)
+            pullop.repo.obsstore.mergemarkers(tr, obsdata)
+            new = len(pullop.repo.obsstore._all) - old
+            ui.status("OBSEXC: %i markers added\n" % new)
+        else:
+            ui.status("OBSEXC: no unknown remote markers\n")
+        ui.status("OBSEXC: DONE\n")
+    if new:
+        pullop.repo.invalidatevolatilesets()
     return tr
 
 def _getobsmarkersstream(repo, heads=None, common=None):
@@ -2615,14 +2672,14 @@
         caps += ' _evoext_pushobsmarkers_0'
         caps += ' _evoext_pullobsmarkers_0'
         caps += ' _evoext_obshash_0'
-        caps += ' _evoext_b2x_obsmarkers_0_pushonly'
+        caps += ' _evoext_b2x_obsmarkers_0'
     return caps
 
 
 @eh.extsetup
 def _installwireprotocol(ui):
     localrepo.moderncaps.add('_evoext_pullobsmarkers_0')
-    localrepo.moderncaps.add('_evoext_b2x_obsmarkers_0_pushonly')
+    localrepo.moderncaps.add('_evoext_b2x_obsmarkers_0')
     hgweb_mod.perms['evoext_pushobsmarkers_0'] = 'push'
     hgweb_mod.perms['evoext_pullobsmarkers_0'] = 'pull'
     hgweb_mod.perms['evoext_obshash'] = 'pull'