From 466f7669b490b400cf94874ec66edaf26dc37584 Mon Sep 17 00:00:00 2001
From: Nico Andrade <nico@nmac.com.ar>
Date: Fri, 2 Oct 2020 18:15:51 -0400
Subject: [PATCH] Add calculation timeout in MongoDB handler to avoid deadlocks

---
 cachier/mongo_core.py    |  2 +-
 tests/test_mongo_core.py | 50 ++++++++++++++++++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/cachier/mongo_core.py b/cachier/mongo_core.py
index 933c084..b501be4 100644
--- a/cachier/mongo_core.py
+++ b/cachier/mongo_core.py
@@ -135,7 +135,7 @@ class _MongoCore(_BaseCore):
         time_spent = 0
         while True:
             time.sleep(MONGO_SLEEP_DURATION_IN_SEC)
-            time_spent += 1
+            time_spent += MONGO_SLEEP_DURATION_IN_SEC
             key, entry = self.get_entry_by_key(key)
             if entry is None:
                 raise RecalculationNeeded()
diff --git a/tests/test_mongo_core.py b/tests/test_mongo_core.py
index b6fd845..9f0906f 100644
--- a/tests/test_mongo_core.py
+++ b/tests/test_mongo_core.py
@@ -99,7 +99,6 @@ 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,53 @@ 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 test_mongo_wait_for_calc_timeout_ok():
+    """Testing MongoDB core handling of waiting for results."""
+    _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
+
+
+@cachier(mongetter=_test_mongetter, stale_after=MONGO_DELTA, next_time=False, wait_for_calc_timeout=2)
+def _wait_for_calc_timeout_mongo_slow(arg_1, arg_2):
+    """Some function."""
+    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 MongoDB core handling of waiting for results."""
+    _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(3)
+    thread1.join()
+    thread2.join()    
+    assert res_queue.qsize() == 2
+    res1 = res_queue.get()
+    res2 = res_queue.get()
+    assert res1 != res2
+
+
 class _BadMongoCollection:
 
     def __init__(self, mongetter):
@@ -188,7 +234,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):