Merge "Make @memoize thread-aware"

This commit is contained in:
Jenkins 2017-09-11 21:30:00 +00:00 committed by Gerrit Code Review
commit 4277b078b9

View File

@ -12,7 +12,9 @@
# License for the specific language governing permissions and limitations # License for the specific language governing permissions and limitations
# under the License. # under the License.
import collections
import functools import functools
import threading
import warnings import warnings
import weakref import weakref
@ -58,6 +60,7 @@ def memoized(func):
# instance for every decorated function, and it's stored in a closure of # instance for every decorated function, and it's stored in a closure of
# the wrapped function. # the wrapped function.
cache = {} cache = {}
locks = collections.defaultdict(threading.Lock)
@functools.wraps(func) @functools.wraps(func)
def wrapped(*args, **kwargs): def wrapped(*args, **kwargs):
@ -71,6 +74,7 @@ def memoized(func):
try: try:
# The key here is from closure, and is calculated later. # The key here is from closure, and is calculated later.
del cache[key] del cache[key]
del locks[key]
except KeyError: except KeyError:
# Some other weak reference might have already removed that # Some other weak reference might have already removed that
# key -- in that case we don't need to do anything. # key -- in that case we don't need to do anything.
@ -78,22 +82,25 @@ def memoized(func):
key = _get_key(args, kwargs, remove) key = _get_key(args, kwargs, remove)
try: try:
# We want cache hit to be as fast as possible, and don't really with locks[key]:
# care much about the speed of a cache miss, because it will only try:
# happen once and likely calls some external API, database, or # We want cache hit to be as fast as possible, and don't
# some other slow thing. That's why the hit is in straightforward # really care much about the speed of a cache miss, because
# code, and the miss is in an exception. # it will only happen once and likely calls some external
value = cache[key] # API, database, or some other slow thing. That's why the
except KeyError: # hit is in straightforward code, and the miss is in an
value = cache[key] = func(*args, **kwargs) # exception.
value = cache[key]
except KeyError:
value = cache[key] = func(*args, **kwargs)
except TypeError: except TypeError:
# The calculated key may be unhashable when an unhashable object, # The calculated key may be unhashable when an unhashable
# such as a list, is passed as one of the arguments. In that case, # object, such as a list, is passed as one of the arguments. In
# we can't cache anything and simply always call the decorated # that case, we can't cache anything and simply always call the
# function. # decorated function.
warnings.warn( warnings.warn(
"The key %r is not hashable and cannot be memoized." % (key,), "The key %r is not hashable and cannot be memoized."
UnhashableKeyWarning, 2) % (key,), UnhashableKeyWarning, 2)
value = func(*args, **kwargs) value = func(*args, **kwargs)
return value return value
return wrapped return wrapped