Replaced event in ResponseFuture for lock

beltran · beltran · commit 302a36e5badc · 2018-02-26T12:51:18.000-05:00
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -32,7 +32,7 @@
 import socket
 import sys
 import time
-from threading import Lock, RLock, Thread, Event
+from threading import Lock, RLock, Thread, Event, ThreadError
 
 import weakref
 from weakref import WeakValueDictionary
@@ -2170,7 +2170,8 @@ def execute_async(self, query, parameters=None, trace=False, custom_payload=None
         """
         future = self._create_response_future(query, parameters, trace, custom_payload, timeout, execution_profile, paging_state)
         future._protocol_handler = self.client_protocol_handler
-        self._on_request(future)
+        if self._request_init_callbacks:
+            self._on_request(future)
         future.send_request()
         return future
 
@@ -3474,7 +3475,9 @@ def __init__(self, session, message, query, timeout, metrics=None, prepared_stat
         self._start_time = start_time or time.time()
         self._spec_execution_plan = speculative_execution_plan or self._spec_execution_plan
         self._make_query_plan()
-        self._event = Event()
+        self._lock_event = Lock()
+        self._lock_event.acquire()
+        self._is_result_set = False
         self._errors = {}
         self._callbacks = []
         self._errbacks = []
@@ -3547,7 +3550,7 @@ def _on_timeout(self, _attempts=0):
 
     def _on_speculative_execute(self):
         self._timer = None
-        if not self._event.is_set():
+        if not self._is_result_set:
 
             # PYTHON-836, the speculative queries must be after
             # the query is sent from the main thread, otherwise the
@@ -3660,7 +3663,7 @@ def warnings(self):
         Otherwise it may throw if the response has not been received.
         """
         # TODO: When timers are introduced, just make this wait
-        if not self._event.is_set():
+        if not self._is_result_set:
             raise DriverException("warnings cannot be retrieved before ResponseFuture is finalized")
         return self._warnings
 
@@ -3678,7 +3681,7 @@ def custom_payload(self):
         :return: :ref:`custom_payload`.
         """
         # TODO: When timers are introduced, just make this wait
-        if not self._event.is_set():
+        if not self._is_result_set:
             raise DriverException("custom_payload cannot be retrieved before ResponseFuture is finalized")
         return self._custom_payload
 
@@ -3697,7 +3700,7 @@ def start_fetching_next_page(self):
 
         self._make_query_plan()
         self.message.paging_state = self._paging_state
-        self._event.clear()
+        self._lock_event.acquire()
         self._final_result = _NOT_SET
         self._final_exception = None
         self._start_timer()
@@ -3941,7 +3944,13 @@ def _set_final_result(self, response):
                 for (fn, args, kwargs) in self._callbacks
             )
 
-        self._event.set()
+        self._is_result_set = True
+        try:
+            self._lock_event.release()
+        # This can happen in speculative executions
+        # _set_result is called several times
+        except ThreadError:
+            pass
 
         # apply each callback
         for callback_partial in to_call:
@@ -3962,7 +3971,14 @@ def _set_final_exception(self, response):
                 partial(fn, response, *args, **kwargs)
                 for (fn, args, kwargs) in self._errbacks
             )
-        self._event.set()
+
+        self._is_result_set = True
+        try:
+            self._lock_event.release()
+        # This can happen in speculative executions
+        # _set_result is called several times
+        except ThreadError:
+            pass
 
         # apply each callback
         for callback_partial in to_call:
@@ -4019,7 +4035,13 @@ def result(self):
             ...     log.exception("Operation failed:")
 
         """
-        self._event.wait()
+        self._lock_event.acquire()
+        try:
+            self._lock_event.release()
+        # This can happen in speculative executions
+        # _set_result is called several times
+        except ThreadError:
+            pass
         if self._final_result is not _NOT_SET:
             return ResultSet(self, self._final_result)
         else:
diff --git a/tests/integration/simulacron/test_connection.py b/tests/integration/simulacron/test_connection.py
@@ -129,7 +129,7 @@ def test_heart_beat_timeout(self):
             futures.append(future)
 
         for f in futures:
-            f._event.wait()
+            f._lock_event.acquire()
             self.assertIsInstance(f._final_exception, OperationTimedOut)
 
         prime_request(PrimeOptions(then=NO_THEN))
@@ -333,7 +333,7 @@ def test_host_is_not_set_to_down_after_query_oto(self):
             futures.append(future)
 
         for f in futures:
-            f._event.wait()
+            f._lock_event.acquire()
             self.assertIsInstance(f._final_exception, OperationTimedOut)
 
         self.assertEqual(listener.hosts_marked_down, [])
diff --git a/tests/integration/simulacron/test_policies.py b/tests/integration/simulacron/test_policies.py
@@ -162,7 +162,7 @@ def test_speculative_and_timeout(self):
         # we have to directly wait for the event
         response_future = self.session.execute_async(statement, execution_profile='spec_ep_brr_lim',
                                                      timeout=14)
-        response_future._event.wait(16)
+        response_future._lock_event.acquire(16)
         self.assertIsInstance(response_future._final_exception, OperationTimedOut)
 
         # This is because 14 / 4 + 1 = 4
diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py
@@ -303,10 +303,11 @@ def test_first_pool_shutdown(self):
         pool_shutdown = self.make_pool()
         pool_shutdown.is_shutdown = True
         pool_ok = self.make_pool()
-        pool_ok.is_shutdown = True
+        pool_ok.is_shutdown = False
         session._pools.get.side_effect = [pool_shutdown, pool_ok]
 
         rf = self.make_response_future(session)
+        rf.query_plan = ['ip1', 'ip2']
         rf.send_request()
 
         rf._set_result(None, None, None, self.make_mock_response([{'col': 'val'}]))
@@ -420,6 +421,7 @@ def test_multiple_errbacks(self):
         retry_policy = Mock()
         retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None)
         rf = ResponseFuture(session, message, query, 1, retry_policy=retry_policy)
+        rf._query_retries = 1
         rf.send_request()
 
         callback = Mock()
@@ -436,7 +438,6 @@ def test_multiple_errbacks(self):
         result = Mock(spec=UnavailableErrorMessage, info={"required_replicas":2, "alive_replicas": 1, "consistency": 1})
         result.to_exception.return_value = expected_exception
         rf._set_result(None, None, None, result)
-        rf._event.set()
         self.assertRaises(Exception, rf.result)
 
         callback.assert_called_once_with(expected_exception, arg, **kwargs)
