def fetch_all_metadata(self):
cluster_md = ClusterMetadata(
metadata_max_age_ms=self._metadata_max_age_ms)
updated = yield from self._metadata_update(cluster_md, [])
if not updated:
raise KafkaError(
'Unable to get cluster metadata over all known brokers')
return cluster_md
def __init__(self, **configs):
self.config = copy.copy(self.DEFAULT_CONFIG)
for key in self.config:
if key in configs:
self.config[key] = configs[key]
self.cluster = ClusterMetadata(**self.config)
self._topics = set() # empty set will fetch all topic metadata
self._metadata_refresh_in_progress = False
self._selector = self.config['selector']()
self._conns = Dict() # object to support weakrefs
self._api_versions = None
self._connecting = set()
self._refresh_on_disconnects = True
self._last_bootstrap = 0
self._bootstrap_fails = 0
self._wake_r, self._wake_w = socket.socketpair()
self._wake_r.setblocking(False)
self._wake_w.settimeout(self.config['wakeup_timeout_ms'] / 1000.0)
self._wake_lock = threading.Lock()
self._lock = threading.RLock()
# when requests complete, they are transferred to this queue prior to
# invocation. The purpose is to avoid invoking them while holding the
# lock above.
self._pending_completion = collections.deque()
self._selector.register(self._wake_r, selectors.EVENT_READ)
self._idle_expiry_manager = IdleConnectionManager(
self.config['connections_max_idle_ms'])
self._closed = False
self._sensors = None
if self.config['metrics']:
self._sensors = KafkaClientMetrics(
self.config['metrics'], self.config['metric_group_prefix'],
weakref.proxy(self._conns))
self._num_bootstrap_hosts = len(
collect_hosts(self.config['bootstrap_servers']))
# Check Broker Version if not set explicitly
if self.config['api_version'] is None:
check_timeout = self.config['api_version_auto_timeout_ms'] / 1000
self.config['api_version'] = self.check_version(
timeout=check_timeout)
def __init__(self, *, loop, bootstrap_servers='localhost',
client_id='aiokafka-' + __version__,
metadata_max_age_ms=300000,
request_timeout_ms=40000,
api_version='auto'):
"""Initialize an asynchronous kafka client
Keyword Arguments:
bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
strings) that the consumer should contact to bootstrap initial
cluster metadata. This does not have to be the full node list.
It just needs to have at least one broker that will respond to
Metadata API Request. Default port is 9092. If no servers are
specified, will default to localhost:9092.
client_id (str): a name for this client. This string is passed in
each request to servers and can be used to identify specific
server-side log entries that correspond to this client. Also
submitted to GroupCoordinator for logging with respect to
consumer group administration. Default: 'aiokafka-{ver}'
request_timeout_ms (int): Client request timeout in milliseconds.
Default: 40000.
metadata_max_age_ms (int): The period of time in milliseconds after
which we force a refresh of metadata even if we haven't seen
any partition leadership changes to proactively discover any
new brokers or partitions. Default: 300000
api_version (str): specify which kafka API version to use.
AIOKafka supports Kafka API versions >=0.9 only.
If set to 'auto', will attempt to infer the broker version by
probing various APIs. Default: auto
"""
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
self._api_version = api_version
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 = asyncio.Future(loop=self._loop)
self._md_update_waiter = asyncio.Future(loop=self._loop)
self._get_conn_lock = asyncio.Lock(loop=loop)
def test_empty_broker_list():
cluster = ClusterMetadata()
assert len(cluster.brokers()) == 0
cluster.update_metadata(MetadataResponse[0](
[(0, 'foo', 12), (1, 'bar', 34)], []))
assert len(cluster.brokers()) == 2
# empty broker list response should be ignored
cluster.update_metadata(MetadataResponse[0](
[], # empty brokers
[(17, 'foo', []), (17, 'bar', [])])) # topics w/ error
assert len(cluster.brokers()) == 2
def __init__(self,
*,
loop,
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):
if security_protocol not in ('SSL', 'PLAINTEXT'):
raise ValueError("`security_protocol` should be SSL or PLAINTEXT")
if security_protocol == "SSL" and ssl_context is None:
raise ValueError(
"`ssl_context` is mandatory if security_protocol=='SSL'")
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.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 = create_future(loop=self._loop)
self._get_conn_lock = asyncio.Lock(loop=loop)
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, loop=self.loop)
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
self.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,
client_id: str,
bootstrap_servers: Union[str, List[str]] = None,
**kwargs) -> None:
""" KafkaClient constructor
Args:
bootstrap_servers (Union[str, List[str]]): ‘host[:port]’ string (or list of ‘host[:port]’ strings) that
the producer should contact to bootstrap initial cluster metadata.
This does not have to be the full node list. It just needs to have at
least one broker that will respond to a Metadata API Request.
Default port is 9092. If no servers are specified, will default
to localhost:9092.
client_id (str): A name for this client. This string is passed in each request to servers and can be
used to identify specific server-side log entries that correspond to this client
cur_instance: Current service instance
nb_replica: Number of service replica
Raises:
BadArgumentKafkaClient: raised if argument are not valid
CurrentInstanceOutOfRange: raised if current instance is highest than replica number
KafkaAdminConfigurationError: raised if KafkaAdminClient argument are not valid
KafkaClientConnectionErrors: raised if KafkaClient can't connect on kafka or no brokers are available
"""
super().__init__(**kwargs)
if bootstrap_servers is None:
self.bootstrap_servers = 'localhost:9092'
else:
self.bootstrap_servers = bootstrap_servers
if isinstance(client_id, str):
self.client_id = client_id
else:
raise BadArgumentKafkaClient
try:
self._kafka_admin_client = KafkaAdminClient(
bootstrap_servers=self.bootstrap_servers,
client_id=f'waiter-{self.cur_instance}')
self._cluster_metadata = ClusterMetadata(
bootstrap_servers=self.bootstrap_servers)
except KafkaConfigurationError:
raise KafkaAdminConfigurationError
except (KafkaConnectionError, NoBrokersAvailable):
raise KafkaClientConnectionErrors
def _get_copartitioned_groups(
cls, topics: Set[str], cluster: ClusterMetadata,
subscriptions: MemberSubscriptionMapping) -> CopartitionedGroups:
topics_by_partitions: MutableMapping[int, Set] = defaultdict(set)
for topic in topics:
num_partitions = len(cluster.partitions_for_topic(topic) or set())
if num_partitions == 0:
logger.warning('Ignoring missing topic: %r', topic)
continue
topics_by_partitions[num_partitions].add(topic)
# We group copartitioned topics by subscribed clients such that
# a group of co-subscribed topics with the same number of partitions
# are copartitioned
copart_grouped = {
num_partitions: cls._group_co_subscribed(topics, subscriptions)
for num_partitions, topics in topics_by_partitions.items()
}
return copart_grouped
def __init__(self, **configs):
self.config = copy.copy(self.DEFAULT_CONFIG)
for key in self.config:
if key in configs:
self.config[key] = configs[key]
self.cluster = ClusterMetadata(**self.config)
self._topics = set() # empty set will fetch all topic metadata
self._metadata_refresh_in_progress = False
self._selector = self.config['selector']()
self._conns = Dict() # object to support weakrefs
self._connecting = set()
self._refresh_on_disconnects = True
self._last_bootstrap = 0
self._bootstrap_fails = 0
self._wake_r, self._wake_w = socket.socketpair()
self._wake_r.setblocking(False)
self._wake_lock = threading.Lock()
self._lock = threading.RLock()
# when requests complete, they are transferred to this queue prior to
# invocation. The purpose is to avoid invoking them while holding the
# lock above.
self._pending_completion = collections.deque()
self._selector.register(self._wake_r, selectors.EVENT_READ)
self._idle_expiry_manager = IdleConnectionManager(self.config['connections_max_idle_ms'])
self._closed = False
self._sensors = None
if self.config['metrics']:
self._sensors = KafkaClientMetrics(self.config['metrics'],
self.config['metric_group_prefix'],
weakref.proxy(self._conns))
self._bootstrap(collect_hosts(self.config['bootstrap_servers']))
# Check Broker Version if not set explicitly
if self.config['api_version'] is None:
check_timeout = self.config['api_version_auto_timeout_ms'] / 1000
self.config['api_version'] = self.check_version(timeout=check_timeout)
def __init__(self, *, loop, bootstrap_servers='localhost',
client_id='aiokafka-'+__version__, metadata_max_age_ms=300000,
request_timeout_ms=40000):
"""Initialize an asynchronous kafka client
Keyword Arguments:
bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
strings) that the consumer should contact to bootstrap initial
cluster metadata. This does not have to be the full node list.
It just needs to have at least one broker that will respond to
Metadata API Request. Default port is 9092. If no servers are
specified, will default to localhost:9092.
client_id (str): a name for this client. This string is passed in
each request to servers and can be used to identify specific
server-side log entries that correspond to this client. Also
submitted to GroupCoordinator for logging with respect to
consumer group administration. Default: 'aiokafka-{ver}'
request_timeout_ms (int): Client request timeout in milliseconds.
Default: 40000.
metadata_max_age_ms (int): The period of time in milliseconds after
which we force a refresh of metadata even if we haven't seen
any partition leadership changes to proactively discover any
new brokers or partitions. Default: 300000
"""
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
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 = asyncio.Future(loop=self._loop)
self._md_update_waiter = asyncio.Future(loop=self._loop)
self._get_conn_lock = asyncio.Lock(loop=loop)
#For URL handling
from urllib2 import urlopen
import socket
#For encryption
# from Crypto.PublicKey import RSA
# from Crypto import Random
myIP = urlopen('http://ip.42.pl/raw').read()
myPrivateIP = socket.gethostbyname(socket.gethostname())
print("myPrivateIP", myPrivateIP)
producer = KafkaProducer(
bootstrap_servers=[myPrivateIP + ':9090', myPrivateIP + ':9091'],
api_version=(0, 10))
clusterMetadata = ClusterMetadata(brokers=producer)
print("Metadata", clusterMetadata)
print("All brokers Metadata", clusterMetadata.brokers())
print("Broker 0 Metadata", clusterMetadata.broker_metadata(0))
print("Broker 1 Metadata", clusterMetadata.broker_metadata(1))
print("Known Topics Metadata", clusterMetadata.topics())
#Assignment of arguments
try:
topic = argv[1]
print("Topic found as arg ", topic)
except IndexError:
print(IndexError)
quit()
#Creation of message structure
from kafka.client import KafkaClient
from kafka.cluster import ClusterMetadata
client = KafkaClient(bootstrap_servers='localhost:9092',
client_id='test_store_builder')
response_metadata = client.poll(future=client.cluster.request_update())
cluster_metadata = ClusterMetadata(bootstrap_servers='localhost:9092')
cluster_metadata.update_metadata(response_metadata[0])
cluster_metadata.partitions_for_topic('test-assignor')
class AIOKafkaClient:
"""This class implements interface for interact with Kafka cluster"""
def __init__(self, *, loop, bootstrap_servers='localhost',
client_id='aiokafka-' + __version__,
metadata_max_age_ms=300000,
request_timeout_ms=40000,
api_version='auto'):
"""Initialize an asynchronous kafka client
Keyword Arguments:
bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
strings) that the consumer should contact to bootstrap initial
cluster metadata. This does not have to be the full node list.
It just needs to have at least one broker that will respond to
Metadata API Request. Default port is 9092. If no servers are
specified, will default to localhost:9092.
client_id (str): a name for this client. This string is passed in
each request to servers and can be used to identify specific
server-side log entries that correspond to this client. Also
submitted to GroupCoordinator for logging with respect to
consumer group administration. Default: 'aiokafka-{ver}'
request_timeout_ms (int): Client request timeout in milliseconds.
Default: 40000.
metadata_max_age_ms (int): The period of time in milliseconds after
which we force a refresh of metadata even if we haven't seen
any partition leadership changes to proactively discover any
new brokers or partitions. Default: 300000
api_version (str): specify which kafka API version to use.
AIOKafka supports Kafka API versions >=0.9 only.
If set to 'auto', will attempt to infer the broker version by
probing various APIs. Default: auto
"""
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
self._api_version = api_version
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 = asyncio.Future(loop=self._loop)
self._md_update_waiter = asyncio.Future(loop=self._loop)
self._get_conn_lock = asyncio.Lock(loop=loop)
def __repr__(self):
return '<AIOKafkaClient client_id=%s>' % self._client_id
@property
def api_version(self):
if type(self._api_version) is tuple:
return self._api_version
# unknown api version, return minimal supported version
return (0, 9)
@property
def hosts(self):
return collect_hosts(self._bootstrap_servers)
@asyncio.coroutine
def close(self):
if self._sync_task:
self._sync_task.cancel()
try:
yield from self._sync_task
except asyncio.CancelledError:
pass
self._sync_task = None
for conn in self._conns.values():
conn.close()
@asyncio.coroutine
def bootstrap(self):
"""Try to to bootstrap initial cluster metadata"""
# using request v0 for bootstap (bcs api version is not detected yet)
metadata_request = MetadataRequest[0]([])
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))
# detect api version if need
if self._api_version == 'auto':
self._api_version = yield from self.check_version()
if type(self._api_version) is not tuple:
self._api_version = tuple(
map(int, self._api_version.split('.')))
if self._sync_task is None:
# starting metadata synchronizer task
self._sync_task = ensure_future(
self._md_synchronizer(), loop=self._loop)
@asyncio.coroutine
def _md_synchronizer(self):
"""routine (async task) for synchronize cluster metadata every
`metadata_max_age_ms` milliseconds"""
while True:
yield from asyncio.wait(
[self._md_update_waiter],
timeout=self._metadata_max_age_ms / 1000,
loop=self._loop)
self._md_update_waiter = asyncio.Future(loop=self._loop)
ret = yield from self._metadata_update(self.cluster, self._topics)
self._md_update_fut.set_result(ret)
self._md_update_fut = asyncio.Future(loop=self._loop)
def get_random_node(self):
"""choice random node from known cluster brokers
Returns:
nodeId - identifier of broker
"""
nodeids = [b.nodeId for b in self.cluster.brokers()]
if not nodeids:
return None
return random.choice(nodeids)
@asyncio.coroutine
def _metadata_update(self, cluster_metadata, topics):
assert isinstance(cluster_metadata, ClusterMetadata)
topics = list(topics)
version_id = 0 if self.api_version < (0, 10) else 1
if version_id == 1 and not topics:
topics = None
metadata_request = MetadataRequest[version_id](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 force_metadata_update(self):
"""Update cluster metadata
Returns:
True/False - metadata updated or not
"""
# Wake up the `_md_synchronizer` task
if not self._md_update_waiter.done():
self._md_update_waiter.set_result(None)
# Metadata will be updated in the background by syncronizer
return self._md_update_fut
@asyncio.coroutine
def fetch_all_metadata(self):
cluster_md = ClusterMetadata(
metadata_max_age_ms=self._metadata_max_age_ms)
updated = yield from self._metadata_update(cluster_md, [])
if not updated:
raise KafkaError(
'Unable to get cluster metadata over all known brokers')
return cluster_md
def add_topic(self, topic):
"""Add a topic to the list of topics tracked via metadata.
Arguments:
topic (str): topic to track
"""
if topic in self._topics:
return
self._topics.add(topic)
def set_topics(self, topics):
"""Set specific topics to track for metadata.
Arguments:
topics (list of str): topics to track
"""
if set(topics).difference(self._topics):
self._topics = set(topics)
# update metadata in async manner
self.force_metadata_update()
@asyncio.coroutine
def _get_conn(self, node_id):
"Get or create a connection to a broker using host and port"
if node_id in self._conns:
conn = self._conns[node_id]
if not conn.connected():
del self._conns[node_id]
else:
return conn
try:
broker = self.cluster.broker_metadata(node_id)
assert broker, 'Broker id %s not in current metadata' % node_id
log.debug("Initiating connection to node %s at %s:%s",
node_id, broker.host, broker.port)
with (yield from self._get_conn_lock):
if node_id in self._conns:
return self._conns[node_id]
self._conns[node_id] = yield from create_conn(
broker.host, broker.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 node with id %s: %s', node_id, err)
return None
else:
return self._conns[node_id]
@asyncio.coroutine
def ready(self, node_id):
conn = yield from self._get_conn(node_id)
if conn is None:
return False
return True
@asyncio.coroutine
def send(self, node_id, request):
"""Send a request to a specific node.
Arguments:
node_id (int): destination node
request (Struct): request object (not-encoded)
Raises:
kafka.common.KafkaTimeoutError
kafka.common.NodeNotReadyError
kafka.commom.ConnectionError
kafka.common.CorrelationIdError
Returns:
Future: resolves to Response struct
"""
if not (yield from self.ready(node_id)):
raise NodeNotReadyError(
"Attempt to send a request to node"
" which is not ready (node id {}).".format(node_id))
# Every request gets a response, except one special case:
expect_response = True
if isinstance(request, tuple(ProduceRequest)) and \
request.required_acks == 0:
expect_response = False
future = self._conns[node_id].send(
request, expect_response=expect_response)
try:
result = yield from future
except asyncio.TimeoutError:
raise KafkaTimeoutError()
else:
return result
@asyncio.coroutine
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_v0, ApiVersionRequest_v0)
from kafka.protocol.commit import (
OffsetFetchRequest_v0, GroupCoordinatorRequest_v0)
from kafka.protocol.metadata import MetadataRequest_v0
test_cases = [
('0.10', ApiVersionRequest_v0()),
('0.9', ListGroupsRequest_v0()),
('0.8.2', GroupCoordinatorRequest_v0('aiokafka-default-group')),
('0.8.1', OffsetFetchRequest_v0('aiokafka-default-group', [])),
('0.8.0', MetadataRequest_v0([])),
]
# 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)
# request can be ignored by Kafka broker,
# so we send metadata request and wait response
task = self._loop.create_task(conn.send(request))
yield from asyncio.wait([task], timeout=0.1, loop=self._loop)
yield from self.fetch_all_metadata()
yield from task
except (KafkaError, asyncio.CancelledError):
continue
else:
return version
raise UnrecognizedBrokerVersion()
def client(mocker):
_cli = mocker.Mock(spec=KafkaClient(bootstrap_servers=[]))
_cli.cluster = mocker.Mock(spec=ClusterMetadata())
return _cli
def assign(self, cluster: ClusterMetadata, members: Dict[str, ConsumerProtocolMemberMetadata]) \
-> Dict[str, ConsumerProtocolMemberAssignment]:
"""Assign function was call by aiokafka for assign consumer on right topic partition.
Args:
cluster (ClusterMetadata): Kafka-python cluster metadata (more detail in kafka-python documentation)
members (Dict[str, ConsumerProtocolMemberMetadata]): members dict which contains
ConsumerProtocolMemberMetadata
(more detail in kafka-python documentation)
Returns:
Dict[str, ConsumerProtocolMemberAssignment]: dict which contain members and assignment protocol (more detail
in kafka-python documentation)
"""
self.logger.info('Statefulset Partition Assignor')
self.logger.debug('Cluster = %s\nMembers = %s', cluster, members)
# Get all topic
all_topics: Set = set()
for key, metadata in members.items():
self.logger.debug('Key = %s\nMetadata = %s', key, metadata)
all_topics.update(metadata.subscription)
# Get all partitions by topic name
all_topic_partitions = []
for topic in all_topics:
partitions = cluster.partitions_for_topic(topic)
if partitions is None:
self.logger.warning('No partition metadata for topic %s', topic)
continue
for partition in partitions:
all_topic_partitions.append(TopicPartition(topic, partition))
# Sort partition
all_topic_partitions.sort()
# Create default dict with lambda
assignment: DefaultDict[str, Any] = collections.defaultdict(lambda: collections.defaultdict(list))
advanced_assignor_dict = self.get_advanced_assignor_dict(all_topic_partitions)
for topic, partitions in advanced_assignor_dict.items():
for member_id, member_data in members.items():
# Loads member assignors data
user_data = json.loads(member_data.user_data)
# Get number of partitions by topic name
topic_number_partitions = len(partitions)
# Logic assignors if nb_replica as same as topic_numbers_partitions (used by StoreBuilder for
# assign each partitions to right instance, in this case nb_replica is same as topic_number_partitions)
if user_data['nb_replica'] == topic_number_partitions:
if user_data['assignor_policy'] == 'all':
for partition in partitions:
assignment[member_id][topic].append(partition)
elif user_data['assignor_policy'] == 'only_own':
if user_data['instance'] in partitions:
assignment[member_id][topic].append(partitions[user_data['instance']])
else:
raise BadAssignorPolicy
else:
raise NotImplementedError
self.logger.debug('Assignment = %s', assignment)
protocol_assignment = {}
for member_id in members:
protocol_assignment[member_id] = ConsumerProtocolMemberAssignment(self.version,
sorted(assignment[member_id].items()),
members[member_id].user_data)
self.logger.debug('Protocol Assignment = %s', protocol_assignment)
return protocol_assignment
class KafkaClient(object):
"""
A network client for asynchronous request/response network I/O.
This is an internal class used to implement the user-facing producer and
consumer clients.
This class is not thread-safe!
Attributes:
cluster (:any:`ClusterMetadata`): Local cache of cluster metadata, retrieved
via MetadataRequests during :meth:`~kafka.KafkaClient.poll`.
Keyword Arguments:
bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
strings) that the consumer should contact to bootstrap initial
cluster metadata. This does not have to be the full node list.
It just needs to have at least one broker that will respond to a
Metadata API Request. Default port is 9092. If no servers are
specified, will default to localhost:9092.
client_id (str): a name for this client. This string is passed in
each request to servers and can be used to identify specific
server-side log entries that correspond to this client. Also
submitted to GroupCoordinator for logging with respect to
consumer group administration. Default: 'kafka-python-{version}'
reconnect_backoff_ms (int): The amount of time in milliseconds to
wait before attempting to reconnect to a given host.
Default: 50.
reconnect_backoff_max_ms (int): The maximum amount of time in
milliseconds to wait when reconnecting to a broker that has
repeatedly failed to connect. If provided, the backoff per host
will increase exponentially for each consecutive connection
failure, up to this maximum. To avoid connection storms, a
randomization factor of 0.2 will be applied to the backoff
resulting in a random range between 20% below and 20% above
the computed value. Default: 1000.
request_timeout_ms (int): Client request timeout in milliseconds.
Default: 30000.
connections_max_idle_ms: Close idle connections after the number of
milliseconds specified by this config. The broker closes idle
connections after connections.max.idle.ms, so this avoids hitting
unexpected socket disconnected errors on the client.
Default: 540000
retry_backoff_ms (int): Milliseconds to backoff when retrying on
errors. Default: 100.
max_in_flight_requests_per_connection (int): Requests are pipelined
to kafka brokers up to this number of maximum requests per
broker connection. Default: 5.
receive_buffer_bytes (int): The size of the TCP receive buffer
(SO_RCVBUF) to use when reading data. Default: None (relies on
system defaults). Java client defaults to 32768.
send_buffer_bytes (int): The size of the TCP send buffer
(SO_SNDBUF) to use when sending data. Default: None (relies on
system defaults). Java client defaults to 131072.
socket_options (list): List of tuple-arguments to socket.setsockopt
to apply to broker connection sockets. Default:
[(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)]
metadata_max_age_ms (int): The period of time in milliseconds after
which we force a refresh of metadata even if we haven't seen any
partition leadership changes to proactively discover any new
brokers or partitions. Default: 300000
security_protocol (str): Protocol used to communicate with brokers.
Valid values are: PLAINTEXT, SSL. Default: PLAINTEXT.
ssl_context (ssl.SSLContext): Pre-configured SSLContext for wrapping
socket connections. If provided, all other ssl_* configurations
will be ignored. Default: None.
ssl_check_hostname (bool): Flag to configure whether SSL handshake
should verify that the certificate matches the broker's hostname.
Default: True.
ssl_cafile (str): Optional filename of CA file to use in certificate
veriication. Default: None.
ssl_certfile (str): Optional filename of file in PEM format containing
the client certificate, as well as any CA certificates needed to
establish the certificate's authenticity. Default: None.
ssl_keyfile (str): Optional filename containing the client private key.
Default: None.
ssl_password (str): Optional password to be used when loading the
certificate chain. Default: None.
ssl_crlfile (str): Optional filename containing the CRL to check for
certificate expiration. By default, no CRL check is done. When
providing a file, only the leaf certificate will be checked against
this CRL. The CRL can only be checked with Python 3.4+ or 2.7.9+.
Default: None.
api_version (tuple): Specify which Kafka API version to use. If set
to None, KafkaClient will attempt to infer the broker version by
probing various APIs. Example: (0, 10, 2). Default: None
api_version_auto_timeout_ms (int): number of milliseconds to throw a
timeout exception from the constructor when checking the broker
api version. Only applies if api_version is None
selector (selectors.BaseSelector): Provide a specific selector
implementation to use for I/O multiplexing.
Default: selectors.DefaultSelector
metrics (kafka.metrics.Metrics): Optionally provide a metrics
instance for capturing network IO stats. Default: None.
metric_group_prefix (str): Prefix for metric names. Default: ''
sasl_mechanism (str): string picking sasl mechanism when security_protocol
is SASL_PLAINTEXT or SASL_SSL. Currently only PLAIN is supported.
Default: None
sasl_plain_username (str): username for sasl PLAIN authentication.
Default: None
sasl_plain_password (str): password for sasl PLAIN authentication.
Default: None
sasl_kerberos_service_name (str): Service name to include in GSSAPI
sasl mechanism handshake. Default: 'kafka'
"""
DEFAULT_CONFIG = {
'bootstrap_servers': 'localhost',
'client_id': 'kafka-python-' + __version__,
'request_timeout_ms': 30000,
'connections_max_idle_ms': 9 * 60 * 1000,
'reconnect_backoff_ms': 50,
'reconnect_backoff_max_ms': 1000,
'max_in_flight_requests_per_connection': 5,
'receive_buffer_bytes': None,
'send_buffer_bytes': None,
'socket_options': [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)],
'sock_chunk_bytes': 4096, # undocumented experimental option
'sock_chunk_buffer_count': 1000, # undocumented experimental option
'retry_backoff_ms': 100,
'metadata_max_age_ms': 300000,
'security_protocol': 'PLAINTEXT',
'ssl_context': None,
'ssl_check_hostname': True,
'ssl_cafile': None,
'ssl_certfile': None,
'ssl_keyfile': None,
'ssl_password': None,
'ssl_crlfile': None,
'api_version': None,
'api_version_auto_timeout_ms': 2000,
'selector': selectors.DefaultSelector,
'metrics': None,
'metric_group_prefix': '',
'sasl_mechanism': None,
'sasl_plain_username': None,
'sasl_plain_password': None,
'sasl_kerberos_service_name': 'kafka',
}
def __init__(self, **configs):
self.config = copy.copy(self.DEFAULT_CONFIG)
for key in self.config:
if key in configs:
self.config[key] = configs[key]
self.cluster = ClusterMetadata(**self.config)
self._topics = set() # empty set will fetch all topic metadata
self._metadata_refresh_in_progress = False
self._selector = self.config['selector']()
self._conns = Dict() # object to support weakrefs
self._connecting = set()
self._refresh_on_disconnects = True
self._last_bootstrap = 0
self._bootstrap_fails = 0
self._wake_r, self._wake_w = socket.socketpair()
self._wake_r.setblocking(False)
self._wake_lock = threading.Lock()
self._lock = threading.RLock()
# when requests complete, they are transferred to this queue prior to
# invocation. The purpose is to avoid invoking them while holding the
# lock above.
self._pending_completion = collections.deque()
self._selector.register(self._wake_r, selectors.EVENT_READ)
self._idle_expiry_manager = IdleConnectionManager(self.config['connections_max_idle_ms'])
self._closed = False
self._sensors = None
if self.config['metrics']:
self._sensors = KafkaClientMetrics(self.config['metrics'],
self.config['metric_group_prefix'],
weakref.proxy(self._conns))
self._bootstrap(collect_hosts(self.config['bootstrap_servers']))
# Check Broker Version if not set explicitly
if self.config['api_version'] is None:
check_timeout = self.config['api_version_auto_timeout_ms'] / 1000
self.config['api_version'] = self.check_version(timeout=check_timeout)
def _bootstrap(self, hosts):
log.info('Bootstrapping cluster metadata from %s', hosts)
# Exponential backoff if bootstrap fails
backoff_ms = self.config['reconnect_backoff_ms'] * 2 ** self._bootstrap_fails
next_at = self._last_bootstrap + backoff_ms / 1000.0
self._refresh_on_disconnects = False
now = time.time()
if next_at > now:
log.debug("Sleeping %0.4f before bootstrapping again", next_at - now)
time.sleep(next_at - now)
self._last_bootstrap = time.time()
if self.config['api_version'] is None or self.config['api_version'] < (0, 10):
metadata_request = MetadataRequest[0]([])
else:
metadata_request = MetadataRequest[1](None)
for host, port, afi in hosts:
log.debug("Attempting to bootstrap via node at %s:%s", host, port)
cb = functools.partial(WeakMethod(self._conn_state_change), 'bootstrap')
bootstrap = BrokerConnection(host, port, afi,
state_change_callback=cb,
node_id='bootstrap',
**self.config)
if not bootstrap.connect_blocking():
bootstrap.close()
continue
future = bootstrap.send(metadata_request)
while not future.is_done:
self._selector.select(1)
for r, f in bootstrap.recv():
f.success(r)
if future.failed():
bootstrap.close()
continue
self.cluster.update_metadata(future.value)
log.info('Bootstrap succeeded: found %d brokers and %d topics.',
len(self.cluster.brokers()), len(self.cluster.topics()))
# 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
else:
bootstrap.close()
self._bootstrap_fails = 0
break
# No bootstrap found...
else:
log.error('Unable to bootstrap from %s', hosts)
# Max exponential backoff is 2^12, x4000 (50ms -> 200s)
self._bootstrap_fails = min(self._bootstrap_fails + 1, 12)
self._refresh_on_disconnects = True
def _can_connect(self, node_id):
if node_id not in self._conns:
if self.cluster.broker_metadata(node_id):
return True
return False
conn = self._conns[node_id]
return conn.disconnected() and not conn.blacked_out()
def _conn_state_change(self, node_id, conn):
with self._lock:
if conn.connecting():
# SSL connections can enter this state 2x (second during Handshake)
if node_id not in self._connecting:
self._connecting.add(node_id)
self._selector.register(conn._sock, selectors.EVENT_WRITE)
elif conn.connected():
log.debug("Node %s connected", node_id)
if node_id in self._connecting:
self._connecting.remove(node_id)
try:
self._selector.unregister(conn._sock)
except KeyError:
pass
self._selector.register(conn._sock, selectors.EVENT_READ, conn)
if self._sensors:
self._sensors.connection_created.record()
self._idle_expiry_manager.update(node_id)
if 'bootstrap' in self._conns and node_id != 'bootstrap':
bootstrap = self._conns.pop('bootstrap')
# XXX: make conn.close() require error to cause refresh
self._refresh_on_disconnects = False
bootstrap.close()
self._refresh_on_disconnects = True
# Connection failures imply that our metadata is stale, so let's refresh
elif conn.state is ConnectionStates.DISCONNECTING:
if node_id in self._connecting:
self._connecting.remove(node_id)
try:
self._selector.unregister(conn._sock)
except KeyError:
pass
if self._sensors:
self._sensors.connection_closed.record()
idle_disconnect = False
if self._idle_expiry_manager.is_expired(node_id):
idle_disconnect = True
self._idle_expiry_manager.remove(node_id)
if self._refresh_on_disconnects and not self._closed and not idle_disconnect:
log.warning("Node %s connection failed -- refreshing metadata", node_id)
self.cluster.request_update()
def _maybe_connect(self, node_id):
"""Idempotent non-blocking connection attempt to the given node id."""
with self._lock:
broker = self.cluster.broker_metadata(node_id)
conn = self._conns.get(node_id)
if conn is None:
assert broker, 'Broker id %s not in current metadata' % node_id
log.debug("Initiating connection to node %s at %s:%s",
node_id, broker.host, broker.port)
host, port, afi = get_ip_port_afi(broker.host)
cb = functools.partial(WeakMethod(self._conn_state_change), node_id)
conn = BrokerConnection(host, broker.port, afi,
state_change_callback=cb,
node_id=node_id,
**self.config)
self._conns[node_id] = conn
# Check if existing connection should be recreated because host/port changed
elif conn.disconnected() and broker is not None:
host, _, __ = get_ip_port_afi(broker.host)
if conn.host != host or conn.port != broker.port:
log.info("Broker metadata change detected for node %s"
" from %s:%s to %s:%s", node_id, conn.host, conn.port,
broker.host, broker.port)
# Drop old connection object.
# It will be recreated on next _maybe_connect
self._conns.pop(node_id)
return False
elif conn.connected():
return True
conn.connect()
return conn.connected()
def ready(self, node_id, metadata_priority=True):
"""Check whether a node is connected and ok to send more requests.
Arguments:
node_id (int): the id of the node to check
metadata_priority (bool): Mark node as not-ready if a metadata
refresh is required. Default: True
Returns:
bool: True if we are ready to send to the given node
"""
self._maybe_connect(node_id)
return self.is_ready(node_id, metadata_priority=metadata_priority)
def connected(self, node_id):
"""Return True iff the node_id is connected."""
with self._lock:
if node_id not in self._conns:
return False
return self._conns[node_id].connected()
def _close(self):
if not self._closed:
self._closed = True
self._wake_r.close()
self._wake_w.close()
self._selector.close()
def close(self, node_id=None):
"""Close one or all broker connections.
Arguments:
node_id (int, optional): the id of the node to close
"""
with self._lock:
if node_id is None:
self._close()
for conn in self._conns.values():
conn.close()
elif node_id in self._conns:
self._conns[node_id].close()
else:
log.warning("Node %s not found in current connection list; skipping", node_id)
return
def __del__(self):
self._close()
def is_disconnected(self, node_id):
"""Check whether the node connection has been disconnected or failed.
A disconnected node has either been closed or has failed. Connection
failures are usually transient and can be resumed in the next ready()
call, but there are cases where transient failures need to be caught
and re-acted upon.
Arguments:
node_id (int): the id of the node to check
Returns:
bool: True iff the node exists and is disconnected
"""
with self._lock:
if node_id not in self._conns:
return False
return self._conns[node_id].disconnected()
def connection_delay(self, node_id):
"""
Return the number of milliseconds to wait, based on the connection
state, before attempting to send data. When disconnected, this respects
the reconnect backoff time. When connecting, returns 0 to allow
non-blocking connect to finish. When connected, returns a very large
number to handle slow/stalled connections.
Arguments:
node_id (int): The id of the node to check
Returns:
int: The number of milliseconds to wait.
"""
with self._lock:
if node_id not in self._conns:
return 0
return self._conns[node_id].connection_delay()
def is_ready(self, node_id, metadata_priority=True):
"""Check whether a node is ready to send more requests.
In addition to connection-level checks, this method also is used to
block additional requests from being sent during a metadata refresh.
Arguments:
node_id (int): id of the node to check
metadata_priority (bool): Mark node as not-ready if a metadata
refresh is required. Default: True
Returns:
bool: True if the node is ready and metadata is not refreshing
"""
if not self._can_send_request(node_id):
return False
# if we need to update our metadata now declare all requests unready to
# make metadata requests first priority
if metadata_priority:
if self._metadata_refresh_in_progress:
return False
if self.cluster.ttl() == 0:
return False
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()
def send(self, node_id, request):
"""Send a request to a specific node.
Arguments:
node_id (int): destination node
request (Struct): request object (not-encoded)
Raises:
AssertionError: if node_id is not in current cluster metadata
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))
return self._conns[node_id].send(request)
def poll(self, timeout_ms=None, future=None):
"""Try to read and write to sockets.
This method will also attempt to complete node connections, refresh
stale metadata, and run previously-scheduled tasks.
Arguments:
timeout_ms (int, optional): maximum amount of time to wait (in ms)
for at least one response. Must be non-negative. The actual
timeout will be the minimum of timeout, request timeout and
metadata timeout. Default: request_timeout_ms
future (Future, optional): if provided, blocks until future.is_done
Returns:
list: responses received (can be empty)
"""
if future is not None:
timeout_ms = 100
elif timeout_ms is None:
timeout_ms = self.config['request_timeout_ms']
elif not isinstance(timeout_ms, (int, float)):
raise RuntimeError('Invalid type for timeout: %s' % type(timeout_ms))
# Loop for futures, break after first loop if None
responses = []
while True:
with self._lock:
# Attempt to complete pending connections
for node_id in list(self._connecting):
self._maybe_connect(node_id)
# Send a metadata request if needed
metadata_timeout_ms = self._maybe_refresh_metadata()
# If we got a future that is already done, don't block in _poll
if future is not None and future.is_done:
timeout = 0
else:
idle_connection_timeout_ms = self._idle_expiry_manager.next_check_ms()
timeout = min(
timeout_ms,
metadata_timeout_ms,
idle_connection_timeout_ms,
self.config['request_timeout_ms'])
timeout = max(0, timeout / 1000) # avoid negative timeouts
self._poll(timeout)
# called without the lock to avoid deadlock potential
# if handlers need to acquire locks
responses.extend(self._fire_pending_completed_requests())
# If all we had was a timeout (future is None) - only do one poll
# If we do have a future, we keep looping until it is done
if future is None or future.is_done:
break
return responses
def _poll(self, timeout):
"""Returns list of (response, future) tuples"""
processed = set()
start_select = time.time()
ready = self._selector.select(timeout)
end_select = time.time()
if self._sensors:
self._sensors.select_time.record((end_select - start_select) * 1000000000)
for key, events in ready:
if key.fileobj is self._wake_r:
self._clear_wake_fd()
continue
elif not (events & selectors.EVENT_READ):
continue
conn = key.data
processed.add(conn)
if not conn.in_flight_requests:
# if we got an EVENT_READ but there were no in-flight requests, one of
# two things has happened:
#
# 1. The remote end closed the connection (because it died, or because
# a firewall timed out, or whatever)
# 2. The protocol is out of sync.
#
# either way, we can no longer safely use this connection
#
# Do a 1-byte read to check protocol didnt get out of sync, and then close the conn
try:
unexpected_data = key.fileobj.recv(1)
if unexpected_data: # anything other than a 0-byte read means protocol issues
log.warning('Protocol out of sync on %r, closing', conn)
except socket.error:
pass
conn.close(Errors.KafkaConnectionError('Socket EVENT_READ without in-flight-requests'))
continue
self._idle_expiry_manager.update(conn.node_id)
self._pending_completion.extend(conn.recv())
# Check for additional pending SSL bytes
if self.config['security_protocol'] in ('SSL', 'SASL_SSL'):
# TODO: optimize
for conn in self._conns.values():
if conn not in processed and conn.connected() and conn._sock.pending():
self._pending_completion.extend(conn.recv())
for conn in six.itervalues(self._conns):
if conn.requests_timed_out():
log.warning('%s timed out after %s ms. Closing connection.',
conn, conn.config['request_timeout_ms'])
conn.close(error=Errors.RequestTimedOutError(
'Request timed out after %s ms' %
conn.config['request_timeout_ms']))
if self._sensors:
self._sensors.io_time.record((time.time() - end_select) * 1000000000)
self._maybe_close_oldest_connection()
def in_flight_request_count(self, node_id=None):
"""Get the number of in-flight requests for a node or all nodes.
Arguments:
node_id (int, optional): a specific node to check. If unspecified,
return the total for all nodes
Returns:
int: pending in-flight requests for the node, or all nodes if None
"""
with self._lock:
if node_id is not None:
if node_id not in self._conns:
return 0
return len(self._conns[node_id].in_flight_requests)
else:
return sum([len(conn.in_flight_requests) for conn in self._conns.values()])
def _fire_pending_completed_requests(self):
responses = []
while True:
try:
# We rely on deque.popleft remaining threadsafe
# to allow both the heartbeat thread and the main thread
# to process responses
response, future = self._pending_completion.popleft()
except IndexError:
break
future.success(response)
responses.append(response)
return responses
def least_loaded_node(self):
"""Choose the node with fewest outstanding requests, with fallbacks.
This method will prefer a node with an existing connection and no
in-flight-requests. If no such node is found, a node will be chosen
randomly from disconnected nodes that are not "blacked out" (i.e.,
are not subject to a reconnect backoff).
Returns:
node_id or None if no suitable node was found
"""
with self._lock:
nodes = [broker.nodeId for broker in self.cluster.brokers()]
random.shuffle(nodes)
inflight = float('inf')
found = None
for node_id in nodes:
conn = self._conns.get(node_id)
connected = conn is not None and conn.connected()
blacked_out = conn is not None and conn.blacked_out()
curr_inflight = len(conn.in_flight_requests) if conn is not None else 0
if connected and curr_inflight == 0:
# if we find an established connection
# with no in-flight requests, we can stop right away
return node_id
elif not blacked_out and curr_inflight < inflight:
# otherwise if this is the best we have found so far, record that
inflight = curr_inflight
found = node_id
if found is not None:
return found
# some broker versions return an empty list of broker metadata
# if there are no topics created yet. the bootstrap process
# should detect this and keep a 'bootstrap' node alive until
# a non-bootstrap node is connected and non-empty broker
# metadata is available
elif 'bootstrap' in self._conns:
return 'bootstrap'
return None
def set_topics(self, topics):
"""Set specific topics to track for metadata.
Arguments:
topics (list of str): topics to check for metadata
Returns:
Future: resolves after metadata request/response
"""
if set(topics).difference(self._topics):
future = self.cluster.request_update()
else:
future = Future().success(set(topics))
self._topics = set(topics)
return future
def add_topic(self, topic):
"""Add a topic to the list of topics tracked via metadata.
Arguments:
topic (str): topic to track
Returns:
Future: resolves after metadata request/response
"""
if topic in self._topics:
return Future().success(set(self._topics))
self._topics.add(topic)
return self.cluster.request_update()
# This method should be locked when running multi-threaded
def _maybe_refresh_metadata(self):
"""Send a metadata request if needed.
Returns:
int: milliseconds until next refresh
"""
ttl = self.cluster.ttl()
wait_for_in_progress_ms = self.config['request_timeout_ms'] if self._metadata_refresh_in_progress else 0
metadata_timeout = max(ttl, wait_for_in_progress_ms)
if metadata_timeout > 0:
return metadata_timeout
# Beware that the behavior of this method and the computation of
# timeouts for poll() are highly dependent on the behavior of
# least_loaded_node()
node_id = self.least_loaded_node()
if node_id is None:
log.debug("Give up sending metadata request since no node is available");
return self.config['reconnect_backoff_ms']
if self._can_send_request(node_id):
topics = list(self._topics)
if self.cluster.need_all_topic_metadata or not topics:
topics = [] if self.config['api_version'] < (0, 10) else None
api_version = 0 if self.config['api_version'] < (0, 10) else 1
request = MetadataRequest[api_version](topics)
log.debug("Sending metadata request %s to node %s", request, node_id)
future = self.send(node_id, request)
future.add_callback(self.cluster.update_metadata)
future.add_errback(self.cluster.failed_update)
self._metadata_refresh_in_progress = True
def refresh_done(val_or_error):
self._metadata_refresh_in_progress = False
future.add_callback(refresh_done)
future.add_errback(refresh_done)
return self.config['request_timeout_ms']
# If there's any connection establishment underway, wait until it completes. This prevents
# the client from unnecessarily connecting to additional nodes while a previous connection
# attempt has not been completed.
if self._connecting:
# Strictly the timeout we should return here is "connect timeout", but as we don't
# have such application level configuration, using request timeout instead.
return self.config['request_timeout_ms']
if self._can_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
# In either case we just need to wait for a network event
# to let us know the selected connection might be usable again.
return float('inf')
def check_version(self, node_id=None, timeout=2, strict=False):
"""Attempt to guess the version of a Kafka broker.
Note: It is possible that this method blocks longer than the
specified timeout. This can happen if the entire cluster
is down and the client enters a bootstrap backoff sleep.
This is only possible if node_id is None.
Returns: version tuple, i.e. (0, 10), (0, 9), (0, 8, 2), ...
Raises:
NodeNotReadyError (if node_id is provided)
NoBrokersAvailable (if node_id is None)
UnrecognizedBrokerVersion: please file bug if seen!
AssertionError (if strict=True): please file bug if seen!
"""
end = time.time() + timeout
while time.time() < end:
# It is possible that least_loaded_node falls back to bootstrap,
# which can block for an increasing backoff period
try_node = node_id or self.least_loaded_node()
if try_node is None:
raise Errors.NoBrokersAvailable()
self._maybe_connect(try_node)
conn = self._conns[try_node]
# We will intentionally cause socket failures
# These should not trigger metadata refresh
self._refresh_on_disconnects = False
try:
remaining = end - time.time()
version = conn.check_version(timeout=remaining, strict=strict)
return version
except Errors.NodeNotReadyError:
# Only raise to user if this is a node-specific request
if node_id is not None:
raise
finally:
self._refresh_on_disconnects = True
# Timeout
else:
raise Errors.NoBrokersAvailable()
def wakeup(self):
with self._wake_lock:
try:
self._wake_w.sendall(b'x')
except socket.error:
log.warning('Unable to send to wakeup socket!')
def _clear_wake_fd(self):
# reading from wake socket should only happen in a single thread
while True:
try:
self._wake_r.recv(1024)
except socket.error:
break
def _maybe_close_oldest_connection(self):
expired_connection = self._idle_expiry_manager.poll_expired_connection()
if expired_connection:
conn_id, ts = expired_connection
idle_ms = (time.time() - ts) * 1000
log.info('Closing idle connection %s, last active %d ms ago', conn_id, idle_ms)
self.close(node_id=conn_id)
class KafkaClient(object):
"""
A network client for asynchronous request/response network I/O.
This is an internal class used to implement the user-facing producer and
consumer clients.
This class is not thread-safe!
Attributes:
cluster (:any:`ClusterMetadata`): Local cache of cluster metadata, retrieved
via MetadataRequests during :meth:`~kafka.KafkaClient.poll`.
Keyword Arguments:
bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
strings) that the client should contact to bootstrap initial
cluster metadata. This does not have to be the full node list.
It just needs to have at least one broker that will respond to a
Metadata API Request. Default port is 9092. If no servers are
specified, will default to localhost:9092.
client_id (str): a name for this client. This string is passed in
each request to servers and can be used to identify specific
server-side log entries that correspond to this client. Also
submitted to GroupCoordinator for logging with respect to
consumer group administration. Default: 'kafka-python-{version}'
reconnect_backoff_ms (int): The amount of time in milliseconds to
wait before attempting to reconnect to a given host.
Default: 50.
reconnect_backoff_max_ms (int): The maximum amount of time in
milliseconds to backoff/wait when reconnecting to a broker that has
repeatedly failed to connect. If provided, the backoff per host
will increase exponentially for each consecutive connection
failure, up to this maximum. Once the maximum is reached,
reconnection attempts will continue periodically with this fixed
rate. To avoid connection storms, a randomization factor of 0.2
will be applied to the backoff resulting in a random range between
20% below and 20% above the computed value. Default: 1000.
request_timeout_ms (int): Client request timeout in milliseconds.
Default: 30000.
connections_max_idle_ms: Close idle connections after the number of
milliseconds specified by this config. The broker closes idle
connections after connections.max.idle.ms, so this avoids hitting
unexpected socket disconnected errors on the client.
Default: 540000
retry_backoff_ms (int): Milliseconds to backoff when retrying on
errors. Default: 100.
max_in_flight_requests_per_connection (int): Requests are pipelined
to kafka brokers up to this number of maximum requests per
broker connection. Default: 5.
receive_buffer_bytes (int): The size of the TCP receive buffer
(SO_RCVBUF) to use when reading data. Default: None (relies on
system defaults). Java client defaults to 32768.
send_buffer_bytes (int): The size of the TCP send buffer
(SO_SNDBUF) to use when sending data. Default: None (relies on
system defaults). Java client defaults to 131072.
socket_options (list): List of tuple-arguments to socket.setsockopt
to apply to broker connection sockets. Default:
[(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)]
metadata_max_age_ms (int): The period of time in milliseconds after
which we force a refresh of metadata even if we haven't seen any
partition leadership changes to proactively discover any new
brokers or partitions. Default: 300000
security_protocol (str): Protocol used to communicate with brokers.
Valid values are: PLAINTEXT, SSL, SASL_PLAINTEXT, SASL_SSL.
Default: PLAINTEXT.
ssl_context (ssl.SSLContext): Pre-configured SSLContext for wrapping
socket connections. If provided, all other ssl_* configurations
will be ignored. Default: None.
ssl_check_hostname (bool): Flag to configure whether SSL handshake
should verify that the certificate matches the broker's hostname.
Default: True.
ssl_cafile (str): Optional filename of CA file to use in certificate
verification. Default: None.
ssl_certfile (str): Optional filename of file in PEM format containing
the client certificate, as well as any CA certificates needed to
establish the certificate's authenticity. Default: None.
ssl_keyfile (str): Optional filename containing the client private key.
Default: None.
ssl_password (str): Optional password to be used when loading the
certificate chain. Default: None.
ssl_crlfile (str): Optional filename containing the CRL to check for
certificate expiration. By default, no CRL check is done. When
providing a file, only the leaf certificate will be checked against
this CRL. The CRL can only be checked with Python 3.4+ or 2.7.9+.
Default: None.
ssl_ciphers (str): optionally set the available ciphers for ssl
connections. It should be a string in the OpenSSL cipher list
format. If no cipher can be selected (because compile-time options
or other configuration forbids use of all the specified ciphers),
an ssl.SSLError will be raised. See ssl.SSLContext.set_ciphers
api_version (tuple): Specify which Kafka API version to use. If set
to None, KafkaClient will attempt to infer the broker version by
probing various APIs. Example: (0, 10, 2). Default: None
api_version_auto_timeout_ms (int): number of milliseconds to throw a
timeout exception from the constructor when checking the broker
api version. Only applies if api_version is None
selector (selectors.BaseSelector): Provide a specific selector
implementation to use for I/O multiplexing.
Default: selectors.DefaultSelector
metrics (kafka.metrics.Metrics): Optionally provide a metrics
instance for capturing network IO stats. Default: None.
metric_group_prefix (str): Prefix for metric names. Default: ''
sasl_mechanism (str): Authentication mechanism when security_protocol
is configured for SASL_PLAINTEXT or SASL_SSL. Valid values are:
PLAIN, GSSAPI, OAUTHBEARER, SCRAM-SHA-256, SCRAM-SHA-512.
sasl_plain_username (str): username for sasl PLAIN and SCRAM authentication.
Required if sasl_mechanism is PLAIN or one of the SCRAM mechanisms.
sasl_plain_password (str): password for sasl PLAIN and SCRAM authentication.
Required if sasl_mechanism is PLAIN or one of the SCRAM mechanisms.
sasl_kerberos_service_name (str): Service name to include in GSSAPI
sasl mechanism handshake. Default: 'kafka'
sasl_kerberos_domain_name (str): kerberos domain name to use in GSSAPI
sasl mechanism handshake. Default: one of bootstrap servers
sasl_oauth_token_provider (AbstractTokenProvider): OAuthBearer token provider
instance. (See kafka.oauth.abstract). Default: None
"""
DEFAULT_CONFIG = {
"bootstrap_servers": "localhost",
"bootstrap_topics_filter": set(),
"client_id": "kafka-python-" + __version__,
"request_timeout_ms": 30000,
"wakeup_timeout_ms": 3000,
"connections_max_idle_ms": 9 * 60 * 1000,
"reconnect_backoff_ms": 50,
"reconnect_backoff_max_ms": 1000,
"max_in_flight_requests_per_connection": 5,
"receive_buffer_bytes": None,
"send_buffer_bytes": None,
"socket_options": [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)],
"sock_chunk_bytes": 4096, # undocumented experimental option
"sock_chunk_buffer_count": 1000, # undocumented experimental option
"retry_backoff_ms": 100,
"metadata_max_age_ms": 300000,
"security_protocol": "PLAINTEXT",
"ssl_context": None,
"ssl_check_hostname": True,
"ssl_cafile": None,
"ssl_certfile": None,
"ssl_keyfile": None,
"ssl_password": None,
"ssl_crlfile": None,
"ssl_ciphers": None,
"api_version": None,
"api_version_auto_timeout_ms": 2000,
"selector": selectors.DefaultSelector,
"metrics": None,
"metric_group_prefix": "",
"sasl_mechanism": None,
"sasl_plain_username": None,
"sasl_plain_password": None,
"sasl_kerberos_service_name": "kafka",
"sasl_kerberos_domain_name": None,
"sasl_oauth_token_provider": None,
}
def __init__(self, **configs):
self.config = copy.copy(self.DEFAULT_CONFIG)
for key in self.config:
if key in configs:
self.config[key] = configs[key]
# these properties need to be set on top of the initialization pipeline
# because they are used when __del__ method is called
self._closed = False
self._wake_r, self._wake_w = socket.socketpair()
self._selector = self.config["selector"]()
self.cluster = ClusterMetadata(**self.config)
self._topics = set() # empty set will fetch all topic metadata
self._metadata_refresh_in_progress = False
self._conns = Dict() # object to support weakrefs
self._api_versions = None
self._connecting = set()
self._sending = set()
self._refresh_on_disconnects = True
self._last_bootstrap = 0
self._bootstrap_fails = 0
self._wake_r.setblocking(False)
self._wake_w.settimeout(self.config["wakeup_timeout_ms"] / 1000.0)
self._wake_lock = threading.Lock()
self._lock = threading.RLock()
# when requests complete, they are transferred to this queue prior to
# invocation. The purpose is to avoid invoking them while holding the
# lock above.
self._pending_completion = collections.deque()
self._selector.register(self._wake_r, selectors.EVENT_READ)
self._idle_expiry_manager = IdleConnectionManager(
self.config["connections_max_idle_ms"])
self._sensors = None
if self.config["metrics"]:
self._sensors = KafkaClientMetrics(
self.config["metrics"],
self.config["metric_group_prefix"],
weakref.proxy(self._conns),
)
self._num_bootstrap_hosts = len(
collect_hosts(self.config["bootstrap_servers"]))
# Check Broker Version if not set explicitly
if self.config["api_version"] is None:
check_timeout = self.config["api_version_auto_timeout_ms"] / 1000
self.config["api_version"] = self.check_version(
timeout=check_timeout)
def _can_bootstrap(self):
effective_failures = self._bootstrap_fails // self._num_bootstrap_hosts
backoff_factor = 2**effective_failures
backoff_ms = min(
self.config["reconnect_backoff_ms"] * backoff_factor,
self.config["reconnect_backoff_max_ms"],
)
backoff_ms *= random.uniform(0.8, 1.2)
next_at = self._last_bootstrap + backoff_ms / 1000.0
now = time.time()
if next_at > now:
return False
return True
def _can_connect(self, node_id):
if node_id not in self._conns:
if self.cluster.broker_metadata(node_id):
return True
return False
conn = self._conns[node_id]
return conn.disconnected() and not conn.blacked_out()
def _conn_state_change(self, node_id, sock, conn):
with self._lock:
if conn.connecting():
# 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(sock, selectors.EVENT_WRITE, conn)
except KeyError:
self._selector.modify(sock, selectors.EVENT_WRITE, conn)
if self.cluster.is_bootstrap(node_id):
self._last_bootstrap = time.time()
elif conn.connected():
log.debug("Node %s connected", node_id)
if node_id in self._connecting:
self._connecting.remove(node_id)
try:
self._selector.modify(sock, selectors.EVENT_READ, conn)
except KeyError:
self._selector.register(sock, selectors.EVENT_READ, conn)
if self._sensors:
self._sensors.connection_created.record()
self._idle_expiry_manager.update(node_id)
if self.cluster.is_bootstrap(node_id):
self._bootstrap_fails = 0
else:
for node_id in list(self._conns.keys()):
if self.cluster.is_bootstrap(node_id):
self._conns.pop(node_id).close()
# Connection failures imply that our metadata is stale, so let's refresh
elif conn.state is ConnectionStates.DISCONNECTED:
if node_id in self._connecting:
self._connecting.remove(node_id)
try:
self._selector.unregister(sock)
except KeyError:
pass
if self._sensors:
self._sensors.connection_closed.record()
idle_disconnect = False
if self._idle_expiry_manager.is_expired(node_id):
idle_disconnect = True
self._idle_expiry_manager.remove(node_id)
# If the connection has already by popped from self._conns,
# we can assume the disconnect was intentional and not a failure
if node_id not in self._conns:
pass
elif self.cluster.is_bootstrap(node_id):
self._bootstrap_fails += 1
elif (self._refresh_on_disconnects and not self._closed
and not idle_disconnect):
log.warning(
"Node %s connection failed -- refreshing metadata",
node_id)
self.cluster.request_update()
def maybe_connect(self, node_id, wakeup=True):
"""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().
if wakeup:
self.wakeup()
return True
return False
def _should_recycle_connection(self, conn):
# Never recycle unless disconnected
if not conn.disconnected():
return False
# Otherwise, only recycle when broker metadata has changed
broker = self.cluster.broker_metadata(conn.node_id)
if broker is None:
return False
host, _, afi = get_ip_port_afi(broker.host)
if conn.host != host or conn.port != broker.port:
log.info(
"Broker metadata change detected for node %s"
" from %s:%s to %s:%s",
conn.node_id,
conn.host,
conn.port,
broker.host,
broker.port,
)
return True
return False
def _maybe_connect(self, node_id):
"""Idempotent non-blocking connection attempt to the given node id."""
with self._lock:
conn = self._conns.get(node_id)
if conn is None:
broker = self.cluster.broker_metadata(node_id)
assert broker, "Broker id %s not in current metadata" % (
node_id, )
log.debug(
"Initiating connection to node %s at %s:%s",
node_id,
broker.host,
broker.port,
)
host, port, afi = get_ip_port_afi(broker.host)
cb = WeakMethod(self._conn_state_change)
conn = BrokerConnection(host,
broker.port,
afi,
state_change_callback=cb,
node_id=node_id,
**self.config)
self._conns[node_id] = conn
# Check if existing connection should be recreated because host/port changed
elif self._should_recycle_connection(conn):
self._conns.pop(node_id)
return False
elif conn.connected():
return True
conn.connect()
return conn.connected()
def ready(self, node_id, metadata_priority=True):
"""Check whether a node is connected and ok to send more requests.
Arguments:
node_id (int): the id of the node to check
metadata_priority (bool): Mark node as not-ready if a metadata
refresh is required. Default: True
Returns:
bool: True if we are ready to send to the given node
"""
self.maybe_connect(node_id)
return self.is_ready(node_id, metadata_priority=metadata_priority)
def connected(self, node_id):
"""Return True iff the node_id is connected."""
conn = self._conns.get(node_id)
if conn is None:
return False
return conn.connected()
def _close(self):
if not self._closed:
self._closed = True
self._wake_r.close()
self._wake_w.close()
self._selector.close()
def close(self, node_id=None):
"""Close one or all broker connections.
Arguments:
node_id (int, optional): the id of the node to close
"""
with self._lock:
if node_id is None:
self._close()
conns = list(self._conns.values())
self._conns.clear()
for conn in conns:
conn.close()
elif node_id in self._conns:
self._conns.pop(node_id).close()
else:
log.warning(
"Node %s not found in current connection list; skipping",
node_id)
return
def __del__(self):
self._close()
def is_disconnected(self, node_id):
"""Check whether the node connection has been disconnected or failed.
A disconnected node has either been closed or has failed. Connection
failures are usually transient and can be resumed in the next ready()
call, but there are cases where transient failures need to be caught
and re-acted upon.
Arguments:
node_id (int): the id of the node to check
Returns:
bool: True iff the node exists and is disconnected
"""
conn = self._conns.get(node_id)
if conn is None:
return False
return conn.disconnected()
def connection_delay(self, node_id):
"""
Return the number of milliseconds to wait, based on the connection
state, before attempting to send data. When disconnected, this respects
the reconnect backoff time. When connecting, returns 0 to allow
non-blocking connect to finish. When connected, returns a very large
number to handle slow/stalled connections.
Arguments:
node_id (int): The id of the node to check
Returns:
int: The number of milliseconds to wait.
"""
conn = self._conns.get(node_id)
if conn is None:
return 0
return conn.connection_delay()
def is_ready(self, node_id, metadata_priority=True):
"""Check whether a node is ready to send more requests.
In addition to connection-level checks, this method also is used to
block additional requests from being sent during a metadata refresh.
Arguments:
node_id (int): id of the node to check
metadata_priority (bool): Mark node as not-ready if a metadata
refresh is required. Default: True
Returns:
bool: True if the node is ready and metadata is not refreshing
"""
if not self._can_send_request(node_id):
return False
# if we need to update our metadata now declare all requests unready to
# make metadata requests first priority
if metadata_priority:
if self._metadata_refresh_in_progress:
return False
if self.cluster.ttl() == 0:
return False
return True
def _can_send_request(self, node_id):
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, wakeup=True):
"""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
request (Struct): request object (not-encoded)
wakeup (bool): optional flag to disable thread-wakeup
Raises:
AssertionError: if node_id is not in current cluster metadata
Returns:
Future: resolves to Response struct or Error
"""
conn = self._conns.get(node_id)
if not conn or not self._can_send_request(node_id):
self.maybe_connect(node_id, wakeup=wakeup)
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 = conn.send(request, blocking=False)
self._sending.add(conn)
# Wakeup signal is useful in case another thread is
# blocked waiting for incoming network traffic while holding
# the client lock in poll().
if wakeup:
self.wakeup()
return future
def poll(self, timeout_ms=None, future=None):
"""Try to read and write to sockets.
This method will also attempt to complete node connections, refresh
stale metadata, and run previously-scheduled tasks.
Arguments:
timeout_ms (int, optional): maximum amount of time to wait (in ms)
for at least one response. Must be non-negative. The actual
timeout will be the minimum of timeout, request timeout and
metadata timeout. Default: request_timeout_ms
future (Future, optional): if provided, blocks until future.is_done
Returns:
list: responses received (can be empty)
"""
if future is not None:
timeout_ms = 100
elif timeout_ms is None:
timeout_ms = self.config["request_timeout_ms"]
elif not isinstance(timeout_ms, (int, float)):
raise TypeError("Invalid type for timeout: %s" % type(timeout_ms))
# Loop for futures, break after first loop if None
responses = []
while True:
with self._lock:
if self._closed:
break
# Attempt to complete pending connections
for node_id in list(self._connecting):
self._maybe_connect(node_id)
# Send a metadata request if needed
metadata_timeout_ms = self._maybe_refresh_metadata()
# If we got a future that is already done, don't block in _poll
if future is not None and future.is_done:
timeout = 0
else:
idle_connection_timeout_ms = (
self._idle_expiry_manager.next_check_ms())
timeout = min(
timeout_ms,
metadata_timeout_ms,
idle_connection_timeout_ms,
self.config["request_timeout_ms"],
)
# if there are no requests in flight, do not block longer than the retry backoff
if self.in_flight_request_count() == 0:
timeout = min(timeout, self.config["retry_backoff_ms"])
timeout = max(0, timeout) # avoid negative timeouts
self._poll(timeout / 1000)
# called without the lock to avoid deadlock potential
# if handlers need to acquire locks
responses.extend(self._fire_pending_completed_requests())
# If all we had was a timeout (future is None) - only do one poll
# If we do have a future, we keep looping until it is done
if future is None or future.is_done:
break
return responses
def _register_send_sockets(self):
while self._sending:
conn = self._sending.pop()
try:
key = self._selector.get_key(conn._sock)
events = key.events | selectors.EVENT_WRITE
self._selector.modify(key.fileobj, events, key.data)
except KeyError:
self._selector.register(conn._sock, selectors.EVENT_WRITE,
conn)
def _poll(self, timeout):
# This needs to be locked, but since it is only called from within the
# locked section of poll(), there is no additional lock acquisition here
processed = set()
# Send pending requests first, before polling for responses
self._register_send_sockets()
start_select = time.time()
ready = self._selector.select(timeout)
end_select = time.time()
if self._sensors:
self._sensors.select_time.record(
(end_select - start_select) * 1000000000)
for key, events in ready:
if key.fileobj is self._wake_r:
self._clear_wake_fd()
continue
# Send pending requests if socket is ready to write
if events & selectors.EVENT_WRITE:
conn = key.data
if conn.connecting():
conn.connect()
else:
if conn.send_pending_requests_v2():
# If send is complete, we dont need to track write readiness
# for this socket anymore
if key.events ^ selectors.EVENT_WRITE:
self._selector.modify(
key.fileobj,
key.events ^ selectors.EVENT_WRITE,
key.data,
)
else:
self._selector.unregister(key.fileobj)
if not (events & selectors.EVENT_READ):
continue
conn = key.data
processed.add(conn)
if not conn.in_flight_requests:
# if we got an EVENT_READ but there were no in-flight requests, one of
# two things has happened:
#
# 1. The remote end closed the connection (because it died, or because
# a firewall timed out, or whatever)
# 2. The protocol is out of sync.
#
# either way, we can no longer safely use this connection
#
# Do a 1-byte read to check protocol didnt get out of sync, and then close the conn
try:
unexpected_data = key.fileobj.recv(1)
if (
unexpected_data
): # anything other than a 0-byte read means protocol issues
log.warning("Protocol out of sync on %r, closing",
conn)
except socket.error:
pass
conn.close(
Errors.KafkaConnectionError(
"Socket EVENT_READ without in-flight-requests"))
continue
self._idle_expiry_manager.update(conn.node_id)
self._pending_completion.extend(conn.recv())
# Check for additional pending SSL bytes
if self.config["security_protocol"] in ("SSL", "SASL_SSL"):
# TODO: optimize
for conn in self._conns.values():
if conn not in processed and conn.connected(
) and conn._sock.pending():
self._pending_completion.extend(conn.recv())
for conn in six.itervalues(self._conns):
if conn.requests_timed_out():
log.warning(
"%s timed out after %s ms. Closing connection.",
conn,
conn.config["request_timeout_ms"],
)
conn.close(error=Errors.RequestTimedOutError(
"Request timed out after %s ms" %
conn.config["request_timeout_ms"]))
if self._sensors:
self._sensors.io_time.record(
(time.time() - end_select) * 1000000000)
self._maybe_close_oldest_connection()
def in_flight_request_count(self, node_id=None):
"""Get the number of in-flight requests for a node or all nodes.
Arguments:
node_id (int, optional): a specific node to check. If unspecified,
return the total for all nodes
Returns:
int: pending in-flight requests for the node, or all nodes if None
"""
if node_id is not None:
conn = self._conns.get(node_id)
if conn is None:
return 0
return len(conn.in_flight_requests)
else:
return sum([
len(conn.in_flight_requests)
for conn in list(self._conns.values())
])
def _fire_pending_completed_requests(self):
responses = []
while True:
try:
# We rely on deque.popleft remaining threadsafe
# to allow both the heartbeat thread and the main thread
# to process responses
response, future = self._pending_completion.popleft()
except IndexError:
break
future.success(response)
responses.append(response)
return responses
def least_loaded_node(self):
"""Choose the node with fewest outstanding requests, with fallbacks.
This method will prefer a node with an existing connection and no
in-flight-requests. If no such node is found, a node will be chosen
randomly from disconnected nodes that are not "blacked out" (i.e.,
are not subject to a reconnect backoff). If no node metadata has been
obtained, will return a bootstrap node (subject to exponential backoff).
Returns:
node_id or None if no suitable node was found
"""
nodes = [broker.nodeId for broker in self.cluster.brokers()]
random.shuffle(nodes)
inflight = float("inf")
found = None
for node_id in nodes:
conn = self._conns.get(node_id)
connected = conn is not None and conn.connected()
blacked_out = conn is not None and conn.blacked_out()
curr_inflight = len(
conn.in_flight_requests) if conn is not None else 0
if connected and curr_inflight == 0:
# if we find an established connection
# with no in-flight requests, we can stop right away
return node_id
elif not blacked_out and curr_inflight < inflight:
# otherwise if this is the best we have found so far, record that
inflight = curr_inflight
found = node_id
return found
def set_topics(self, topics):
"""Set specific topics to track for metadata.
Arguments:
topics (list of str): topics to check for metadata
Returns:
Future: resolves after metadata request/response
"""
if set(topics).difference(self._topics):
future = self.cluster.request_update()
else:
future = Future().success(set(topics))
self._topics = set(topics)
return future
def add_topic(self, topic):
"""Add a topic to the list of topics tracked via metadata.
Arguments:
topic (str): topic to track
Returns:
Future: resolves after metadata request/response
"""
if topic in self._topics:
return Future().success(set(self._topics))
self._topics.add(topic)
return self.cluster.request_update()
# This method should be locked when running multi-threaded
def _maybe_refresh_metadata(self, wakeup=False):
"""Send a metadata request if needed.
Returns:
int: milliseconds until next refresh
"""
ttl = self.cluster.ttl()
wait_for_in_progress_ms = (self.config["request_timeout_ms"] if
self._metadata_refresh_in_progress else 0)
metadata_timeout = max(ttl, wait_for_in_progress_ms)
if metadata_timeout > 0:
return metadata_timeout
# Beware that the behavior of this method and the computation of
# timeouts for poll() are highly dependent on the behavior of
# least_loaded_node()
node_id = self.least_loaded_node()
if node_id is None:
log.debug(
"Give up sending metadata request since no node is available")
return self.config["reconnect_backoff_ms"]
if self._can_send_request(node_id):
topics = list(self._topics)
if not topics and self.cluster.is_bootstrap(node_id):
topics = list(self.config["bootstrap_topics_filter"])
if self.cluster.need_all_topic_metadata or not topics:
topics = [] if self.config["api_version"] < (0, 10) else None
api_version = 0 if self.config["api_version"] < (0, 10) else 1
request = MetadataRequest[api_version](topics)
log.debug("Sending metadata request %s to node %s", request,
node_id)
future = self.send(node_id, request, wakeup=wakeup)
future.add_callback(self.cluster.update_metadata)
future.add_errback(self.cluster.failed_update)
self._metadata_refresh_in_progress = True
def refresh_done(val_or_error):
self._metadata_refresh_in_progress = False
future.add_callback(refresh_done)
future.add_errback(refresh_done)
return self.config["request_timeout_ms"]
# If there's any connection establishment underway, wait until it completes. This prevents
# the client from unnecessarily connecting to additional nodes while a previous connection
# attempt has not been completed.
if self._connecting:
return self.config["reconnect_backoff_ms"]
if self.maybe_connect(node_id, wakeup=wakeup):
log.debug(
"Initializing connection to node %s for metadata request",
node_id)
return self.config["reconnect_backoff_ms"]
# connected but can't send more, OR connecting
# In either case we just need to wait for a network event
# to let us know the selected connection might be usable again.
return float("inf")
def get_api_versions(self):
"""Return the ApiVersions map, if available.
Note: A call to check_version must previously have succeeded and returned
version 0.10.0 or later
Returns: a map of dict mapping {api_key : (min_version, max_version)},
or None if ApiVersion is not supported by the kafka cluster.
"""
return self._api_versions
def check_version(self, node_id=None, timeout=2, strict=False):
"""Attempt to guess the version of a Kafka broker.
Note: It is possible that this method blocks longer than the
specified timeout. This can happen if the entire cluster
is down and the client enters a bootstrap backoff sleep.
This is only possible if node_id is None.
Returns: version tuple, i.e. (0, 10), (0, 9), (0, 8, 2), ...
Raises:
NodeNotReadyError (if node_id is provided)
NoBrokersAvailable (if node_id is None)
UnrecognizedBrokerVersion: please file bug if seen!
AssertionError (if strict=True): please file bug if seen!
"""
self._lock.acquire()
end = time.time() + timeout
while time.time() < end:
# It is possible that least_loaded_node falls back to bootstrap,
# which can block for an increasing backoff period
try_node = node_id or self.least_loaded_node()
if try_node is None:
self._lock.release()
raise Errors.NoBrokersAvailable()
self._maybe_connect(try_node)
conn = self._conns[try_node]
# We will intentionally cause socket failures
# These should not trigger metadata refresh
self._refresh_on_disconnects = False
try:
remaining = end - time.time()
version = conn.check_version(
timeout=remaining,
strict=strict,
topics=list(self.config["bootstrap_topics_filter"]),
)
if version >= (0, 10, 0):
# cache the api versions map if it's available (starting
# in 0.10 cluster version)
self._api_versions = conn.get_api_versions()
self._lock.release()
return version
except Errors.NodeNotReadyError:
# Only raise to user if this is a node-specific request
if node_id is not None:
self._lock.release()
raise
finally:
self._refresh_on_disconnects = True
# Timeout
else:
self._lock.release()
raise Errors.NoBrokersAvailable()
def wakeup(self):
with self._wake_lock:
try:
self._wake_w.sendall(b"x")
except socket.timeout:
log.warning("Timeout to send to wakeup socket!")
raise Errors.KafkaTimeoutError()
except socket.error:
log.warning("Unable to send to wakeup socket!")
def _clear_wake_fd(self):
# reading from wake socket should only happen in a single thread
while True:
try:
self._wake_r.recv(1024)
except socket.error:
break
def _maybe_close_oldest_connection(self):
expired_connection = self._idle_expiry_manager.poll_expired_connection(
)
if expired_connection:
conn_id, ts = expired_connection
idle_ms = (time.time() - ts) * 1000
log.info("Closing idle connection %s, last active %d ms ago",
conn_id, idle_ms)
self.close(node_id=conn_id)
def bootstrap_connected(self):
"""Return True if a bootstrap node is connected"""
for node_id in self._conns:
if not self.cluster.is_bootstrap(node_id):
continue
if self._conns[node_id].connected():
return True
else:
return False
class KafkaClient(object):
"""
A network client for asynchronous request/response network I/O.
This is an internal class used to implement the user-facing producer and
consumer clients.
This class is not thread-safe!
Attributes:
cluster (:any:`ClusterMetadata`): Local cache of cluster metadata, retrieved
via MetadataRequests during :meth:`~kafka.KafkaClient.poll`.
Keyword Arguments:
bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
strings) that the consumer should contact to bootstrap initial
cluster metadata. This does not have to be the full node list.
It just needs to have at least one broker that will respond to a
Metadata API Request. Default port is 9092. If no servers are
specified, will default to localhost:9092.
client_id (str): a name for this client. This string is passed in
each request to servers and can be used to identify specific
server-side log entries that correspond to this client. Also
submitted to GroupCoordinator for logging with respect to
consumer group administration. Default: 'kafka-python-{version}'
reconnect_backoff_ms (int): The amount of time in milliseconds to
wait before attempting to reconnect to a given host.
Default: 50.
reconnect_backoff_max_ms (int): The maximum amount of time in
milliseconds to wait when reconnecting to a broker that has
repeatedly failed to connect. If provided, the backoff per host
will increase exponentially for each consecutive connection
failure, up to this maximum. To avoid connection storms, a
randomization factor of 0.2 will be applied to the backoff
resulting in a random range between 20% below and 20% above
the computed value. Default: 1000.
request_timeout_ms (int): Client request timeout in milliseconds.
Default: 30000.
retry_backoff_ms (int): Milliseconds to backoff when retrying on
errors. Default: 100.
max_in_flight_requests_per_connection (int): Requests are pipelined
to kafka brokers up to this number of maximum requests per
broker connection. Default: 5.
receive_buffer_bytes (int): The size of the TCP receive buffer
(SO_RCVBUF) to use when reading data. Default: None (relies on
system defaults). Java client defaults to 32768.
send_buffer_bytes (int): The size of the TCP send buffer
(SO_SNDBUF) to use when sending data. Default: None (relies on
system defaults). Java client defaults to 131072.
socket_options (list): List of tuple-arguments to socket.setsockopt
to apply to broker connection sockets. Default:
[(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)]
metadata_max_age_ms (int): The period of time in milliseconds after
which we force a refresh of metadata even if we haven't seen any
partition leadership changes to proactively discover any new
brokers or partitions. Default: 300000
security_protocol (str): Protocol used to communicate with brokers.
Valid values are: PLAINTEXT, SSL. Default: PLAINTEXT.
ssl_context (ssl.SSLContext): pre-configured SSLContext for wrapping
socket connections. If provided, all other ssl_* configurations
will be ignored. Default: None.
ssl_check_hostname (bool): flag to configure whether ssl handshake
should verify that the certificate matches the brokers hostname.
default: true.
ssl_cafile (str): optional filename of ca file to use in certificate
veriication. default: none.
ssl_certfile (str): optional filename of file in pem format containing
the client certificate, as well as any ca certificates needed to
establish the certificate's authenticity. default: none.
ssl_keyfile (str): optional filename containing the client private key.
default: none.
ssl_password (str): optional password to be used when loading the
certificate chain. default: none.
ssl_crlfile (str): optional filename containing the CRL to check for
certificate expiration. By default, no CRL check is done. When
providing a file, only the leaf certificate will be checked against
this CRL. The CRL can only be checked with Python 3.4+ or 2.7.9+.
default: none.
api_version (tuple): Specify which Kafka API version to use. If set
to None, KafkaClient will attempt to infer the broker version by
probing various APIs. Example: (0, 10, 2). Default: None
api_version_auto_timeout_ms (int): number of milliseconds to throw a
timeout exception from the constructor when checking the broker
api version. Only applies if api_version is None
selector (selectors.BaseSelector): Provide a specific selector
implementation to use for I/O multiplexing.
Default: selectors.DefaultSelector
metrics (kafka.metrics.Metrics): Optionally provide a metrics
instance for capturing network IO stats. Default: None.
metric_group_prefix (str): Prefix for metric names. Default: ''
sasl_mechanism (str): string picking sasl mechanism when security_protocol
is SASL_PLAINTEXT or SASL_SSL. Currently only PLAIN is supported.
Default: None
sasl_plain_username (str): username for sasl PLAIN authentication.
Default: None
sasl_plain_password (str): password for sasl PLAIN authentication.
Default: None
sasl_kerberos_service_name (str): Service name to include in GSSAPI
sasl mechanism handshake. Default: 'kafka'
"""
DEFAULT_CONFIG = {
'bootstrap_servers': 'localhost',
'client_id': 'kafka-python-' + __version__,
'request_timeout_ms': 30000,
'connections_max_idle_ms': 9 * 60 * 1000,
'reconnect_backoff_ms': 50,
'reconnect_backoff_max_ms': 1000,
'max_in_flight_requests_per_connection': 5,
'receive_buffer_bytes': None,
'send_buffer_bytes': None,
'socket_options': [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)],
'sock_chunk_bytes': 4096, # undocumented experimental option
'sock_chunk_buffer_count': 1000, # undocumented experimental option
'retry_backoff_ms': 100,
'metadata_max_age_ms': 300000,
'security_protocol': 'PLAINTEXT',
'ssl_context': None,
'ssl_check_hostname': True,
'ssl_cafile': None,
'ssl_certfile': None,
'ssl_keyfile': None,
'ssl_password': None,
'ssl_crlfile': None,
'api_version': None,
'api_version_auto_timeout_ms': 2000,
'selector': selectors.DefaultSelector,
'metrics': None,
'metric_group_prefix': '',
'sasl_mechanism': None,
'sasl_plain_username': None,
'sasl_plain_password': None,
'sasl_kerberos_service_name': 'kafka',
}
def __init__(self, **configs):
self.config = copy.copy(self.DEFAULT_CONFIG)
for key in self.config:
if key in configs:
self.config[key] = configs[key]
self.cluster = ClusterMetadata(**self.config)
self._topics = set() # empty set will fetch all topic metadata
self._metadata_refresh_in_progress = False
self._selector = self.config['selector']()
self._conns = Dict() # object to support weakrefs
self._connecting = set()
self._refresh_on_disconnects = True
self._last_bootstrap = 0
self._bootstrap_fails = 0
self._wake_r, self._wake_w = socket.socketpair()
self._wake_r.setblocking(False)
self._wake_lock = threading.Lock()
self._lock = threading.RLock()
# when requests complete, they are transferred to this queue prior to
# invocation. The purpose is to avoid invoking them while holding the
# lock above.
self._pending_completion = collections.deque()
self._selector.register(self._wake_r, selectors.EVENT_READ)
self._idle_expiry_manager = IdleConnectionManager(
self.config['connections_max_idle_ms'])
self._closed = False
self._sensors = None
if self.config['metrics']:
self._sensors = KafkaClientMetrics(
self.config['metrics'], self.config['metric_group_prefix'],
weakref.proxy(self._conns))
self._bootstrap(collect_hosts(self.config['bootstrap_servers']))
# Check Broker Version if not set explicitly
if self.config['api_version'] is None:
check_timeout = self.config['api_version_auto_timeout_ms'] / 1000
self.config['api_version'] = self.check_version(
timeout=check_timeout)
def _bootstrap(self, hosts):
log.info('Bootstrapping cluster metadata from %s', hosts)
# Exponential backoff if bootstrap fails
backoff_ms = self.config[
'reconnect_backoff_ms'] * 2**self._bootstrap_fails
next_at = self._last_bootstrap + backoff_ms / 1000.0
self._refresh_on_disconnects = False
now = time.time()
if next_at > now:
log.debug("Sleeping %0.4f before bootstrapping again",
next_at - now)
time.sleep(next_at - now)
self._last_bootstrap = time.time()
if self.config['api_version'] is None or self.config['api_version'] < (
0, 10):
metadata_request = MetadataRequest[0]([])
else:
metadata_request = MetadataRequest[1](None)
for host, port, afi in hosts:
log.debug("Attempting to bootstrap via node at %s:%s", host, port)
cb = functools.partial(WeakMethod(self._conn_state_change),
'bootstrap')
bootstrap = BrokerConnection(host,
port,
afi,
state_change_callback=cb,
node_id='bootstrap',
**self.config)
if not bootstrap.connect_blocking():
bootstrap.close()
continue
future = bootstrap.send(metadata_request)
while not future.is_done:
self._selector.select(1)
for r, f in bootstrap.recv():
f.success(r)
if future.failed():
bootstrap.close()
continue
self.cluster.update_metadata(future.value)
log.info('Bootstrap succeeded: found %d brokers and %d topics.',
len(self.cluster.brokers()), len(self.cluster.topics()))
# 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
else:
bootstrap.close()
self._bootstrap_fails = 0
break
# No bootstrap found...
else:
log.error('Unable to bootstrap from %s', hosts)
# Max exponential backoff is 2^12, x4000 (50ms -> 200s)
self._bootstrap_fails = min(self._bootstrap_fails + 1, 12)
self._refresh_on_disconnects = True
def _can_connect(self, node_id):
if node_id not in self._conns:
if self.cluster.broker_metadata(node_id):
return True
return False
conn = self._conns[node_id]
return conn.disconnected() and not conn.blacked_out()
def _conn_state_change(self, node_id, conn):
with self._lock:
if conn.connecting():
# SSL connections can enter this state 2x (second during Handshake)
if node_id not in self._connecting:
self._connecting.add(node_id)
self._selector.register(conn._sock, selectors.EVENT_WRITE)
elif conn.connected():
log.debug("Node %s connected", node_id)
if node_id in self._connecting:
self._connecting.remove(node_id)
try:
self._selector.unregister(conn._sock)
except KeyError:
pass
self._selector.register(conn._sock, selectors.EVENT_READ, conn)
if self._sensors:
self._sensors.connection_created.record()
self._idle_expiry_manager.update(node_id)
if 'bootstrap' in self._conns and node_id != 'bootstrap':
bootstrap = self._conns.pop('bootstrap')
# XXX: make conn.close() require error to cause refresh
self._refresh_on_disconnects = False
bootstrap.close()
self._refresh_on_disconnects = True
# Connection failures imply that our metadata is stale, so let's refresh
elif conn.state is ConnectionStates.DISCONNECTING:
if node_id in self._connecting:
self._connecting.remove(node_id)
try:
self._selector.unregister(conn._sock)
except KeyError:
pass
if self._sensors:
self._sensors.connection_closed.record()
idle_disconnect = False
if self._idle_expiry_manager.is_expired(node_id):
idle_disconnect = True
self._idle_expiry_manager.remove(node_id)
if self._refresh_on_disconnects and not self._closed and not idle_disconnect:
log.warning(
"Node %s connection failed -- refreshing metadata",
node_id)
self.cluster.request_update()
def _maybe_connect(self, node_id):
"""Idempotent non-blocking connection attempt to the given node id."""
with self._lock:
broker = self.cluster.broker_metadata(node_id)
conn = self._conns.get(node_id)
if conn is None:
assert broker, 'Broker id %s not in current metadata' % node_id
log.debug("Initiating connection to node %s at %s:%s", node_id,
broker.host, broker.port)
host, port, afi = get_ip_port_afi(broker.host)
cb = functools.partial(WeakMethod(self._conn_state_change),
node_id)
conn = BrokerConnection(host,
broker.port,
afi,
state_change_callback=cb,
node_id=node_id,
**self.config)
self._conns[node_id] = conn
# Check if existing connection should be recreated because host/port changed
elif conn.disconnected() and broker is not None:
host, _, __ = get_ip_port_afi(broker.host)
if conn.host != host or conn.port != broker.port:
log.info(
"Broker metadata change detected for node %s"
" from %s:%s to %s:%s", node_id, conn.host, conn.port,
broker.host, broker.port)
# Drop old connection object.
# It will be recreated on next _maybe_connect
self._conns.pop(node_id)
return False
elif conn.connected():
return True
conn.connect()
return conn.connected()
def ready(self, node_id, metadata_priority=True):
"""Check whether a node is connected and ok to send more requests.
Arguments:
node_id (int): the id of the node to check
metadata_priority (bool): Mark node as not-ready if a metadata
refresh is required. Default: True
Returns:
bool: True if we are ready to send to the given node
"""
self._maybe_connect(node_id)
return self.is_ready(node_id, metadata_priority=metadata_priority)
def connected(self, node_id):
"""Return True iff the node_id is connected."""
with self._lock:
if node_id not in self._conns:
return False
return self._conns[node_id].connected()
def _close(self):
if not self._closed:
self._closed = True
self._wake_r.close()
self._wake_w.close()
self._selector.close()
def close(self, node_id=None):
"""Close one or all broker connections.
Arguments:
node_id (int, optional): the id of the node to close
"""
with self._lock:
if node_id is None:
self._close()
for conn in self._conns.values():
conn.close()
elif node_id in self._conns:
self._conns[node_id].close()
else:
log.warning(
"Node %s not found in current connection list; skipping",
node_id)
return
def __del__(self):
self._close()
def is_disconnected(self, node_id):
"""Check whether the node connection has been disconnected or failed.
A disconnected node has either been closed or has failed. Connection
failures are usually transient and can be resumed in the next ready()
call, but there are cases where transient failures need to be caught
and re-acted upon.
Arguments:
node_id (int): the id of the node to check
Returns:
bool: True iff the node exists and is disconnected
"""
with self._lock:
if node_id not in self._conns:
return False
return self._conns[node_id].disconnected()
def connection_delay(self, node_id):
"""
Return the number of milliseconds to wait, based on the connection
state, before attempting to send data. When disconnected, this respects
the reconnect backoff time. When connecting, returns 0 to allow
non-blocking connect to finish. When connected, returns a very large
number to handle slow/stalled connections.
Arguments:
node_id (int): The id of the node to check
Returns:
int: The number of milliseconds to wait.
"""
with self._lock:
if node_id not in self._conns:
return 0
return self._conns[node_id].connection_delay()
def is_ready(self, node_id, metadata_priority=True):
"""Check whether a node is ready to send more requests.
In addition to connection-level checks, this method also is used to
block additional requests from being sent during a metadata refresh.
Arguments:
node_id (int): id of the node to check
metadata_priority (bool): Mark node as not-ready if a metadata
refresh is required. Default: True
Returns:
bool: True if the node is ready and metadata is not refreshing
"""
if not self._can_send_request(node_id):
return False
# if we need to update our metadata now declare all requests unready to
# make metadata requests first priority
if metadata_priority:
if self._metadata_refresh_in_progress:
return False
if self.cluster.ttl() == 0:
return False
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()
def send(self, node_id, request):
"""Send a request to a specific node.
Arguments:
node_id (int): destination node
request (Struct): request object (not-encoded)
Raises:
AssertionError: if node_id is not in current cluster metadata
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))
return self._conns[node_id].send(request)
def poll(self, timeout_ms=None, future=None):
"""Try to read and write to sockets.
This method will also attempt to complete node connections, refresh
stale metadata, and run previously-scheduled tasks.
Arguments:
timeout_ms (int, optional): maximum amount of time to wait (in ms)
for at least one response. Must be non-negative. The actual
timeout will be the minimum of timeout, request timeout and
metadata timeout. Default: request_timeout_ms
future (Future, optional): if provided, blocks until future.is_done
Returns:
list: responses received (can be empty)
"""
if future is not None:
timeout_ms = 100
elif timeout_ms is None:
timeout_ms = self.config['request_timeout_ms']
elif not isinstance(timeout_ms, (int, float)):
raise RuntimeError('Invalid type for timeout: %s' %
type(timeout_ms))
# Loop for futures, break after first loop if None
responses = []
while True:
with self._lock:
# Attempt to complete pending connections
for node_id in list(self._connecting):
self._maybe_connect(node_id)
# Send a metadata request if needed
metadata_timeout_ms = self._maybe_refresh_metadata()
# If we got a future that is already done, don't block in _poll
if future is not None and future.is_done:
timeout = 0
else:
idle_connection_timeout_ms = self._idle_expiry_manager.next_check_ms(
)
timeout = min(timeout_ms, metadata_timeout_ms,
idle_connection_timeout_ms,
self.config['request_timeout_ms'])
timeout = max(0, timeout / 1000) # avoid negative timeouts
self._poll(timeout)
# called without the lock to avoid deadlock potential
# if handlers need to acquire locks
responses.extend(self._fire_pending_completed_requests())
# If all we had was a timeout (future is None) - only do one poll
# If we do have a future, we keep looping until it is done
if future is None or future.is_done:
break
return responses
def _poll(self, timeout):
"""Returns list of (response, future) tuples"""
processed = set()
start_select = time.time()
ready = self._selector.select(timeout)
end_select = time.time()
if self._sensors:
self._sensors.select_time.record(
(end_select - start_select) * 1000000000)
for key, events in ready:
if key.fileobj is self._wake_r:
self._clear_wake_fd()
continue
elif not (events & selectors.EVENT_READ):
continue
conn = key.data
processed.add(conn)
if not conn.in_flight_requests:
# if we got an EVENT_READ but there were no in-flight requests, one of
# two things has happened:
#
# 1. The remote end closed the connection (because it died, or because
# a firewall timed out, or whatever)
# 2. The protocol is out of sync.
#
# either way, we can no longer safely use this connection
#
# Do a 1-byte read to check protocol didnt get out of sync, and then close the conn
try:
unexpected_data = key.fileobj.recv(1)
if unexpected_data: # anything other than a 0-byte read means protocol issues
log.warning('Protocol out of sync on %r, closing',
conn)
except socket.error:
pass
conn.close(
Errors.ConnectionError(
'Socket EVENT_READ without in-flight-requests'))
continue
self._idle_expiry_manager.update(conn.node_id)
self._pending_completion.extend(conn.recv())
# Check for additional pending SSL bytes
if self.config['security_protocol'] in ('SSL', 'SASL_SSL'):
# TODO: optimize
for conn in self._conns.values():
if conn not in processed and conn.connected(
) and conn._sock.pending():
self._pending_completion.extend(conn.recv())
for conn in six.itervalues(self._conns):
if conn.requests_timed_out():
log.warning('%s timed out after %s ms. Closing connection.',
conn, conn.config['request_timeout_ms'])
conn.close(error=Errors.RequestTimedOutError(
'Request timed out after %s ms' %
conn.config['request_timeout_ms']))
if self._sensors:
self._sensors.io_time.record(
(time.time() - end_select) * 1000000000)
self._maybe_close_oldest_connection()
def in_flight_request_count(self, node_id=None):
"""Get the number of in-flight requests for a node or all nodes.
Arguments:
node_id (int, optional): a specific node to check. If unspecified,
return the total for all nodes
Returns:
int: pending in-flight requests for the node, or all nodes if None
"""
with self._lock:
if node_id is not None:
if node_id not in self._conns:
return 0
return len(self._conns[node_id].in_flight_requests)
else:
return sum([
len(conn.in_flight_requests)
for conn in self._conns.values()
])
def _fire_pending_completed_requests(self):
responses = []
while True:
try:
# We rely on deque.popleft remaining threadsafe
# to allow both the heartbeat thread and the main thread
# to process responses
response, future = self._pending_completion.popleft()
except IndexError:
break
future.success(response)
responses.append(response)
return responses
def least_loaded_node(self):
"""Choose the node with fewest outstanding requests, with fallbacks.
This method will prefer a node with an existing connection and no
in-flight-requests. If no such node is found, a node will be chosen
randomly from disconnected nodes that are not "blacked out" (i.e.,
are not subject to a reconnect backoff).
Returns:
node_id or None if no suitable node was found
"""
with self._lock:
nodes = [broker.nodeId for broker in self.cluster.brokers()]
random.shuffle(nodes)
inflight = float('inf')
found = None
for node_id in nodes:
conn = self._conns.get(node_id)
connected = conn is not None and conn.connected()
blacked_out = conn is not None and conn.blacked_out()
curr_inflight = len(
conn.in_flight_requests) if conn is not None else 0
if connected and curr_inflight == 0:
# if we find an established connection
# with no in-flight requests, we can stop right away
return node_id
elif not blacked_out and curr_inflight < inflight:
# otherwise if this is the best we have found so far, record that
inflight = curr_inflight
found = node_id
if found is not None:
return found
# some broker versions return an empty list of broker metadata
# if there are no topics created yet. the bootstrap process
# should detect this and keep a 'bootstrap' node alive until
# a non-bootstrap node is connected and non-empty broker
# metadata is available
elif 'bootstrap' in self._conns:
return 'bootstrap'
return None
def set_topics(self, topics):
"""Set specific topics to track for metadata.
Arguments:
topics (list of str): topics to check for metadata
Returns:
Future: resolves after metadata request/response
"""
if set(topics).difference(self._topics):
future = self.cluster.request_update()
else:
future = Future().success(set(topics))
self._topics = set(topics)
return future
def add_topic(self, topic):
"""Add a topic to the list of topics tracked via metadata.
Arguments:
topic (str): topic to track
Returns:
Future: resolves after metadata request/response
"""
if topic in self._topics:
return Future().success(set(self._topics))
self._topics.add(topic)
return self.cluster.request_update()
# This method should be locked when running multi-threaded
def _maybe_refresh_metadata(self):
"""Send a metadata request if needed.
Returns:
int: milliseconds until next refresh
"""
ttl = self.cluster.ttl()
wait_for_in_progress_ms = self.config[
'request_timeout_ms'] if self._metadata_refresh_in_progress else 0
metadata_timeout = max(ttl, wait_for_in_progress_ms)
if metadata_timeout > 0:
return metadata_timeout
# Beware that the behavior of this method and the computation of
# timeouts for poll() are highly dependent on the behavior of
# least_loaded_node()
node_id = self.least_loaded_node()
if node_id is None:
log.debug(
"Give up sending metadata request since no node is available")
return self.config['reconnect_backoff_ms']
if self._can_send_request(node_id):
topics = list(self._topics)
if self.cluster.need_all_topic_metadata or not topics:
topics = [] if self.config['api_version'] < (0, 10) else None
api_version = 0 if self.config['api_version'] < (0, 10) else 1
request = MetadataRequest[api_version](topics)
log.debug("Sending metadata request %s to node %s", request,
node_id)
future = self.send(node_id, request)
future.add_callback(self.cluster.update_metadata)
future.add_errback(self.cluster.failed_update)
self._metadata_refresh_in_progress = True
def refresh_done(val_or_error):
self._metadata_refresh_in_progress = False
future.add_callback(refresh_done)
future.add_errback(refresh_done)
return self.config['request_timeout_ms']
# If there's any connection establishment underway, wait until it completes. This prevents
# the client from unnecessarily connecting to additional nodes while a previous connection
# attempt has not been completed.
if self._connecting:
# Strictly the timeout we should return here is "connect timeout", but as we don't
# have such application level configuration, using request timeout instead.
return self.config['request_timeout_ms']
if self._can_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
# In either case we just need to wait for a network event
# to let us know the selected connection might be usable again.
return float('inf')
def check_version(self, node_id=None, timeout=2, strict=False):
"""Attempt to guess the version of a Kafka broker.
Note: It is possible that this method blocks longer than the
specified timeout. This can happen if the entire cluster
is down and the client enters a bootstrap backoff sleep.
This is only possible if node_id is None.
Returns: version tuple, i.e. (0, 10), (0, 9), (0, 8, 2), ...
Raises:
NodeNotReadyError (if node_id is provided)
NoBrokersAvailable (if node_id is None)
UnrecognizedBrokerVersion: please file bug if seen!
AssertionError (if strict=True): please file bug if seen!
"""
end = time.time() + timeout
while time.time() < end:
# It is possible that least_loaded_node falls back to bootstrap,
# which can block for an increasing backoff period
try_node = node_id or self.least_loaded_node()
if try_node is None:
raise Errors.NoBrokersAvailable()
self._maybe_connect(try_node)
conn = self._conns[try_node]
# We will intentionally cause socket failures
# These should not trigger metadata refresh
self._refresh_on_disconnects = False
try:
remaining = end - time.time()
version = conn.check_version(timeout=remaining, strict=strict)
return version
except Errors.NodeNotReadyError:
# Only raise to user if this is a node-specific request
if node_id is not None:
raise
finally:
self._refresh_on_disconnects = True
# Timeout
else:
raise Errors.NoBrokersAvailable()
def wakeup(self):
with self._wake_lock:
try:
self._wake_w.sendall(b'x')
except socket.error:
log.warning('Unable to send to wakeup socket!')
def _clear_wake_fd(self):
# reading from wake socket should only happen in a single thread
while True:
try:
self._wake_r.recv(1024)
except socket.error:
break
def _maybe_close_oldest_connection(self):
expired_connection = self._idle_expiry_manager.poll_expired_connection(
)
if expired_connection:
conn_id, ts = expired_connection
idle_ms = (time.time() - ts) * 1000
log.info('Closing idle connection %s, last active %d ms ago',
conn_id, idle_ms)
self.close(node_id=conn_id)
class AIOKafkaClient:
"""This class implements interface for interact with Kafka cluster"""
def __init__(self, *, loop, bootstrap_servers='localhost',
client_id='aiokafka-'+__version__, metadata_max_age_ms=300000,
request_timeout_ms=40000):
"""Initialize an asynchronous kafka client
Keyword Arguments:
bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
strings) that the consumer should contact to bootstrap initial
cluster metadata. This does not have to be the full node list.
It just needs to have at least one broker that will respond to
Metadata API Request. Default port is 9092. If no servers are
specified, will default to localhost:9092.
client_id (str): a name for this client. This string is passed in
each request to servers and can be used to identify specific
server-side log entries that correspond to this client. Also
submitted to GroupCoordinator for logging with respect to
consumer group administration. Default: 'aiokafka-{ver}'
request_timeout_ms (int): Client request timeout in milliseconds.
Default: 40000.
metadata_max_age_ms (int): The period of time in milliseconds after
which we force a refresh of metadata even if we haven't seen
any partition leadership changes to proactively discover any
new brokers or partitions. Default: 300000
"""
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
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 = asyncio.Future(loop=self._loop)
self._md_update_waiter = asyncio.Future(loop=self._loop)
self._get_conn_lock = asyncio.Lock(loop=loop)
def __repr__(self):
return '<AIOKafkaClient client_id=%s>' % self._client_id
@property
def hosts(self):
return collect_hosts(self._bootstrap_servers)
@asyncio.coroutine
def close(self):
if self._sync_task:
self._sync_task.cancel()
try:
yield from self._sync_task
except asyncio.CancelledError:
pass
self._sync_task = None
for conn in self._conns.values():
conn.close()
@asyncio.coroutine
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)
@asyncio.coroutine
def _md_synchronizer(self):
"""routine (async task) for synchronize cluster metadata every
`metadata_max_age_ms` milliseconds"""
while True:
yield from asyncio.wait(
[self._md_update_waiter],
timeout=self._metadata_max_age_ms / 1000,
loop=self._loop)
self._md_update_waiter = asyncio.Future(loop=self._loop)
ret = yield from self._metadata_update(self.cluster, self._topics)
self._md_update_fut.set_result(ret)
self._md_update_fut = asyncio.Future(loop=self._loop)
def get_random_node(self):
"""choice random node from known cluster brokers
Returns:
nodeId - identifier of broker
"""
nodeids = [b.nodeId for b in self.cluster.brokers()]
if not nodeids:
return None
return random.choice(nodeids)
@asyncio.coroutine
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 force_metadata_update(self):
"""Update cluster metadata
Returns:
True/False - metadata updated or not
"""
# Wake up the `_md_synchronizer` task
if not self._md_update_waiter.done():
self._md_update_waiter.set_result(None)
# Metadata will be updated in the background by syncronizer
return self._md_update_fut
@asyncio.coroutine
def fetch_all_metadata(self):
cluster_md = ClusterMetadata(
metadata_max_age_ms=self._metadata_max_age_ms)
updated = yield from self._metadata_update(cluster_md, [])
if not updated:
raise KafkaError(
'Unable to get cluster metadata over all known brokers')
return cluster_md
def add_topic(self, topic):
"""Add a topic to the list of topics tracked via metadata.
Arguments:
topic (str): topic to track
"""
if topic in self._topics:
return
self._topics.add(topic)
def set_topics(self, topics):
"""Set specific topics to track for metadata.
Arguments:
topics (list of str): topics to track
"""
if set(topics).difference(self._topics):
self._topics = set(topics)
# update metadata in async manner
self.force_metadata_update()
@asyncio.coroutine
def _get_conn(self, node_id):
"Get or create a connection to a broker using host and port"
if node_id in self._conns:
conn = self._conns[node_id]
if not conn.connected():
del self._conns[node_id]
else:
return conn
try:
broker = self.cluster.broker_metadata(node_id)
assert broker, 'Broker id %s not in current metadata' % node_id
log.debug("Initiating connection to node %s at %s:%s",
node_id, broker.host, broker.port)
with (yield from self._get_conn_lock):
if node_id in self._conns:
return self._conns[node_id]
self._conns[node_id] = yield from create_conn(
broker.host, broker.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 node with id %s: %s', node_id, err)
return None
else:
return self._conns[node_id]
@asyncio.coroutine
def ready(self, node_id):
conn = yield from self._get_conn(node_id)
if conn is None:
return False
return True
@asyncio.coroutine
def send(self, node_id, request):
"""Send a request to a specific node.
Arguments:
node_id (int): destination node
request (Struct): request object (not-encoded)
Raises:
kafka.common.KafkaTimeoutError
kafka.common.NodeNotReadyError
kafka.commom.ConnectionError
kafka.common.CorrelationIdError
Returns:
Future: resolves to Response struct
"""
if not (yield from self.ready(node_id)):
raise NodeNotReadyError(
"Attempt to send a request to node"
" which is not ready (node id {}).".format(node_id))
# Every request gets a response, except one special case:
expect_response = True
if isinstance(request, tuple(ProduceRequest)) and \
request.required_acks == 0:
expect_response = False
future = self._conns[node_id].send(
request, expect_response=expect_response)
try:
result = yield from future
except asyncio.TimeoutError:
raise KafkaTimeoutError()
else:
return result
@asyncio.coroutine
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_v0
from kafka.protocol.commit import (
OffsetFetchRequest_v0, GroupCoordinatorRequest_v0)
from kafka.protocol.metadata import MetadataRequest_v0
test_cases = [
('0.9', ListGroupsRequest_v0()),
('0.8.2', GroupCoordinatorRequest_v0('aiokafka-default-group')),
('0.8.1', OffsetFetchRequest_v0('aiokafka-default-group', [])),
('0.8.0', MetadataRequest_v0([])),
]
# 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_batch_pending_batch_list(self):
# In message accumulator we have _pending_batches list, that stores
# batches when those are delivered to node. We must be sure we never
# lose a batch during retries and that we don't produce duplicate batch
# links in the process
tp0 = TopicPartition("test-topic", 0)
def mocked_leader_for_partition(tp):
if tp == tp0:
return 0
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)
fut1 = await ma.add_message(tp0, b'key', b'value', timeout=2)
# Drain and Reenqueu
batches, _ = ma.drain_by_nodes(ignore_nodes=[])
batch = batches[0][tp0]
self.assertIn(batch, ma._pending_batches)
self.assertFalse(ma._batches)
self.assertFalse(fut1.done())
ma.reenqueue(batch)
self.assertEqual(batch.retry_count, 1)
self.assertFalse(ma._pending_batches)
self.assertIn(batch, ma._batches[tp0])
self.assertFalse(fut1.done())
# Drain and Reenqueu again. We check for repeated call
batches, _ = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(batches[0][tp0], batch)
self.assertEqual(batch.retry_count, 2)
self.assertIn(batch, ma._pending_batches)
self.assertFalse(ma._batches)
self.assertFalse(fut1.done())
ma.reenqueue(batch)
self.assertEqual(batch.retry_count, 2)
self.assertFalse(ma._pending_batches)
self.assertIn(batch, ma._batches[tp0])
self.assertFalse(fut1.done())
# Drain and mark as done. Check that no link to batch remained
batches, _ = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(batches[0][tp0], batch)
self.assertEqual(batch.retry_count, 3)
self.assertIn(batch, ma._pending_batches)
self.assertFalse(ma._batches)
self.assertFalse(fut1.done())
if hasattr(batch.future, "_callbacks"): # Vanilla asyncio
self.assertEqual(len(batch.future._callbacks), 1)
batch.done_noack()
await asyncio.sleep(0.01)
self.assertEqual(batch.retry_count, 3)
self.assertFalse(ma._pending_batches)
self.assertFalse(ma._batches)
def client(mocker):
_cli = mocker.Mock(
spec=KafkaClient(bootstrap_servers=(), api_version=(0, 9)))
_cli.cluster = mocker.Mock(spec=ClusterMetadata())
return _cli
import signal
import pprint
from kafka.admin import KafkaAdminClient, NewTopic, ConfigResourceType, ConfigResource
from kafka.cluster import ClusterMetadata
from config import bootstrap_servers, ssl_cafile, ssl_certfile, ssl_keyfile
adminClient = KafkaAdminClient(bootstrap_servers=bootstrap_servers,
security_protocol="SSL",
ssl_cafile=ssl_cafile,
ssl_certfile=ssl_certfile,
ssl_keyfile=ssl_keyfile)
clusterMetadata = ClusterMetadata(bootstrap_servers=bootstrap_servers, )
def createTopic():
try:
topic = str(input("Please enter a topic name:")).strip()
topic_partitions = int(
input("Please enter a number of partition [1]:") or 1)
if str(
input(
"Confirm the creation of topic '{}' with {} partitions? [Y/N]"
.format(topic, topic_partitions))) == 'Y':
print("Creating topic {}...".format(topic))
newTopic_list = [
NewTopic(name=topic,
num_partitions=topic_partitions,
replication_factor=1)
]
adminClient.create_topics(newTopic_list)
def test_batch_done(self):
tp0 = TopicPartition("test-topic", 0)
tp1 = TopicPartition("test-topic", 1)
tp2 = TopicPartition("test-topic", 2)
tp3 = TopicPartition("test-topic", 3)
def mocked_leader_for_partition(tp):
if tp == tp0:
return 0
if tp == tp1:
return 1
if tp == tp2:
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, None, 1, self.loop)
fut1 = yield from ma.add_message(
tp2, None, b'msg for tp@2', timeout=2)
fut2 = yield from ma.add_message(
tp3, None, b'msg for tp@3', timeout=2)
yield from ma.add_message(tp1, None, b'0123456789'*70, timeout=2)
with self.assertRaises(KafkaTimeoutError):
yield from ma.add_message(tp1, None, b'0123456789'*70, timeout=2)
batches, _ = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(batches[1][tp1].expired(), True)
with self.assertRaises(LeaderNotAvailableError):
yield from fut1
with self.assertRaises(NotLeaderForPartitionError):
yield from fut2
fut01 = yield from ma.add_message(
tp0, b'key0', b'value#0', timeout=2)
fut02 = yield from ma.add_message(
tp0, b'key1', b'value#1', timeout=2)
fut10 = yield from ma.add_message(
tp1, None, b'0123456789'*70, timeout=2)
batches, _ = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(batches[0][tp0].expired(), False)
self.assertEqual(batches[1][tp1].expired(), False)
batch_data = batches[0][tp0].get_data_buffer()
self.assertEqual(type(batch_data), io.BytesIO)
batches[0][tp0].done(base_offset=10)
class TestException(Exception):
pass
batches[1][tp1].done(exception=TestException())
res = yield from fut01
self.assertEqual(res.topic, "test-topic")
self.assertEqual(res.partition, 0)
self.assertEqual(res.offset, 10)
res = yield from fut02
self.assertEqual(res.topic, "test-topic")
self.assertEqual(res.partition, 0)
self.assertEqual(res.offset, 11)
with self.assertRaises(TestException):
yield from fut10
fut01 = yield from ma.add_message(
tp0, b'key0', b'value#0', timeout=2)
batches, _ = ma.drain_by_nodes(ignore_nodes=[])
batches[0][tp0].done(base_offset=None)
res = yield from fut01
self.assertEqual(res, None)
# cancelling future
fut01 = yield from ma.add_message(
tp0, b'key0', b'value#2', timeout=2)
batches, _ = ma.drain_by_nodes(ignore_nodes=[])
fut01.cancel()
batches[0][tp0].done(base_offset=21) # no error in this case
def test_basic(self):
cluster = ClusterMetadata(metadata_max_age_ms=10000)
ma = MessageAccumulator(cluster, 1000, None, 30, self.loop)
data_waiter = asyncio.async(ma.data_waiter(), loop=self.loop)
done, _ = yield from asyncio.wait(
[data_waiter], timeout=0.2, loop=self.loop)
self.assertFalse(bool(done)) # no data in accumulator yet...
tp0 = TopicPartition("test-topic", 0)
tp1 = TopicPartition("test-topic", 1)
yield from ma.add_message(tp0, b'key', b'value', timeout=2)
yield from ma.add_message(tp1, None, b'value without key', timeout=2)
done, _ = yield from asyncio.wait(
[data_waiter], timeout=0.2, loop=self.loop)
self.assertTrue(bool(done))
batches, unknown_leaders_exist = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(batches, {})
self.assertEqual(unknown_leaders_exist, True)
def mocked_leader_for_partition(tp):
if tp == tp0:
return 0
if tp == tp1:
return 1
return -1
cluster.leader_for_partition = mock.MagicMock()
cluster.leader_for_partition.side_effect = mocked_leader_for_partition
batches, unknown_leaders_exist = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(len(batches), 2)
self.assertEqual(unknown_leaders_exist, False)
m_set0 = batches[0].get(tp0)
self.assertEqual(type(m_set0), MessageBatch)
m_set1 = batches[1].get(tp1)
self.assertEqual(type(m_set1), MessageBatch)
self.assertEqual(m_set0.expired(), False)
data_waiter = asyncio.async(ma.data_waiter(), loop=self.loop)
done, _ = yield from asyncio.wait(
[data_waiter], timeout=0.2, loop=self.loop)
self.assertFalse(bool(done)) # no data in accumulator again...
# testing batch overflow
tp2 = TopicPartition("test-topic", 2)
yield from ma.add_message(
tp0, None, b'some short message', timeout=2)
yield from ma.add_message(
tp0, None, b'some other short message', timeout=2)
yield from ma.add_message(
tp1, None, b'0123456789'*70, timeout=2)
yield from ma.add_message(
tp2, None, b'message to unknown leader', timeout=2)
# next we try to add message with len=500,
# as we have buffer_size=1000 coroutine will block until data will be
# drained
add_task = asyncio.async(
ma.add_message(tp1, None, b'0123456789'*50, timeout=2),
loop=self.loop)
done, _ = yield from asyncio.wait(
[add_task], timeout=0.2, loop=self.loop)
self.assertFalse(bool(done))
batches, unknown_leaders_exist = ma.drain_by_nodes(ignore_nodes=[1, 2])
self.assertEqual(unknown_leaders_exist, True)
m_set0 = batches[0].get(tp0)
self.assertEqual(m_set0._relative_offset, 2)
m_set1 = batches[1].get(tp1)
self.assertEqual(m_set1, None)
done, _ = yield from asyncio.wait(
[add_task], timeout=0.1, loop=self.loop)
self.assertFalse(bool(done)) # we stil not drained data for tp1
batches, unknown_leaders_exist = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(unknown_leaders_exist, True)
m_set0 = batches[0].get(tp0)
self.assertEqual(m_set0, None)
m_set1 = batches[1].get(tp1)
self.assertEqual(m_set1._relative_offset, 1)
done, _ = yield from asyncio.wait(
[add_task], timeout=0.2, loop=self.loop)
self.assertTrue(bool(done))
batches, unknown_leaders_exist = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(unknown_leaders_exist, True)
m_set1 = batches[1].get(tp1)
self.assertEqual(m_set1._relative_offset, 1)
class AIOKafkaClient:
"""Initialize an asynchronous kafka client
Keyword Arguments:
bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
strings) that the consumer should contact to bootstrap initial
cluster metadata. This does not have to be the full node list.
It just needs to have at least one broker that will respond to
Metadata API Request. Default port is 9092. If no servers are
specified, will default to localhost:9092.
client_id (str): a name for this client. This string is passed in
each request to servers and can be used to identify specific
server-side log entries that correspond to this client. Also
submitted to GroupCoordinator for logging with respect to
consumer group administration. Default: 'aiokafka-{ver}'
request_timeout_ms (int): Client request timeout in milliseconds.
Default: 40000.
metadata_max_age_ms (int): The period of time in milliseconds after
which we force a refresh of metadata even if we haven't seen
any partition leadership changes to proactively discover any
new brokers or partitions. Default: 300000
retry_backoff_ms (int): Milliseconds to backoff when retrying on
errors. Default: 100.
api_version (str): specify which kafka API version to use.
AIOKafka supports Kafka API versions >=0.9 only.
If set to 'auto', will attempt to infer the broker version by
probing various APIs. Default: auto
security_protocol (str): Protocol used to communicate with brokers.
Valid values are: PLAINTEXT, SSL. Default: PLAINTEXT.
ssl_context (ssl.SSLContext): pre-configured SSLContext for wrapping
socket connections. For more information see :ref:`ssl_auth`.
Default: None.
connections_max_idle_ms (int): Close idle connections after the number
of milliseconds specified by this config. Specifying `None` will
disable idle checks. Default: 540000 (9hours).
"""
def __init__(self, *, loop, 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):
if security_protocol not in ('SSL', 'PLAINTEXT'):
raise ValueError("`security_protocol` should be SSL or PLAINTEXT")
if security_protocol == "SSL" and ssl_context is None:
raise ValueError(
"`ssl_context` is mandatory if security_protocol=='SSL'")
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.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 = create_future(loop=self._loop)
self._get_conn_lock = asyncio.Lock(loop=loop)
def __repr__(self):
return '<AIOKafkaClient client_id=%s>' % self._client_id
@property
def api_version(self):
if type(self._api_version) is tuple:
return self._api_version
# unknown api version, return minimal supported version
return (0, 9, 0)
@property
def hosts(self):
return collect_hosts(self._bootstrap_servers)
@asyncio.coroutine
def close(self):
if self._sync_task:
self._sync_task.cancel()
try:
yield from self._sync_task
except asyncio.CancelledError:
pass
self._sync_task = None
# Be careful to wait for graceful closure of all connections, so we
# process all pending buffers.
futs = []
for conn in self._conns.values():
futs.append(conn.close(reason=CloseReason.SHUTDOWN))
if futs:
yield from asyncio.gather(*futs, loop=self._loop)
@asyncio.coroutine
def bootstrap(self):
"""Try to to bootstrap initial cluster metadata"""
# using request v0 for bootstap (bcs api version is not detected yet)
metadata_request = MetadataRequest[0]([])
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,
ssl_context=self._ssl_context,
security_protocol=self._security_protocol,
max_idle_ms=self._connections_max_idle_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 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 ConnectionError(
'Unable to bootstrap from {}'.format(self.hosts))
# detect api version if need
if self._api_version == 'auto':
self._api_version = yield from self.check_version()
if self._sync_task is None:
# starting metadata synchronizer task
self._sync_task = ensure_future(
self._md_synchronizer(), loop=self._loop)
@asyncio.coroutine
def _md_synchronizer(self):
"""routine (async task) for synchronize cluster metadata every
`metadata_max_age_ms` milliseconds"""
while True:
yield from asyncio.wait(
[self._md_update_waiter],
timeout=self._metadata_max_age_ms / 1000,
loop=self._loop)
topics = self._topics
if self._md_update_fut is None:
self._md_update_fut = create_future(loop=self._loop)
ret = yield from self._metadata_update(self.cluster, topics)
# If list of topics changed during metadata update we must update
# it again right away.
if topics != self._topics:
continue
# Earlier this waiter was set before sending metadata_request,
# but that was to avoid topic list changes being unnoticed, which
# is handled explicitly now.
self._md_update_waiter = create_future(loop=self._loop)
self._md_update_fut.set_result(ret)
self._md_update_fut = None
def get_random_node(self):
"""choice random node from known cluster brokers
Returns:
nodeId - identifier of broker
"""
nodeids = [b.nodeId for b in self.cluster.brokers()]
if not nodeids:
return None
return random.choice(nodeids)
@asyncio.coroutine
def _metadata_update(self, cluster_metadata, topics):
assert isinstance(cluster_metadata, ClusterMetadata)
topics = list(topics)
version_id = 0 if self.api_version < (0, 10) else 1
if version_id == 1 and not topics:
topics = None
metadata_request = MetadataRequest[version_id](topics)
nodeids = [b.nodeId for b in self.cluster.brokers()]
bootstrap_id = ('bootstrap', ConnectionGroup.DEFAULT)
if bootstrap_id 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 node %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
# don't update the cluster if there are no valid nodes...the topic
# we want may still be in the process of being created which means
# we will get errors and no nodes until it exists
if not metadata.brokers:
return False
cluster_metadata.update_metadata(metadata)
# We only keep bootstrap connection to update metadata until
# proper cluster layout is available.
if bootstrap_id in self._conns and len(self.cluster.brokers()):
conn = self._conns.pop(bootstrap_id)
conn.close()
break
else:
log.error('Unable to update metadata from %s', nodeids)
cluster_metadata.failed_update(None)
return False
return True
def force_metadata_update(self):
"""Update cluster metadata
Returns:
True/False - metadata updated or not
"""
if self._md_update_fut is None:
# Wake up the `_md_synchronizer` task
if not self._md_update_waiter.done():
self._md_update_waiter.set_result(None)
self._md_update_fut = create_future(loop=self._loop)
# Metadata will be updated in the background by syncronizer
return asyncio.shield(self._md_update_fut, loop=self._loop)
@asyncio.coroutine
def fetch_all_metadata(self):
cluster_md = ClusterMetadata(
metadata_max_age_ms=self._metadata_max_age_ms)
updated = yield from self._metadata_update(cluster_md, [])
if not updated:
raise KafkaError(
'Unable to get cluster metadata over all known brokers')
return cluster_md
def add_topic(self, topic):
"""Add a topic to the list of topics tracked via metadata.
Arguments:
topic (str): topic to track
"""
if topic in self._topics:
res = create_future(loop=self._loop)
res.set_result(True)
else:
res = self.force_metadata_update()
self._topics.add(topic)
return res
def set_topics(self, topics):
"""Set specific topics to track for metadata.
Arguments:
topics (list of str): topics to track
"""
assert not isinstance(topics, str)
if not topics or set(topics).difference(self._topics):
res = self.force_metadata_update()
else:
res = create_future(loop=self._loop)
res.set_result(True)
self._topics = set(topics)
return res
def _on_connection_closed(self, conn, reason):
""" Callback called when connection is closed
"""
# Connection failures imply that our metadata is stale, so let's
# refresh
if reason == CloseReason.CONNECTION_BROKEN or \
reason == CloseReason.CONNECTION_TIMEOUT:
self.force_metadata_update()
@asyncio.coroutine
def _get_conn(self, node_id, *, group=ConnectionGroup.DEFAULT):
"Get or create a connection to a broker using host and port"
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
try:
broker = self.cluster.broker_metadata(node_id)
assert broker, 'Broker id %s not in current metadata' % node_id
log.debug("Initiating connection to node %s at %s:%s",
node_id, broker.host, broker.port)
with (yield from self._get_conn_lock):
if conn_id in self._conns:
return self._conns[conn_id]
self._conns[conn_id] = yield from create_conn(
broker.host, broker.port, loop=self._loop,
client_id=self._client_id,
request_timeout_ms=self._request_timeout_ms,
ssl_context=self._ssl_context,
security_protocol=self._security_protocol,
on_close=self._on_connection_closed,
max_idle_ms=self._connections_max_idle_ms)
except (OSError, asyncio.TimeoutError) as err:
log.error('Unable connect to node with id %s: %s', node_id, err)
# Connection failures imply that our metadata is stale, so let's
# refresh
self.force_metadata_update()
return None
else:
return self._conns[conn_id]
@asyncio.coroutine
def ready(self, node_id, *, group=ConnectionGroup.DEFAULT):
conn = yield from self._get_conn(node_id, group=group)
if conn is None:
return False
return True
@asyncio.coroutine
def send(self, node_id, request, *, group=ConnectionGroup.DEFAULT):
"""Send a request to a specific node.
Arguments:
node_id (int): destination node
request (Struct): request object (not-encoded)
Raises:
kafka.common.RequestTimedOutError
kafka.common.NodeNotReadyError
kafka.common.ConnectionError
kafka.common.CorrelationIdError
Returns:
Future: resolves to Response struct
"""
if not (yield from self.ready(node_id, group=group)):
raise NodeNotReadyError(
"Attempt to send a request to node"
" which is not ready (node id {}).".format(node_id))
# Every request gets a response, except one special case:
expect_response = True
if isinstance(request, tuple(ProduceRequest)) and \
request.required_acks == 0:
expect_response = False
future = self._conns[(node_id, group)].send(
request, expect_response=expect_response)
try:
result = yield from future
except asyncio.TimeoutError:
# close connection so it is renewed in next request
self._conns[(node_id, group)].close(
reason=CloseReason.CONNECTION_TIMEOUT)
raise RequestTimedOutError()
else:
return result
@asyncio.coroutine
def check_version(self, node_id=None):
"""Attempt to guess the broker version"""
if node_id is None:
default_group_conns = [
n_id for (n_id, group) in self._conns.keys()
if group == ConnectionGroup.DEFAULT
]
if default_group_conns:
node_id = default_group_conns[0]
else:
assert self.cluster.brokers(), 'no brokers in metadata'
node_id = list(self.cluster.brokers())[0].nodeId
from kafka.protocol.admin import (
ListGroupsRequest_v0, ApiVersionRequest_v0)
from kafka.protocol.commit import (
OffsetFetchRequest_v0, GroupCoordinatorRequest_v0)
from kafka.protocol.metadata import MetadataRequest_v0
test_cases = [
((0, 10), ApiVersionRequest_v0()),
((0, 9), ListGroupsRequest_v0()),
((0, 8, 2), GroupCoordinatorRequest_v0('aiokafka-default-group')),
((0, 8, 1), OffsetFetchRequest_v0('aiokafka-default-group', [])),
((0, 8, 0), MetadataRequest_v0([])),
]
# 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)
# request can be ignored by Kafka broker,
# so we send metadata request and wait response
task = self._loop.create_task(conn.send(request))
yield from asyncio.wait([task], timeout=0.1, loop=self._loop)
try:
yield from conn.send(MetadataRequest_v0([]))
except KafkaError:
# metadata request can be cancelled in case
# of invalid correlationIds order
pass
response = yield from task
except KafkaError:
continue
else:
# To avoid having a connection in undefined state
if node_id != "bootstrap" and conn.connected():
conn.close()
if isinstance(request, ApiVersionRequest_v0):
# Starting from 0.10 kafka broker we determine version
# by looking at ApiVersionResponse
version = self._check_api_version_response(response)
return version
raise UnrecognizedBrokerVersion()
def _check_api_version_response(self, response):
# The logic here is to check the list of supported request versions
# in descending order. As soon as we find one that works, return it
test_cases = [
# format (<broker verion>, <needed struct>)
((1, 0, 0), MetadataRequest[0].API_KEY, 5),
((0, 11, 0), MetadataRequest[0].API_KEY, 4),
((0, 10, 2), OffsetFetchRequest[0].API_KEY, 2),
((0, 10, 1), MetadataRequest[0].API_KEY, 2),
]
error_type = Errors.for_code(response.error_code)
assert error_type is Errors.NoError, "API version check failed"
max_versions = dict([
(api_key, max_version)
for api_key, _, max_version in response.api_versions
])
# Get the best match of test cases
for broker_version, api_key, version in test_cases:
if max_versions.get(api_key, -1) >= version:
return broker_version
# We know that ApiVersionResponse is only supported in 0.10+
# so if all else fails, choose that
return (0, 10, 0)
@asyncio.coroutine
def _wait_on_metadata(self, topic):
"""
Wait for cluster metadata including partitions for the given topic to
be available.
Arguments:
topic (str): topic we want metadata for
Returns:
set: partition ids for the topic
Raises:
UnknownTopicOrPartitionError: if no topic or partitions found
in cluster metadata
"""
if topic in self.cluster.topics():
return self.cluster.partitions_for_topic(topic)
# add topic to metadata topic list if it is not there already.
self.add_topic(topic)
t0 = self._loop.time()
while True:
yield from self.force_metadata_update()
if topic in self.cluster.topics():
break
if (self._loop.time() - t0) > (self._request_timeout_ms / 1000):
raise UnknownTopicOrPartitionError()
yield from asyncio.sleep(self._retry_backoff, loop=self._loop)
return self.cluster.partitions_for_topic(topic)
@asyncio.coroutine
def _maybe_wait_metadata(self):
if self._md_update_fut is not None:
yield from asyncio.shield(
self._md_update_fut, loop=self._loop)
from kafka.cluster import ClusterMetadata
topic_name = "test"
cm = ClusterMetadata(bootstrap_servers='localhost:9092')
for i in cm.topics():
print(i)
asyncio.set_event_loop(waiter_app['loop'])
waiter_app['serializer'] = AvroSerializer(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
'examples/coffee_bar/avro_schemas'))
# Creates & registers local store memory / global store memory
waiter_app['local_store'] = LocalStoreMemory(
name=f'waiter-{cur_instance}-local-memory')
waiter_app['global_store'] = GlobalStoreMemory(
name=f'waiter-{cur_instance}-global-memory')
cluster_admin = KafkaAdminClient(bootstrap_servers='localhost:9092',
client_id=f'waiter-{cur_instance}')
cluster_metadata = ClusterMetadata(bootstrap_servers='localhost:9092')
# Creates & registers store builder
waiter_app['store_builder'] = StoreBuilder(
name=f'waiter-{cur_instance}-store-builder',
current_instance=cur_instance,
nb_replica=nb_replica,
topic_store='waiter-stores',
serializer=waiter_app['serializer'],
local_store=waiter_app['local_store'],
global_store=waiter_app['global_store'],
bootstrap_server='localhost:9092',
cluster_metadata=cluster_metadata,
cluster_admin=cluster_admin,
loop=waiter_app['loop'],
rebuild=True,
def assign(self, cluster: ClusterMetadata, members: Dict[str, ConsumerProtocolMemberMetadata]) \
-> Dict[str, ConsumerProtocolMemberAssignment]:
logger.info('Statefulset Partition Assignor')
logger.debug(f'Cluster = {cluster}\nMembers = {members}')
# Get all topic
all_topics: Set = set()
for key, metadata in members.items():
all_topics.update(metadata.subscription)
# Get all partitions by topic name
all_topic_partitions = []
for topic in all_topics:
partitions = cluster.partitions_for_topic(topic)
if partitions is None:
logger.warning('No partition metadata for topic %s', topic)
continue
for partition in partitions:
all_topic_partitions.append(TopicPartition(topic, partition))
# Sort partition
all_topic_partitions.sort()
# Create default dict with lambda
assignment: DefaultDict[str, Any] = collections.defaultdict(
lambda: collections.defaultdict(list))
advanced_assignor_dict = self.get_advanced_assignor_dict(
all_topic_partitions)
for topic, partitions in advanced_assignor_dict.items():
for member_id, member_data in members.items():
# Loads member assignors data
user_data = json.loads(member_data.user_data)
# Get number of partitions by topic name
topic_number_partitions = len(partitions)
# Logic assignors if nb_replica as same as topic_numbers_partitions (used by StoreBuilder for
# assign each partitions to right instance, in this case nb_replica is same as topic_number_partitions)
if user_data['nb_replica'] == topic_number_partitions:
if user_data['assignor_policy'] == 'all':
for partition in partitions:
assignment[member_id][topic].append(partition)
elif user_data['assignor_policy'] == 'only_own':
if user_data['instance'] in partitions:
assignment[member_id][topic].append(
partitions[user_data['instance']])
else:
raise BadAssignorPolicy
else:
# Todo Add repartition
raise NotImplementedError
logger.debug(f'Assignment = {assignment}')
protocol_assignment = {}
for member_id in members:
protocol_assignment[member_id] = ConsumerProtocolMemberAssignment(
self.version, sorted(assignment[member_id].items()),
members[member_id].user_data)
logger.debug(f'Protocol Assignment = {protocol_assignment}')
return protocol_assignment
def test_basic(self):
cluster = ClusterMetadata(metadata_max_age_ms=10000)
ma = MessageAccumulator(cluster, 1000, None, 30, self.loop)
data_waiter = ma.data_waiter()
done, _ = yield from asyncio.wait(
[data_waiter], timeout=0.2, loop=self.loop)
self.assertFalse(bool(done)) # no data in accumulator yet...
tp0 = TopicPartition("test-topic", 0)
tp1 = TopicPartition("test-topic", 1)
yield from ma.add_message(tp0, b'key', b'value', timeout=2)
yield from ma.add_message(tp1, None, b'value without key', timeout=2)
done, _ = yield from asyncio.wait(
[data_waiter], timeout=0.2, loop=self.loop)
self.assertTrue(bool(done))
batches, unknown_leaders_exist = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(batches, {})
self.assertEqual(unknown_leaders_exist, True)
def mocked_leader_for_partition(tp):
if tp == tp0:
return 0
if tp == tp1:
return 1
return -1
cluster.leader_for_partition = mock.MagicMock()
cluster.leader_for_partition.side_effect = mocked_leader_for_partition
batches, unknown_leaders_exist = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(len(batches), 2)
self.assertEqual(unknown_leaders_exist, False)
m_set0 = batches[0].get(tp0)
self.assertEqual(type(m_set0), MessageBatch)
m_set1 = batches[1].get(tp1)
self.assertEqual(type(m_set1), MessageBatch)
self.assertEqual(m_set0.expired(), False)
data_waiter = ensure_future(ma.data_waiter(), loop=self.loop)
done, _ = yield from asyncio.wait(
[data_waiter], timeout=0.2, loop=self.loop)
self.assertFalse(bool(done)) # no data in accumulator again...
# testing batch overflow
tp2 = TopicPartition("test-topic", 2)
yield from ma.add_message(
tp0, None, b'some short message', timeout=2)
yield from ma.add_message(
tp0, None, b'some other short message', timeout=2)
yield from ma.add_message(
tp1, None, b'0123456789' * 70, timeout=2)
yield from ma.add_message(
tp2, None, b'message to unknown leader', timeout=2)
# next we try to add message with len=500,
# as we have buffer_size=1000 coroutine will block until data will be
# drained
add_task = ensure_future(
ma.add_message(tp1, None, b'0123456789' * 50, timeout=2),
loop=self.loop)
done, _ = yield from asyncio.wait(
[add_task], timeout=0.2, loop=self.loop)
self.assertFalse(bool(done))
batches, unknown_leaders_exist = ma.drain_by_nodes(ignore_nodes=[1, 2])
self.assertEqual(unknown_leaders_exist, True)
m_set0 = batches[0].get(tp0)
self.assertEqual(m_set0._builder._relative_offset, 2)
m_set1 = batches[1].get(tp1)
self.assertEqual(m_set1, None)
done, _ = yield from asyncio.wait(
[add_task], timeout=0.1, loop=self.loop)
self.assertFalse(bool(done)) # we stil not drained data for tp1
batches, unknown_leaders_exist = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(unknown_leaders_exist, True)
m_set0 = batches[0].get(tp0)
self.assertEqual(m_set0, None)
m_set1 = batches[1].get(tp1)
self.assertEqual(m_set1._builder._relative_offset, 1)
done, _ = yield from asyncio.wait(
[add_task], timeout=0.2, loop=self.loop)
self.assertTrue(bool(done))
batches, unknown_leaders_exist = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(unknown_leaders_exist, True)
m_set1 = batches[1].get(tp1)
self.assertEqual(m_set1._builder._relative_offset, 1)
def test_batch_done(self):
tp0 = TopicPartition("test-topic", 0)
tp1 = TopicPartition("test-topic", 1)
tp2 = TopicPartition("test-topic", 2)
tp3 = TopicPartition("test-topic", 3)
def mocked_leader_for_partition(tp):
if tp == tp0:
return 0
if tp == tp1:
return 1
if tp == tp2:
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, None, 1, self.loop)
fut1 = yield from ma.add_message(
tp2, None, b'msg for tp@2', timeout=2)
fut2 = yield from ma.add_message(
tp3, None, b'msg for tp@3', timeout=2)
yield from ma.add_message(tp1, None, b'0123456789'*70, timeout=2)
with self.assertRaises(KafkaTimeoutError):
yield from ma.add_message(tp1, None, b'0123456789'*70, timeout=2)
batches, _ = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(batches[1][tp1].expired(), True)
with self.assertRaises(LeaderNotAvailableError):
yield from fut1
with self.assertRaises(NotLeaderForPartitionError):
yield from fut2
fut01 = yield from ma.add_message(
tp0, b'key0', b'value#0', timeout=2)
fut02 = yield from ma.add_message(
tp0, b'key1', b'value#1', timeout=2)
fut10 = yield from ma.add_message(
tp1, None, b'0123456789'*70, timeout=2)
batches, _ = ma.drain_by_nodes(ignore_nodes=[])
self.assertEqual(batches[0][tp0].expired(), False)
self.assertEqual(batches[1][tp1].expired(), False)
batch_data = batches[0][tp0].data()
self.assertEqual(type(batch_data), io.BytesIO)
batches[0][tp0].done(base_offset=10)
class TestException(Exception):
pass
batches[1][tp1].done(exception=TestException())
res = yield from fut01
self.assertEqual(res.topic, "test-topic")
self.assertEqual(res.partition, 0)
self.assertEqual(res.offset, 10)
res = yield from fut02
self.assertEqual(res.topic, "test-topic")
self.assertEqual(res.partition, 0)
self.assertEqual(res.offset, 11)
with self.assertRaises(TestException):
yield from fut10
fut01 = yield from ma.add_message(
tp0, b'key0', b'value#0', timeout=2)
batches, _ = ma.drain_by_nodes(ignore_nodes=[])
batches[0][tp0].done(base_offset=None)
res = yield from fut01
self.assertEqual(res, None)
