annotate hgext/simple4server.py @ 875:10867a8e27c6

exchange: properly wrap wireprotocol capabilities

Just wrapping the module function only works for ssh. Ssh gets capability
through the `hello` command. We also need to wrap the function stored in the
command dict. This fix discovery for http peer.

'''enable experimental obsolescence feature of Mercurial

OBSOLESCENCE IS AN EXPERIMENTAL FEATURE MAKE SURE YOU UNDERSTOOD THE INVOLVED
CONCEPT BEFORE USING IT.

/!\ THIS EXTENSION IS INTENDED FOR SERVER SIDE ONLY USAGE /!\

For client side usages it is recommended to use the evolve extension for
improved user interface.'''

testedwith = '2.9.1'
buglink = 'https://bitbucket.org/marmoute/mutable-history/issues'

import mercurial.obsolete
mercurial.obsolete._enabled = True

import struct
from mercurial import util
from mercurial import wireproto
from mercurial import extensions
from mercurial import obsolete
from cStringIO import StringIO
from mercurial import node
from mercurial.hgweb import hgweb_mod
_pack = struct.pack

# from evolve extension: 1a23c7c52a43
def srv_pushobsmarkers(repo, proto):
    """That receives a stream of markers and apply then to the repo"""
    fp = StringIO()
    proto.redirect()
    proto.getfile(fp)
    data = fp.getvalue()
    fp.close()
    lock = repo.lock()
    try:
        tr = repo.transaction('pushkey: obsolete markers')
        try:
            repo.obsstore.mergemarkers(tr, data)
            tr.close()
        finally:
            tr.release()
    finally:
        lock.release()
    return wireproto.pushres(0)

# from mercurial.obsolete: 19e9478c1a22
def _encodemarkersstream(fp, markers):
    """write a binary version of a set of markers

    Includes the initial version number"""
    fp.write(_pack('>B', 0))
    for mark in markers:
        fp.write(obsolete._encodeonemarker(mark))

# from evolve extension: 1a23c7c52a43
def _getobsmarkersstream(repo, heads=None, common=None):
    """Get a binary stream for all markers relevant to `::<heads> - ::<common>`
    """
    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

# from evolve extension: 1a23c7c52a43
class pruneobsstore(obsolete.obsstore):
    """And extended obsstore class that read parent information from v1 format

    Evolve extension adds parent information in prune marker. We use it to make
    markers relevant to pushed changeset."""

    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

# from evolve extension: 1a23c7c52a43
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 is 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

# from evolve extension: 1a23c7c52a43
def srv_pullobsmarkers(repo, proto, others):
    """serves a binary stream of markers.

    Serves relevant to changeset between heads and common. The stream is prefix
    by a -string- representation of an integer. This integer is the size of the
    stream."""
    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())


# from evolve extension: 1a23c7c52a43
def _obsrelsethashtree(repo):
    """Build an obshash for every node in a repo

    return a [(node), (obshash)] list. in revision order."""
    cache = []
    unfi = repo.unfiltered()
    for i in unfi:
        ctx = unfi[i]
        entry = 0
        sha = util.sha1()
        # add data from p1
        for p in ctx.parents():
            p = p.rev()
            if p < 0:
                p = node.nullid
            else:
                p = cache[p][1]
            if p != node.nullid:
                entry += 1
                sha.update(p)
        tmarkers = repo.obsstore.relevantmarkers([ctx.node()])
        if tmarkers:
            bmarkers = [obsolete._encodeonemarker(m) for m in tmarkers]
            bmarkers.sort()
            for m in bmarkers:
                entry += 1
                sha.update(m)
        if entry:
            cache.append((ctx.node(), sha.digest()))
        else:
            cache.append((ctx.node(), node.nullid))
    return cache

# from evolve extension: 1a23c7c52a43
def _obshash(repo, nodes):
    """hash of binary version of relevant markers + obsparent

    (special case so that all empty are hashed as nullid)"""
    hashs = _obsrelsethashtree(repo)
    nm = repo.changelog.nodemap
    return  [hashs[nm.get(n)][1] for n in nodes]

# from evolve extension: 1a23c7c52a43
def srv_obshash(repo, proto, nodes):
    """give the obshash of a a set of node

    Used for markes discovery"""
    return wireproto.encodelist(_obshash(repo, wireproto.decodelist(nodes)))

# from evolve extension: 1a23c7c52a43
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'
        caps += ' _evoext_obshash_0'
    return caps

# from evolve extension: 1a23c7c52a43
# heavily modified
def extsetup(ui):
    obsolete.obsstore = pruneobsstore
    obsolete.obsstore.relevantmarkers = relevantmarkers
    hgweb_mod.perms['evoext_pushobsmarkers_0'] = 'push'
    hgweb_mod.perms['evoext_pullobsmarkers_0'] = 'pull'
    hgweb_mod.perms['evoext_obshash'] = 'pull'
    wireproto.commands['evoext_pushobsmarkers_0'] = (srv_pushobsmarkers, '')
    wireproto.commands['evoext_pullobsmarkers_0'] = (srv_pullobsmarkers, '*')
    # wrap module content
    extensions.wrapfunction(wireproto, 'capabilities', capabilities)
    # wrap command content
    oldcap, args = wireproto.commands['capabilities']
    def newcap(repo, proto):
        return capabilities(oldcap, repo, proto)
    wireproto.commands['capabilities'] = (newcap, args)
    wireproto.commands['evoext_obshash'] = (srv_obshash, 'nodes')