diff --git a/cachier/core.py b/cachier/core.py index 1e106d4..d0d811b 100644 --- a/cachier/core.py +++ b/cachier/core.py @@ -80,6 +80,7 @@ def cachier( mongetter=None, cache_dir=None, hash_params=None, + wait_for_calc_timeout=0 ): """A persistent, stale-free memoization decorator. @@ -118,6 +119,12 @@ def cachier( and returns a hash key for them. This parameter can be used to enable the use of cachier with functions that get arguments that are not automatically hashable by Python. + wait_for_calc_timeout: int, optional, for MongoDB only + The maximum time to wait for an ongoing calculation. When a + process started to calculate the value setting being_calculated to + True, any process trying to read the same entry will wait a maximum of + seconds specified in this parameter. 0 means wait forever. + Once the timeout expires the calculation will be triggered. """ # print('Inside the wrapper maker') # print('mongetter={}'.format(mongetter)) @@ -125,7 +132,7 @@ def cachier( # print('next_time={}'.format(next_time)) if mongetter: - core = _MongoCore(mongetter, stale_after, next_time) + core = _MongoCore(mongetter, stale_after, next_time, wait_for_calc_timeout) else: core = _PickleCore( # pylint: disable=R0204 stale_after=stale_after, diff --git a/cachier/mongo_core.py b/cachier/mongo_core.py index e05b6da..933c084 100644 --- a/cachier/mongo_core.py +++ b/cachier/mongo_core.py @@ -37,7 +37,7 @@ class _MongoCore(_BaseCore): _INDEX_NAME = 'func_1_key_1' - def __init__(self, mongetter, stale_after, next_time): + def __init__(self, mongetter, stale_after, next_time, wait_for_calc_timeout): if 'pymongo' not in sys.modules: warnings.warn(( "Cachier warning: pymongo was not found. " @@ -45,6 +45,7 @@ class _MongoCore(_BaseCore): _BaseCore.__init__(self, stale_after, next_time) self.mongetter = mongetter self.mongo_collection = self.mongetter() + self.wait_for_calc_timeout = wait_for_calc_timeout index_inf = self.mongo_collection.index_information() if _MongoCore._INDEX_NAME not in index_inf: func1key1 = IndexModel( @@ -131,13 +132,21 @@ class _MongoCore(_BaseCore): pass # don't care in this case def wait_on_entry_calc(self, key): + time_spent = 0 while True: time.sleep(MONGO_SLEEP_DURATION_IN_SEC) + time_spent += 1 key, entry = self.get_entry_by_key(key) if entry is None: raise RecalculationNeeded() - if entry is not None and not entry['being_calculated']: - return entry['value'] + + if entry is not None: + if not entry['being_calculated']: + return entry['value'] + + if self.wait_for_calc_timeout > 0 and time_spent >= self.wait_for_calc_timeout: + raise RecalculationNeeded() + def clear_cache(self): self.mongo_collection.delete_many(