Fix race conditions with conn.in_flight_requests

kafka/client_async.py

@@ -612,7 +612,7 @@ def _poll(self, timeout):
             conn = key.data
             processed.add(conn)
 
-            if not conn.in_flight_requests:
+            if not conn.has_in_flight_requests():
                 # if we got an EVENT_READ but there were no in-flight requests, one of
                 # two things has happened:
                 #

kafka/conn.py

@@ -271,6 +271,8 @@ def __init__(self, host, port, afi, **configs):
         # including tracking request futures and timestamps, we
         # can use a simple dictionary of correlation_id => request data
         self.in_flight_requests = dict()
+        # Ensures that all in_flight_requests futures are properly closed on disconnect.
+        self._ifr_lock = threading.Lock()
 
         self._protocol = KafkaProtocol(
             client_id=self.config['client_id'],
@@ -741,15 +743,19 @@ def close(self, error=None):
         self.config['state_change_callback'](self)
         self._update_reconnect_backoff()
         self._close_socket()
-        self.state = ConnectionStates.DISCONNECTED
-        self._sasl_auth_future = None
-        self._protocol = KafkaProtocol(
-            client_id=self.config['client_id'],
-            api_version=self.config['api_version'])
+
+        with self._ifr_lock:
+            self.state = ConnectionStates.DISCONNECTED
+            self._sasl_auth_future = None
+            self._protocol = KafkaProtocol(
+                client_id=self.config['client_id'],
+                api_version=self.config['api_version'])
+            fail_ifrs = dict(self.in_flight_requests)
+            self.in_flight_requests.clear()
+
         if error is None:
             error = Errors.Cancelled(str(self))
-        while self.in_flight_requests:
-            (_correlation_id, (future, _timestamp)) = self.in_flight_requests.popitem()
+        for future, _timestamp in fail_ifrs.values():
             future.failure(error)
         self.config['state_change_callback'](self)
 
@@ -770,13 +776,18 @@ def _send(self, request, blocking=True):
         with self._lock:
             correlation_id = self._protocol.send_request(request)
 
-        log.debug('%s Request %d: %s', self, correlation_id, request)
-        if request.expect_response():
-            sent_time = time.time()
-            assert correlation_id not in self.in_flight_requests, 'Correlation ID already in-flight!'
-            self.in_flight_requests[correlation_id] = (future, sent_time)
-        else:
-            future.success(None)
+            log.debug('%s Request %d: %s', self, correlation_id, request)
+            if request.expect_response():
+                with self._ifr_lock:
+                    if self.disconnected():
+                        log.debug("%s: Connection already closed.", self)
+                        future.failure(Errors.Cancelled(str(self)))
+                    else:
+                        sent_time = time.time()
+                        assert correlation_id not in self.in_flight_requests, 'Correlation ID already in-flight!'
+                        self.in_flight_requests[correlation_id] = (future, sent_time)
+            else:
+                future.success(None)
 
         # Attempt to replicate behavior from prior to introduction of
         # send_pending_requests() / async sends
@@ -808,22 +819,27 @@ def send_pending_requests(self):
     def can_send_more(self):
         """Return True unless there are max_in_flight_requests_per_connection."""
         max_ifrs = self.config['max_in_flight_requests_per_connection']
-        return len(self.in_flight_requests) < max_ifrs
+        with self._ifr_lock:
+            return len(self.in_flight_requests) < max_ifrs
 
     def recv(self):
         """Non-blocking network receive.
 
         Return list of (response, future) tuples
         """
-        if not self.connected() and not self.state is ConnectionStates.AUTHENTICATING:
+        with self._ifr_lock:
+            state_ok = self.connected() or self.state is ConnectionStates.AUTHENTICATING
+            has_ifr = bool(self.in_flight_requests)
+
+        if not state_ok:
             log.warning('%s cannot recv: socket not connected', self)
             # If requests are pending, we should close the socket and
             # fail all the pending request futures
-            if self.in_flight_requests:
+            if has_ifr:
                 self.close(Errors.KafkaConnectionError('Socket not connected during recv with in-flight-requests'))
             return ()
 
-        elif not self.in_flight_requests:
+        elif not has_ifr:
             log.warning('%s: No in-flight-requests to recv', self)
             return ()
 
@@ -839,9 +855,14 @@ def recv(self):
         # augment respones w/ correlation_id, future, and timestamp
         for i, (correlation_id, response) in enumerate(responses):
             try:
-                (future, timestamp) = self.in_flight_requests.pop(correlation_id)
+                with self._ifr_lock:
+                    (future, timestamp) = self.in_flight_requests.pop(correlation_id)
             except KeyError:
-                self.close(Errors.KafkaConnectionError('Received unrecognized correlation id'))
+                if self.disconnected():
+                    log.warning('%s: Received response %s after the connection had been closed.',
+                                self, response.__class__.__name__)
+                else:
+                    self.close(Errors.KafkaConnectionError('Received unrecognized correlation id'))
                 return ()
             latency_ms = (time.time() - timestamp) * 1000
             if self._sensors:
@@ -895,15 +916,20 @@ def _recv(self):
         else:
             return responses
 
+    def has_in_flight_requests(self):
+        with self._ifr_lock:
+            return bool(self.in_flight_requests)
+
     def requests_timed_out(self):
-        if self.in_flight_requests:
-            get_timestamp = lambda v: v[1]
-            oldest_at = min(map(get_timestamp,
-                                self.in_flight_requests.values()))
-            timeout = self.config['request_timeout_ms'] / 1000.0
-            if time.time() >= oldest_at + timeout:
-                return True
-        return False
+        with self._ifr_lock:
+            if self.in_flight_requests:
+                get_timestamp = lambda v: v[1]
+                oldest_at = min(map(get_timestamp,
+                                    self.in_flight_requests.values()))
+                timeout = self.config['request_timeout_ms'] / 1000.0
+                if time.time() >= oldest_at + timeout:
+                    return True
+            return False
 
     def _handle_api_version_response(self, response):
         error_type = Errors.for_code(response.error_code)