1 # wireprotosimplecache.py - Extension providing in-memory wire protocol cache

     2 #

     3 # Copyright 2018 Gregory Szorc <gregory.szorc@gmail.com>

     4 #

     5 # This software may be used and distributed according to the terms of the

     6 # GNU General Public License version 2 or any later version.

     7 

     8 from __future__ import absolute_import

     9 

    10 from mercurial import (

    11     extensions,

    12     registrar,

    13     repository,

    14     util,

    15     wireprototypes,

    16     wireprotov2server,

    17 )

    18 from mercurial.utils import (

    19     interfaceutil,

    20 )

    21 

    22 CACHE = None

    23 

    24 configtable = {}

    25 configitem = registrar.configitem(configtable)

    26 

    27 configitem('simplecache', 'cacheobjects',

    28            default=False)

    29 

    30 @interfaceutil.implementer(repository.iwireprotocolcommandcacher)

    31 class memorycacher(object):

    32     def __init__(self, ui, command, encodefn):

    33         self.ui = ui

    34         self.encodefn = encodefn

    35         self.key = None

    36         self.cacheobjects = ui.configbool('simplecache', 'cacheobjects')

    37         self.buffered = []

    38 

    39         ui.log('simplecache', 'cacher constructed for %s\n', command)

    40 

    41     def __enter__(self):

    42         return self

    43 

    44     def __exit__(self, exctype, excvalue, exctb):

    45         if exctype:

    46             self.ui.log('simplecache', 'cacher exiting due to error\n')

    47 

    48     def adjustcachekeystate(self, state):

    49         # Needed in order to make tests deterministic. Don't copy this

    50         # pattern for production caches!

    51         del state[b'repo']

    52 

    53     def setcachekey(self, key):

    54         self.key = key

    55         return True

    56 

    57     def lookup(self):

    58         if self.key not in CACHE:

    59             self.ui.log('simplecache', 'cache miss for %s\n', self.key)

    60             return None

    61 

    62         entry = CACHE[self.key]

    63         self.ui.log('simplecache', 'cache hit for %s\n', self.key)

    64 

    65         if self.cacheobjects:

    66             return {

    67                 'objs': entry,

    68             }

    69         else:

    70             return {

    71                 'objs': [wireprototypes.encodedresponse(entry)],

    72             }

    73 

    74     def onobject(self, obj):

    75         if self.cacheobjects:

    76             self.buffered.append(obj)

    77         else:

    78             self.buffered.extend(self.encodefn(obj))

    79 

    80         yield obj

    81 

    82     def onfinished(self):

    83         self.ui.log('simplecache', 'storing cache entry for %s\n', self.key)

    84         if self.cacheobjects:

    85             CACHE[self.key] = self.buffered

    86         else:

    87             CACHE[self.key] = b''.join(self.buffered)

    88 

    89         return []

    90 

    91 def makeresponsecacher(orig, repo, proto, command, args, objencoderfn):

    92     return memorycacher(repo.ui, command, objencoderfn)

    93 

    94 def extsetup(ui):

    95     global CACHE

    96 

    97     CACHE = util.lrucachedict(10000)

    98 

    99     extensions.wrapfunction(wireprotov2server, 'makeresponsecacher',

   100                             makeresponsecacher)