tests/wireprotosimplecache.py

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/tests/wireprotosimplecache.py	Wed Sep 26 17:16:56 2018 -0700
@@ -0,0 +1,100 @@
+# wireprotosimplecache.py - Extension providing in-memory wire protocol cache
+#
+# Copyright 2018 Gregory Szorc <[e-mail]>
+#
+# 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,
+    repository,
+    util,
+    wireprototypes,
+    wireprotov2server,
+)
+from mercurial.utils import (
+    interfaceutil,
+)
+
+CACHE = None
+
+configtable = {}
+configitem = registrar.configitem(configtable)
+
+configitem('simplecache', 'cacheobjects',
+           default=False)
+
[e-mail](repository.iwireprotocolcommandcacher)
+class memorycacher(object):
+    def __init__(self, ui, command, encodefn):
+        self.ui = ui
+        self.encodefn = encodefn
+        self.key = None
+        self.cacheobjects = ui.configbool('simplecache', 'cacheobjects')
+        self.buffered = []
+
+        ui.log('simplecache', 'cacher constructed for %s\n', command)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exctype, excvalue, exctb):
+        if exctype:
+            self.ui.log('simplecache', '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('simplecache', 'cache miss for %s\n', self.key)
+            return None
+
+        entry = CACHE[self.key]
+        self.ui.log('simplecache', 'cache hit for %s\n', self.key)
+
+        if self.cacheobjects:
+            return {
+                'objs': entry,
+            }
+        else:
+            return {
+                '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('simplecache', '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):
+    return memorycacher(repo.ui, command, objencoderfn)
+
+def extsetup(ui):
+    global CACHE
+
+    CACHE = util.lrucachedict(10000)
+
+    extensions.wrapfunction(wireprotov2server, 'makeresponsecacher',
+                            makeresponsecacher)