# HG changeset patch
# User Pierre-Yves David <pierre-yves.david@fb.com>
# Date 1409321599 -7200
# Node ID 580a2d838996a8d1942e52643f72c168669c4bbd
# Parent  661f3a112f7a90b21c6e8089035ce38496a7471d
pull: use discovery to pull less obsmarkers through bundle2

This hooks in the official API for bundle2

diff -r 661f3a112f7a -r 580a2d838996 hgext/evolve.py
--- a/hgext/evolve.py	Fri Aug 29 16:13:01 2014 +0200
+++ b/hgext/evolve.py	Fri Aug 29 16:13:19 2014 +0200
@@ -2494,7 +2494,33 @@
 
 @eh.uisetup
 def addgetbundleargs(self):
-    gboptsmap['evo_obscommon'] = 'plain'
+    gboptsmap['evo_obscommon'] = 'nodes'
+
+@eh.wrapfunction(exchange, '_pullbundle2extraprepare')
+def _addobscommontob2pull(orig, pullop, kwargs):
+    ret = orig(pullop, kwargs)
+    if 'obsmarkers' in kwargs and pullop.remote.capable('_evoext_getbundle_obscommon'):
+        boundaries = _buildpullobsmarkersboundaries(pullop)
+        common = boundaries['common']
+        if common != [nullid]:
+            kwargs['evo_obscommon'] = common
+    return ret
+
+@eh.wrapfunction(exchange, '_getbundleobsmarkerpart')
+def _getbundleobsmarkerpart(orig, bundler, repo, source, heads=None, common=None,
+                            bundlecaps=None, **kwargs):
+    if 'evo_obscommon' not in kwargs:
+        return orig(bundler, repo, source, heads, common, bundlecaps, **kwargs)
+
+    if kwargs.get('obsmarkers', False):
+        if heads is None:
+            heads = repo.heads()
+        obscommon = kwargs.get('evo_obscommon', ())
+        obsset = repo.set('::%ln - ::%ln', heads, obscommon)
+        subset = [c.node() for c in obsset]
+        markers = repo.obsstore.relevantmarkers(subset)
+        exchange.buildobsmarkerspart(bundler, markers)
+
 
 @eh.wrapfunction(exchange, '_pullobsolete')
 def _pullobsolete(orig, pullop):
diff -r 661f3a112f7a -r 580a2d838996 hgext/simple4server.py
--- a/hgext/simple4server.py	Fri Aug 29 16:13:01 2014 +0200
+++ b/hgext/simple4server.py	Fri Aug 29 16:13:19 2014 +0200
@@ -237,22 +237,28 @@
         caps += ' _evoext_pushobsmarkers_0'
         caps += ' _evoext_pullobsmarkers_0'
         caps += ' _evoext_obshash_0'
-        caps += ' _evoext_b2x_obsmarkers_0'
+        caps += ' _evoext_getbundle_obscommon'
     return caps
 
+def _getbundleobsmarkerpart(orig, bundler, repo, source, heads=None, common=None,
+                            bundlecaps=None, **kwargs):
+    if 'evo_obscommon' not in kwargs:
+        return orig(bundler, repo, source, heads, common, bundlecaps, **kwargs)
+
+    if kwargs.get('obsmarkers', False):
+        if heads is None:
+            heads = repo.heads()
+        obscommon = kwargs.get('evo_obscommon', ())
+        obsset = repo.set('::%ln - ::%ln', heads, obscommon)
+        subset = [c.node() for c in obsset]
+        markers = repo.obsstore.relevantmarkers(subset)
+        exchange.buildobsmarkerspart(bundler, markers)
 
 # from evolve extension: 10867a8e27c6
 # heavily modified
 def extsetup(ui):
     localrepo.moderncaps.add('_evoext_b2x_obsmarkers_0')
-    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')
+    gboptsmap['evo_obscommon'] = 'nodes'
     if not util.safehasattr(obsolete.obsstore, 'relevantmarkers'):
         obsolete.obsstore = pruneobsstore
         obsolete.obsstore.relevantmarkers = relevantmarkers
@@ -262,7 +268,7 @@
     wireproto.commands['evoext_pushobsmarkers_0'] = (srv_pushobsmarkers, '')
     wireproto.commands['evoext_pullobsmarkers_0'] = (srv_pullobsmarkers, '*')
     # wrap module content
-    extensions.wrapfunction(exchange, '_getbundleextrapart', _getbundleextrapart)
+    extensions.wrapfunction(exchange, '_pullbundle2extraprepare', _getbundleobsmarkerpart)
     extensions.wrapfunction(wireproto, 'capabilities', capabilities)
     # wrap command content
     oldcap, args = wireproto.commands['capabilities']
@@ -274,47 +280,3 @@
     extensions.wrapfunction(pushkey, '_nslist', _nslist)
     pushkey._namespaces['namespaces'] = (lambda *x: False, pushkey._nslist)
 
-
-#from evolve extension
-@bundle2.parthandler('evolve:b2x:obsmarkerv1')
-def handleobsmarkerv1(op, inpart):
-    """add a stream of obsmarker to the repo"""
-    tr = op.gettransaction()
-    advparams = dict(inpart.advisoryparams)
-    length = advparams.get('totalbytes')
-    if length is None:
-        obsdata = inpart.read()
-    else:
-        length = int(length)
-        data = StringIO()
-        current = 0
-        op.ui.progress('OBSEXC', current, unit="bytes", total=length)
-        while current < length:
-            readsize = min(length-current, 4096)
-            data.write(inpart.read(readsize))
-            current += readsize
-            op.ui.progress('OBSEXC', current, unit="bytes", total=length)
-        op.ui.progress('OBSEXC', None)
-        obsdata = data.getvalue()
-    totalsize = len(obsdata)
-    old = len(op.repo.obsstore._all)
-    op.repo.obsstore.mergemarkers(tr, obsdata)
-    new = len(op.repo.obsstore._all) - old
-    op.records.add('evo_obsmarkers', {'new': new, 'bytes': totalsize})
-    tr.hookargs['evolve_new_obsmarkers'] = str(new)
-
-#from evolve extension
-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:
-            advparams = [('totalbytes', str(len(obsdata.getvalue())))]
-            obspart = bundle2.bundlepart('EVOLVE:B2X:OBSMARKERV1',
-                                         advisoryparams=advparams,
-                                         data=obsdata)
-            bundler.addpart(obspart)
-    orig(bundler, repo, source)
diff -r 661f3a112f7a -r 580a2d838996 tests/test-simple4server.t
--- a/tests/test-simple4server.t	Fri Aug 29 16:13:01 2014 +0200
+++ b/tests/test-simple4server.t	Fri Aug 29 16:13:19 2014 +0200
@@ -73,9 +73,9 @@
 ===================
 
   $ curl --silent http://localhost:$HGPORT/?cmd=hello
-  capabilities: lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_b2x_obsmarkers_0
+  capabilities: lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_getbundle_obscommon
   $ curl --silent http://localhost:$HGPORT/?cmd=capabilities
-  lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_b2x_obsmarkers_0 (no-eol)
+  lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_getbundle_obscommon (no-eol)
 
   $ curl --silent "http://localhost:$HGPORT/?cmd=listkeys&namespace=namespaces" | sort
   bookmarks	
@@ -136,9 +136,9 @@
   obsolete	
   phases	
   $ curl --silent http://localhost:$HGPORT/?cmd=hello
-  capabilities: lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_b2x_obsmarkers_0
+  capabilities: lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_getbundle_obscommon
   $ curl --silent http://localhost:$HGPORT/?cmd=capabilities
-  lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_b2x_obsmarkers_0 (no-eol)
+  lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_getbundle_obscommon (no-eol)
 
   $ echo '[__temporary__]' >> server/.hg/hgrc
   $ echo 'advertiseobsolete=False' >> server/.hg/hgrc
@@ -166,6 +166,6 @@
   obsolete	
   phases	
   $ curl --silent http://localhost:$HGPORT/?cmd=hello
-  capabilities: lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_b2x_obsmarkers_0
+  capabilities: lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_getbundle_obscommon
   $ curl --silent http://localhost:$HGPORT/?cmd=capabilities
-  lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_b2x_obsmarkers_0 (no-eol)
+  lookup changegroupsubset branchmap pushkey known getbundle unbundlehash batch stream unbundle=HG10GZ,HG10BZ,HG10UN httpheader=1024 _evoext_pushobsmarkers_0 _evoext_pullobsmarkers_0 _evoext_obshash_0 _evoext_getbundle_obscommon (no-eol)