diff --git a/djblets/util/misc.py b/djblets/util/misc.py index df40bb8d0876a1523d4178ade85b08520c5549f2..bc1d52ef4843c90ed2ca7498b381a3a99d5f4b08 100644 --- a/djblets/util/misc.py +++ b/djblets/util/misc.py @@ -66,7 +66,7 @@ class MissingChunkError(Exception): pass -def _cache_fetch_large_data(cache, key): +def _cache_fetch_large_data(cache, key, compress_large_data): chunk_count = cache.get(key) data = [] @@ -81,7 +81,9 @@ def _cache_fetch_large_data(cache, key): data = ''.join(data) - data = zlib.decompress(data) + if compress_large_data: + data = zlib.decompress(data) + try: unpickler = pickle.Unpickler(StringIO(data)) data = unpickler.load() @@ -92,7 +94,7 @@ def _cache_fetch_large_data(cache, key): return data -def _cache_store_large_data(cache, key, data, expiration): +def _cache_store_large_data(cache, key, data, expiration, compress_large_data): # We store large data in the cache broken into chunks that are 1M in size. # To do this easily, we first pickle the data and compress it with zlib. # This gives us a string which can be chunked easily. These are then stored @@ -103,13 +105,15 @@ def _cache_store_large_data(cache, key, data, expiration): pickler = pickle.Pickler(file) pickler.dump(data) data = file.getvalue() - data = zlib.compress(data) + + if compress_large_data: + data = zlib.compress(data) i = 0 while len(data) > CACHE_CHUNK_SIZE: chunk = data[0:CACHE_CHUNK_SIZE] data = data[CACHE_CHUNK_SIZE:] - cache.set('%s-%d' % (key, i), chunk, expiration) + cache.set('%s-%d' % (key, i), [chunk], expiration) i += 1 cache.set('%s-%d' % (key, i), [data], expiration) @@ -120,18 +124,21 @@ def cache_memoize(key, lookup_callable, expiration=getattr(settings, "CACHE_EXPIRATION_TIME", DEFAULT_EXPIRATION_TIME), force_overwrite=False, - large_data=False): + large_data=False, + compress_large_data=True): """Memoize the results of a callable inside the configured cache. Keyword arguments: - expiration -- The expiration time for the key. - force_overwrite -- If True, the value will always be computed and stored - regardless of whether it exists in the cache already. - large_data -- If True, the resulting data will be pickled, gzipped, - and (potentially) split up into megabyte-sized chunks. - This is useful for very large, computationally - intensive hunks of data which we don't want to store - in a database due to the way things are accessed. + expiration -- The expiration time for the key. + force_overwrite -- If True, the value will always be computed and stored + regardless of whether it exists in the cache already. + large_data -- If True, the resulting data will be pickled, gzipped, + and (potentially) split up into megabyte-sized chunks. + This is useful for very large, computationally + intensive hunks of data which we don't want to store + in a database due to the way things are accessed. + compress_large_data -- Compresses the data with zlib compression when + large_data is True. """ try: site = Site.objects.get_current() @@ -156,7 +163,7 @@ def cache_memoize(key, lookup_callable, if large_data: if not force_overwrite and cache.has_key(key): try: - data = _cache_fetch_large_data(cache, key) + data = _cache_fetch_large_data(cache, key, compress_large_data) return data except Exception, e: logging.warning('Failed to fetch large data from cache for key %s: %s.' % (key, e)) @@ -164,7 +171,8 @@ def cache_memoize(key, lookup_callable, logging.info('Cache miss for key %s.' % key) data = lookup_callable() - _cache_store_large_data(cache, key, data, expiration) + _cache_store_large_data(cache, key, data, expiration, + compress_large_data) return data else: diff --git a/djblets/util/tests.py b/djblets/util/tests.py index bb51e4a4076fea9506a1cdbbd69452d491a6ea3c..106833d2b84a15f5b5f25e03c912ad2c18d36cb3 100644 --- a/djblets/util/tests.py +++ b/djblets/util/tests.py @@ -27,13 +27,15 @@ import datetime import unittest +from django.contrib.sites.models import Site +from django.core.cache import cache from django.http import HttpRequest from django.template import Token, TOKEN_TEXT, TemplateSyntaxError from django.utils.html import strip_spaces_between_tags from djblets.util.http import get_http_accept_lists, \ get_http_requested_mimetype -from djblets.util.misc import cache_memoize +from djblets.util.misc import cache_memoize, CACHE_CHUNK_SIZE from djblets.util.testing import TestCase, TagTest from djblets.util.templatetags import djblets_deco from djblets.util.templatetags import djblets_email @@ -45,14 +47,14 @@ def normalize_html(s): class CacheTest(TestCase): - def testCacheMemoize(self): + def test_cache_memoize(self): """Testing cache_memoize""" cacheKey = "abc123" testStr = "Test 123" - def cacheFunc(cacheCalled=False): - self.assert_(not cacheCalled) - cacheCalled = True + def cacheFunc(cacheCalled=[]): + self.assertTrue(not cacheCalled) + cacheCalled.append(True) return testStr result = cache_memoize(cacheKey, cacheFunc) @@ -62,6 +64,34 @@ class CacheTest(TestCase): result = cache_memoize(cacheKey, cacheFunc) self.assertEqual(result, testStr) + def test_cache_memoize_large_files(self): + """Testing cache_memoize with large files""" + cacheKey = "abc123" + + # This takes into account the size of the pickle data, and will + # get us to exactly 2 chunks of data in cache. + data = 'x' * (CACHE_CHUNK_SIZE * 2 - 8) + + def cacheFunc(cacheCalled=[]): + self.assertTrue(not cacheCalled) + cacheCalled.append(True) + return data + + result = cache_memoize(cacheKey, cacheFunc, large_data=True, + compress_large_data=False) + self.assertEqual(result, data) + + site = Site.objects.get_current() + full_key = '%s:%s' % (site.domain, cacheKey) + self.assertTrue(cache.has_key(full_key)) + self.assertTrue(cache.has_key('%s-0' % full_key)) + self.assertTrue(cache.has_key('%s-1' % full_key)) + self.assertFalse(cache.has_key('%s-2' % full_key)) + + result = cache_memoize(cacheKey, cacheFunc, large_data=True, + compress_large_data=False) + self.assertEqual(result, data) + class BoxTest(TagTest): def testPlain(self):