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

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/hgext/pushexperiment.py	Mon Jun 03 14:53:24 2013 +0200
@@ -0,0 +1,77 @@
+"""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)
+
+