hgext/evolve.py
changeset 823 9aa20585e158
parent 821 202376586cf6
child 824 fed090e07621
--- a/hgext/evolve.py	Fri Feb 28 00:55:34 2014 -0800
+++ b/hgext/evolve.py	Fri Feb 28 13:25:34 2014 -0800
@@ -63,6 +63,7 @@
 from mercurial.commands import walkopts, commitopts, commitopts2
 from mercurial.node import nullid
 from mercurial import wireproto
+from mercurial import localrepo
 
 _pack = struct.pack
 
@@ -2183,14 +2184,92 @@
         lock.release()
     return wireproto.pushres(0)
 
+@eh.wrapfunction(exchange, '_pullobsolete')
+def _pullobsolete(orig, pullop):
+    if not obsolete._enabled:
+        return None
+    if not pullop.remote.capable('_evoext_pullobsmarkers_0'):
+        return orig(pullop)
+    tr = None
+    ui = pullop.repo.ui
+    ui.status("OBSEXC: pull obsolescence markers\n")
+    obsdata = pullop.remote.evoext_pullobsmarkers_0()
+    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()
+    ui.status("OBSEXC: DONE\n")
+    return tr
+
+def _getobsmarkersstream(repo, heads=None, common=None):
+    revset = ''
+    args = []
+    repo = repo.unfiltered()
+    if heads is None:
+        revset = 'all()'
+    elif heads:
+        revset += "(::%ln)"
+        args.append(heads)
+    else:
+        assert False, 'pulling no heads?'
+    if common:
+        revset += ' - (::%ln)'
+        args.append(common)
+    nodes = [c.node() for c in repo.set(revset, *args)]
+    markers = repo.obsstore.relevantmarkers(nodes)
+    obsdata = StringIO()
+    _encodemarkersstream(obsdata, markers)
+    obsdata.seek(0)
+    return obsdata
+
+@eh.addattr(wireproto.wirepeer, 'evoext_pullobsmarkers_0')
+def client_pullobsmarkers(self, heads=None, common=None):
+    self.requirecap('_evoext_pullobsmarkers_0', _('look up remote obsmarkers'))
+    opts = {}
+    if heads is not None:
+        opts['heads'] = wireproto.encodelist(heads)
+    if common is not None:
+        opts['common'] = wireproto.encodelist(common)
+    f = self._callstream("evoext_pullobsmarkers_0", **opts)
+    f = self._decompress(f)
+    length= int(f.read(20))
+    return StringIO(f.read(length))
+
+@eh.addattr(localrepo.localpeer, 'evoext_pullobsmarkers_0')
+def local_pullobsmarkers(self, heads=None, common=None):
+    return _getobsmarkersstream(self._repo, heads=heads, common=common)
+
+def srv_pullobsmarkers(repo, proto, others):
+    opts = wireproto.options('', ['heads', 'common'], others)
+    for k, v in opts.iteritems():
+        if k in ('heads', 'common'):
+            opts[k] = wireproto.decodelist(v)
+    obsdata = _getobsmarkersstream(repo, **opts)
+    length = '%20i' % len(obsdata.getvalue())
+    def data():
+        yield length
+        for c in proto.groupchunks(obsdata):
+            yield c
+    return wireproto.streamres(data())
+
 @eh.wrapfunction(wireproto, 'capabilities')
 def capabilities(orig, repo, proto):
     """wrapper to advertise new capability"""
     caps = orig(repo, proto)
     if obsolete._enabled:
         caps += ' _evoext_pushobsmarkers_0'
+        caps += ' _evoext_pullobsmarkers_0'
     return caps
 
 @eh.extsetup
 def _installwireprotocol(ui):
+    localrepo.MODERNCAPS.add('_evoext_pullobsmarkers_0')
     wireproto.commands['evoext_pushobsmarkers_0'] = (srv_pushobsmarkers, '')
+    wireproto.commands['evoext_pullobsmarkers_0'] = (srv_pullobsmarkers, '*')