1 """Small extension altering some push behavior

     2 

     3 - Add a new wire protocol command to exchange obsolescence markers. Sending the

     4   raw file as a binary instead of using pushkey hack.

     5 - Add a "push done" notification

     6 - Push obsolescence marker before anything else (This works around the lack

     7 of global transaction)

     8 

     9 """

    10 

    11 import errno

    12 from StringIO import StringIO

    13 

    14 from mercurial.i18n import _

    15 from mercurial import error

    16 from mercurial import extensions

    17 from mercurial import wireproto

    18 from mercurial import obsolete

    19 from mercurial import localrepo

    20 

    21 

    22 def client_pushobsmarkers(self, obsfile):

    23     """wireprotocol peer method"""

    24     self.requirecap('_push_experiment_pushobsmarkers_0',

    25                     _('push obsolete markers faster'))

    26     ret, output = self._callpush('push_experiment_pushobsmarkers_0', obsfile)

    27     for l in output.splitlines(True):

    28         self.ui.status(_('remote: '), l)

    29     return ret

    30 

    31 

    32 def srv_pushobsmarkers(repo, proto):

    33     """wireprotocol command"""

    34     fp = StringIO()

    35     proto.redirect()

    36     proto.getfile(fp)

    37     data = fp.getvalue()

    38     fp.close()

    39     lock = repo.lock()

    40     try:

    41         tr = repo.transaction('pushkey: obsolete markers')

    42         try:

    43             repo.obsstore.mergemarkers(tr, data)

    44             tr.close()

    45         finally:

    46             tr.release()

    47     finally:

    48         lock.release()

    49     return wireproto.pushres(0)

    50 

    51 

    52 def syncpush(orig, repo, remote):

    53     """wraper for obsolete.syncpush to use the fast way if possible"""

    54     if not (obsolete.isenabled(repo, obsolete.exchangeopt) and

    55             repo.obsstore):

    56         return

    57     if remote.capable('_push_experiment_pushobsmarkers_0'):

    58         return # already pushed before changeset

    59         remote.push_experiment_pushobsmarkers_0(obsfp)

    60         return

    61     return orig(repo, remote)

    62 

    63 

    64 def client_notifypushend(self):

    65     """wire peer  command to notify a push is done"""

    66     self.requirecap('_push_experiment_notifypushend_0',

    67                     _('hook once push is all done'))

    68     return self._call('push_experiment_notifypushend_0')

    69 

    70 

    71 def srv_notifypushend(repo, proto):

    72     """wire protocol command to notify a push is done"""

    73     proto.redirect()

    74     repo.hook('notifypushend')

    75     return wireproto.pushres(0)

    76 

    77 

    78 def augmented_push(orig, repo, remote, *args, **kwargs):

    79     """push wrapped that call the wire protocol command"""

    80     if not remote.canpush():

    81         raise error.Abort(_("destination does not support push"))

    82     if (obsolete.isenabled(repo, obsolete.exchangeopt) and repo.obsstore

    83         and remote.capable('_push_experiment_pushobsmarkers_0')):

    84         # push marker early to limit damage of pushing too early.

    85         try:

    86             obsfp = repo.svfs('obsstore')

    87         except IOError as e:

    88             if e.errno != errno.ENOENT:

    89                 raise

    90         else:

    91             remote.push_experiment_pushobsmarkers_0(obsfp)

    92     ret = orig(repo, remote, *args, **kwargs)

    93     if remote.capable('_push_experiment_notifypushend_0'):

    94         remote.push_experiment_notifypushend_0()

    95     return ret

    96 

    97 

    98 def capabilities(orig, repo, proto):

    99     """wrapper to advertise new capability"""

   100     caps = orig(repo, proto)

   101     if obsolete.isenabled(repo, obsolete.exchangeopt):

   102         caps += ' _push_experiment_pushobsmarkers_0'

   103     caps += ' _push_experiment_notifypushend_0'

   104     return caps

   105 

   106 

   107 def extsetup(ui):

   108     wireproto.wirepeer.push_experiment_pushobsmarkers_0 = client_pushobsmarkers

   109     wireproto.wirepeer.push_experiment_notifypushend_0 = client_notifypushend

   110     wireproto.commands['push_experiment_pushobsmarkers_0'] = \

   111         (srv_pushobsmarkers, '')

   112     wireproto.commands['push_experiment_notifypushend_0'] = \

   113         (srv_notifypushend, '')

   114     extensions.wrapfunction(wireproto, 'capabilities', capabilities)

   115     extensions.wrapfunction(obsolete, 'syncpush', syncpush)

   116     extensions.wrapfunction(localrepo.localrepository, 'push', augmented_push)

   117 

   118