hgext/fastannotate/protocol.py
author Augie Fackler <augie@google.com>
Sun, 06 Oct 2019 09:45:02 -0400
changeset 43076 2372284d9457
parent 41365 876494fd967d
child 43077 687b865b95ad
permissions -rw-r--r--
formatting: blacken the codebase This is using my patch to black (https://github.com/psf/black/pull/826) so we don't un-wrap collection literals. Done with: hg files 'set:**.py - mercurial/thirdparty/** - "contrib/python-zstandard/**"' | xargs black -S # skip-blame mass-reformatting only # no-check-commit reformats foo_bar functions Differential Revision: https://phab.mercurial-scm.org/D6971
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
     1
# Copyright 2016-present Facebook. All Rights Reserved.
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
     2
#
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
     3
# protocol: logic for a server providing fastannotate support
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
     4
#
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
     5
# This software may be used and distributed according to the terms of the
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
     6
# GNU General Public License version 2 or any later version.
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
     7
from __future__ import absolute_import
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
     8
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
     9
import contextlib
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    10
import os
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    11
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    12
from mercurial.i18n import _
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    13
from mercurial import (
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    14
    error,
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    15
    extensions,
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    16
    hg,
39251
bb2b462f81da fastannotate: pconvert paths from the server for Windows
Matt Harbison <matt_harbison@yahoo.com>
parents: 39215
diff changeset
    17
    util,
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    18
    wireprotov1peer,
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    19
    wireprotov1server,
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    20
)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    21
from . import context
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    22
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    23
# common
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    24
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    25
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    26
def _getmaster(ui):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    27
    """get the mainbranch, and enforce it is set"""
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    28
    master = ui.config('fastannotate', 'mainbranch')
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    29
    if not master:
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    30
        raise error.Abort(
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    31
            _(
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    32
                'fastannotate.mainbranch is required '
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    33
                'for both the client and the server'
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    34
            )
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    35
        )
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    36
    return master
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    37
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    38
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    39
# server-side
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    40
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    41
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    42
def _capabilities(orig, repo, proto):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    43
    result = orig(repo, proto)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    44
    result.append('getannotate')
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    45
    return result
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    46
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    47
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    48
def _getannotate(repo, proto, path, lastnode):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    49
    # output:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    50
    #   FILE := vfspath + '\0' + str(size) + '\0' + content
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    51
    #   OUTPUT := '' | FILE + OUTPUT
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    52
    result = ''
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    53
    buildondemand = repo.ui.configbool(
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    54
        'fastannotate', 'serverbuildondemand', True
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    55
    )
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    56
    with context.annotatecontext(repo, path) as actx:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    57
        if buildondemand:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    58
            # update before responding to the client
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    59
            master = _getmaster(repo.ui)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    60
            try:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    61
                if not actx.isuptodate(master):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    62
                    actx.annotate(master, master)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    63
            except Exception:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    64
                # non-fast-forward move or corrupted. rebuild automically.
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    65
                actx.rebuild()
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    66
                try:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    67
                    actx.annotate(master, master)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    68
                except Exception:
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    69
                    actx.rebuild()  # delete files
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    70
            finally:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    71
                # although the "with" context will also do a close/flush, we
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    72
                # need to do it early so we can send the correct respond to
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    73
                # client.
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    74
                actx.close()
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    75
        # send back the full content of revmap and linelog, in the future we
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    76
        # may want to do some rsync-like fancy updating.
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    77
        # the lastnode check is not necessary if the client and the server
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    78
        # agree where the main branch is.
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    79
        if actx.lastnode != lastnode:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    80
            for p in [actx.revmappath, actx.linelogpath]:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    81
                if not os.path.exists(p):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    82
                    continue
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    83
                with open(p, 'rb') as f:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    84
                    content = f.read()
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    85
                vfsbaselen = len(repo.vfs.base + '/')
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    86
                relpath = p[vfsbaselen:]
39620
b220851999b5 py3: use '%d' for integers instead of '%s'
Pulkit Goyal <pulkit@yandex-team.ru>
parents: 39543
diff changeset
    87
                result += '%s\0%d\0%s' % (relpath, len(content), content)
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    88
    return result
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    89
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    90
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    91
def _registerwireprotocommand():
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    92
    if 'getannotate' in wireprotov1server.commands:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    93
        return
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    94
    wireprotov1server.wireprotocommand('getannotate', 'path lastnode')(
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    95
        _getannotate
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    96
    )
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
    97
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    98
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
    99
def serveruisetup(ui):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   100
    _registerwireprotocommand()
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   101
    extensions.wrapfunction(wireprotov1server, '_capabilities', _capabilities)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   102
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   103
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   104
# client-side
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   105
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   106
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   107
def _parseresponse(payload):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   108
    result = {}
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   109
    i = 0
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   110
    l = len(payload) - 1
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   111
    state = 0  # 0: vfspath, 1: size
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   112
    vfspath = size = ''
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   113
    while i < l:
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   114
        ch = payload[i : i + 1]
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   115
        if ch == '\0':
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   116
            if state == 1:
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   117
                result[vfspath] = payload[i + 1 : i + 1 + int(size)]
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   118
                i += int(size)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   119
                state = 0
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   120
                vfspath = size = ''
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   121
            elif state == 0:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   122
                state = 1
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   123
        else:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   124
            if state == 1:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   125
                size += ch
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   126
            elif state == 0:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   127
                vfspath += ch
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   128
        i += 1
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   129
    return result
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   130
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   131
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   132
def peersetup(ui, peer):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   133
    class fastannotatepeer(peer.__class__):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   134
        @wireprotov1peer.batchable
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   135
        def getannotate(self, path, lastnode=None):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   136
            if not self.capable('getannotate'):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   137
                ui.warn(_('remote peer cannot provide annotate cache\n'))
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   138
                yield None, None
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   139
            else:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   140
                args = {'path': path, 'lastnode': lastnode or ''}
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   141
                f = wireprotov1peer.future()
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   142
                yield args, f
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   143
                yield _parseresponse(f.value)
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   144
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   145
    peer.__class__ = fastannotatepeer
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   146
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   147
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   148
@contextlib.contextmanager
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   149
def annotatepeer(repo):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   150
    ui = repo.ui
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   151
39213
303dae0136b0 fastannotate: rip out specialized support for remotefilelog
Augie Fackler <augie@google.com>
parents: 39211
diff changeset
   152
    remotepath = ui.expandpath(
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   153
        ui.config('fastannotate', 'remotepath', 'default')
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   154
    )
39213
303dae0136b0 fastannotate: rip out specialized support for remotefilelog
Augie Fackler <augie@google.com>
parents: 39211
diff changeset
   155
    peer = hg.peer(ui, {}, remotepath)
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   156
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   157
    try:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   158
        yield peer
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   159
    finally:
39213
303dae0136b0 fastannotate: rip out specialized support for remotefilelog
Augie Fackler <augie@google.com>
parents: 39211
diff changeset
   160
        peer.close()
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   161
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   162
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   163
def clientfetch(repo, paths, lastnodemap=None, peer=None):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   164
    """download annotate cache from the server for paths"""
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   165
    if not paths:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   166
        return
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   167
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   168
    if peer is None:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   169
        with annotatepeer(repo) as peer:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   170
            return clientfetch(repo, paths, lastnodemap, peer)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   171
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   172
    if lastnodemap is None:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   173
        lastnodemap = {}
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   174
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   175
    ui = repo.ui
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   176
    results = []
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   177
    with peer.commandexecutor() as batcher:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   178
        ui.debug('fastannotate: requesting %d files\n' % len(paths))
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   179
        for p in paths:
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   180
            results.append(
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   181
                batcher.callcommand(
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   182
                    'getannotate', {'path': p, 'lastnode': lastnodemap.get(p)}
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   183
                )
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   184
            )
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   185
39715
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   186
        for result in results:
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   187
            r = result.result()
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   188
            # TODO: pconvert these paths on the server?
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   189
            r = {util.pconvert(p): v for p, v in r.iteritems()}
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   190
            for path in sorted(r):
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   191
                # ignore malicious paths
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   192
                if not path.startswith('fastannotate/') or '/../' in (
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   193
                    path + '/'
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   194
                ):
39715
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   195
                    ui.debug('fastannotate: ignored malicious path %s\n' % path)
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   196
                    continue
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   197
                content = r[path]
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   198
                if ui.debugflag:
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   199
                    ui.debug(
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   200
                        'fastannotate: writing %d bytes to %s\n'
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   201
                        % (len(content), path)
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   202
                    )
39715
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   203
                repo.vfs.makedirs(os.path.dirname(path))
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   204
                with repo.vfs(path, 'wb') as f:
d8a7690ccc74 fastannotate: process files as they arrive
Martin von Zweigbergk <martinvonz@google.com>
parents: 39620
diff changeset
   205
                    f.write(content)
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   206
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   207
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   208
def _filterfetchpaths(repo, paths):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   209
    """return a subset of paths whose history is long and need to fetch linelog
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   210
    from the server. works with remotefilelog and non-remotefilelog repos.
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   211
    """
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   212
    threshold = repo.ui.configint('fastannotate', 'clientfetchthreshold', 10)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   213
    if threshold <= 0:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   214
        return paths
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   215
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   216
    result = []
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   217
    for path in paths:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   218
        try:
39213
303dae0136b0 fastannotate: rip out specialized support for remotefilelog
Augie Fackler <augie@google.com>
parents: 39211
diff changeset
   219
            if len(repo.file(path)) >= threshold:
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   220
                result.append(path)
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   221
        except Exception:  # file not found etc.
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   222
            result.append(path)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   223
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   224
    return result
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   225
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   226
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   227
def localreposetup(ui, repo):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   228
    class fastannotaterepo(repo.__class__):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   229
        def prefetchfastannotate(self, paths, peer=None):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   230
            master = _getmaster(self.ui)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   231
            needupdatepaths = []
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   232
            lastnodemap = {}
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   233
            try:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   234
                for path in _filterfetchpaths(self, paths):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   235
                    with context.annotatecontext(self, path) as actx:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   236
                        if not actx.isuptodate(master, strict=False):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   237
                            needupdatepaths.append(path)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   238
                            lastnodemap[path] = actx.lastnode
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   239
                if needupdatepaths:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   240
                    clientfetch(self, needupdatepaths, lastnodemap, peer)
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   241
            except Exception as ex:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   242
                # could be directory not writable or so, not fatal
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   243
                self.ui.debug('fastannotate: prefetch failed: %r\n' % ex)
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   244
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   245
    repo.__class__ = fastannotaterepo
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   246
43076
2372284d9457 formatting: blacken the codebase
Augie Fackler <augie@google.com>
parents: 41365
diff changeset
   247
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   248
def clientreposetup(ui, repo):
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   249
    _registerwireprotocommand()
39543
2182e67ea912 fastannotate: use repo.local()
Gregory Szorc <gregory.szorc@gmail.com>
parents: 39251
diff changeset
   250
    if repo.local():
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   251
        localreposetup(ui, repo)
39214
8da20fc9fc07 fastannotate: move some global state mutation to extsetup()
Augie Fackler <augie@google.com>
parents: 39213
diff changeset
   252
    # TODO: this mutates global state, but only if at least one repo
8da20fc9fc07 fastannotate: move some global state mutation to extsetup()
Augie Fackler <augie@google.com>
parents: 39213
diff changeset
   253
    # has the extension enabled. This is probably bad for hgweb.
39210
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   254
    if peersetup not in hg.wirepeersetupfuncs:
1ddb296e0dee fastannotate: initial import from Facebook's hg-experimental
Augie Fackler <augie@google.com>
parents:
diff changeset
   255
        hg.wirepeersetupfuncs.append(peersetup)