Merge pull request apache#873 from datastax/python-836

beltran · web-flow · commit aaccf2fad406 · 2017-12-18T12:23:48.000-05:00
PYTHON-836: Added lock to query plan for speculative execution
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -3460,12 +3460,12 @@ def __init__(self, session, message, query, timeout, metrics=None, prepared_stat
         self.prepared_statement = prepared_statement
         self._callback_lock = Lock()
         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._errors = {}
         self._callbacks = []
         self._errbacks = []
-        self._spec_execution_plan = speculative_execution_plan or self._spec_execution_plan
         self.attempted_hosts = []
         self._start_timer()
 
@@ -3534,6 +3534,10 @@ def _make_query_plan(self):
         # calls to send_request (which retries may do) will resume where
         # they last left off
         self.query_plan = iter(self._load_balancer.make_query_plan(self.session.keyspace, self.query))
+        # Make iterator thread safe when there can be a speculative delay since it could
+        # from different threads
+        if isinstance(self._spec_execution_plan, NoSpeculativeExecutionPlan):
+            self.query_plan = _threadsafe_iter(self.query_plan)
 
     def send_request(self, error_no_hosts=True):
         """ Internal """
@@ -4306,3 +4310,16 @@ def paging_state(self):
         avoid sending this to untrusted parties.
         """
         return self.response_future._paging_state
+
+
+class _threadsafe_iter(six.Iterator):
+    def __init__(self, it):
+        self.it = it
+        self.lock = Lock()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        with self.lock:
+            return next(self.it)
diff --git a/tests/integration/standard/test_policies.py b/tests/integration/standard/test_policies.py
@@ -104,7 +104,7 @@ def test_delay_can_be_0(self):
         @expected_result all the queries are executed immediately
         @test_category policy
         """
-        number_of_requests = 6
+        number_of_requests = 4
         spec = ExecutionProfile(speculative_execution_policy=ConstantSpeculativeExecutionPolicy(0, number_of_requests))
 
         cluster = Cluster()
