Permalink
Browse files

Add a memcache client + environment var. handling.

  • Loading branch information...
1 parent 9691632 commit 172b16125ff9a96f1475bfcb9c29f4aa785a1dbe @almet almet committed Apr 5, 2012
Showing with 44 additions and 6 deletions.
  1. +42 −4 tokenserver/crypto/pyworker.py
  2. +2 −2 tokenserver/tests/test_crypto_pyworker.py
View
46 tokenserver/crypto/pyworker.py
@@ -1,3 +1,4 @@
+import os
import json
import sys
import traceback
@@ -16,6 +17,19 @@
import pylibmc
+class MemcacheClient(pylibmc.Client):
+
+ def __init__(self, *args, **kwargs):
+ self.ttl = kwargs.pop('ttl', 0)
+ super(MemcacheClient, self).__init__(*args, **kwargs)
+
+ def set(self, key, value, time=None, *args, **kwargs):
+ if time is None:
+ time = self.ttl
+ return super(MemcacheClient, self).set(str(key), value, time, *args,
+ **kwargs)
+
+
class ExpiredValue(KeyError):
pass
@@ -63,14 +77,15 @@ def __init__(self, memory=None, memcache=None, loadtest_mode=False):
setting :param memory: or :param memcache: to False will cause the
certificates manager to not use them
+
+ :param memory: the dict to use for in-memory cache
+ :param memcache: the memcache instance, already configured.
"""
if memory is None:
memory = TTLedDict(60 * 10) # TTL of 10 minutes for the certs
if memcache is None:
memcache = False
- else:
- memcache = pylibmc.Client((memcache, ))
self.memory = memory
self.memcache = memcache
@@ -86,6 +101,7 @@ def __getitem__(self, hostname):
shared memcache. If it's not in the the memcache, download it and store
it both in memory and in the memcache.
"""
+ hostname = str(hostname)
if self.memory and hostname in self.memory:
return self.memory[hostname]
else:
@@ -130,11 +146,33 @@ def get_certificate(algo, filename=None, key=None):
return cls(obj=obj)
+def get_certificate_manager(loadtest_mode=False):
+ """look at the environment variables and return an appropritatly configured
+ certificate manager.
+ """
+ mc_host = os.environ.get('MEMCACHE_HOST', None)
+ if mc_host is not None:
+ mc_ttl = os.environ.get('MEMCACHE_TTL', 60 * 30)
+ memcache = MemcacheClient((mc_host,), ttl=mc_ttl)
+ else:
+ memcache = False
+
+ mem_ttl = os.environ.get('MEMORY_TTL', 60 * 10)
+ memory = TTLedDict(ttl=mem_ttl)
+
+ return CertificatesManagerWithCache(
+ loadtest_mode=loadtest_mode,
+ memory=memory,
+ memcache=memcache)
+
+
class CryptoWorker(object):
- def __init__(self, loadtest_mode=False):
+ def __init__(self, loadtest_mode=False, certs=None):
logger.info('starting a crypto worker')
- self.certs = CertificatesManagerWithCache(loadtest_mode=loadtest_mode)
+ if certs is None:
+ certs = get_certificate_manager(loadtest_mode=loadtest_mode)
+ self.certs = certs
def __call__(self, job):
"""proxy to the functions exposed by the worker"""
View
4 tokenserver/tests/test_crypto_pyworker.py
@@ -1,7 +1,8 @@
from unittest import TestCase
import time
-from tokenserver.crypto.pyworker import CryptoWorker, TTLedDict, ExpiredValue
+from tokenserver.crypto.pyworker import (CryptoWorker, TTLedDict, ExpiredValue,
+ CertificatesManagerWithCache)
from tokenserver.tests.mockworker import MockCryptoWorker
from tokenserver.tests.support import (
sign_data,
@@ -76,7 +77,6 @@ def test_key_derivation(self):
class TestTTledDict(TestCase):
def test_ttled_dict(self):
-
# setup a dict with an expiration of 100ms.
cache = TTLedDict(1)
# asking for something not defined raises an exception

0 comments on commit 172b161

Please sign in to comment.