1 files changed, 50 insertions, 4 deletions
diff --git a/dogpile/cache/util.py b/dogpile/cache/util.py
index 47f7bbe..8122565 100644
--- a/dogpile/cache/util.py
+++ b/dogpile/cache/util.py
@@ -2,6 +2,7 @@ from hashlib import sha1
 import inspect
 import sys
 import re
+import collections
 
 try:
     import threading
@@ -73,13 +74,13 @@ def function_key_generator(namespace, fn):
     """Return a function that generates a string
     key, based on a given function as well as
     arguments to the returned function itself.
-    
+
     This is used by :meth:`.CacheRegion.cache_on_arguments`
     to generate a cache key from a decorated function.
-    
+
     It can be replaced using the ``function_key_generator``
     argument passed to :func:`.make_region`.
-    
+
     """
 
     if namespace is None:
@@ -107,7 +108,7 @@ def sha1_mangle_key(key):
 def length_conditional_mangler(length, mangler):
     """a key mangler that mangles if the length of the key is
     past a certain threshold.
-    
+
     """
     def mangle(key):
         if len(key) >= length:
@@ -137,3 +138,48 @@ def to_list(x, default=None):
         return [x]
     else:
         return x
+
+
+class KeyReentrantMutex(object):
+
+    def __init__(self, key, mutex, keys):
+        self.key = key
+        self.mutex = mutex
+        self.keys = keys
+
+    @classmethod
+    def factory(cls, mutex):
+        # this collection holds zero or one
+        # thread idents as the key; a set of
+        # keynames held as the value.
+        keystore = collections.defaultdict(set)
+        def fac(key):
+            return KeyReentrantMutex(key, mutex, keystore)
+        return fac
+
+    def acquire(self, wait=True):
+        current_thread = threading.current_thread().ident
+        keys = self.keys.get(current_thread)
+        if keys is not None and \
+            self.key not in keys:
+            # current lockholder, new key. add it in
+            keys.add(self.key)
+            return True
+        elif self.mutex.acquire(wait=wait):
+            # after acquire, create new set and add our key
+            self.keys[current_thread].add(self.key)
+            return True
+        else:
+            return False
+
+    def release(self):
+        current_thread = threading.current_thread().ident
+        keys = self.keys.get(current_thread)
+        assert keys is not None, "this thread didn't do the acquire"
+        assert self.key in keys, "No acquire held for key '%s'" % self.key
+        keys.remove(self.key)
+        if not keys:
+            # when list of keys empty, remove
+            # the thread ident and unlock.
+            del self.keys[current_thread]
+            self.mutex.release()