Commit 911cc952 authored by Matt Robenolt's avatar Matt Robenolt Committed by Sean Bleier

Implement cache.add with SETNX so cache.add can be atomic like it's expected

The way it is now, doing a cache.add has a slight race condition
between checking if the key exists already before setting.

This rewrite uses the native SETNX command to only set the key if one
doesn't exist yet.
parent b3121178
...@@ -167,10 +167,7 @@ class CacheClass(BaseCache): ...@@ -167,10 +167,7 @@ class CacheClass(BaseCache):
Returns ``True`` if the object was added, ``False`` if not. Returns ``True`` if the object was added, ``False`` if not.
""" """
key = self.make_key(key, version=version) return self.set(key, value, timeout, _add_only=True)
if self._client.exists(key):
return False
return self.set(key, value, timeout)
def get(self, key, default=None, version=None): def get(self, key, default=None, version=None):
""" """
...@@ -188,15 +185,22 @@ class CacheClass(BaseCache): ...@@ -188,15 +185,22 @@ class CacheClass(BaseCache):
result = self.unpickle(value) result = self.unpickle(value)
return result return result
def _set(self, key, value, timeout, client): def _set(self, key, value, timeout, client, _add_only=False):
if timeout == 0: if timeout == 0:
if _add_only:
return client.setnx(key, value)
return client.set(key, value) return client.set(key, value)
elif timeout > 0: elif timeout > 0:
return client.setex(key, value, int(timeout)) if _add_only:
added = client.setnx(key, value)
if added:
client.expire(key, timeout)
return added
return client.setex(key, value, timeout)
else: else:
return False return False
def set(self, key, value, timeout=None, version=None, client=None): def set(self, key, value, timeout=None, version=None, client=None, _add_only=False):
""" """
Persist a value to the cache, and set an optional expiration time. Persist a value to the cache, and set an optional expiration time.
""" """
...@@ -211,9 +215,9 @@ class CacheClass(BaseCache): ...@@ -211,9 +215,9 @@ class CacheClass(BaseCache):
if int(value) != value: if int(value) != value:
raise TypeError raise TypeError
except (ValueError, TypeError): except (ValueError, TypeError):
result = self._set(key, pickle.dumps(value), int(timeout), client) result = self._set(key, pickle.dumps(value), int(timeout), client, _add_only)
else: else:
result = self._set(key, int(value), int(timeout), client) result = self._set(key, int(value), int(timeout), client, _add_only)
# result is a boolean # result is a boolean
return result return result
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment