1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220
|
# 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,
wireprotoserver,
wireprototypes,
wireprotov2server,
)
from mercurial.interfaces import (
repository,
util as interfaceutil,
)
from mercurial.utils import stringutil
CACHE = None
configtable = {}
configitem = registrar.configitem(configtable)
configitem(b'simplecache', b'cacheapi', default=False)
configitem(b'simplecache', b'cacheobjects', default=False)
configitem(b'simplecache', b'redirectsfile', default=None)
# API handler that makes cached keys available.
def handlecacherequest(rctx, req, res, checkperm, urlparts):
if rctx.repo.ui.configbool(b'simplecache', b'cacheobjects'):
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])
def cachedescriptor(req, repo):
return {}
wireprotoserver.API_HANDLERS[b'simplecache'] = {
b'config': (b'simplecache', b'cacheapi'),
b'handler': handlecacherequest,
b'apidescriptor': cachedescriptor,
}
@interfaceutil.implementer(repository.iwireprotocolcommandcacher)
class memorycacher(object):
def __init__(
self, ui, command, encodefn, redirecttargets, redirecthashes, req
):
self.ui = ui
self.encodefn = encodefn
self.redirecttargets = redirecttargets
self.redirecthashes = redirecthashes
self.req = req
self.key = None
self.cacheobjects = ui.configbool(b'simplecache', b'cacheobjects')
self.cacheapi = ui.configbool(b'simplecache', b'cacheapi')
self.buffered = []
ui.log(b'simplecache', b'cacher constructed for %s\n', command)
def __enter__(self):
return self
def __exit__(self, exctype, excvalue, exctb):
if exctype:
self.ui.log(b'simplecache', b'cacher exiting due to error\n')
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:
self.ui.log(b'simplecache', b'cache miss for %s\n', self.key)
return None
entry = CACHE[self.key]
self.ui.log(b'simplecache', b'cache hit for %s\n', self.key)
redirectable = True
if not self.cacheapi:
redirectable = False
elif not self.redirecttargets:
redirectable = False
else:
clienttargets = set(self.redirecttargets)
ourtargets = {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)
url = b'%s/%s' % (self.req.baseurl, b'/'.join(paths))
# 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,
)
response = wireprototypes.alternatelocationresponse(
url=url, mediatype=b'application/mercurial-cbor'
)
return {b'objs': [response]}
if self.cacheobjects:
return {
b'objs': entry,
}
else:
return {
b'objs': [wireprototypes.encodedresponse(entry)],
}
def onobject(self, obj):
if self.cacheobjects:
self.buffered.append(obj)
else:
self.buffered.extend(self.encodefn(obj))
yield obj
def onfinished(self):
self.ui.log(b'simplecache', b'storing cache entry for %s\n', self.key)
if self.cacheobjects:
CACHE[self.key] = self.buffered
else:
CACHE[self.key] = b''.join(self.buffered)
return []
def makeresponsecacher(
orig,
repo,
proto,
command,
args,
objencoderfn,
redirecttargets,
redirecthashes,
):
return memorycacher(
repo.ui,
command,
objencoderfn,
redirecttargets,
redirecthashes,
proto._req,
)
def loadredirecttargets(ui):
path = ui.config(b'simplecache', b'redirectsfile')
if not path:
return []
with open(path, 'rb') as fh:
s = fh.read()
return stringutil.evalpythonliteral(s)
def getadvertisedredirecttargets(orig, repo, proto):
return loadredirecttargets(repo.ui)
def extsetup(ui):
global CACHE
CACHE = util.lrucachedict(10000)
extensions.wrapfunction(
wireprotov2server, b'makeresponsecacher', makeresponsecacher
)
extensions.wrapfunction(
wireprotov2server,
b'getadvertisedredirecttargets',
getadvertisedredirecttargets,
)
|