# HG changeset patch
# User Pierre-Yves David <pierre-yves.david@fb.com>
# Date 1400621953 25200
# Node ID 0c1142059a824cbb6154f1cbc3f609e365eb10d1
# Parent  d39a5c8c82ade2748abeec579d556c14ea8d34f5
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.

diff -r d39a5c8c82ad -r 0c1142059a82 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'