diff options
Diffstat (limited to 'kafka')
| -rw-r--r-- | kafka/client_async.py | 59 | ||||
| -rw-r--r-- | kafka/conn.py | 49 | ||||
| -rw-r--r-- | kafka/consumer/group.py | 13 | ||||
| -rw-r--r-- | kafka/coordinator/base.py | 4 |
4 files changed, 76 insertions, 49 deletions
diff --git a/kafka/client_async.py b/kafka/client_async.py index e2bdda9..d608e6a 100644 --- a/kafka/client_async.py +++ b/kafka/client_async.py @@ -304,7 +304,10 @@ class KafkaClient(object): # SSL connections can enter this state 2x (second during Handshake) if node_id not in self._connecting: self._connecting.add(node_id) + try: self._selector.register(conn._sock, selectors.EVENT_WRITE) + except KeyError: + self._selector.modify(conn._sock, selectors.EVENT_WRITE) elif conn.connected(): log.debug("Node %s connected", node_id) @@ -312,10 +315,10 @@ class KafkaClient(object): self._connecting.remove(node_id) try: - self._selector.unregister(conn._sock) + self._selector.modify(conn._sock, selectors.EVENT_READ, conn) except KeyError: - pass - self._selector.register(conn._sock, selectors.EVENT_READ, conn) + self._selector.register(conn._sock, selectors.EVENT_READ, conn) + if self._sensors: self._sensors.connection_created.record() @@ -336,6 +339,7 @@ class KafkaClient(object): self._selector.unregister(conn._sock) except KeyError: pass + if self._sensors: self._sensors.connection_closed.record() @@ -348,6 +352,17 @@ class KafkaClient(object): log.warning("Node %s connection failed -- refreshing metadata", node_id) self.cluster.request_update() + def maybe_connect(self, node_id): + """Queues a node for asynchronous connection during the next .poll()""" + if self._can_connect(node_id): + self._connecting.add(node_id) + # Wakeup signal is useful in case another thread is + # blocked waiting for incoming network traffic while holding + # the client lock in poll(). + self.wakeup() + return True + return False + def _maybe_connect(self, node_id): """Idempotent non-blocking connection attempt to the given node id.""" with self._lock: @@ -397,7 +412,7 @@ class KafkaClient(object): Returns: bool: True if we are ready to send to the given node """ - self._maybe_connect(node_id) + self.maybe_connect(node_id) return self.is_ready(node_id, metadata_priority=metadata_priority) def connected(self, node_id): @@ -499,14 +514,15 @@ class KafkaClient(object): return True def _can_send_request(self, node_id): - with self._lock: - if node_id not in self._conns: - return False - conn = self._conns[node_id] - return conn.connected() and conn.can_send_more() + conn = self._conns.get(node_id) + if not conn: + return False + return conn.connected() and conn.can_send_more() def send(self, node_id, request): - """Send a request to a specific node. + """Send a request to a specific node. Bytes are placed on an + internal per-connection send-queue. Actual network I/O will be + triggered in a subsequent call to .poll() Arguments: node_id (int): destination node @@ -518,11 +534,21 @@ class KafkaClient(object): Returns: Future: resolves to Response struct or Error """ - with self._lock: - if not self._maybe_connect(node_id): - return Future().failure(Errors.NodeNotReadyError(node_id)) + if not self._can_send_request(node_id): + self.maybe_connect(node_id) + return Future().failure(Errors.NodeNotReadyError(node_id)) + + # conn.send will queue the request internally + # we will need to call send_pending_requests() + # to trigger network I/O + future = self._conns[node_id].send(request, blocking=False) - return self._conns[node_id].send(request) + # Wakeup signal is useful in case another thread is + # blocked waiting for incoming network traffic while holding + # the client lock in poll(). + self.wakeup() + + return future def poll(self, timeout_ms=None, future=None): """Try to read and write to sockets. @@ -640,6 +666,8 @@ class KafkaClient(object): conn.close(error=Errors.RequestTimedOutError( 'Request timed out after %s ms' % conn.config['request_timeout_ms'])) + else: + conn.send_pending_requests() if self._sensors: self._sensors.io_time.record((time.time() - end_select) * 1000000000) @@ -801,9 +829,8 @@ class KafkaClient(object): # have such application level configuration, using request timeout instead. return self.config['request_timeout_ms'] - if self._can_connect(node_id): + if self.maybe_connect(node_id): log.debug("Initializing connection to node %s for metadata request", node_id) - self._maybe_connect(node_id) return self.config['reconnect_backoff_ms'] # connected but can't send more, OR connecting diff --git a/kafka/conn.py b/kafka/conn.py index 7dfc8bd..6b5aff9 100644 --- a/kafka/conn.py +++ b/kafka/conn.py @@ -733,11 +733,8 @@ class BrokerConnection(object): future.failure(error) self.config['state_change_callback'](self) - def send(self, request): - """send request, return Future() - - Can block on network if request is larger than send_buffer_bytes - """ + def send(self, request, blocking=True): + """Queue request for async network send, return Future()""" future = Future() if self.connecting(): return future.failure(Errors.NodeNotReadyError(str(self))) @@ -745,35 +742,49 @@ class BrokerConnection(object): return future.failure(Errors.KafkaConnectionError(str(self))) elif not self.can_send_more(): return future.failure(Errors.TooManyInFlightRequests(str(self))) - return self._send(request) + return self._send(request, blocking=blocking) - def _send(self, request): + def _send(self, request, blocking=True): assert self.state in (ConnectionStates.AUTHENTICATING, ConnectionStates.CONNECTED) future = Future() correlation_id = self._protocol.send_request(request) + + # Attempt to replicate behavior from prior to introduction of + # send_pending_requests() / async sends + if blocking: + error = self.send_pending_requests() + if isinstance(error, Exception): + future.failure(error) + return future + + log.debug('%s Request %d: %s', self, correlation_id, request) + if request.expect_response(): + sent_time = time.time() + ifr = (correlation_id, future, sent_time) + self.in_flight_requests.append(ifr) + else: + future.success(None) + return future + + def send_pending_requests(self): + """Can block on network if request is larger than send_buffer_bytes""" + if self.state not in (ConnectionStates.AUTHENTICATING, + ConnectionStates.CONNECTED): + return Errors.NodeNotReadyError(str(self)) data = self._protocol.send_bytes() try: # In the future we might manage an internal write buffer # and send bytes asynchronously. For now, just block # sending each request payload - sent_time = time.time() total_bytes = self._send_bytes_blocking(data) if self._sensors: self._sensors.bytes_sent.record(total_bytes) + return total_bytes except ConnectionError as e: - log.exception("Error sending %s to %s", request, self) + log.exception("Error sending request data to %s", self) error = Errors.KafkaConnectionError("%s: %s" % (self, e)) self.close(error=error) - return future.failure(error) - log.debug('%s Request %d: %s', self, correlation_id, request) - - if request.expect_response(): - ifr = (correlation_id, future, sent_time) - self.in_flight_requests.append(ifr) - else: - future.success(None) - - return future + return error def can_send_more(self): """Return True unless there are max_in_flight_requests_per_connection.""" diff --git a/kafka/consumer/group.py b/kafka/consumer/group.py index 8d2c65e..531c107 100644 --- a/kafka/consumer/group.py +++ b/kafka/consumer/group.py @@ -1070,16 +1070,6 @@ class KafkaConsumer(six.Iterator): # like heartbeats, auto-commits, and metadata refreshes timeout_at = self._next_timeout() - # Because the consumer client poll does not sleep unless blocking on - # network IO, we need to explicitly sleep when we know we are idle - # because we haven't been assigned any partitions to fetch / consume - if self._use_consumer_group() and not self.assignment(): - sleep_time = max(timeout_at - time.time(), 0) - if sleep_time > 0 and not self._client.in_flight_request_count(): - log.debug('No partitions assigned; sleeping for %s', sleep_time) - time.sleep(sleep_time) - continue - # Short-circuit the fetch iterator if we are already timed out # to avoid any unintentional interaction with fetcher setup if time.time() > timeout_at: @@ -1090,8 +1080,7 @@ class KafkaConsumer(six.Iterator): if time.time() > timeout_at: log.debug("internal iterator timeout - breaking for poll") break - if self._client.in_flight_request_count(): - self._client.poll(timeout_ms=0) + self._client.poll(timeout_ms=0) # An else block on a for loop only executes if there was no break # so this should only be called on a StopIteration from the fetcher diff --git a/kafka/coordinator/base.py b/kafka/coordinator/base.py index 1435183..664e8d2 100644 --- a/kafka/coordinator/base.py +++ b/kafka/coordinator/base.py @@ -252,7 +252,7 @@ class BaseCoordinator(object): if self.config['api_version'] < (0, 8, 2): self.coordinator_id = self._client.least_loaded_node() if self.coordinator_id is not None: - self._client.ready(self.coordinator_id) + self._client.maybe_connect(self.coordinator_id) continue future = self.lookup_coordinator() @@ -686,7 +686,7 @@ class BaseCoordinator(object): self.coordinator_id = response.coordinator_id log.info("Discovered coordinator %s for group %s", self.coordinator_id, self.group_id) - self._client.ready(self.coordinator_id) + self._client.maybe_connect(self.coordinator_id) self.heartbeat.reset_timeouts() future.success(self.coordinator_id) |