def test_relookup_on_failure():
hostname = 'example.org'
port = 9092
conn = BrokerConnection(hostname, port, socket.AF_UNSPEC)
assert conn.host == hostname
mock_return1 = []
with mock.patch("socket.getaddrinfo", return_value=mock_return1) as m:
last_attempt = conn.last_attempt
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
assert conn.disconnected()
assert conn.last_attempt > last_attempt
afi2 = socket.AF_INET
sockaddr2 = ('127.0.0.2', 9092)
mock_return2 = [
(afi2, socket.SOCK_STREAM, 6, '', sockaddr2),
]
with mock.patch("socket.getaddrinfo", return_value=mock_return2) as m:
conn.last_attempt = 0
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
assert conn._sock_afi == afi2
assert conn._sock_addr == sockaddr2
conn.close()
def test_lookup_on_connect():
hostname = 'example.org'
port = 9092
conn = BrokerConnection(hostname, port, socket.AF_UNSPEC)
assert conn.host == hostname
assert conn.port == port
assert conn.afi == socket.AF_UNSPEC
afi1 = socket.AF_INET
sockaddr1 = ('127.0.0.1', 9092)
mock_return1 = [
(afi1, socket.SOCK_STREAM, 6, '', sockaddr1),
]
with mock.patch("socket.getaddrinfo", return_value=mock_return1) as m:
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
assert conn._sock_afi == afi1
assert conn._sock_addr == sockaddr1
conn.close()
afi2 = socket.AF_INET6
sockaddr2 = ('::1', 9092, 0, 0)
mock_return2 = [
(afi2, socket.SOCK_STREAM, 6, '', sockaddr2),
]
with mock.patch("socket.getaddrinfo", return_value=mock_return2) as m:
conn.last_attempt = 0
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
assert conn._sock_afi == afi2
assert conn._sock_addr == sockaddr2
conn.close()
def test_lookup_on_connect():
hostname = 'example.org'
port = 9092
conn = BrokerConnection(hostname, port, socket.AF_UNSPEC)
assert conn.host == conn.hostname == hostname
ip1 = '127.0.0.1'
mock_return1 = [
(2, 2, 17, '', (ip1, 9092)),
]
with mock.patch("socket.getaddrinfo", return_value=mock_return1) as m:
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
conn.close()
assert conn.host == ip1
ip2 = '127.0.0.2'
mock_return2 = [
(2, 2, 17, '', (ip2, 9092)),
]
with mock.patch("socket.getaddrinfo", return_value=mock_return2) as m:
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
conn.close()
assert conn.host == ip2
def test_relookup_on_failure():
hostname = 'example.org'
port = 9092
conn = BrokerConnection(hostname, port, socket.AF_UNSPEC)
assert conn.host == hostname
mock_return1 = []
with mock.patch("socket.getaddrinfo", return_value=mock_return1) as m:
last_attempt = conn.last_attempt
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
assert conn.disconnected()
assert conn.last_attempt > last_attempt
afi2 = socket.AF_INET
sockaddr2 = ('127.0.0.2', 9092)
mock_return2 = [
(afi2, socket.SOCK_STREAM, 6, '', sockaddr2),
]
with mock.patch("socket.getaddrinfo", return_value=mock_return2) as m:
conn.last_attempt = 0
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
assert conn._sock_afi == afi2
assert conn._sock_addr == sockaddr2
conn.close()
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 _get_conn(self, host, port, afi, node_id='bootstrap'):
"""Get or create a connection to a broker using host and port"""
host_key = (host, port)
if host_key not in self._conns:
self._conns[host_key] = BrokerConnection(
host, port, afi,
request_timeout_ms=self.timeout * 1000,
client_id=self.client_id,
metrics=self._metrics_registry,
metric_group_prefix='simple-client',
node_id=node_id,
)
conn = self._conns[host_key]
conn.connect()
if conn.connected():
return conn
timeout = time.time() + self.timeout
while time.time() < timeout and conn.connecting():
if conn.connect() is ConnectionStates.CONNECTED:
break
else:
time.sleep(0.05)
else:
conn.close()
raise ConnectionError("%s:%s (%s)" % (host, port, afi))
return conn
def _get_conn(self, host, port, afi):
"""Get or create a connection to a broker using host and port"""
host_key = (host, port)
if host_key not in self._conns:
self._conns[host_key] = BrokerConnection(
host,
port,
afi,
request_timeout_ms=self.timeout * 1000,
client_id=self.client_id)
conn = self._conns[host_key]
conn.connect()
if conn.connected():
return conn
timeout = time.time() + self.timeout
while time.time() < timeout and conn.connecting():
if conn.connect() is ConnectionStates.CONNECTED:
break
else:
time.sleep(0.05)
else:
conn.close()
raise ConnectionError("%s:%s (%s)" % (host, port, afi))
return conn
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 = 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 self._should_recycle_connection(conn):
self._conns.pop(node_id)
return False
elif conn.connected():
return True
conn.connect()
return conn.connected()
def _get_conn(self, host, port):
"""Get or create a connection to a broker using host and port"""
host_key = (host, port)
if host_key not in self._conns:
self._conns[host_key] = BrokerConnection(
host,
port,
request_timeout_ms=self.timeout * 1000,
client_id=self.client_id)
conn = self._conns[host_key]
while conn.connect() == ConnectionStates.CONNECTING:
pass
return conn
def _get_conn(self, host, port, afi):
"""Get or create a connection to a broker using host and port"""
host_key = (host, port)
if host_key not in self._conns:
self._conns[host_key] = BrokerConnection(
host, port, afi,
request_timeout_ms=self.timeout * 1000,
client_id=self.client_id
)
conn = self._conns[host_key]
if not conn.connect_blocking(self.timeout):
conn.close()
raise ConnectionError("%s:%s (%s)" % (host, port, afi))
return conn
def test_relookup_on_failure():
hostname = 'example.org'
port = 9092
conn = BrokerConnection(hostname, port, socket.AF_UNSPEC)
assert conn.host == conn.hostname == hostname
mock_return1 = []
with mock.patch("socket.getaddrinfo", return_value=mock_return1) as m:
last_attempt = conn.last_attempt
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
assert conn.disconnected()
assert conn.last_attempt > last_attempt
ip2 = '127.0.0.2'
mock_return2 = [
(2, 2, 17, '', (ip2, 9092)),
]
with mock.patch("socket.getaddrinfo", return_value=mock_return2) as m:
conn.last_attempt = 0
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
conn.close()
assert conn.host == ip2
def _get_conn(self, host, port, afi, node_id='bootstrap'):
"""Get or create a connection to a broker using host and port"""
host_key = (host, port)
if host_key not in self._conns:
self._conns[host_key] = BrokerConnection(
host,
port,
afi,
request_timeout_ms=self.timeout * 1000,
client_id=self.client_id,
metrics=self._metrics_registry,
metric_group_prefix='simple-client',
node_id=node_id,
)
conn = self._conns[host_key]
if not conn.connect_blocking(self.timeout):
conn.close()
raise KafkaConnectionError("%s:%s (%s)" % (host, port, afi))
return conn
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 _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 _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:
# Note that when bootstrapping, each call to broker_metadata may
# return a different host/port. So we need to be careful to only
# call when necessary to avoid skipping some possible bootstrap
# source.
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 = 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 self._should_recycle_connection(conn):
self._conns.pop(node_id)
return False
elif conn.connected():
return True
conn.connect()
return conn.connected()
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 conn(_socket):
conn = BrokerConnection('localhost', 9092, socket.AF_INET)
return conn
def test_recv_disconnected():
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind(('127.0.0.1', 0))
port = sock.getsockname()[1]
sock.listen(5)
conn = BrokerConnection('127.0.0.1', port, socket.AF_INET)
timeout = time.time() + 1
while time.time() < timeout:
conn.connect()
if conn.connected():
break
else:
assert False, 'Connection attempt to local socket timed-out ?'
conn.send(MetadataRequest[0]([]))
# Disconnect server socket
sock.close()
# Attempt to receive should mark connection as disconnected
assert conn.connected()
conn.recv()
assert conn.disconnected()
def test_lookup_on_connect():
hostname = 'example.org'
port = 9092
conn = BrokerConnection(hostname, port, socket.AF_UNSPEC)
assert conn.host == hostname
assert conn.port == port
assert conn.afi == socket.AF_UNSPEC
afi1 = socket.AF_INET
sockaddr1 = ('127.0.0.1', 9092)
mock_return1 = [
(afi1, socket.SOCK_STREAM, 6, '', sockaddr1),
]
with mock.patch("socket.getaddrinfo", return_value=mock_return1) as m:
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
assert conn._sock_afi == afi1
assert conn._sock_addr == sockaddr1
conn.close()
afi2 = socket.AF_INET6
sockaddr2 = ('::1', 9092, 0, 0)
mock_return2 = [
(afi2, socket.SOCK_STREAM, 6, '', sockaddr2),
]
with mock.patch("socket.getaddrinfo", return_value=mock_return2) as m:
conn.last_attempt = 0
conn.connect()
m.assert_called_once_with(hostname, port, 0, 1)
assert conn._sock_afi == afi2
assert conn._sock_addr == sockaddr2
conn.close()
def test_correlation_id_rollover(kafka_broker):
logging.getLogger('kafka.conn').setLevel(logging.ERROR)
from kafka.protocol.metadata import MetadataRequest
conn = BrokerConnection('localhost',
kafka_broker.port,
receive_buffer_bytes=131072,
max_in_flight_requests_per_connection=100)
req = MetadataRequest([])
while not conn.connected():
conn.connect()
futures = collections.deque()
start = time.time()
done = 0
for i in six.moves.xrange(2**13):
if not conn.can_send_more():
conn.recv(timeout=None)
futures.append(conn.send(req))
conn.recv()
while futures and futures[0].is_done:
f = futures.popleft()
if not f.succeeded():
raise f.exception
done += 1
if time.time() > start + 10:
print("%d done" % done)
start = time.time()
while futures:
conn.recv()
if futures[0].is_done:
f = futures.popleft()
if not f.succeeded():
raise f.exception
def test_correlation_id_rollover(kafka_broker):
logging.getLogger('kafka.conn').setLevel(logging.ERROR)
from kafka.protocol.metadata import MetadataRequest
conn = BrokerConnection('localhost', kafka_broker.port,
receive_buffer_bytes=131072,
max_in_flight_requests_per_connection=100)
req = MetadataRequest([])
while not conn.connected():
conn.connect()
futures = collections.deque()
start = time.time()
done = 0
for i in six.moves.xrange(2**13):
if not conn.can_send_more():
conn.recv(timeout=None)
futures.append(conn.send(req))
conn.recv()
while futures and futures[0].is_done:
f = futures.popleft()
if not f.succeeded():
raise f.exception
done += 1
if time.time() > start + 10:
print ("%d done" % done)
start = time.time()
while futures:
conn.recv()
if futures[0].is_done:
f = futures.popleft()
if not f.succeeded():
raise f.exception
def conn(socket):
from socket import AF_INET
conn = BrokerConnection('localhost', 9092, AF_INET)
return conn
def test_connection(self):
conn = BrokerConnection('localhost', 9092, socket.AF_INET)
assert conn.connected() is False
conn.state = ConnectionStates.CONNECTED
assert conn.connected() is True
assert conn.connecting() is False
conn.state = ConnectionStates.CONNECTING
assert conn.connecting() is True
conn.state = ConnectionStates.CONNECTED
assert conn.connecting() is False
conn.state = ConnectionStates.DISCONNECTED
assert conn.connected() is False