wireprotosimplecache.py
220 lines
| 5.8 KiB
| text/x-python
|
PythonLexer
/ tests / wireprotosimplecache.py
|
r40057 | # wireprotosimplecache.py - Extension providing in-memory wire protocol cache | ||
# | ||||
# Copyright 2018 Gregory Szorc <gregory.szorc@gmail.com> | ||||
# | ||||
# This software may be used and distributed according to the terms of the | ||||
# GNU General Public License version 2 or any later version. | ||||
from __future__ import absolute_import | ||||
from mercurial import ( | ||||
extensions, | ||||
registrar, | ||||
util, | ||||
|
r40061 | wireprotoserver, | ||
|
r40057 | wireprototypes, | ||
wireprotov2server, | ||||
) | ||||
|
r43078 | from mercurial.interfaces import ( | ||
repository, | ||||
|
r43079 | util as interfaceutil, | ||
|
r43078 | ) | ||
|
r43346 | from mercurial.utils import stringutil | ||
|
r40057 | |||
CACHE = None | ||||
configtable = {} | ||||
configitem = registrar.configitem(configtable) | ||||
|
r43346 | configitem(b'simplecache', b'cacheapi', default=False) | ||
configitem(b'simplecache', b'cacheobjects', default=False) | ||||
configitem(b'simplecache', b'redirectsfile', default=None) | ||||
|
r40057 | |||
|
r40061 | # API handler that makes cached keys available. | ||
def handlecacherequest(rctx, req, res, checkperm, urlparts): | ||||
|
r40254 | if rctx.repo.ui.configbool(b'simplecache', b'cacheobjects'): | ||
|
r40061 | res.status = b'500 Internal Server Error' | ||
res.setbodybytes(b'cacheobjects not supported for api server') | ||||
return | ||||
if not urlparts: | ||||
res.status = b'200 OK' | ||||
res.headers[b'Content-Type'] = b'text/plain' | ||||
res.setbodybytes(b'simple cache server') | ||||
return | ||||
key = b'/'.join(urlparts) | ||||
if key not in CACHE: | ||||
res.status = b'404 Not Found' | ||||
res.headers[b'Content-Type'] = b'text/plain' | ||||
res.setbodybytes(b'key not found in cache') | ||||
return | ||||
res.status = b'200 OK' | ||||
res.headers[b'Content-Type'] = b'application/mercurial-cbor' | ||||
res.setbodybytes(CACHE[key]) | ||||
|
r43346 | |||
|
r40061 | def cachedescriptor(req, repo): | ||
return {} | ||||
|
r43346 | |||
|
r40061 | wireprotoserver.API_HANDLERS[b'simplecache'] = { | ||
|
r40254 | b'config': (b'simplecache', b'cacheapi'), | ||
b'handler': handlecacherequest, | ||||
b'apidescriptor': cachedescriptor, | ||||
|
r40061 | } | ||
|
r43346 | |||
|
r40057 | @interfaceutil.implementer(repository.iwireprotocolcommandcacher) | ||
class memorycacher(object): | ||||
|
r43346 | def __init__( | ||
self, ui, command, encodefn, redirecttargets, redirecthashes, req | ||||
): | ||||
|
r40057 | self.ui = ui | ||
self.encodefn = encodefn | ||||
|
r40061 | self.redirecttargets = redirecttargets | ||
self.redirecthashes = redirecthashes | ||||
self.req = req | ||||
|
r40057 | self.key = None | ||
|
r40254 | self.cacheobjects = ui.configbool(b'simplecache', b'cacheobjects') | ||
self.cacheapi = ui.configbool(b'simplecache', b'cacheapi') | ||||
|
r40057 | self.buffered = [] | ||
|
r40254 | ui.log(b'simplecache', b'cacher constructed for %s\n', command) | ||
|
r40057 | |||
def __enter__(self): | ||||
return self | ||||
def __exit__(self, exctype, excvalue, exctb): | ||||
if exctype: | ||||
|
r40254 | self.ui.log(b'simplecache', b'cacher exiting due to error\n') | ||
|
r40057 | |||
def adjustcachekeystate(self, state): | ||||
# Needed in order to make tests deterministic. Don't copy this | ||||
# pattern for production caches! | ||||
del state[b'repo'] | ||||
def setcachekey(self, key): | ||||
self.key = key | ||||
return True | ||||
def lookup(self): | ||||
if self.key not in CACHE: | ||||
|
r40254 | self.ui.log(b'simplecache', b'cache miss for %s\n', self.key) | ||
|
r40057 | return None | ||
entry = CACHE[self.key] | ||||
|
r40254 | self.ui.log(b'simplecache', b'cache hit for %s\n', self.key) | ||
|
r40057 | |||
|
r40061 | redirectable = True | ||
if not self.cacheapi: | ||||
redirectable = False | ||||
elif not self.redirecttargets: | ||||
redirectable = False | ||||
else: | ||||
clienttargets = set(self.redirecttargets) | ||||
ourtargets = set(t[b'name'] for t in loadredirecttargets(self.ui)) | ||||
# We only ever redirect to a single target (for now). So we don't | ||||
# need to store which target matched. | ||||
if not clienttargets & ourtargets: | ||||
redirectable = False | ||||
if redirectable: | ||||
paths = self.req.dispatchparts[:-3] | ||||
paths.append(b'simplecache') | ||||
paths.append(self.key) | ||||
|
r40206 | url = b'%s/%s' % (self.req.baseurl, b'/'.join(paths)) | ||
|
r40061 | |||
|
r43346 | # url = b'http://example.com/%s' % self.key | ||
self.ui.log( | ||||
b'simplecache', | ||||
b'sending content redirect for %s to ' b'%s\n', | ||||
self.key, | ||||
url, | ||||
) | ||||
|
r40061 | response = wireprototypes.alternatelocationresponse( | ||
|
r43346 | url=url, mediatype=b'application/mercurial-cbor' | ||
) | ||||
|
r40061 | |||
|
r40254 | return {b'objs': [response]} | ||
|
r40061 | |||
|
r40057 | if self.cacheobjects: | ||
return { | ||||
|
r40254 | b'objs': entry, | ||
|
r40057 | } | ||
else: | ||||
return { | ||||
|
r40254 | b'objs': [wireprototypes.encodedresponse(entry)], | ||
|
r40057 | } | ||
def onobject(self, obj): | ||||
if self.cacheobjects: | ||||
self.buffered.append(obj) | ||||
else: | ||||
self.buffered.extend(self.encodefn(obj)) | ||||
yield obj | ||||
def onfinished(self): | ||||
|
r40254 | self.ui.log(b'simplecache', b'storing cache entry for %s\n', self.key) | ||
|
r40057 | if self.cacheobjects: | ||
CACHE[self.key] = self.buffered | ||||
else: | ||||
CACHE[self.key] = b''.join(self.buffered) | ||||
return [] | ||||
|
r43346 | |||
def makeresponsecacher( | ||||
orig, | ||||
repo, | ||||
proto, | ||||
command, | ||||
args, | ||||
objencoderfn, | ||||
redirecttargets, | ||||
redirecthashes, | ||||
): | ||||
return memorycacher( | ||||
repo.ui, | ||||
command, | ||||
objencoderfn, | ||||
redirecttargets, | ||||
redirecthashes, | ||||
proto._req, | ||||
) | ||||
|
r40057 | |||
|
r40059 | def loadredirecttargets(ui): | ||
|
r40254 | path = ui.config(b'simplecache', b'redirectsfile') | ||
|
r40059 | if not path: | ||
return [] | ||||
with open(path, 'rb') as fh: | ||||
s = fh.read() | ||||
return stringutil.evalpythonliteral(s) | ||||
|
r43346 | |||
|
r40059 | def getadvertisedredirecttargets(orig, repo, proto): | ||
return loadredirecttargets(repo.ui) | ||||
|
r43346 | |||
|
r40057 | def extsetup(ui): | ||
global CACHE | ||||
CACHE = util.lrucachedict(10000) | ||||
|
r43346 | extensions.wrapfunction( | ||
wireprotov2server, b'makeresponsecacher', makeresponsecacher | ||||
) | ||||
extensions.wrapfunction( | ||||
wireprotov2server, | ||||
b'getadvertisedredirecttargets', | ||||
getadvertisedredirecttargets, | ||||
) | ||||