def _send_broker_unaware_request(self, payloads, encoder_fn, decoder_fn):
"""
Attempt to send a broker-agnostic request to one of the available
brokers. Keep trying until you succeed.
"""
for (host, port) in self.hosts:
requestId = self._next_id()
log.debug('Request %s: %s', requestId, payloads)
try:
conn = self._get_conn(host, port)
request = encoder_fn(client_id=self.client_id,
correlation_id=requestId,
payloads=payloads)
conn.send(requestId, request)
response = conn.recv(requestId)
decoded = decoder_fn(response)
log.debug('Response %s: %s', requestId, decoded)
return decoded
except Exception:
log.exception(
'Error sending request [%s] to server %s:%s, '
'trying next server', requestId, host, port)
raise KafkaUnavailableError('All servers failed to process request')
def _send_broker_unaware_request(self, payloads, encoder_fn, decoder_fn):
"""
Attempt to send a broker-agnostic request to one of the available
brokers. Keep trying until you succeed.
"""
hosts = set()
for broker in self.brokers.values():
host, port, afi = get_ip_port_afi(broker.host)
hosts.add((host, broker.port, afi))
hosts.update(self.hosts)
hosts = list(hosts)
random.shuffle(hosts)
for (host, port, afi) in hosts:
conn = self._get_conn(host, port, afi)
if not conn.connected():
log.warning("Skipping unconnected connection: %s", conn)
continue
request = encoder_fn(payloads=payloads)
future = conn.send(request)
# Block
while not future.is_done:
conn.recv()
if future.failed():
log.error("Request failed: %s", future.exception)
continue
return decoder_fn(future.value)
raise KafkaUnavailableError(
'All servers failed to process request: %s' % hosts)
def test_get_kafka_topics_error():
expected = {'topic1': [0, 1, 2, 3], 'topic2': [0, 1]}
mock_client = mock.Mock()
mock_client.topic_partitions = expected
mock_client.load_metadata_for_topics.side_effect = KafkaUnavailableError(
'Boom!')
with pytest.raises(KafkaUnavailableError):
utils.get_kafka_topics(mock_client)
def test_get_kafka_topics_recover_from_error():
expected = {'topic1': [0, 1, 2, 3], 'topic2': [0, 1]}
mock_client = mock.Mock()
mock_client.topic_partitions = expected
mock_client.load_metadata_for_topics.side_effect = [
KafkaUnavailableError(), None
]
actual = utils.get_kafka_topics(mock_client)
assert expected == actual
def test_get_metadata_kafka_error(self, kafka_client_mock):
with mock.patch.object(MyKafkaToolClient,
'load_metadata_for_topics',
side_effect=KafkaUnavailableError("Boom!"),
autospec=True) as mock_func:
with pytest.raises(KafkaUnavailableError):
get_consumer_offsets_metadata(
kafka_client_mock,
self.group,
{'topic1': [99]},
)
assert mock_func.call_count == 2
def _send_broker_unaware_request(self, requestId, request):
"""
Attempt to send a broker-agnostic request to one of the available
brokers. Keep trying until you succeed.
"""
for (host, port) in self.hosts:
try:
conn = self._get_conn(host, port)
conn.send(requestId, request)
response = conn.recv(requestId)
return response
except Exception as e:
log.warning("Could not send request [%r] to server %s:%i, "
"trying next server: %s" % (request, host, port, e))
raise KafkaUnavailableError("All servers failed to process request")
def get_kafka_connection(cls, timeout_seconds=15):
"""Returns a kafka connection, waiting timeout_seconds for the container
to come up.
Args:
timeout_seconds: Retry time (seconds) to get a kafka connection
"""
end_time = time.time() + timeout_seconds
logger.info("Getting connection to Kafka container on yocalhost")
while end_time > time.time():
try:
return KafkaClient(get_config().cluster_config.broker_list)
except KafkaUnavailableError:
logger.info("Kafka not yet available, waiting...")
time.sleep(0.1)
raise KafkaUnavailableError()
def _send_broker_unaware_request(self, payloads, encoder_fn, decoder_fn):
"""
Attempt to send a broker-agnostic request to one of the available
brokers. Keep trying until you succeed.
"""
for (host, port) in self.hosts:
requestId = self._next_id()
try:
conn = self._get_conn(host, port)
request = encoder_fn(client_id=self.client_id,
correlation_id=requestId,
payloads=payloads)
conn.send(requestId, request)
response = conn.recv(requestId)
return decoder_fn(response)
except Exception as e:
log.warning("Could not send request [%r] to server %s:%i, "
"trying next server: %s" % (requestId, host, port, e))
raise KafkaUnavailableError("All servers failed to process request")
class KafkaClient(object):
CLIENT_ID = "kafka-python"
ID_GEN = count()
# NOTE: The timeout given to the client should always be greater than the
# one passed to SimpleConsumer.get_message(), otherwise you can get a
# socket timeout.
def __init__(self,
hosts,
client_id=CLIENT_ID,
timeout=DEFAULT_SOCKET_TIMEOUT_SECONDS):
# We need one connection to bootstrap
self.client_id = client_id
self.timeout = timeout
self.hosts = collect_hosts(hosts)
# create connections only when we need them
self.conns = {}
self.brokers = {} # broker_id -> BrokerMetadata
self.topics_to_brokers = {} # topic_id -> broker_id
self.topic_partitions = {} # topic_id -> [0, 1, 2, ...]
self.load_metadata_for_topics() # bootstrap with all metadata
##################
# Private API #
##################
def _get_conn(self, host, port):
"Get or create a connection to a broker using host and port"
host_key = (host, port)
if host_key not in self.conns:
self.conns[host_key] = KafkaConnection(host, port)
return self.conns[host_key]
def _get_conn_for_broker(self, broker):
"""
Get or create a connection to a broker
"""
if (broker.host, broker.port) not in self.conns:
self.conns[(broker.host, broker.port)] = \
KafkaConnection(broker.host, broker.port, timeout=self.timeout)
return self._get_conn(broker.host, broker.port)
def _get_leader_for_partition(self, topic, partition):
"""
Returns the leader for a partition or None if the partition exists
but has no leader.
PartitionUnavailableError will be raised if the topic or partition
is not part of the metadata.
"""
key = TopicAndPartition(topic, partition)
# reload metadata whether the partition is not available
# or has no leader (broker is None)
if self.topics_to_brokers.get(key) is None:
self.load_metadata_for_topics(topic)
if key not in self.topics_to_brokers:
raise PartitionUnavailableError("%s not available" % str(key))
return self.topics_to_brokers[key]
def _next_id(self):
"""
Generate a new correlation id
"""
return KafkaClient.ID_GEN.next()
def _send_broker_unaware_request(self, requestId, request):
"""
Attempt to send a broker-agnostic request to one of the available
brokers. Keep trying until you succeed.
"""
for (host, port) in self.hosts:
try:
conn = self._get_conn(host, port)
conn.send(requestId, request)
response = conn.recv(requestId)
return response
except Exception, e:
log.warning("Could not send request [%r] to server %s:%i, "
"trying next server: %s" %
(request, host, port, e))
continue
raise KafkaUnavailableError("All servers failed to process request")
class KafkaClient(object):
CLIENT_ID = "kafka-python"
ID_GEN = count()
# NOTE: The timeout given to the client should always be greater than the
# one passed to SimpleConsumer.get_message(), otherwise you can get a
# socket timeout.
def __init__(self,
host,
port,
client_id=CLIENT_ID,
timeout=DEFAULT_SOCKET_TIMEOUT_SECONDS):
# We need one connection to bootstrap
self.client_id = client_id
self.timeout = timeout
self.conns = { # (host, port) -> KafkaConnection
(host, port): KafkaConnection(host, port, timeout=timeout)
}
self.brokers = {} # broker_id -> BrokerMetadata
self.topics_to_brokers = {} # topic_id -> broker_id
self.topic_partitions = {} # topic_id -> [0, 1, 2, ...]
self.load_metadata_for_topics() # bootstrap with all metadata
##################
# Private API #
##################
def _get_conn_for_broker(self, broker):
"""
Get or create a connection to a broker
"""
if (broker.host, broker.port) not in self.conns:
self.conns[(broker.host, broker.port)] = \
KafkaConnection(broker.host, broker.port, timeout=self.timeout)
return self.conns[(broker.host, broker.port)]
def _get_leader_for_partition(self, topic, partition):
key = TopicAndPartition(topic, partition)
if key not in self.topics_to_brokers:
self.load_metadata_for_topics(topic)
if key not in self.topics_to_brokers:
raise KafkaRequestError("Partition does not exist: %s" % str(key))
return self.topics_to_brokers[key]
def _next_id(self):
"""
Generate a new correlation id
"""
return KafkaClient.ID_GEN.next()
def _send_broker_unaware_request(self, requestId, request):
"""
Attempt to send a broker-agnostic request to one of the available
brokers. Keep trying until you succeed.
"""
for conn in self.conns.values():
try:
conn.send(requestId, request)
response = conn.recv(requestId)
return response
except Exception, e:
log.warning("Could not send request [%r] to server %s, "
"trying next server: %s" % (request, conn, e))
continue
raise KafkaUnavailableError("All servers failed to process request")
