view hgext/simple4server.py @ 852:aa722de36179

exchange: fix pulling from simple for server

There is far too much code duplication down there.

'''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.'''

import mercurial.obsolete
mercurial.obsolete._enabled = True

import struct
from mercurial import wireproto
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"""
    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)

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():
        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())

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

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)