Merge pull request #36 from non-senses/master

Calculation timeout for MongoDB Backend

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,

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,14 +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 += MONGO_SLEEP_DURATION_IN_SEC
             key, entry = self.get_entry_by_key(key)
             if entry is None:
                 raise RecalculationNeeded()
-            if entry is not None and not entry['being_calculated']:
+
+            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(
             filter={'func': _MongoCore._get_func_str(self.func)}

tests/test_mongo_core.py

@@ -92,14 +92,13 @@ def test_mongo_core():
 
 
 MONGO_DELTA = timedelta(seconds=3)
-
+MONGO_DELTA_LONG = timedelta(seconds=10)
 
 @cachier(mongetter=_test_mongetter, stale_after=MONGO_DELTA, next_time=False)
 def _stale_after_mongo(arg_1, arg_2):
     """Some function."""
     return random() + arg_1 + arg_2
 
-
 def test_mongo_stale_after():
     """Testing MongoDB core stale_after functionality."""
     _stale_after_mongo.clear_cache()
@@ -142,6 +141,75 @@ def test_mongo_being_calculated():
     assert res1 == res2
 
 
+@cachier(mongetter=_test_mongetter, stale_after=MONGO_DELTA, next_time=False, wait_for_calc_timeout=2)
+def _wait_for_calc_timeout_mongo_fast(arg_1, arg_2):
+    """Some function."""
+    sleep(1)
+    return random() + arg_1 + arg_2
+
+def _calls_wait_for_calc_timeout_mongo_fast(res_queue):
+    res = _wait_for_calc_timeout_mongo_fast(1, 2)
+    res_queue.put(res)
+
+
+def test_mongo_wait_for_calc_timeout_ok():
+    """ Testing calls that avoid timeouts store the values in cache. """
+    _wait_for_calc_timeout_mongo_fast.clear_cache()
+    val1 = _wait_for_calc_timeout_mongo_fast(1, 2)
+    val2 = _wait_for_calc_timeout_mongo_fast(1, 2)
+    assert val1 == val2
+
+    res_queue = queue.Queue()
+    thread1 = threading.Thread(
+        target=_calls_wait_for_calc_timeout_mongo_fast, kwargs={'res_queue': res_queue})
+    thread2 = threading.Thread(
+        target=_calls_wait_for_calc_timeout_mongo_fast, kwargs={'res_queue': res_queue})
+
+    thread1.start()
+    thread2.start()
+    sleep(2)
+    thread1.join()
+    thread2.join()    
+    assert res_queue.qsize() == 2
+    res1 = res_queue.get()
+    res2 = res_queue.get()
+    assert res1 == res2 # Timeout did not kick in, a single call was done
+
+@cachier(mongetter=_test_mongetter, stale_after=MONGO_DELTA_LONG, next_time=False, wait_for_calc_timeout=2)
+def _wait_for_calc_timeout_mongo_slow(arg_1, arg_2):
+    sleep(3)
+    return random() + arg_1 + arg_2
+
+
+def _calls_wait_for_calc_timeout_mongo_slow(res_queue):
+    res = _wait_for_calc_timeout_mongo_slow(1, 2)
+    res_queue.put(res)
+
+
+def test_mongo_wait_for_calc_timeout_slow():
+    """Testing for calls timing out to be performed twice when needed."""
+    _wait_for_calc_timeout_mongo_slow.clear_cache()
+    res_queue = queue.Queue()
+    thread1 = threading.Thread(
+        target=_calls_wait_for_calc_timeout_mongo_slow, kwargs={'res_queue': res_queue})
+    thread2 = threading.Thread(
+        target=_calls_wait_for_calc_timeout_mongo_slow, kwargs={'res_queue': res_queue})
+
+    thread1.start()
+    thread2.start()
+    sleep(1)
+    res3 = _wait_for_calc_timeout_mongo_slow(1, 2)
+    sleep(4)
+    thread1.join()
+    thread2.join()
+    assert res_queue.qsize() == 2
+    res1 = res_queue.get()
+    res2 = res_queue.get()
+    assert res1 != res2 # Timeout kicked in.  Two calls were done
+    res4 = _wait_for_calc_timeout_mongo_slow(1, 2)
+    assert res1 == res4 or res2 == res4 or res3 == res4 # One of the cached values is returned
+
+
 class _BadMongoCollection:
 
     def __init__(self, mongetter):
@@ -188,7 +256,7 @@ def test_stalled_mongo_db_cache():
     @cachier(mongetter=_test_mongetter)
     def _stalled_func():
         return 1
-    core = _MongoCore(_test_mongetter, None, False)
+    core = _MongoCore(_test_mongetter, None, False, 0)
     core.set_func(_stalled_func)
     core.clear_cache()
     with pytest.raises(RecalculationNeeded):