async def test_connections_max_idle_ms(self):
host, port = self.kafka_host, self.kafka_port
conn = await create_conn(host, port, max_idle_ms=200)
self.assertEqual(conn.connected(), True)
await asyncio.sleep(0.1)
# Do some work
request = MetadataRequest([])
await conn.send(request)
await asyncio.sleep(0.15)
# Check if we're still connected after 250ms, as we were not idle
self.assertEqual(conn.connected(), True)
# It shouldn't break if we have a long running call either
readexactly = conn._reader.readexactly
with mock.patch.object(conn._reader, 'readexactly') as mocked:
async def long_read(n):
await asyncio.sleep(0.2)
return (await readexactly(n))
mocked.side_effect = long_read
await conn.send(MetadataRequest([]))
self.assertEqual(conn.connected(), True)
await asyncio.sleep(0.2)
self.assertEqual(conn.connected(), False)
def test_connections_max_idle_ms(self):
host, port = self.kafka_host, self.kafka_port
conn = yield from create_conn(host,
port,
loop=self.loop,
max_idle_ms=200)
self.assertEqual(conn.connected(), True)
yield from asyncio.sleep(0.1, loop=self.loop)
# Do some work
request = MetadataRequest([])
yield from conn.send(request)
yield from asyncio.sleep(0.15, loop=self.loop)
# Check if we're stil connected after 250ms, as we were not idle
self.assertEqual(conn.connected(), True)
# It shouldn't break if we have a long running call either
readexactly = conn._reader.readexactly
with mock.patch.object(conn._reader, 'readexactly') as mocked:
@asyncio.coroutine
def long_read(n):
yield from asyncio.sleep(0.2, loop=self.loop)
return (yield from readexactly(n))
mocked.side_effect = long_read
yield from conn.send(MetadataRequest([]))
self.assertEqual(conn.connected(), True)
yield from asyncio.sleep(0.2, loop=self.loop)
self.assertEqual(conn.connected(), False)
def test_send_timeout_deletes_connection(self):
correct_response = MetadataResponse([], [])
@asyncio.coroutine
def send_exception(*args, **kwargs):
raise asyncio.TimeoutError()
@asyncio.coroutine
def send(*args, **kwargs):
return correct_response
@asyncio.coroutine
def get_conn(self, node_id, *, group=0):
conn_id = (node_id, group)
if conn_id in self._conns:
conn = self._conns[conn_id]
if not conn.connected():
del self._conns[conn_id]
else:
return conn
conn = mock.MagicMock()
conn.send.side_effect = send
self._conns[conn_id] = conn
return conn
node_id = 0
conn = mock.MagicMock()
conn.send.side_effect = send_exception
conn.connected.return_value = True
mocked_conns = {(node_id, 0): conn}
client = AIOKafkaClient(loop=self.loop,
bootstrap_servers=['broker_1:4567'])
client._conns = mocked_conns
client._get_conn = types.MethodType(get_conn, client)
# first send timeouts
with self.assertRaises(RequestTimedOutError):
yield from client.send(0, MetadataRequest([]))
conn.close.assert_called_once_with(
reason=CloseReason.CONNECTION_TIMEOUT)
# this happens because conn was closed
conn.connected.return_value = False
# second send gets new connection and obtains result
response = yield from client.send(0, MetadataRequest([]))
self.assertEqual(response, correct_response)
self.assertNotEqual(conn, client._conns[(node_id, 0)])
async def test_invalid_correlation_id(self):
host, port = self.kafka_host, self.kafka_port
request = MetadataRequest([])
# setup connection with mocked reader and writer
conn = AIOKafkaConnection(host=host, port=port)
# setup reader
reader = mock.MagicMock()
int32 = struct.Struct('>i')
resp = MetadataResponse(brokers=[], topics=[])
resp = resp.encode()
resp = int32.pack(999) + resp # set invalid correlation id
async def first_resp(*args: Any, **kw: Any):
return int32.pack(len(resp))
async def second_resp(*args: Any, **kw: Any):
return resp
reader.readexactly.side_effect = [first_resp(), second_resp()]
writer = mock.MagicMock()
conn._reader = reader
conn._writer = writer
# invoke reader task
conn._read_task = conn._create_reader_task()
with self.assertRaises(CorrelationIdError):
await conn.send(request)
def test_correlation_id_rollover(kafka_broker):
logging.getLogger('kafka.conn').setLevel(logging.ERROR)
from kafka.protocol.metadata import MetadataRequest
conn = BrokerConnection('localhost',
kafka_broker.port,
receive_buffer_bytes=131072,
max_in_flight_requests_per_connection=100)
req = MetadataRequest([])
while not conn.connected():
conn.connect()
futures = collections.deque()
start = time.time()
done = 0
for i in six.moves.xrange(2**13):
if not conn.can_send_more():
conn.recv(timeout=None)
futures.append(conn.send(req))
conn.recv()
while futures and futures[0].is_done:
f = futures.popleft()
if not f.succeeded():
raise f.exception
done += 1
if time.time() > start + 10:
print("%d done" % done)
start = time.time()
while futures:
conn.recv()
if futures[0].is_done:
f = futures.popleft()
if not f.succeeded():
raise f.exception
async def test_no_concurrent_send_on_connection(self):
client = AIOKafkaClient(bootstrap_servers=self.hosts,
metadata_max_age_ms=10000)
await client.bootstrap()
self.add_cleanup(client.close)
await self.wait_topic(client, self.topic)
node_id = client.get_random_node()
wait_request = FetchRequest_v0(
-1, # replica_id
500, # max_wait_ms
1024 * 1024, # min_bytes
[(self.topic, [(0, 0, 1024)])])
vanila_request = MetadataRequest([])
loop = get_running_loop()
send_time = loop.time()
long_task = create_task(client.send(node_id, wait_request))
await asyncio.sleep(0.0001)
self.assertFalse(long_task.done())
await client.send(node_id, vanila_request)
resp_time = loop.time()
fetch_resp = await long_task
# Check error code like resp->topics[0]->partitions[0]->error_code
self.assertEqual(fetch_resp.topics[0][1][0][1], 0)
# Check that vanila request actually executed after wait request
self.assertGreaterEqual(resp_time - send_time, 0.5)
def _metadata_update(self, cluster_metadata, topics):
assert isinstance(cluster_metadata, ClusterMetadata)
metadata_request = MetadataRequest(list(topics))
nodeids = [b.nodeId for b in self.cluster.brokers()]
if 'bootstrap' in self._conns:
nodeids.append('bootstrap')
random.shuffle(nodeids)
for node_id in nodeids:
conn = yield from self._get_conn(node_id)
if conn is None:
continue
log.debug("Sending metadata request %s to %s", metadata_request,
node_id)
try:
metadata = yield from conn.send(metadata_request)
except KafkaError as err:
log.error(
'Unable to request metadata from node with id %s: %s',
node_id, err)
continue
cluster_metadata.update_metadata(metadata)
break
else:
log.error('Unable to update metadata from %s', nodeids)
cluster_metadata.failed_update(None)
return False
return True
def test_invalid_correlation_id(self):
host, port = self.kafka_host, self.kafka_port
request = MetadataRequest([])
# setup connection with mocked reader and writer
conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)
# setup reader
reader = mock.MagicMock()
int32 = struct.Struct('>i')
resp = MetadataResponse(brokers=[], topics=[]).encode()
resp = int32.pack(999) + resp # set invalid correlation id
reader.readexactly.side_effect = [
asyncio.coroutine(lambda *a, **kw: int32.pack(len(resp)))(),
asyncio.coroutine(lambda *a, **kw: resp)()
]
writer = mock.MagicMock()
conn._reader = reader
conn._writer = writer
# invoke reader task
conn._read_task = asyncio. async (conn._read(), loop=self.loop)
with self.assertRaises(CorrelationIdError):
yield from conn.send(request)
def test_concurrent_send_on_different_connection_groups(self):
client = AIOKafkaClient(loop=self.loop,
bootstrap_servers=self.hosts,
metadata_max_age_ms=10000)
yield from client.bootstrap()
self.add_cleanup(client.close)
yield from self.wait_topic(client, self.topic)
node_id = client.get_random_node()
wait_request = FetchRequest_v0(
-1, # replica_id
500, # max_wait_ms
1024 * 1024, # min_bytes
[(self.topic, [(0, 0, 1024)])])
vanila_request = MetadataRequest([])
send_time = self.loop.time()
long_task = self.loop.create_task(client.send(node_id, wait_request))
yield from asyncio.sleep(0.0001, loop=self.loop)
self.assertFalse(long_task.done())
yield from client.send(node_id,
vanila_request,
group=ConnectionGroup.COORDINATION)
resp_time = self.loop.time()
self.assertFalse(long_task.done())
fetch_resp = yield from long_task
# Check error code like resp->topics[0]->partitions[0]->error_code
self.assertEqual(fetch_resp.topics[0][1][0][1], 0)
# Check that vanila request actually executed after wait request
self.assertLess(resp_time - send_time, 0.5)
def test_osserror_in_reader_task(self):
host, port = self.kafka_host, self.kafka_port
@asyncio.coroutine
def invoke_osserror(*a, **kw):
yield from asyncio.sleep(0.1, loop=self.loop)
raise OSError('test oserror')
request = MetadataRequest([])
# setup connection with mocked reader and writer
conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)
# setup reader
reader = mock.MagicMock()
reader.readexactly.return_value = invoke_osserror()
writer = mock.MagicMock()
conn._reader = reader
conn._writer = writer
# invoke reader task
conn._read_task = asyncio. async (conn._read(), loop=self.loop)
with self.assertRaises(ConnectionError):
yield from conn.send(request)
self.assertEqual(conn.connected(), False)
def test_send_request(self):
client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
yield from client.bootstrap()
node_id = client.get_random_node()
resp = yield from client.send(node_id, MetadataRequest([]))
self.assertTrue(isinstance(resp, MetadataResponse))
yield from client.close()
async def test_send_request(self):
client = AIOKafkaClient(bootstrap_servers=self.hosts)
await client.bootstrap()
node_id = client.get_random_node()
resp = await client.send(node_id, MetadataRequest([]))
self.assertTrue(isinstance(resp, MetadataResponse))
await client.close()
async def test_basic_connection_load_meta(self):
host, port = self.kafka_host, self.kafka_port
conn = await create_conn(host, port)
self.assertEqual(conn.connected(), True)
request = MetadataRequest([])
response = await conn.send(request)
conn.close()
self.assertIsInstance(response, MetadataResponse)
def test_bootstrap_success(conn):
conn.state = ConnectionStates.CONNECTED
cli = KafkaClient()
conn.assert_called_once_with('localhost', 9092, **cli.config)
conn.connect.assert_called_with()
conn.send.assert_called_once_with(MetadataRequest([]))
assert cli._bootstrap_fails == 0
assert cli.cluster.brokers() == set([BrokerMetadata(0, 'foo', 12),
BrokerMetadata(1, 'bar', 34)])
async def test_send_to_closed(self):
host, port = self.kafka_host, self.kafka_port
conn = AIOKafkaConnection(host=host, port=port)
request = MetadataRequest([])
with self.assertRaises(KafkaConnectionError):
await conn.send(request)
conn._writer = mock.MagicMock()
conn._writer.write.side_effect = OSError('mocked writer is closed')
with self.assertRaises(KafkaConnectionError):
await conn.send(request)
def test_send_to_closed(self):
host, port = self.kafka_host, self.kafka_port
conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)
request = MetadataRequest([])
with self.assertRaises(ConnectionError):
yield from conn.send(request)
@asyncio.coroutine
def invoke_osserror(*a, **kw):
yield from asyncio.sleep(0.1, loop=self.loop)
raise OSError('mocked writer is closed')
conn._writer = mock.MagicMock()
conn._writer.write.side_effect = OSError('mocked writer is closed')
with self.assertRaises(ConnectionError):
yield from conn.send(request)
def bootstrap(self):
"""Try to to bootstrap initial cluster metadata"""
metadata_request = MetadataRequest([])
for host, port, _ in self.hosts:
log.debug("Attempting to bootstrap via node at %s:%s", host, port)
try:
bootstrap_conn = yield from create_conn(
host,
port,
loop=self._loop,
client_id=self._client_id,
request_timeout_ms=self._request_timeout_ms)
except (OSError, asyncio.TimeoutError) as err:
log.error('Unable connect to "%s:%s": %s', host, port, err)
continue
try:
metadata = yield from bootstrap_conn.send(metadata_request)
except KafkaError as err:
log.warning('Unable to request metadata from "%s:%s": %s',
host, port, err)
bootstrap_conn.close()
continue
self.cluster.update_metadata(metadata)
# A cluster with no topics can return no broker metadata
# in that case, we should keep the bootstrap connection
if not len(self.cluster.brokers()):
self._conns['bootstrap'] = bootstrap_conn
else:
bootstrap_conn.close()
log.debug('Received cluster metadata: %s', self.cluster)
break
else:
raise ConnectionError('Unable to bootstrap from {}'.format(
self.hosts))
if self._sync_task is None:
# starting metadata synchronizer task
self._sync_task = ensure_future(self._md_synchronizer(),
loop=self._loop)
def test_send(conn):
cli = KafkaClient()
try:
cli.send(2, None)
except Errors.NodeNotReadyError:
pass
else:
assert False, 'NodeNotReadyError not raised'
cli._initiate_connect(0)
# ProduceRequest w/ 0 required_acks -> no response
request = ProduceRequest(0, 0, [])
ret = cli.send(0, request)
assert conn.send.called_with(request, expect_response=False)
assert isinstance(ret, Future)
request = MetadataRequest([])
cli.send(0, request)
assert conn.send.called_with(request, expect_response=True)
async def test_concurrent_send_on_different_connection_groups(self):
client = AIOKafkaClient(bootstrap_servers=self.hosts,
metadata_max_age_ms=10000)
await client.bootstrap()
self.add_cleanup(client.close)
await self.wait_topic(client, self.topic)
node_id = client.get_random_node()
broker = client.cluster.broker_metadata(node_id)
client.cluster.add_coordinator(node_id,
broker.host,
broker.port,
rack=None,
purpose=(CoordinationType.GROUP, ""))
wait_request = FetchRequest_v0(
-1, # replica_id
500, # max_wait_ms
1024 * 1024, # min_bytes
[(self.topic, [(0, 0, 1024)])])
vanila_request = MetadataRequest([])
loop = get_running_loop()
send_time = loop.time()
long_task = create_task(client.send(node_id, wait_request))
await asyncio.sleep(0.0001)
self.assertFalse(long_task.done())
await client.send(node_id,
vanila_request,
group=ConnectionGroup.COORDINATION)
resp_time = loop.time()
self.assertFalse(long_task.done())
fetch_resp = await long_task
# Check error code like resp->topics[0]->partitions[0]->error_code
self.assertEqual(fetch_resp.topics[0][1][0][1], 0)
# Check that vanila request actually executed after wait request
self.assertLess(resp_time - send_time, 0.5)
def check_version(self, node_id=None):
"""Attempt to guess the broker version"""
if node_id is None:
if self._conns:
node_id = list(self._conns.keys())[0]
else:
assert self.cluster.brokers(), 'no brokers in metadata'
node_id = list(self.cluster.brokers())[0].nodeId
from kafka.protocol.admin import ListGroupsRequest
from kafka.protocol.commit import (OffsetFetchRequest_v0,
GroupCoordinatorRequest)
from kafka.protocol.metadata import MetadataRequest
test_cases = [
('0.9', ListGroupsRequest()),
('0.8.2', GroupCoordinatorRequest('kafka-python-default-group')),
('0.8.1', OffsetFetchRequest_v0('kafka-python-default-group', [])),
('0.8.0', MetadataRequest([])),
]
# kafka kills the connection when it doesnt recognize an API request
# so we can send a test request and then follow immediately with a
# vanilla MetadataRequest. If the server did not recognize the first
# request, both will be failed with a ConnectionError that wraps
# socket.error (32, 54, or 104)
conn = yield from self._get_conn(node_id)
if conn is None:
raise ConnectionError(
"No connection to node with id {}".format(node_id))
for version, request in test_cases:
try:
if not conn.connected():
yield from conn.connect()
assert conn, 'no connection to node with id {}'.format(node_id)
yield from conn.send(request)
except KafkaError:
continue
else:
return version
raise UnrecognizedBrokerVersion()
async def test_metadata_updated_on_socket_disconnect(self):
# Related to issue 176. A disconnect means that either we lost
# connection to the node, or we have a node failure. In both cases
# there's a high probability that Leader distribution will also change.
client = AIOKafkaClient(bootstrap_servers=self.hosts,
metadata_max_age_ms=10000)
await client.bootstrap()
self.add_cleanup(client.close)
# Init a clonnection
node_id = client.get_random_node()
assert node_id is not None
req = MetadataRequest([])
await client.send(node_id, req)
# No metadata update pending atm
self.assertFalse(client._md_update_waiter.done())
# Connection disconnect should trigger an update
conn = await client._get_conn(node_id)
conn.close(reason=CloseReason.CONNECTION_BROKEN)
self.assertTrue(client._md_update_waiter.done())