annotate hgext/pushexperiment.py @ 727:cb907cf3b556

Add a new pushexperiment extension

This extension introduce a faster way to push obsolescence marker. Send the
obsstore directly instead of inefficiently going through pushkey's
inefficiency.

"""Small extension altering some push behavior

- Add a new wire protocol command to exchange obsolescence marker. Sending the
  raw file as a binary instead of using pushkey hack.

"""

import errno
from StringIO import StringIO

from mercurial.i18n import _
from mercurial import extensions
from mercurial import wireproto
from mercurial import obsolete


def client_pushobsmarkers(self, obsfile):
    """wireprotocol peer method"""
    self.requirecap('_push_experiment_pushobsmarkers_0',
                    _('push obsolete markers faster'))
    ret, output = self._callpush('push_experiment_pushobsmarkers_0', obsfile)
    for l in output.splitlines(True):
        self.ui.status(_('remote: '), l)
    return ret


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 syncpush(orig, repo, remote):
    """wraper for obsolete.syncpush to use the fast way if possible"""
    if not (obsolete._enabled and repo.obsstore):
        return
    if remote.capable('_push_experiment_pushobsmarkers_0'):
        try:
            obsfp = repo.sopener('obsstore')
        except IOError as e:
            if e.errno != errno.ENOENT:
                raise
            return
        remote.push_experiment_pushobsmarkers_0(obsfp)
        return
    return orig(repo, remote)


def capabilities(orig, repo, proto):
    """wrapper to advertise new capability"""
    caps = orig(repo, proto)
    if obsolete._enabled:
        caps += ' _push_experiment_pushobsmarkers_0'
    return caps


def extsetup(ui):
    wireproto.wirepeer.push_experiment_pushobsmarkers_0 = client_pushobsmarkers
    wireproto.commands['push_experiment_pushobsmarkers_0'] = (srv_pushobsmarkers, '')
    extensions.wrapfunction(wireproto, 'capabilities', capabilities)
    extensions.wrapfunction(obsolete, 'syncpush', syncpush)