From 04920bb89f9d73e4028dbd487719975c65954592 Mon Sep 17 00:00:00 2001 From: Enrico Canzonieri Date: Mon, 12 Oct 2015 00:11:18 -0700 Subject: [PATCH 1/2] Unblocking broker aware request --- kafka/client.py | 48 ++++++++++++++++++++++++++++++------------------ kafka/conn.py | 5 +++++ 2 files changed, 35 insertions(+), 18 deletions(-) diff --git a/kafka/client.py b/kafka/client.py index 13777a449..68277ed2a 100644 --- a/kafka/client.py +++ b/kafka/client.py @@ -2,6 +2,7 @@ import copy import functools import logging +import select import time import kafka.common @@ -177,6 +178,10 @@ def _send_broker_aware_request(self, payloads, encoder_fn, decoder_fn): # For each broker, send the list of request payloads # and collect the responses and errors broker_failures = [] + + # For each KafkaConnection we store the real socket so that we can use + # a select to perform unblocking I/O + socket_connection = {} for broker, payloads in payloads_by_broker.items(): requestId = self._next_id() log.debug('Request %s to %s: %s', requestId, broker, payloads) @@ -210,27 +215,34 @@ def _send_broker_aware_request(self, payloads, encoder_fn, decoder_fn): topic_partition = (payload.topic, payload.partition) responses[topic_partition] = None continue + else: + socket_connection[conn.get_connected_socket()] = (conn, broker) - try: - response = conn.recv(requestId) - except ConnectionError as e: - broker_failures.append(broker) - log.warning('ConnectionError attempting to receive a ' - 'response to request %s from server %s: %s', - requestId, broker, e) + conn = None + while socket_connection: + sockets = socket_connection.keys() + rlist, _, _ = select.select(sockets, [], [], None) + conn, broker = socket_connection.pop(rlist[0]) + try: + response = conn.recv(requestId) + except ConnectionError as e: + broker_failures.append(broker) + log.warning('ConnectionError attempting to receive a ' + 'response to request %s from server %s: %s', + requestId, broker, e) - for payload in payloads: - topic_partition = (payload.topic, payload.partition) - responses[topic_partition] = FailedPayloadsError(payload) + for payload in payloads: + topic_partition = (payload.topic, payload.partition) + responses[topic_partition] = FailedPayloadsError(payload) - else: - _resps = [] - for payload_response in decoder_fn(response): - topic_partition = (payload_response.topic, - payload_response.partition) - responses[topic_partition] = payload_response - _resps.append(payload_response) - log.debug('Response %s: %s', requestId, _resps) + else: + _resps = [] + for payload_response in decoder_fn(response): + topic_partition = (payload_response.topic, + payload_response.partition) + responses[topic_partition] = payload_response + _resps.append(payload_response) + log.debug('Response %s: %s', requestId, _resps) # Connection errors generally mean stale metadata # although sometimes it means incorrect api request diff --git a/kafka/conn.py b/kafka/conn.py index 432e10b0c..f1a12dc9b 100644 --- a/kafka/conn.py +++ b/kafka/conn.py @@ -118,6 +118,11 @@ def _read_bytes(self, num_bytes): # TODO multiplex socket communication to allow for multi-threaded clients + def get_connected_socket(self): + if not self._sock: + self.reinit() + return self._sock + def send(self, request_id, payload): """ Send a request to Kafka From c2adeeab057b825c8cccae67aac822be02293211 Mon Sep 17 00:00:00 2001 From: Enrico Canzonieri Date: Sat, 24 Oct 2015 16:50:46 -0700 Subject: [PATCH 2/2] Add tests. Bug fix. Rename socket_conn dict. --- kafka/client.py | 14 +++++++------- test/test_conn.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/kafka/client.py b/kafka/client.py index 68277ed2a..6603a4743 100644 --- a/kafka/client.py +++ b/kafka/client.py @@ -179,9 +179,9 @@ def _send_broker_aware_request(self, payloads, encoder_fn, decoder_fn): # and collect the responses and errors broker_failures = [] - # For each KafkaConnection we store the real socket so that we can use + # For each KafkaConnection keep the real socket so that we can use # a select to perform unblocking I/O - socket_connection = {} + connections_by_socket = {} for broker, payloads in payloads_by_broker.items(): requestId = self._next_id() log.debug('Request %s to %s: %s', requestId, broker, payloads) @@ -216,13 +216,13 @@ def _send_broker_aware_request(self, payloads, encoder_fn, decoder_fn): responses[topic_partition] = None continue else: - socket_connection[conn.get_connected_socket()] = (conn, broker) + connections_by_socket[conn.get_connected_socket()] = (conn, broker) conn = None - while socket_connection: - sockets = socket_connection.keys() + while connections_by_socket: + sockets = connections_by_socket.keys() rlist, _, _ = select.select(sockets, [], [], None) - conn, broker = socket_connection.pop(rlist[0]) + conn, broker = connections_by_socket.pop(rlist[0]) try: response = conn.recv(requestId) except ConnectionError as e: @@ -231,7 +231,7 @@ def _send_broker_aware_request(self, payloads, encoder_fn, decoder_fn): 'response to request %s from server %s: %s', requestId, broker, e) - for payload in payloads: + for payload in payloads_by_broker[broker]: topic_partition = (payload.topic, payload.partition) responses[topic_partition] = FailedPayloadsError(payload) diff --git a/test/test_conn.py b/test/test_conn.py index 2b7034461..1bdfc1eb0 100644 --- a/test/test_conn.py +++ b/test/test_conn.py @@ -165,6 +165,23 @@ def test_recv__doesnt_consume_extra_data_in_stream(self): self.assertEqual(self.conn.recv(self.config['request_id']), self.config['payload']) self.assertEqual(self.conn.recv(self.config['request_id']), self.config['payload2']) + def test_get_connected_socket(self): + s = self.conn.get_connected_socket() + + self.assertEqual(s, self.MockCreateConn()) + + def test_get_connected_socket_on_dirty_conn(self): + # Dirty the connection + try: + self.conn._raise_connection_error() + except ConnectionError: + pass + + # Test that get_connected_socket tries to connect + self.assertEqual(self.MockCreateConn.call_count, 0) + self.conn.get_connected_socket() + self.assertEqual(self.MockCreateConn.call_count, 1) + def test_close__object_is_reusable(self): # test that sending to a closed connection