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 __init__(self, topics: Iterable[str], loop=None):
if loop is None:
loop = get_running_loop()
self._topics = frozenset(topics) # type: FrozenSet[str]
self._assignment = None # type: Assignment
self.unsubscribe_future = loop.create_future() # type: Future
self._reassignment_in_progress = True
async def test_global_loop_for_create_conn(self):
loop = get_running_loop()
host, port = self.kafka_host, self.kafka_port
conn = await create_conn(host, port)
self.assertIs(conn._loop, loop)
conn.close()
# make sure second closing does nothing and we have full coverage
# of *if self._reader:* condition
conn.close()
def __init__(self, loop=None):
if loop is None:
loop = get_running_loop()
self._loop = loop
self._subscription_waiters = [] # type: List[Future]
self._assignment_waiters = [] # type: List[Future]
# Fetch contexts
self._fetch_count = 0
self._last_fetch_ended = time.monotonic()
async def test_add_batch_builder(self):
tp0 = TopicPartition("test-topic", 0)
tp1 = TopicPartition("test-topic", 1)
def mocked_leader_for_partition(tp):
if tp == tp0:
return 0
if tp == tp1:
return 1
return None
cluster = ClusterMetadata(metadata_max_age_ms=10000)
cluster.leader_for_partition = mock.MagicMock()
cluster.leader_for_partition.side_effect = mocked_leader_for_partition
ma = MessageAccumulator(cluster, 1000, 0, 1)
builder0 = ma.create_builder()
builder1_1 = ma.create_builder()
builder1_2 = ma.create_builder()
# batches may queued one-per-TP
self.assertFalse(ma._wait_data_future.done())
await ma.add_batch(builder0, tp0, 1)
self.assertTrue(ma._wait_data_future.done())
self.assertEqual(len(ma._batches[tp0]), 1)
await ma.add_batch(builder1_1, tp1, 1)
self.assertEqual(len(ma._batches[tp1]), 1)
with self.assertRaises(KafkaTimeoutError):
await ma.add_batch(builder1_2, tp1, 0.1)
self.assertTrue(ma._wait_data_future.done())
self.assertEqual(len(ma._batches[tp1]), 1)
# second batch gets added once the others are cleared out
get_running_loop().call_later(0.1, ma.drain_by_nodes, [])
await ma.add_batch(builder1_2, tp1, 1)
self.assertTrue(ma._wait_data_future.done())
self.assertEqual(len(ma._batches[tp0]), 0)
self.assertEqual(len(ma._batches[tp1]), 1)
def __init__(
self, cluster, batch_size, compression_type, batch_ttl, *,
txn_manager=None, loop=None):
if loop is None:
loop = get_running_loop()
self._loop = loop
self._batches = collections.defaultdict(collections.deque)
self._pending_batches = set([])
self._cluster = cluster
self._batch_size = batch_size
self._compression_type = compression_type
self._batch_ttl = batch_ttl
self._wait_data_future = loop.create_future()
self._closed = False
self._api_version = (0, 9)
self._txn_manager = txn_manager
self._exception = None # Critical exception
async def start(self):
"""Connect to Kafka cluster and check server version"""
assert self._loop is get_running_loop(), (
"Please create objects with the same loop as running with")
log.debug("Starting the Kafka producer") # trace
await self.client.bootstrap()
if self._compression_type == 'lz4':
assert self.client.api_version >= (0, 8, 2), \
'LZ4 Requires >= Kafka 0.8.2 Brokers'
if self._txn_manager is not None and self.client.api_version < (0, 11):
raise UnsupportedVersionError(
"Idempotent producer available only for Broker version 0.11"
" and above")
await self._sender.start()
self._message_accumulator.set_api_version(self.client.api_version)
self._producer_magic = 0 if self.client.api_version < (0, 10) else 1
log.debug("Kafka producer started")
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)
async def test_sender__handler_base_do(self):
sender = await self._setup_sender()
class MockHandler(BaseHandler):
def create_request(self):
return MetadataRequest[0]([])
mock_handler = MockHandler(sender)
mock_handler.handle_response = mock.Mock(return_value=0.1)
success = await mock_handler.do(node_id=0)
self.assertFalse(success)
MockHandler.return_value = None
mock_handler.handle_response = mock.Mock(return_value=None)
success = await mock_handler.do(node_id=0)
self.assertTrue(success)
loop = get_running_loop()
time = loop.time()
sender.client.send = mock.Mock(side_effect=UnknownError())
success = await mock_handler.do(node_id=0)
self.assertFalse(success)
self.assertAlmostEqual(loop.time() - time, 0.1, 1)
async def test_fetcher__update_fetch_positions(self):
client = AIOKafkaClient(bootstrap_servers=[])
subscriptions = SubscriptionState()
fetcher = Fetcher(client, subscriptions)
self.add_cleanup(fetcher.close)
# Disable background task
fetcher._fetch_task.cancel()
try:
await fetcher._fetch_task
except asyncio.CancelledError:
pass
fetcher._fetch_task = create_task(asyncio.sleep(1000000))
partition = TopicPartition('test', 0)
offsets = {partition: OffsetAndTimestamp(12, -1)}
async def _proc_offset_request(node_id, topic_data):
return offsets
fetcher._proc_offset_request = mock.Mock()
fetcher._proc_offset_request.side_effect = _proc_offset_request
def reset_assignment():
subscriptions.assign_from_user({partition})
assignment = subscriptions.subscription.assignment
tp_state = assignment.state_value(partition)
return assignment, tp_state
assignment, tp_state = reset_assignment()
self.assertIsNone(tp_state._position)
# CASE: reset from committed
# In basic case we will need to wait for committed
update_task = create_task(
fetcher._update_fetch_positions(assignment, 0, [partition]), )
await asyncio.sleep(0.1)
self.assertFalse(update_task.done())
# Will continue only after committed is resolved
tp_state.update_committed(OffsetAndMetadata(4, ""))
needs_wakeup = await update_task
self.assertFalse(needs_wakeup)
self.assertEqual(tp_state._position, 4)
self.assertEqual(fetcher._proc_offset_request.call_count, 0)
# CASE: will not query committed if position already present
await fetcher._update_fetch_positions(assignment, 0, [partition])
self.assertEqual(tp_state._position, 4)
self.assertEqual(fetcher._proc_offset_request.call_count, 0)
# CASE: awaiting_reset for the partition
tp_state.await_reset(OffsetResetStrategy.LATEST)
self.assertIsNone(tp_state._position)
await fetcher._update_fetch_positions(assignment, 0, [partition])
self.assertEqual(tp_state._position, 12)
self.assertEqual(fetcher._proc_offset_request.call_count, 1)
# CASE: seeked while waiting for committed to be resolved
assignment, tp_state = reset_assignment()
update_task = create_task(
fetcher._update_fetch_positions(assignment, 0, [partition]), )
await asyncio.sleep(0.1)
self.assertFalse(update_task.done())
tp_state.seek(8)
tp_state.update_committed(OffsetAndMetadata(4, ""))
await update_task
self.assertEqual(tp_state._position, 8)
self.assertEqual(fetcher._proc_offset_request.call_count, 1)
# CASE: awaiting_reset during waiting for committed
assignment, tp_state = reset_assignment()
update_task = create_task(
fetcher._update_fetch_positions(assignment, 0, [partition]), )
await asyncio.sleep(0.1)
self.assertFalse(update_task.done())
tp_state.await_reset(OffsetResetStrategy.LATEST)
tp_state.update_committed(OffsetAndMetadata(4, ""))
await update_task
self.assertEqual(tp_state._position, 12)
self.assertEqual(fetcher._proc_offset_request.call_count, 2)
# CASE: reset using default strategy if committed offset undefined
assignment, tp_state = reset_assignment()
loop = get_running_loop()
loop.call_later(0.01, tp_state.update_committed,
OffsetAndMetadata(-1, ""))
await fetcher._update_fetch_positions(assignment, 0, [partition])
self.assertEqual(tp_state._position, 12)
self.assertEqual(fetcher._records, {})
# CASE: set error if _default_reset_strategy = OffsetResetStrategy.NONE
assignment, tp_state = reset_assignment()
loop.call_later(0.01, tp_state.update_committed,
OffsetAndMetadata(-1, ""))
fetcher._default_reset_strategy = OffsetResetStrategy.NONE
needs_wakeup = await fetcher._update_fetch_positions(
assignment, 0, [partition])
self.assertTrue(needs_wakeup)
self.assertIsNone(tp_state._position)
self.assertIsInstance(fetcher._records[partition], FetchError)
fetcher._records.clear()
# CASE: if _proc_offset_request errored, we will retry on another spin
fetcher._proc_offset_request.side_effect = UnknownError()
assignment, tp_state = reset_assignment()
tp_state.await_reset(OffsetResetStrategy.LATEST)
await fetcher._update_fetch_positions(assignment, 0, [partition])
self.assertIsNone(tp_state._position)
self.assertTrue(tp_state.awaiting_reset)
# CASE: reset 2 partitions separately, 1 will raise, 1 will get
# committed
fetcher._proc_offset_request.side_effect = _proc_offset_request
partition2 = TopicPartition('test', 1)
subscriptions.assign_from_user({partition, partition2})
assignment = subscriptions.subscription.assignment
tp_state = assignment.state_value(partition)
tp_state2 = assignment.state_value(partition2)
tp_state.await_reset(OffsetResetStrategy.LATEST)
loop.call_later(0.01, tp_state2.update_committed,
OffsetAndMetadata(5, ""))
await fetcher._update_fetch_positions(assignment, 0,
[partition, partition2])
self.assertEqual(tp_state.position, 12)
self.assertEqual(tp_state2.position, 5)
def __init__(self,
host,
port,
*,
client_id='aiokafka',
request_timeout_ms=40000,
api_version=(0, 8, 2),
ssl_context=None,
security_protocol='PLAINTEXT',
max_idle_ms=None,
on_close=None,
sasl_mechanism=None,
sasl_plain_password=None,
sasl_plain_username=None,
sasl_kerberos_service_name='kafka',
sasl_kerberos_domain_name=None,
sasl_oauth_token_provider=None,
version_hint=None):
loop = get_running_loop()
if sasl_mechanism == "GSSAPI":
assert gssapi is not None, "gssapi library required"
if sasl_mechanism == "OAUTHBEARER":
if sasl_oauth_token_provider is None or \
not isinstance(
sasl_oauth_token_provider, AbstractTokenProvider):
raise ValueError("sasl_oauth_token_provider needs to be \
provided implementing aiokafka.abc.AbstractTokenProvider")
assert callable(getattr(
sasl_oauth_token_provider, "token", None)), (
'sasl_oauth_token_provider must implement method #token()')
self._loop = loop
self._host = host
self._port = port
self._request_timeout = request_timeout_ms / 1000
self._api_version = api_version
self._client_id = client_id
self._ssl_context = ssl_context
self._security_protocol = security_protocol
self._sasl_mechanism = sasl_mechanism
self._sasl_plain_username = sasl_plain_username
self._sasl_plain_password = sasl_plain_password
self._sasl_kerberos_service_name = sasl_kerberos_service_name
self._sasl_kerberos_domain_name = sasl_kerberos_domain_name
self._sasl_oauth_token_provider = sasl_oauth_token_provider
# Version hint is the version determined by initial client bootstrap
self._version_hint = version_hint
self._version_info = VersionInfo({})
self._reader = self._writer = self._protocol = None
# Even on small size seems to be a bit faster than list.
# ~2x on size of 2 in Python3.6
self._requests = collections.deque()
self._read_task = None
self._correlation_id = 0
self._closed_fut = None
self._max_idle_ms = max_idle_ms
self._last_action = time.monotonic()
self._idle_handle = None
self._on_close_cb = on_close
if loop.get_debug():
self._source_traceback = traceback.extract_stack(sys._getframe(1))
def __init__(self,
*,
loop=None,
bootstrap_servers='localhost',
client_id=None,
metadata_max_age_ms=300000,
request_timeout_ms=40000,
api_version='auto',
acks=_missing,
key_serializer=None,
value_serializer=None,
compression_type=None,
max_batch_size=16384,
partitioner=DefaultPartitioner(),
max_request_size=1048576,
linger_ms=0,
send_backoff_ms=100,
retry_backoff_ms=100,
security_protocol="PLAINTEXT",
ssl_context=None,
connections_max_idle_ms=540000,
enable_idempotence=False,
transactional_id=None,
transaction_timeout_ms=60000,
sasl_mechanism="PLAIN",
sasl_plain_password=None,
sasl_plain_username=None,
sasl_kerberos_service_name='kafka',
sasl_kerberos_domain_name=None):
if loop is None:
loop = get_running_loop()
if acks not in (0, 1, -1, 'all', _missing):
raise ValueError("Invalid ACKS parameter")
if compression_type not in ('gzip', 'snappy', 'lz4', None):
raise ValueError("Invalid compression type!")
if compression_type:
checker, compression_attrs = self._COMPRESSORS[compression_type]
if not checker():
raise RuntimeError(
"Compression library for {} not found".format(
compression_type))
else:
compression_attrs = 0
if transactional_id is not None:
enable_idempotence = True
else:
transaction_timeout_ms = INTEGER_MAX_VALUE
if enable_idempotence:
if acks is _missing:
acks = -1
elif acks not in ('all', -1):
raise ValueError(
"acks={} not supported if enable_idempotence=True".format(
acks))
self._txn_manager = TransactionManager(transactional_id,
transaction_timeout_ms,
loop=loop)
else:
self._txn_manager = None
if acks is _missing:
acks = 1
elif acks == 'all':
acks = -1
AIOKafkaProducer._PRODUCER_CLIENT_ID_SEQUENCE += 1
if client_id is None:
client_id = 'aiokafka-producer-%s' % \
AIOKafkaProducer._PRODUCER_CLIENT_ID_SEQUENCE
self._key_serializer = key_serializer
self._value_serializer = value_serializer
self._compression_type = compression_type
self._partitioner = partitioner
self._max_request_size = max_request_size
self._request_timeout_ms = request_timeout_ms
self.client = AIOKafkaClient(
loop=loop,
bootstrap_servers=bootstrap_servers,
client_id=client_id,
metadata_max_age_ms=metadata_max_age_ms,
request_timeout_ms=request_timeout_ms,
retry_backoff_ms=retry_backoff_ms,
api_version=api_version,
security_protocol=security_protocol,
ssl_context=ssl_context,
connections_max_idle_ms=connections_max_idle_ms,
sasl_mechanism=sasl_mechanism,
sasl_plain_username=sasl_plain_username,
sasl_plain_password=sasl_plain_password,
sasl_kerberos_service_name=sasl_kerberos_service_name,
sasl_kerberos_domain_name=sasl_kerberos_domain_name)
self._metadata = self.client.cluster
self._message_accumulator = MessageAccumulator(
self._metadata,
max_batch_size,
compression_attrs,
self._request_timeout_ms / 1000,
txn_manager=self._txn_manager,
loop=loop)
self._sender = Sender(self.client,
acks=acks,
txn_manager=self._txn_manager,
retry_backoff_ms=retry_backoff_ms,
linger_ms=linger_ms,
message_accumulator=self._message_accumulator,
request_timeout_ms=request_timeout_ms,
loop=loop)
self._loop = loop
if loop.get_debug():
self._source_traceback = traceback.extract_stack(sys._getframe(1))
self._closed = False
def __init__(self,
*,
loop=None,
bootstrap_servers='localhost',
client_id='aiokafka-' + __version__,
metadata_max_age_ms=300000,
request_timeout_ms=40000,
retry_backoff_ms=100,
ssl_context=None,
security_protocol='PLAINTEXT',
api_version='auto',
connections_max_idle_ms=540000,
sasl_mechanism='PLAIN',
sasl_plain_username=None,
sasl_plain_password=None,
sasl_kerberos_service_name='kafka',
sasl_kerberos_domain_name=None,
sasl_oauth_token_provider=None):
if loop is None:
loop = get_running_loop()
if security_protocol not in ('SSL', 'PLAINTEXT', 'SASL_PLAINTEXT',
'SASL_SSL'):
raise ValueError("`security_protocol` should be SSL or PLAINTEXT")
if security_protocol in ["SSL", "SASL_SSL"] and ssl_context is None:
raise ValueError(
"`ssl_context` is mandatory if security_protocol=='SSL'")
if security_protocol in ["SASL_SSL", "SASL_PLAINTEXT"]:
if sasl_mechanism not in ("PLAIN", "GSSAPI", "SCRAM-SHA-256",
"SCRAM-SHA-512", "OAUTHBEARER"):
raise ValueError("only `PLAIN`, `GSSAPI`, `SCRAM-SHA-256`, "
"`SCRAM-SHA-512` and `OAUTHBEARER`"
"sasl_mechanism are supported "
"at the moment")
if sasl_mechanism == "PLAIN" and \
(sasl_plain_username is None or sasl_plain_password is None):
raise ValueError(
"sasl_plain_username and sasl_plain_password required for "
"PLAIN sasl")
self._bootstrap_servers = bootstrap_servers
self._client_id = client_id
self._metadata_max_age_ms = metadata_max_age_ms
self._request_timeout_ms = request_timeout_ms
if api_version != "auto":
api_version = parse_kafka_version(api_version)
self._api_version = api_version
self._security_protocol = security_protocol
self._ssl_context = ssl_context
self._retry_backoff = retry_backoff_ms / 1000
self._connections_max_idle_ms = connections_max_idle_ms
self._sasl_mechanism = sasl_mechanism
self._sasl_plain_username = sasl_plain_username
self._sasl_plain_password = sasl_plain_password
self._sasl_kerberos_service_name = sasl_kerberos_service_name
self._sasl_kerberos_domain_name = sasl_kerberos_domain_name
self._sasl_oauth_token_provider = sasl_oauth_token_provider
self.cluster = ClusterMetadata(metadata_max_age_ms=metadata_max_age_ms)
self._topics = set() # empty set will fetch all topic metadata
self._conns = {}
self._loop = loop
self._sync_task = None
self._md_update_fut = None
self._md_update_waiter = loop.create_future()
self._get_conn_lock_value = None
async def bootstrap(self):
"""Try to to bootstrap initial cluster metadata"""
assert self._loop is get_running_loop(), (
"Please create objects with the same loop as running with")
# using request v0 for bootstrap if not sure v1 is available
if self._api_version == "auto" or self._api_version < (0, 10):
metadata_request = MetadataRequest[0]([])
else:
metadata_request = MetadataRequest[1]([])
version_hint = None
if self._api_version != "auto":
version_hint = self._api_version
for host, port, _ in self.hosts:
log.debug("Attempting to bootstrap via node at %s:%s", host, port)
try:
bootstrap_conn = await create_conn(
host,
port,
client_id=self._client_id,
request_timeout_ms=self._request_timeout_ms,
ssl_context=self._ssl_context,
security_protocol=self._security_protocol,
max_idle_ms=self._connections_max_idle_ms,
sasl_mechanism=self._sasl_mechanism,
sasl_plain_username=self._sasl_plain_username,
sasl_plain_password=self._sasl_plain_password,
sasl_kerberos_service_name=self.
_sasl_kerberos_service_name, # noqa: ignore=E501
sasl_kerberos_domain_name=self._sasl_kerberos_domain_name,
sasl_oauth_token_provider=self._sasl_oauth_token_provider,
version_hint=version_hint)
except (OSError, asyncio.TimeoutError) as err:
log.error('Unable connect to "%s:%s": %s', host, port, err)
continue
try:
metadata = await bootstrap_conn.send(metadata_request)
except (KafkaError, asyncio.TimeoutError) 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 till
# we get a normal cluster layout.
if not len(self.cluster.brokers()):
bootstrap_id = ('bootstrap', ConnectionGroup.DEFAULT)
self._conns[bootstrap_id] = bootstrap_conn
else:
bootstrap_conn.close()
log.debug('Received cluster metadata: %s', self.cluster)
break
else:
raise KafkaConnectionError('Unable to bootstrap from {}'.format(
self.hosts))
# detect api version if need
if self._api_version == 'auto':
self._api_version = await self.check_version()
if self._sync_task is None:
# starting metadata synchronizer task
self._sync_task = create_task(self._md_synchronizer())
