hgext/evolve.py
changeset 960 0c1142059a82
parent 959 d39a5c8c82ad
child 961 8de88b323fb6
--- 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'