diff hgext/simple4server.py @ 852:aa722de36179

exchange: fix pulling from simple for server

There is far too much code duplication down there.

--- a/hgext/simple4server.py	Mon Mar 03 21:23:52 2014 -0800
+++ b/hgext/simple4server.py	Mon Mar 03 21:37:18 2014 -0800
@@ -11,8 +11,13 @@
 import mercurial.obsolete
 mercurial.obsolete._enabled = True
 
+import struct
 from mercurial import wireproto
-from mercurial import extension
+from mercurial import extensions
+from mercurial import obsolete
+from cStringIO import StringIO
+from mercurial import node
+_pack = struct.pack
 
 def srv_pushobsmarkers(repo, proto):
     """wireprotocol command"""
@@ -33,6 +38,82 @@
         lock.release()
     return wireproto.pushres(0)
 
+def _encodemarkersstream(fp, markers):
+    fp.write(_pack('>B', 0))
+    for mark in markers:
+        fp.write(obsolete._encodeonemarker(mark))
+
+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
+
+class pruneobsstore(obsolete.obsstore):
+
+    def __init__(self, *args, **kwargs):
+        self.prunedchildren = {}
+        return super(pruneobsstore, self).__init__(*args, **kwargs)
+
+    def _load(self, markers):
+        markers = self._prunedetectingmarkers(markers)
+        return super(pruneobsstore, self)._load(markers)
+
+
+    def _prunedetectingmarkers(self, markers):
+        for m in markers:
+            if not m[1]: # no successors
+                meta = obsolete.decodemeta(m[3])
+                if 'p1' in meta:
+                    p1 = node.bin(meta['p1'])
+                    self.prunedchildren.setdefault(p1, set()).add(m)
+                if 'p2' in meta:
+                    p2 = node.bin(meta['p2'])
+                    self.prunedchildren.setdefault(p2, set()).add(m)
+            yield m
+
+def relevantmarkers(self, nodes):
+    """return a set of all obsolescence marker relevant to a set of node.
+
+    "relevant" to a set of node mean:
+
+    - marker that use this changeset as successors
+    - prune marker of direct children on this changeset.
+    - recursive application of the two rules on precursors of these markers
+
+    It  a set so you cannot rely on order"""
+    seennodes = set(nodes)
+    seenmarkers = set()
+    pendingnodes = set(nodes)
+    precursorsmarkers = self.precursors
+    prunedchildren = self.prunedchildren
+    while pendingnodes:
+        direct = set()
+        for current in pendingnodes:
+            direct.update(precursorsmarkers.get(current, ()))
+            direct.update(prunedchildren.get(current, ()))
+        direct -= seenmarkers
+        pendingnodes = set([m[0] for m in direct])
+        seenmarkers |= direct
+        pendingnodes -= seennodes
+        seennodes |= pendingnodes
+    return seenmarkers
+
 def srv_pullobsmarkers(repo, proto, others):
     opts = wireproto.options('', ['heads', 'common'], others)
     for k, v in opts.iteritems():
@@ -55,6 +136,8 @@
     return caps
 
 def extsetup(ui):
+    obsolete.obsstore = pruneobsstore
+    obsolete.obsstore.relevantmarkers = relevantmarkers
     wireproto.commands['evoext_pushobsmarkers_0'] = (srv_pushobsmarkers, '')
     wireproto.commands['evoext_pullobsmarkers_0'] = (srv_pullobsmarkers, '*')
     extensions.wrapfunction(wireproto, 'capabilities', capabilities)