PYTHON-1248: Libevreactor: Raise ConnectionBusy if tcp send buffer is full

aboudreault · aboudreault · commit e14a7a097fb6 · 2020-06-02T12:55:20.000-04:00
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -6,6 +6,7 @@ Features
 --------
 * Make geomet an optional dependency at runtime (PYTHON-1237)
 * Add use_default_tempdir cloud config options (PYTHON-1245)
+* Tcp flow control for libevreactor (PYTHON-1248)
 
 Bug Fixes
 ---------
diff --git a/Jenkinsfile b/Jenkinsfile
@@ -601,7 +601,7 @@ pipeline {
     EVENT_LOOP_MANAGER = "${params.EVENT_LOOP_MANAGER.toLowerCase()}"
     EXECUTE_LONG_TESTS = "${params.EXECUTE_LONG_TESTS ? 'True' : 'False'}"
     CCM_ENVIRONMENT_SHELL = '/usr/local/bin/ccm_environment.sh'
-    CCM_MAX_HEAP_SIZE = '1536M'
+    CCM_MAX_HEAP_SIZE = '1280M'
   }
 
   stages {
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -48,7 +48,7 @@
 from cassandra.connection import (ConnectionException, ConnectionShutdown,
                                   ConnectionHeartbeat, ProtocolVersionUnsupported,
                                   EndPoint, DefaultEndPoint, DefaultEndPointFactory,
-                                  ContinuousPagingState, SniEndPointFactory)
+                                  ContinuousPagingState, SniEndPointFactory, ConnectionBusy)
 from cassandra.cqltypes import UserType
 from cassandra.encoder import Encoder
 from cassandra.protocol import (QueryMessage, ResultMessage,
@@ -4445,15 +4445,18 @@ def _query(self, host, message=None, cb=None):
         except NoConnectionsAvailable as exc:
             log.debug("All connections for host %s are at capacity, moving to the next host", host)
             self._errors[host] = exc
-            return None
+        except ConnectionBusy as exc:
+            log.debug("Connection for host %s is busy, moving to the next host", host)
+            self._errors[host] = exc
         except Exception as exc:
             log.debug("Error querying host %s", host, exc_info=True)
             self._errors[host] = exc
             if self._metrics is not None:
                 self._metrics.on_connection_error()
             if connection:
                 pool.return_connection(connection)
-            return None
+
+        return None
 
     @property
     def has_more_pages(self):
diff --git a/cassandra/connection.py b/cassandra/connection.py
@@ -692,6 +692,7 @@ def __init__(self, host='127.0.0.1', port=9042, authenticator=None,
         self._requests = {}
         self._iobuf = io.BytesIO()
         self._continuous_paging_sessions = {}
+        self._socket_writable = True
 
         if ssl_options:
             self._check_hostname = bool(self.ssl_options.pop('check_hostname', False))
@@ -926,6 +927,8 @@ def send_msg(self, msg, request_id, cb, encoder=ProtocolHandler.encode_message,
             raise ConnectionShutdown("Connection to %s is defunct" % self.endpoint)
         elif self.is_closed:
             raise ConnectionShutdown("Connection to %s is closed" % self.endpoint)
+        elif not self._socket_writable:
+            raise ConnectionBusy("Connection %s is overloaded" % self.endpoint)
 
         # queue the decoder function with the request
         # this allows us to inject custom functions per request to encode, decode messages
diff --git a/cassandra/io/libevreactor.py b/cassandra/io/libevreactor.py
@@ -310,13 +310,17 @@ def handle_write(self, watcher, revents, errno=None):
                 with self._deque_lock:
                     next_msg = self.deque.popleft()
             except IndexError:
+                if not self._socket_writable:
+                    self._socket_writable = True
                 return
 
             try:
                 sent = self._socket.send(next_msg)
             except socket.error as err:
                 if (err.args[0] in NONBLOCKING or
                         err.args[0] in (ssl.SSL_ERROR_WANT_READ, ssl.SSL_ERROR_WANT_WRITE)):
+                    if err.args[0] in NONBLOCKING:
+                        self._socket_writable = False
                     with self._deque_lock:
                         self.deque.appendleft(next_msg)
                 else:
diff --git a/tests/integration/simulacron/test_backpressure.py b/tests/integration/simulacron/test_backpressure.py
@@ -45,7 +45,7 @@ def _fill_buffers(self, session, query, expected_blocked=3, **execute_kwargs):
             for pool in session.get_pools():
                 if not pool._connection._socket_writable:
                     total_blocked += 1
-            if total_blocked == expected_blocked:
+            if total_blocked >= expected_blocked:
                 break
         else:
             raise Exception("Unable to fill TCP send buffer on expected number of nodes")

Original file line number	Diff line number	Diff line change
`@@ -601,7 +601,7 @@ pipeline {`
`601`	`601`	`EVENT_LOOP_MANAGER = "${params.EVENT_LOOP_MANAGER.toLowerCase()}"`
`602`	`602`	`EXECUTE_LONG_TESTS = "${params.EXECUTE_LONG_TESTS ? 'True' : 'False'}"`
`603`	`603`	`CCM_ENVIRONMENT_SHELL = '/usr/local/bin/ccm_environment.sh'`
`604`		`- CCM_MAX_HEAP_SIZE = '1536M'`
	`604`	`+ CCM_MAX_HEAP_SIZE = '1280M'`
`605`	`605`	`}`
`606`	`606`
`607`	`607`	`stages {`