Avoid conn leak in prepared statements, better err handling

thobbs · thobbs · commit 4c44321381a7 · 2013-10-29T16:13:10.000-05:00
If a prepared statement was not recognized by the server (for
example, if it was restarted), connections were not being properly
returned to the pool.

Additionally, the re-submission of the original query after the
preparation completed was not borrowing connections properly
(by finding the least loaded conn and incrementing its in_flight
count).

Last, if there was a connection error while preparing an
unrecognized prepared statement, that error was raised directly
instead of retrying the operation on other hosts.
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -1681,11 +1681,17 @@ def _query(self, host):
             self._errors[host] = ConnectionException("Pool is shutdown")
             return None
 
+        return self._borrow_conn_and_send_message(host, pool, self.message, self._set_result)
+
+    def _borrow_conn_and_send_message(self, host, pool, message, cb):
+        if cb is None:
+            cb = self._set_result
+
         connection = None
         try:
             # TODO get connectTimeout from cluster settings
             connection = pool.borrow_connection(timeout=2.0)
-            request_id = connection.send_msg(self.message, cb=self._set_result)
+            request_id = connection.send_msg(message, cb=cb)
             self._current_host = host
             self._current_pool = pool
             self._connection = connection
@@ -1791,9 +1797,11 @@ def _set_result(self, response):
                     prepare_message = PrepareMessage(query=prepared_statement.query_string)
                     # since this might block, run on the executor to avoid hanging
                     # the event loop thread
-                    self.session.submit(self._connection.send_msg,
+                    self.session.submit(self._borrow_conn_and_send_message,
+                                        self._current_host,
+                                        self._current_pool,
                                         prepare_message,
-                                        cb=self._execute_after_prepare)
+                                        self._execute_after_prepare)
                     return
                 else:
                     if hasattr(response, 'to_exception'):
@@ -1846,8 +1854,8 @@ def _execute_after_prepare(self, response):
         Handle the response to our attempt to prepare a statement.
         If it succeeded, run the original query again against the same host.
         """
-        if self._final_exception:
-            return
+        if self._current_pool and self._connection:
+            self._current_pool.return_connection(self._connection)
 
         if isinstance(response, ResultMessage):
             if response.kind == ResultMessage.KIND_PREPARED:
@@ -1860,6 +1868,12 @@ def _execute_after_prepare(self, response):
                     "on host %s: %s" % (self._current_host, response)))
         elif isinstance(response, ErrorMessage):
             self._set_final_exception(response)
+        elif isinstance(response, ConnectionException):
+            log.debug("Connection error when preparing statement on host %s: %s",
+                      self._current_host, response)
+            # try again on a different host, preparing again if necessary
+            self._errors[self._current_host] = response
+            self.send_request()
         else:
             self._set_final_exception(ConnectionException(
                 "Got unexpected response type when preparing "
diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py
@@ -366,8 +366,8 @@ def test_prepared_query_not_found(self):
 
         session.submit.assert_called_once()
         args, kwargs = session.submit.call_args
-        self.assertIsInstance(args[-1], PrepareMessage)
-        self.assertEquals(args[-1].query, "SELECT * FROM foobar")
+        self.assertIsInstance(args[-2], PrepareMessage)
+        self.assertEquals(args[-2].query, "SELECT * FROM foobar")
 
     def test_prepared_query_not_found_bad_keyspace(self):
         session = self.make_session()
