def drain_ready(self):
"""Compress batch to be ready for send"""
memview = self._buffer.getbuffer()
self._drain_waiter.set_result(None)
if self._compression_type:
_, compressor, attrs = self._COMPRESSORS[self._compression_type]
msg = Message(compressor(memview[4:].tobytes()), attributes=attrs,
magic=self._version_id)
encoded = msg.encode()
# if compressed message is longer than original
# we should send it as is (not compressed)
header_size = 16 # 4(all size) + 8(offset) + 4(compressed size)
if len(encoded) + header_size < len(memview):
# write compressed message set (with header) to buffer
# using memory view (for avoid memory copying)
memview[:4] = Int32.encode(len(encoded) + 12)
memview[4:12] = Int64.encode(0) # offset 0
memview[12:16] = Int32.encode(len(encoded))
memview[16:16 + len(encoded)] = encoded
self._buffer.seek(0)
return
# update batch size (first 4 bytes of buffer)
memview[:4] = Int32.encode(self._buffer.tell() - 4)
self._buffer.seek(0)
def _build(self):
if self._closed:
self._buffer.seek(0)
return self._buffer
self._closed = True
memview = self._buffer.getbuffer()
if self._compression_type:
_, compressor, attrs = self._COMPRESSORS[self._compression_type]
msg = Message(compressor(memview[4:].tobytes()),
attributes=attrs,
magic=self._magic)
encoded = msg.encode()
# if compressed message is longer than original
# we should send it as is (not compressed)
header_size = 16 # 4(all size) + 8(offset) + 4(compressed size)
if len(encoded) + header_size < len(memview):
# write compressed message set (with header) to buffer
# using memory view (for avoid memory copying)
memview[:4] = Int32.encode(len(encoded) + 12)
memview[4:12] = Int64.encode(0) # offset 0
memview[12:16] = Int32.encode(len(encoded))
memview[16:16 + len(encoded)] = encoded
memview.release()
self._buffer.seek(16 + len(encoded))
self._buffer.truncate()
self._buffer.seek(0)
return self._buffer
# update batch size (first 4 bytes of buffer)
memview[:4] = Int32.encode(self._buffer.tell() - 4)
self._buffer.seek(0)
return self._buffer
def test_decode_fetch_response_partial():
encoded = b''.join([
Int32.encode(1), # Num Topics (Array)
String('utf-8').encode('foobar'),
Int32.encode(2), # Num Partitions (Array)
Int32.encode(0), # Partition id
Int16.encode(0), # Error Code
Int64.encode(1234), # Highwater offset
Int32.encode(52), # MessageSet size
Int64.encode(0), # Msg Offset
Int32.encode(18), # Msg Size
struct.pack('>i', 1474775406), # CRC
struct.pack('>bb', 0, 0), # Magic, flags
struct.pack('>i', 2), # Length of key
b'k1', # Key
struct.pack('>i', 2), # Length of value
b'v1', # Value
Int64.encode(1), # Msg Offset
struct.pack('>i', 24), # Msg Size (larger than remaining MsgSet size)
struct.pack('>i', -16383415), # CRC
struct.pack('>bb', 0, 0), # Magic, flags
struct.pack('>i', 2), # Length of key
b'k2', # Key
struct.pack('>i', 8), # Length of value
b'ar', # Value (truncated)
Int32.encode(1),
Int16.encode(0),
Int64.encode(2345),
Int32.encode(52), # MessageSet size
Int64.encode(0), # Msg Offset
Int32.encode(18), # Msg Size
struct.pack('>i', 1474775406), # CRC
struct.pack('>bb', 0, 0), # Magic, flags
struct.pack('>i', 2), # Length of key
b'k1', # Key
struct.pack('>i', 2), # Length of value
b'v1', # Value
Int64.encode(1), # Msg Offset
struct.pack('>i', 24), # Msg Size (larger than remaining MsgSet size)
struct.pack('>i', -16383415), # CRC
struct.pack('>bb', 0, 0), # Magic, flags
struct.pack('>i', 2), # Length of key
b'k2', # Key
struct.pack('>i', 8), # Length of value
b'ar', # Value (truncated)
])
resp = FetchResponse[0].decode(io.BytesIO(encoded))
assert len(resp.topics) == 1
topic, partitions = resp.topics[0]
assert topic == 'foobar'
assert len(partitions) == 2
m1 = MessageSet.decode(
partitions[0][3], bytes_to_read=len(partitions[0][3]))
assert len(m1) == 2
assert m1[1] == (None, None, PartialMessage())
def test__unpack_message_set_compressed_v1(fetcher):
fetcher.config['check_crcs'] = False
tp = TopicPartition('foo', 0)
messages = [
(0, None, Message(b'a')),
(1, None, Message(b'b')),
(2, None, Message(b'c')),
]
message_bytes = []
for offset, _, m in messages:
encoded = m.encode()
message_bytes.append(
Int64.encode(offset) + Int32.encode(len(encoded)) + encoded)
compressed_bytes = gzip_encode(b''.join(message_bytes))
compressed_base_offset = 10
compressed_msgs = [(compressed_base_offset, None,
Message(compressed_bytes,
magic=1,
attributes=Message.CODEC_GZIP))]
records = list(fetcher._unpack_message_set(tp, compressed_msgs))
assert len(records) == 3
assert all(map(lambda x: isinstance(x, ConsumerRecord), records))
assert records[0].value == b'a'
assert records[1].value == b'b'
assert records[2].value == b'c'
assert records[0].offset == 8
assert records[1].offset == 9
assert records[2].offset == 10
def _process_response(self, read_buffer):
assert not self._processing, 'Recursion not supported'
self._processing = True
ifr = self.in_flight_requests.popleft()
# verify send/recv correlation ids match
recv_correlation_id = Int32.decode(read_buffer)
# 0.8.2 quirk
if (self.config['api_version'] == (0, 8, 2) and
ifr.response_type is GroupCoordinatorResponse and
recv_correlation_id == 0):
raise Errors.KafkaError(
'Kafka 0.8.2 quirk -- try creating a topic first')
elif ifr.correlation_id != recv_correlation_id:
error = Errors.CorrelationIdError(
'Correlation ids do not match: sent %d, recv %d'
% (ifr.correlation_id, recv_correlation_id))
ifr.future.fail(error)
self.close()
self._processing = False
return None
# decode response
response = ifr.response_type.decode(read_buffer)
log.debug('%s Response %d: %s', self, ifr.correlation_id, response)
ifr.future.success(response)
self._processing = False
return response
def decode(cls, data, bytes_to_read=None):
"""Compressed messages should pass in bytes_to_read (via message size)
otherwise, we decode from data as Int32
"""
if isinstance(data, bytes):
data = io.BytesIO(data)
if bytes_to_read is None:
bytes_to_read = Int32.decode(data)
# if FetchRequest max_bytes is smaller than the available message set
# the server returns partial data for the final message
# So create an internal buffer to avoid over-reading
raw = io.BytesIO(data.read(bytes_to_read))
items = []
while bytes_to_read:
try:
offset = Int64.decode(raw)
msg_bytes = Bytes.decode(raw)
bytes_to_read -= 8 + 4 + len(msg_bytes)
items.append(
(offset, len(msg_bytes), Message.decode(msg_bytes)))
except ValueError:
# PartialMessage to signal that max_bytes may be too small
items.append((None, None, PartialMessage()))
break
return items
def _process_response(self, read_buffer):
assert not self._processing, 'Recursion not supported'
self._processing = True
ifr = self.in_flight_requests.popleft()
# verify send/recv correlation ids match
recv_correlation_id = Int32.decode(read_buffer)
# 0.8.2 quirk
if (self.config['api_version'] == (0, 8, 2)
and ifr.response_type is GroupCoordinatorResponse[0]
and ifr.correlation_id != 0 and recv_correlation_id == 0):
log.warning('Kafka 0.8.2 quirk -- GroupCoordinatorResponse'
' coorelation id does not match request. This'
' should go away once at least one topic has been'
' initialized on the broker')
elif ifr.correlation_id != recv_correlation_id:
error = Errors.CorrelationIdError(
'%s: Correlation ids do not match: sent %d, recv %d' %
(str(self), ifr.correlation_id, recv_correlation_id))
ifr.future.failure(error)
self.close()
self._processing = False
return None
# decode response
response = ifr.response_type.decode(read_buffer)
log.debug('%s Response %d: %s', self, ifr.correlation_id, response)
ifr.future.success(response)
self._processing = False
return response
def _try_authenticate_plain(self, future):
if self.config['security_protocol'] == 'SASL_PLAINTEXT':
log.warning('%s: Sending username and password in the clear', self)
data = b''
# Send PLAIN credentials per RFC-4616
msg = bytes('\0'.join([self.config['sasl_plain_username'],
self.config['sasl_plain_username'],
self.config['sasl_plain_password']]).encode('utf-8'))
size = Int32.encode(len(msg))
try:
self._send_bytes_blocking(size + msg)
# The server will send a zero sized message (that is Int32(0)) on success.
# The connection is closed on failure
self._recv_bytes_blocking(4)
except ConnectionError as e:
log.exception("%s: Error receiving reply from server", self)
error = Errors.ConnectionError("%s: %s" % (self, e))
self.close(error=error)
return future.failure(error)
if data != b'\x00\x00\x00\x00':
error = Errors.AuthenticationFailedError('Unrecognized response during authentication')
return future.failure(error)
log.info('%s: Authenticated as %s via PLAIN', self, self.config['sasl_plain_username'])
return future.success(True)
def send_request(self, request, correlation_id=None):
"""Encode and queue a kafka api request for sending.
Arguments:
request (object): An un-encoded kafka request.
correlation_id (int, optional): Optionally specify an ID to
correlate requests with responses. If not provided, an ID will
be generated automatically.
Returns:
correlation_id
"""
log.debug("Sending request %s", request)
if correlation_id is None:
correlation_id = self._next_correlation_id()
header = RequestHeader(request,
correlation_id=correlation_id,
client_id=self._client_id)
message = b"".join([header.encode(), request.encode()])
size = Int32.encode(len(message))
data = size + message
self.bytes_to_send.append(data)
if request.expect_response():
ifr = (correlation_id, request)
self.in_flight_requests.append(ifr)
return correlation_id
def append(self, key, value, timestamp_ms):
"""Append message (key and value) to batch
Returns:
None if batch is full
or
asyncio.Future that will resolved when message is delivered
"""
if self._is_full(key, value):
return None
# `.encode()` is a weak method for some reason, so we need to save
# reference before calling it.
if self._version_id == 0:
msg_inst = Message(value, key=key, magic=self._version_id)
else:
msg_inst = Message(value,
key=key,
magic=self._version_id,
timestamp=timestamp_ms)
encoded = msg_inst.encode()
msg = Int64.encode(self._relative_offset) + Int32.encode(len(encoded))
msg += encoded
self._buffer.write(msg)
future = asyncio.Future(loop=self._loop)
self._msg_futures.append(future)
self._relative_offset += 1
return future
def decode(cls, data, bytes_to_read=None):
"""Compressed messages should pass in bytes_to_read (via message size)
otherwise, we decode from data as Int32
"""
if isinstance(data, bytes):
data = io.BytesIO(data)
if bytes_to_read is None:
bytes_to_read = Int32.decode(data)
# if FetchRequest max_bytes is smaller than the available message set
# the server returns partial data for the final message
# So create an internal buffer to avoid over-reading
raw = io.BytesIO(data.read(bytes_to_read))
items = []
while bytes_to_read:
try:
offset = Int64.decode(raw)
msg_bytes = Bytes.decode(raw)
bytes_to_read -= 8 + 4 + len(msg_bytes)
items.append((offset, len(msg_bytes), Message.decode(msg_bytes)))
except ValueError:
# PartialMessage to signal that max_bytes may be too small
items.append((None, None, PartialMessage()))
break
return items
def __init__(self, tp, batch_size, compression_type,
ttl, api_version, loop):
# Kafka 0.8/0.9 had a quirky lz4...
version_id = 0 if api_version < (0, 10) else 1
if compression_type == 'lz4' and version_id == 0:
compression_type = 'lz4-old-kafka'
if compression_type:
checker, _, _ = self._COMPRESSORS[compression_type]
assert checker(), 'Compression Libraries Not Found'
self._tp = tp
self._batch_size = batch_size
self._compression_type = compression_type
self._buffer = io.BytesIO()
self._buffer.write(Int32.encode(0)) # first 4 bytes for batch size
self._relative_offset = 0
self._loop = loop
self._ttl = ttl
self._ctime = loop.time()
self._version_id = version_id
# Waiters
# Set when messages are delivered to Kafka based on ACK setting
self._msg_futures = []
# Set when sender takes this batch
self._drain_waiter = asyncio.Future(loop=loop)
def _process_response(self, read_buffer):
assert not self._processing, 'Recursion not supported'
self._processing = True
ifr = self.in_flight_requests.popleft()
# verify send/recv correlation ids match
recv_correlation_id = Int32.decode(read_buffer)
# 0.8.2 quirk
if (self.config['api_version'] == (0, 8, 2) and
ifr.response_type is GroupCoordinatorResponse and
recv_correlation_id == 0):
raise Errors.KafkaError(
'Kafka 0.8.2 quirk -- try creating a topic first')
elif ifr.correlation_id != recv_correlation_id:
error = Errors.CorrelationIdError(
'Correlation ids do not match: sent %d, recv %d'
% (ifr.correlation_id, recv_correlation_id))
ifr.future.fail(error)
self.close()
self._processing = False
return None
# decode response
response = ifr.response_type.decode(read_buffer)
log.debug('%s Response %d: %s', self, ifr.correlation_id, response)
ifr.future.success(response)
self._processing = False
return response
def _process_response(self, read_buffer):
assert not self._processing, 'Recursion not supported'
self._processing = True
ifr = self.in_flight_requests.popleft()
# verify send/recv correlation ids match
recv_correlation_id = Int32.decode(read_buffer)
# 0.8.2 quirk
if (self.config['api_version'] == (0, 8, 2) and
ifr.response_type is GroupCoordinatorResponse[0] and
ifr.correlation_id != 0 and
recv_correlation_id == 0):
log.warning('Kafka 0.8.2 quirk -- GroupCoordinatorResponse'
' coorelation id does not match request. This'
' should go away once at least one topic has been'
' initialized on the broker')
elif ifr.correlation_id != recv_correlation_id:
error = Errors.CorrelationIdError(
'%s: Correlation ids do not match: sent %d, recv %d'
% (str(self), ifr.correlation_id, recv_correlation_id))
ifr.future.failure(error)
self.close()
self._processing = False
return None
# decode response
response = ifr.response_type.decode(read_buffer)
log.debug('%s Response %d: %s', self, ifr.correlation_id, response)
ifr.future.success(response)
self._processing = False
return response
def send_request(self, request, correlation_id=None):
"""Encode and queue a kafka api request for sending.
Arguments:
request (object): An un-encoded kafka request.
correlation_id (int, optional): Optionally specify an ID to
correlate requests with responses. If not provided, an ID will
be generated automatically.
Returns:
correlation_id
"""
log.debug('Sending request %s', request)
if correlation_id is None:
correlation_id = self._next_correlation_id()
header = RequestHeader(request,
correlation_id=correlation_id,
client_id=self._client_id)
message = b''.join([header.encode(), request.encode()])
size = Int32.encode(len(message))
data = size + message
self.bytes_to_send.append(data)
if request.expect_response():
ifr = (correlation_id, request)
self.in_flight_requests.append(ifr)
return correlation_id
def test_decode_fetch_response_partial():
encoded = b''.join([
Int32.encode(1), # Num Topics (Array)
String('utf-8').encode('foobar'),
Int32.encode(2), # Num Partitions (Array)
Int32.encode(0), # Partition id
Int16.encode(0), # Error Code
Int64.encode(1234), # Highwater offset
Int32.encode(52), # MessageSet size
Int64.encode(0), # Msg Offset
Int32.encode(18), # Msg Size
struct.pack('>i', 1474775406), # CRC
struct.pack('>bb', 0, 0), # Magic, flags
struct.pack('>i', 2), # Length of key
b'k1', # Key
struct.pack('>i', 2), # Length of value
b'v1', # Value
Int64.encode(1), # Msg Offset
struct.pack('>i', 24), # Msg Size (larger than remaining MsgSet size)
struct.pack('>i', -16383415), # CRC
struct.pack('>bb', 0, 0), # Magic, flags
struct.pack('>i', 2), # Length of key
b'k2', # Key
struct.pack('>i', 8), # Length of value
b'ar', # Value (truncated)
Int32.encode(1),
Int16.encode(0),
Int64.encode(2345),
Int32.encode(52), # MessageSet size
Int64.encode(0), # Msg Offset
Int32.encode(18), # Msg Size
struct.pack('>i', 1474775406), # CRC
struct.pack('>bb', 0, 0), # Magic, flags
struct.pack('>i', 2), # Length of key
b'k1', # Key
struct.pack('>i', 2), # Length of value
b'v1', # Value
Int64.encode(1), # Msg Offset
struct.pack('>i', 24), # Msg Size (larger than remaining MsgSet size)
struct.pack('>i', -16383415), # CRC
struct.pack('>bb', 0, 0), # Magic, flags
struct.pack('>i', 2), # Length of key
b'k2', # Key
struct.pack('>i', 8), # Length of value
b'ar', # Value (truncated)
])
resp = FetchResponse[0].decode(io.BytesIO(encoded))
assert len(resp.topics) == 1
topic, partitions = resp.topics[0]
assert topic == 'foobar'
assert len(partitions) == 2
m1 = partitions[0][3]
assert len(m1) == 2
assert m1[1] == (None, None, PartialMessage())
def server_thread():
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as ss:
ss.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
ss.bind(('127.0.0.1', PORT))
ss.listen(1)
event.set()
(cs, _) = ss.accept()
with cs:
size = Int32.decode(BytesIO(cs.recv(4, socket.MSG_WAITALL)))
id = Int32.decode(BytesIO(cs.recv(4, socket.MSG_WAITALL)))
cs.recv(size - 4, socket.MSG_WAITALL)
if need_response:
message = b''.join([Int32.encode(id), response.encode()])
cs.send(Int32.encode(len(message)) + message)
event.set()
def receive_bytes(self, data):
"""Process bytes received from the network.
Arguments:
data (bytes): any length bytes received from a network connection
to a kafka broker.
Returns:
responses (list of (correlation_id, response)): any/all completed
responses, decoded from bytes to python objects.
Raises:
KafkaProtocolError: if the bytes received could not be decoded.
CorrelationIdError: if the response does not match the request
correlation id.
"""
i = 0
n = len(data)
responses = []
while i < n:
# Not receiving is the state of reading the payload header
if not self._receiving:
bytes_to_read = min(4 - self._header.tell(), n - i)
self._header.write(data[i:i + bytes_to_read])
i += bytes_to_read
if self._header.tell() == 4:
self._header.seek(0)
nbytes = Int32.decode(self._header)
# reset buffer and switch state to receiving payload bytes
self._rbuffer = KafkaBytes(nbytes)
self._receiving = True
elif self._header.tell() > 4:
raise Errors.KafkaError(
"this should not happen - are you threading?")
if self._receiving:
total_bytes = len(self._rbuffer)
staged_bytes = self._rbuffer.tell()
bytes_to_read = min(total_bytes - staged_bytes, n - i)
self._rbuffer.write(data[i:i + bytes_to_read])
i += bytes_to_read
staged_bytes = self._rbuffer.tell()
if staged_bytes > total_bytes:
raise Errors.KafkaError(
"Receive buffer has more bytes than expected?")
if staged_bytes != total_bytes:
break
self._receiving = False
self._rbuffer.seek(0)
resp = self._process_response(self._rbuffer)
responses.append(resp)
self._reset_buffer()
return responses
def receive_bytes(self, data):
"""Process bytes received from the network.
Arguments:
data (bytes): any length bytes received from a network connection
to a kafka broker.
Returns:
responses (list of (correlation_id, response)): any/all completed
responses, decoded from bytes to python objects.
Raises:
KafkaProtocolError: if the bytes received could not be decoded.
CorrelationIdError: if the response does not match the request
correlation id.
"""
i = 0
n = len(data)
responses = []
while i < n:
# Not receiving is the state of reading the payload header
if not self._receiving:
bytes_to_read = min(4 - self._header.tell(), n - i)
self._header.write(data[i:i+bytes_to_read])
i += bytes_to_read
if self._header.tell() == 4:
self._header.seek(0)
nbytes = Int32.decode(self._header)
# reset buffer and switch state to receiving payload bytes
self._rbuffer = KafkaBytes(nbytes)
self._receiving = True
elif self._header.tell() > 4:
raise Errors.KafkaError('this should not happen - are you threading?')
if self._receiving:
total_bytes = len(self._rbuffer)
staged_bytes = self._rbuffer.tell()
bytes_to_read = min(total_bytes - staged_bytes, n - i)
self._rbuffer.write(data[i:i+bytes_to_read])
i += bytes_to_read
staged_bytes = self._rbuffer.tell()
if staged_bytes > total_bytes:
raise Errors.KafkaError('Receive buffer has more bytes than expected?')
if staged_bytes != total_bytes:
break
self._receiving = False
self._rbuffer.seek(0)
resp = self._process_response(self._rbuffer)
responses.append(resp)
self._reset_buffer()
return responses
def send(request, wait_response=True):
global correlation_id
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.connect(('127.0.0.1', PORT))
header = RequestHeader(
request=request,
correlation_id=correlation_id,
client_id='test'
)
correlation_id += 1
message = b''.join([header.encode(), request.encode()])
size = Int32.encode(len(message))
s.send(size)
s.send(message)
if wait_response:
size = s.recv(4, socket.MSG_WAITALL)
size = Int32.decode(BytesIO(size))
s.recv(size, socket.MSG_WAITALL)
def __init__(self, magic, batch_size, compression_type):
if compression_type == 'lz4' and magic == 0:
compression_type = 'lz4-old-kafka'
if compression_type:
checker, _, _ = self._COMPRESSORS[compression_type]
assert checker(), 'Compression Libraries Not Found'
self._magic = magic
self._batch_size = batch_size
self._compression_type = compression_type
self._relative_offset = 0
self._buffer = io.BytesIO()
self._buffer.write(Int32.encode(0)) # first 4 bytes for batch size
self._closed = False
def _process_response(self, read_buffer):
assert not self._processing, 'Recursion not supported'
self._processing = True
ifr = self.in_flight_requests.popleft()
if self._sensors:
self._sensors.request_time.record(
(time.time() - ifr.timestamp) * 1000)
# verify send/recv correlation ids match
recv_correlation_id = Int32.decode(read_buffer)
# 0.8.2 quirk
if (self.config['api_version'] == (0, 8, 2)
and ifr.response_type is GroupCoordinatorResponse[0]
and ifr.correlation_id != 0 and recv_correlation_id == 0):
log.warning('Kafka 0.8.2 quirk -- GroupCoordinatorResponse'
' Correlation ID does not match request. This'
' should go away once at least one topic has been'
' initialized on the broker.')
elif ifr.correlation_id != recv_correlation_id:
error = Errors.CorrelationIdError(
'%s: Correlation IDs do not match: sent %d, recv %d' %
(self, ifr.correlation_id, recv_correlation_id))
ifr.future.failure(error)
self.close(error)
self._processing = False
return None
# decode response
try:
response = ifr.response_type.decode(read_buffer)
except ValueError:
read_buffer.seek(0)
buf = read_buffer.read()
log.error(
'%s Response %d [ResponseType: %s Request: %s]:'
' Unable to decode %d-byte buffer: %r',
self, ifr.correlation_id, ifr.response_type, ifr.request,
len(buf), buf)
error = Errors.UnknownError('Unable to decode response')
ifr.future.failure(error)
self.close(error)
self._processing = False
return None
log.debug('%s Response %d: %s', self, ifr.correlation_id, response)
ifr.future.success(response)
self._processing = False
return response
def drain_ready(self):
"""Compress batch to be ready for send"""
memview = self._buffer.getbuffer()
self._drain_waiter.set_result(None)
if self._compression_type:
_, compressor, attrs = self._COMPRESSORS[self._compression_type]
msg = Message(compressor(memview[4:].tobytes()), attributes=attrs)
encoded = msg.encode()
# if compressed message is longer than original
# we should send it as is (not compressed)
header_size = 16 # 4(all size) + 8(offset) + 4(compressed size)
if len(encoded) + header_size < len(memview):
# write compressed message set (with header) to buffer
# using memory view (for avoid memory copying)
memview[:4] = Int32.encode(len(encoded) + 12)
memview[4:12] = Int64.encode(0) # offset 0
memview[12:16] = Int32.encode(len(encoded))
memview[16:16+len(encoded)] = encoded
self._buffer.seek(0)
return
# update batch size (first 4 bytes of buffer)
memview[:4] = Int32.encode(self._buffer.tell()-4)
self._buffer.seek(0)
def _process_response(self, read_buffer):
assert not self._processing, 'Recursion not supported'
self._processing = True
ifr = self.in_flight_requests.popleft()
if self._sensors:
self._sensors.request_time.record((time.time() - ifr.timestamp) * 1000)
# verify send/recv correlation ids match
recv_correlation_id = Int32.decode(read_buffer)
# 0.8.2 quirk
if (self.config['api_version'] == (0, 8, 2) and
ifr.response_type is GroupCoordinatorResponse[0] and
ifr.correlation_id != 0 and
recv_correlation_id == 0):
log.warning('Kafka 0.8.2 quirk -- GroupCoordinatorResponse'
' coorelation id does not match request. This'
' should go away once at least one topic has been'
' initialized on the broker')
elif ifr.correlation_id != recv_correlation_id:
error = Errors.CorrelationIdError(
'%s: Correlation ids do not match: sent %d, recv %d'
% (str(self), ifr.correlation_id, recv_correlation_id))
ifr.future.failure(error)
self.close()
self._processing = False
return None
# decode response
try:
response = ifr.response_type.decode(read_buffer)
except ValueError:
read_buffer.seek(0)
buf = read_buffer.read()
log.error('%s Response %d [ResponseType: %s Request: %s]:'
' Unable to decode %d-byte buffer: %r', self,
ifr.correlation_id, ifr.response_type,
ifr.request, len(buf), buf)
ifr.future.failure(Errors.UnknownError('Unable to decode response'))
self.close()
self._processing = False
return None
log.debug('%s Response %d: %s', self, ifr.correlation_id, response)
ifr.future.success(response)
self._processing = False
return response
def send(self, request, expect_response=True):
"""send request, return Future()
Can block on network if request is larger than send_buffer_bytes
"""
future = Future()
if self.connecting():
return future.failure(Errors.NodeNotReadyError(str(self)))
elif not self.connected():
return future.failure(Errors.ConnectionError(str(self)))
elif not self.can_send_more():
return future.failure(Errors.TooManyInFlightRequests(str(self)))
correlation_id = self._next_correlation_id()
header = RequestHeader(request,
correlation_id=correlation_id,
client_id=self.config['client_id'])
message = b''.join([header.encode(), request.encode()])
size = Int32.encode(len(message))
try:
# In the future we might manage an internal write buffer
# and send bytes asynchronously. For now, just block
# sending each request payload
self._sock.setblocking(True)
for data in (size, message):
total_sent = 0
while total_sent < len(data):
sent_bytes = self._sock.send(data[total_sent:])
total_sent += sent_bytes
assert total_sent == len(data)
self._sock.setblocking(False)
except (AssertionError, ConnectionError) as e:
log.exception("Error sending %s to %s", request, self)
error = Errors.ConnectionError("%s: %s" % (str(self), e))
self.close(error=error)
return future.failure(error)
log.debug('%s Request %d: %s', self, correlation_id, request)
if expect_response:
ifr = InFlightRequest(request=request,
correlation_id=correlation_id,
response_type=request.RESPONSE_TYPE,
future=future,
timestamp=time.time())
self.in_flight_requests.append(ifr)
else:
future.success(None)
return future
def send(self, request, expect_response=True):
"""send request, return Future()
Can block on network if request is larger than send_buffer_bytes
"""
future = Future()
if self.connecting():
return future.failure(Errors.NodeNotReadyError(str(self)))
elif not self.connected():
return future.failure(Errors.ConnectionError(str(self)))
elif not self.can_send_more():
return future.failure(Errors.TooManyInFlightRequests(str(self)))
correlation_id = self._next_correlation_id()
header = RequestHeader(request,
correlation_id=correlation_id,
client_id=self.config['client_id'])
message = b''.join([header.encode(), request.encode()])
size = Int32.encode(len(message))
try:
# In the future we might manage an internal write buffer
# and send bytes asynchronously. For now, just block
# sending each request payload
self._sock.setblocking(True)
for data in (size, message):
total_sent = 0
while total_sent < len(data):
sent_bytes = self._sock.send(data[total_sent:])
total_sent += sent_bytes
assert total_sent == len(data)
self._sock.setblocking(False)
except (AssertionError, ConnectionError) as e:
log.exception("Error sending %s to %s", request, self)
error = Errors.ConnectionError("%s: %s" % (str(self), e))
self.close(error=error)
return future.failure(error)
log.debug('%s Request %d: %s', self, correlation_id, request)
if expect_response:
ifr = InFlightRequest(request=request,
correlation_id=correlation_id,
response_type=request.RESPONSE_TYPE,
future=future,
timestamp=time.time())
self.in_flight_requests.append(ifr)
else:
future.success(None)
return future
def _process_response(self, read_buffer):
recv_correlation_id = Int32.decode(read_buffer)
log.debug("Received correlation id: %d", recv_correlation_id)
if not self.in_flight_requests:
raise Errors.CorrelationIdError(
"No in-flight-request found for server response"
" with correlation ID %d" % (recv_correlation_id, ))
(correlation_id, request) = self.in_flight_requests.popleft()
# 0.8.2 quirk
if (recv_correlation_id == 0 and correlation_id != 0
and request.RESPONSE_TYPE is GroupCoordinatorResponse[0] and
(self._api_version == (0, 8, 2) or self._api_version is None)):
log.warning("Kafka 0.8.2 quirk -- GroupCoordinatorResponse"
" Correlation ID does not match request. This"
" should go away once at least one topic has been"
" initialized on the broker.")
elif correlation_id != recv_correlation_id:
# return or raise?
raise Errors.CorrelationIdError(
"Correlation IDs do not match: sent %d, recv %d" %
(correlation_id, recv_correlation_id))
# decode response
log.debug("Processing response %s", request.RESPONSE_TYPE.__name__)
try:
response = request.RESPONSE_TYPE.decode(read_buffer)
except ValueError:
read_buffer.seek(0)
buf = read_buffer.read()
log.error(
"Response %d [ResponseType: %s Request: %s]:"
" Unable to decode %d-byte buffer: %r",
correlation_id,
request.RESPONSE_TYPE,
request,
len(buf),
buf,
)
raise Errors.KafkaProtocolError("Unable to decode response")
return (correlation_id, response)
def close(self):
"""Close the batch to further updates.
Closing the batch before submitting to the producer ensures that no
messages are added via the ``producer.send()`` interface. To gracefully
support both the batch and individual message interfaces, leave the
batch open. For complete control over the batch's contents, close
before submission. Closing a batch has no effect on when it's sent to
the broker.
A batch may not be reopened after it's closed.
"""
if self._closed:
return
self._closed = True
data = self._builder.build()
self._buffer = io.BytesIO(Int32.encode(len(data)) + data)
del self._builder
def _try_authenticate_gssapi(self, future):
gssapi_name = gssapi.Name(
self.config['sasl_kerberos_service_name'] + '@' + self.hostname,
name_type=gssapi.NameType.hostbased_service
).canonicalize(gssapi.MechType.kerberos)
log.debug('%s: GSSAPI name: %s', self, gssapi_name)
# Exchange tokens until authentication either succeeds or fails
client_ctx = gssapi.SecurityContext(name=gssapi_name, usage='initiate')
received_token = None
try:
while not client_ctx.complete:
# calculate an output token from kafka token (or None if first iteration)
output_token = client_ctx.step(received_token)
if output_token is None:
continue
# pass output token to kafka
try:
msg = output_token
size = Int32.encode(len(msg))
self._send_bytes_blocking(size + msg)
# The server will send a token back. Processing of this token either
# establishes a security context, or it needs further token exchange.
# The gssapi will be able to identify the needed next step.
# The connection is closed on failure.
header = self._recv_bytes_blocking(4)
(token_size,) = struct.unpack('>i', header)
received_token = self._recv_bytes_blocking(token_size)
except ConnectionError as e:
log.exception("%s: Error receiving reply from server", self)
error = Errors.ConnectionError("%s: %s" % (self, e))
self.close(error=error)
return future.failure(error)
except Exception as e:
return future.failure(e)
log.info('%s: Authenticated as %s via GSSAPI', self, gssapi_name)
return future.success(True)
def encode(cls, items, prepend_size=True):
# RecordAccumulator encodes messagesets internally
if isinstance(items, (io.BytesIO, KafkaBytes)):
size = Int32.decode(items)
if prepend_size:
# rewind and return all the bytes
items.seek(items.tell() - 4)
size += 4
return items.read(size)
encoded_values = []
for (offset, message) in items:
encoded_values.append(Int64.encode(offset))
encoded_values.append(Bytes.encode(message))
encoded = b''.join(encoded_values)
if prepend_size:
return Bytes.encode(encoded)
else:
return encoded
def encode(cls, items, prepend_size=True):
# RecordAccumulator encodes messagesets internally
if isinstance(items, (io.BytesIO, KafkaBytes)):
size = Int32.decode(items)
if prepend_size:
# rewind and return all the bytes
items.seek(items.tell() - 4)
size += 4
return items.read(size)
encoded_values = []
for (offset, message) in items:
encoded_values.append(Int64.encode(offset))
encoded_values.append(Bytes.encode(message))
encoded = b''.join(encoded_values)
if prepend_size:
return Bytes.encode(encoded)
else:
return encoded
def _send(self, request):
assert self.state in (ConnectionStates.AUTHENTICATING,
ConnectionStates.CONNECTED)
future = Future()
correlation_id = self._next_correlation_id()
header = RequestHeader(request,
correlation_id=correlation_id,
client_id=self.config['client_id'])
message = b''.join([header.encode(), request.encode()])
size = Int32.encode(len(message))
data = size + message
try:
# In the future we might manage an internal write buffer
# and send bytes asynchronously. For now, just block
# sending each request payload
self._sock.setblocking(True)
total_sent = 0
while total_sent < len(data):
sent_bytes = self._sock.send(data[total_sent:])
total_sent += sent_bytes
assert total_sent == len(data)
if self._sensors:
self._sensors.bytes_sent.record(total_sent)
self._sock.setblocking(False)
except (AssertionError, ConnectionError) as e:
log.exception("Error sending %s to %s", request, self)
error = Errors.ConnectionError("%s: %s" % (self, e))
self.close(error=error)
return future.failure(error)
log.debug('%s Request %d: %s', self, correlation_id, request)
if request.expect_response():
ifr = InFlightRequest(request=request,
correlation_id=correlation_id,
response_type=request.RESPONSE_TYPE,
future=future,
timestamp=time.time())
self.in_flight_requests.append(ifr)
else:
future.success(None)
return future
def _process_response(self, read_buffer):
recv_correlation_id = Int32.decode(read_buffer)
log.debug('Received correlation id: %d', recv_correlation_id)
if not self.in_flight_requests:
raise Errors.CorrelationIdError(
'No in-flight-request found for server response'
' with correlation ID %d'
% recv_correlation_id)
(correlation_id, request) = self.in_flight_requests.popleft()
# 0.8.2 quirk
if (self._api_version == (0, 8, 2) and
request.RESPONSE_TYPE is GroupCoordinatorResponse[0] and
correlation_id != 0 and
recv_correlation_id == 0):
log.warning('Kafka 0.8.2 quirk -- GroupCoordinatorResponse'
' Correlation ID does not match request. This'
' should go away once at least one topic has been'
' initialized on the broker.')
elif correlation_id != recv_correlation_id:
# return or raise?
raise Errors.CorrelationIdError(
'Correlation IDs do not match: sent %d, recv %d'
% (correlation_id, recv_correlation_id))
# decode response
log.debug('Processing response %s', request.RESPONSE_TYPE.__name__)
try:
response = request.RESPONSE_TYPE.decode(read_buffer)
except ValueError:
read_buffer.seek(0)
buf = read_buffer.read()
log.error('Response %d [ResponseType: %s Request: %s]:'
' Unable to decode %d-byte buffer: %r',
correlation_id, request.RESPONSE_TYPE,
request, len(buf), buf)
raise Errors.KafkaProtocolError('Unable to decode response')
return (correlation_id, response)
def append(self, *, timestamp, key, value):
if not self._has_room_for(key, value):
return 0
# `.encode()` is a weak method for some reason, so we need to save
# reference before calling it.
if self._magic == 0:
msg_inst = Message(value, key=key, magic=self._magic)
else:
msg_inst = Message(value,
key=key,
magic=self._magic,
timestamp=timestamp)
encoded = msg_inst.encode()
msg = Int64.encode(self._relative_offset) + Int32.encode(len(encoded))
msg += encoded
actual_size = self._buffer.write(msg)
self._relative_offset += 1
return actual_size
def append(self, key, value):
"""Append message (key and value) to batch
Returns:
None if batch is full
or
asyncio.Future that will resolved when message is delivered
"""
if self._is_full(key, value):
return None
encoded = Message(value, key=key).encode()
msg = Int64.encode(self._relative_offset) + Int32.encode(len(encoded))
msg += encoded
self._buffer.write(msg)
future = asyncio.Future(loop=self._loop)
self._msg_futures.append(future)
self._relative_offset += 1
return future
def append(self, key, value):
"""Append message (key and value) to batch
Returns:
None if batch is full
or
asyncio.Future that will resolved when message is delivered
"""
if self._is_full(key, value):
return None
encoded = Message(value, key=key).encode()
msg = Int64.encode(self._relative_offset) + Int32.encode(len(encoded))
msg += encoded
self._buffer.write(msg)
future = asyncio.Future(loop=self._loop)
self._msg_futures.append(future)
self._relative_offset += 1
return future
def __init__(self, tp, batch_size, compression_type, ttl, loop):
if compression_type:
checker, _, _ = self._COMPRESSORS[compression_type]
assert checker(), 'Compression Libraries Not Found'
self._tp = tp
self._batch_size = batch_size
self._compression_type = compression_type
self._buffer = io.BytesIO()
self._buffer.write(Int32.encode(0)) # first 4 bytes for batch size
self._relative_offset = 0
self._loop = loop
self._ttl = ttl
self._ctime = loop.time()
# Waiters
# Set when messages are delivered to Kafka based on ACK setting
self._msg_futures = []
# Set when sender takes this batch
self._drain_waiter = asyncio.Future(loop=loop)
def __init__(self, tp, batch_size, compression_type, ttl, loop):
if compression_type:
checker, _, _ = self._COMPRESSORS[compression_type]
assert checker(), 'Compression Libraries Not Found'
self._tp = tp
self._batch_size = batch_size
self._compression_type = compression_type
self._buffer = io.BytesIO()
self._buffer.write(Int32.encode(0)) # first 4 bytes for batch size
self._relative_offset = 0
self._loop = loop
self._ttl = ttl
self._ctime = loop.time()
# Waiters
# Set when messages are delivered to Kafka based on ACK setting
self._msg_futures = []
# Set when sender takes this batch
self._drain_waiter = asyncio.Future(loop=loop)
def _try_authenticate_plain(self, future):
if self.config['security_protocol'] == 'SASL_PLAINTEXT':
log.warning('%s: Sending username and password in the clear',
str(self))
data = b''
try:
self._sock.setblocking(True)
# Send PLAIN credentials per RFC-4616
msg = bytes('\0'.join([
self.config['sasl_plain_username'],
self.config['sasl_plain_username'],
self.config['sasl_plain_password']
]).encode('utf-8'))
size = Int32.encode(len(msg))
self._sock.sendall(size + msg)
# The server will send a zero sized message (that is Int32(0)) on success.
# The connection is closed on failure
while len(data) < 4:
fragment = self._sock.recv(4 - len(data))
if not fragment:
log.error('%s: Authentication failed for user %s', self,
self.config['sasl_plain_username'])
error = Errors.AuthenticationFailedError(
'Authentication failed for user {0}'.format(
self.config['sasl_plain_username']))
future.failure(error)
raise error
data += fragment
self._sock.setblocking(False)
except (AssertionError, ConnectionError) as e:
log.exception("%s: Error receiving reply from server", self)
error = Errors.ConnectionError("%s: %s" % (str(self), e))
future.failure(error)
self.close(error=error)
if data != b'\x00\x00\x00\x00':
return future.failure(Errors.AuthenticationFailedError())
return future.success(True)
def _send(self, request, expect_response=True):
future = Future()
correlation_id = self._next_correlation_id()
header = RequestHeader(request,
correlation_id=correlation_id,
client_id=self.config['client_id'])
message = b''.join([header.encode(), request.encode()])
size = Int32.encode(len(message))
data = size + message
try:
# In the future we might manage an internal write buffer
# and send bytes asynchronously. For now, just block
# sending each request payload
self._sock.setblocking(True)
total_sent = 0
while total_sent < len(data):
sent_bytes = self._sock.send(data[total_sent:])
total_sent += sent_bytes
assert total_sent == len(data)
if self._sensors:
self._sensors.bytes_sent.record(total_sent)
self._sock.setblocking(False)
except (AssertionError, ConnectionError) as e:
log.exception("Error sending %s to %s", request, self)
error = Errors.ConnectionError("%s: %s" % (str(self), e))
self.close(error=error)
return future.failure(error)
log.debug('%s Request %d: %s', self, correlation_id, request)
if expect_response:
ifr = InFlightRequest(request=request,
correlation_id=correlation_id,
response_type=request.RESPONSE_TYPE,
future=future,
timestamp=time.time())
self.in_flight_requests.append(ifr)
else:
future.success(None)
return future
def _try_authenticate_plain(self, future):
if self.config['security_protocol'] == 'SASL_PLAINTEXT':
log.warning('%s: Sending username and password in the clear', str(self))
data = b''
try:
self._sock.setblocking(True)
# Send PLAIN credentials per RFC-4616
msg = bytes('\0'.join([self.config['sasl_plain_username'],
self.config['sasl_plain_username'],
self.config['sasl_plain_password']]).encode('utf-8'))
size = Int32.encode(len(msg))
self._sock.sendall(size + msg)
# The server will send a zero sized message (that is Int32(0)) on success.
# The connection is closed on failure
while len(data) < 4:
fragment = self._sock.recv(4 - len(data))
if not fragment:
log.error('%s: Authentication failed for user %s', self, self.config['sasl_plain_username'])
error = Errors.AuthenticationFailedError(
'Authentication failed for user {0}'.format(
self.config['sasl_plain_username']))
future.failure(error)
raise error
data += fragment
self._sock.setblocking(False)
except (AssertionError, ConnectionError) as e:
log.exception("%s: Error receiving reply from server", self)
error = Errors.ConnectionError("%s: %s" % (str(self), e))
future.failure(error)
self.close(error=error)
if data != b'\x00\x00\x00\x00':
return future.failure(Errors.AuthenticationFailedError())
return future.success(True)
def recv(self, timeout=0):
"""Non-blocking network receive.
Return response if available
"""
assert not self._processing, 'Recursion not supported'
if not self.connected():
log.warning('%s cannot recv: socket not connected', self)
# If requests are pending, we should close the socket and
# fail all the pending request futures
if self.in_flight_requests:
self.close()
return None
elif not self.in_flight_requests:
log.warning('%s: No in-flight-requests to recv', self)
return None
elif self._requests_timed_out():
log.warning('%s timed out after %s ms. Closing connection.',
self, self.config['request_timeout_ms'])
self.close(error=Errors.RequestTimedOutError(
'Request timed out after %s ms' %
self.config['request_timeout_ms']))
return None
readable, _, _ = select([self._sock], [], [], timeout)
if not readable:
return None
# Not receiving is the state of reading the payload header
if not self._receiving:
try:
# An extremely small, but non-zero, probability that there are
# more than 0 but not yet 4 bytes available to read
self._rbuffer.write(self._sock.recv(4 - self._rbuffer.tell()))
except ConnectionError as e:
if six.PY2 and e.errno == errno.EWOULDBLOCK:
# This shouldn't happen after selecting above
# but just in case
return None
log.exception('%s: Error receiving 4-byte payload header -'
' closing socket', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
if self._rbuffer.tell() == 4:
self._rbuffer.seek(0)
self._next_payload_bytes = Int32.decode(self._rbuffer)
# reset buffer and switch state to receiving payload bytes
self._rbuffer.seek(0)
self._rbuffer.truncate()
self._receiving = True
elif self._rbuffer.tell() > 4:
raise Errors.KafkaError('this should not happen - are you threading?')
if self._receiving:
staged_bytes = self._rbuffer.tell()
try:
self._rbuffer.write(self._sock.recv(self._next_payload_bytes - staged_bytes))
except ConnectionError as e:
# Extremely small chance that we have exactly 4 bytes for a
# header, but nothing to read in the body yet
if six.PY2 and e.errno == errno.EWOULDBLOCK:
return None
log.exception('%s: Error in recv', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
staged_bytes = self._rbuffer.tell()
if staged_bytes > self._next_payload_bytes:
self.close(error=Errors.KafkaError('Receive buffer has more bytes than expected?'))
if staged_bytes != self._next_payload_bytes:
return None
self._receiving = False
self._next_payload_bytes = 0
self._rbuffer.seek(0)
response = self._process_response(self._rbuffer)
self._rbuffer.seek(0)
self._rbuffer.truncate()
return response
def _recv(self):
# Not receiving is the state of reading the payload header
if not self._receiving:
try:
bytes_to_read = 4 - self._rbuffer.tell()
data = self._sock.recv(bytes_to_read)
# We expect socket.recv to raise an exception if there is not
# enough data to read the full bytes_to_read
# but if the socket is disconnected, we will get empty data
# without an exception raised
if not data:
log.error('%s: socket disconnected', self)
self.close(error=Errors.ConnectionError('socket disconnected'))
return None
self._rbuffer.write(data)
except ssl.SSLWantReadError:
return None
except ConnectionError as e:
if six.PY2 and e.errno == errno.EWOULDBLOCK:
return None
log.exception('%s: Error receiving 4-byte payload header -'
' closing socket', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
if self._rbuffer.tell() == 4:
self._rbuffer.seek(0)
self._next_payload_bytes = Int32.decode(self._rbuffer)
# reset buffer and switch state to receiving payload bytes
self._rbuffer.seek(0)
self._rbuffer.truncate()
self._receiving = True
elif self._rbuffer.tell() > 4:
raise Errors.KafkaError('this should not happen - are you threading?')
if self._receiving:
staged_bytes = self._rbuffer.tell()
try:
bytes_to_read = self._next_payload_bytes - staged_bytes
data = self._sock.recv(bytes_to_read)
# We expect socket.recv to raise an exception if there is not
# enough data to read the full bytes_to_read
# but if the socket is disconnected, we will get empty data
# without an exception raised
if bytes_to_read and not data:
log.error('%s: socket disconnected', self)
self.close(error=Errors.ConnectionError('socket disconnected'))
return None
self._rbuffer.write(data)
except ssl.SSLWantReadError:
return None
except ConnectionError as e:
# Extremely small chance that we have exactly 4 bytes for a
# header, but nothing to read in the body yet
if six.PY2 and e.errno == errno.EWOULDBLOCK:
return None
log.exception('%s: Error in recv', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
staged_bytes = self._rbuffer.tell()
if staged_bytes > self._next_payload_bytes:
self.close(error=Errors.KafkaError('Receive buffer has more bytes than expected?'))
if staged_bytes != self._next_payload_bytes:
return None
self._receiving = False
self._next_payload_bytes = 0
if self._sensors:
self._sensors.bytes_received.record(4 + self._rbuffer.tell())
self._rbuffer.seek(0)
response = self._process_response(self._rbuffer)
self._rbuffer.seek(0)
self._rbuffer.truncate()
return response
def recv(self):
"""Non-blocking network receive.
Return response if available
"""
assert not self._processing, 'Recursion not supported'
if not self.connected():
log.warning('%s cannot recv: socket not connected', self)
# If requests are pending, we should close the socket and
# fail all the pending request futures
if self.in_flight_requests:
self.close()
return None
elif not self.in_flight_requests:
log.warning('%s: No in-flight-requests to recv', self)
return None
elif self._requests_timed_out():
log.warning('%s timed out after %s ms. Closing connection.',
self, self.config['request_timeout_ms'])
self.close(error=Errors.RequestTimedOutError(
'Request timed out after %s ms' %
self.config['request_timeout_ms']))
return None
# Not receiving is the state of reading the payload header
if not self._receiving:
try:
bytes_to_read = 4 - self._rbuffer.tell()
data = self._sock.recv(bytes_to_read)
# We expect socket.recv to raise an exception if there is not
# enough data to read the full bytes_to_read
# but if the socket is disconnected, we will get empty data
# without an exception raised
if not data:
log.error('%s: socket disconnected', self)
self.close(error=Errors.ConnectionError('socket disconnected'))
return None
self._rbuffer.write(data)
except ssl.SSLWantReadError:
return None
except ConnectionError as e:
if six.PY2 and e.errno == errno.EWOULDBLOCK:
return None
log.exception('%s: Error receiving 4-byte payload header -'
' closing socket', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
if self._rbuffer.tell() == 4:
self._rbuffer.seek(0)
self._next_payload_bytes = Int32.decode(self._rbuffer)
# reset buffer and switch state to receiving payload bytes
self._rbuffer.seek(0)
self._rbuffer.truncate()
self._receiving = True
elif self._rbuffer.tell() > 4:
raise Errors.KafkaError('this should not happen - are you threading?')
if self._receiving:
staged_bytes = self._rbuffer.tell()
try:
bytes_to_read = self._next_payload_bytes - staged_bytes
data = self._sock.recv(bytes_to_read)
# We expect socket.recv to raise an exception if there is not
# enough data to read the full bytes_to_read
# but if the socket is disconnected, we will get empty data
# without an exception raised
if not data:
log.error('%s: socket disconnected', self)
self.close(error=Errors.ConnectionError('socket disconnected'))
return None
self._rbuffer.write(data)
except ssl.SSLWantReadError:
return None
except ConnectionError as e:
# Extremely small chance that we have exactly 4 bytes for a
# header, but nothing to read in the body yet
if six.PY2 and e.errno == errno.EWOULDBLOCK:
return None
log.exception('%s: Error in recv', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
staged_bytes = self._rbuffer.tell()
if staged_bytes > self._next_payload_bytes:
self.close(error=Errors.KafkaError('Receive buffer has more bytes than expected?'))
if staged_bytes != self._next_payload_bytes:
return None
self._receiving = False
self._next_payload_bytes = 0
self._rbuffer.seek(0)
response = self._process_response(self._rbuffer)
self._rbuffer.seek(0)
self._rbuffer.truncate()
return response
def recv(self, timeout=0):
"""Non-blocking network receive.
Return response if available
"""
assert not self._processing, 'Recursion not supported'
if not self.connected():
log.warning('%s cannot recv: socket not connected', self)
# If requests are pending, we should close the socket and
# fail all the pending request futures
if self.in_flight_requests:
self.close()
return None
elif not self.in_flight_requests:
log.warning('%s: No in-flight-requests to recv', self)
return None
elif self._requests_timed_out():
log.warning('%s timed out after %s ms. Closing connection.',
self, self.config['request_timeout_ms'])
self.close(error=Errors.RequestTimedOutError(
'Request timed out after %s ms' %
self.config['request_timeout_ms']))
return None
readable, _, _ = select([self._sock], [], [], timeout)
if not readable:
return None
# Not receiving is the state of reading the payload header
if not self._receiving:
try:
# An extremely small, but non-zero, probability that there are
# more than 0 but not yet 4 bytes available to read
self._rbuffer.write(self._sock.recv(4 - self._rbuffer.tell()))
except ConnectionError as e:
if six.PY2 and e.errno == errno.EWOULDBLOCK:
# This shouldn't happen after selecting above
# but just in case
return None
log.exception('%s: Error receiving 4-byte payload header -'
' closing socket', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
if self._rbuffer.tell() == 4:
self._rbuffer.seek(0)
self._next_payload_bytes = Int32.decode(self._rbuffer)
# reset buffer and switch state to receiving payload bytes
self._rbuffer.seek(0)
self._rbuffer.truncate()
self._receiving = True
elif self._rbuffer.tell() > 4:
raise Errors.KafkaError('this should not happen - are you threading?')
if self._receiving:
staged_bytes = self._rbuffer.tell()
try:
self._rbuffer.write(self._sock.recv(self._next_payload_bytes - staged_bytes))
except ConnectionError as e:
# Extremely small chance that we have exactly 4 bytes for a
# header, but nothing to read in the body yet
if six.PY2 and e.errno == errno.EWOULDBLOCK:
return None
log.exception('%s: Error in recv', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
staged_bytes = self._rbuffer.tell()
if staged_bytes > self._next_payload_bytes:
self.close(error=Errors.KafkaError('Receive buffer has more bytes than expected?'))
if staged_bytes != self._next_payload_bytes:
return None
self._receiving = False
self._next_payload_bytes = 0
self._rbuffer.seek(0)
response = self._process_response(self._rbuffer)
self._rbuffer.seek(0)
self._rbuffer.truncate()
return response
def recv(self):
"""Non-blocking network receive.
Return response if available
"""
assert not self._processing, 'Recursion not supported'
if not self.connected():
log.warning('%s cannot recv: socket not connected', self)
# If requests are pending, we should close the socket and
# fail all the pending request futures
if self.in_flight_requests:
self.close()
return None
elif not self.in_flight_requests:
log.warning('%s: No in-flight-requests to recv', self)
return None
elif self._requests_timed_out():
log.warning('%s timed out after %s ms. Closing connection.',
self, self.config['request_timeout_ms'])
self.close(error=Errors.RequestTimedOutError(
'Request timed out after %s ms' %
self.config['request_timeout_ms']))
return None
# Not receiving is the state of reading the payload header
if not self._receiving:
try:
bytes_to_read = 4 - self._rbuffer.tell()
data = self._sock.recv(bytes_to_read)
# We expect socket.recv to raise an exception if there is not
# enough data to read the full bytes_to_read
# but if the socket is disconnected, we will get empty data
# without an exception raised
if not data:
log.error('%s: socket disconnected', self)
self.close(error=Errors.ConnectionError('socket disconnected'))
return None
self._rbuffer.write(data)
except ssl.SSLWantReadError:
return None
except ConnectionError as e:
if six.PY2 and e.errno == errno.EWOULDBLOCK:
return None
log.exception('%s: Error receiving 4-byte payload header -'
' closing socket', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
if self._rbuffer.tell() == 4:
self._rbuffer.seek(0)
self._next_payload_bytes = Int32.decode(self._rbuffer)
# reset buffer and switch state to receiving payload bytes
self._rbuffer.seek(0)
self._rbuffer.truncate()
self._receiving = True
elif self._rbuffer.tell() > 4:
raise Errors.KafkaError('this should not happen - are you threading?')
if self._receiving:
staged_bytes = self._rbuffer.tell()
try:
bytes_to_read = self._next_payload_bytes - staged_bytes
data = self._sock.recv(bytes_to_read)
# We expect socket.recv to raise an exception if there is not
# enough data to read the full bytes_to_read
# but if the socket is disconnected, we will get empty data
# without an exception raised
if not data:
log.error('%s: socket disconnected', self)
self.close(error=Errors.ConnectionError('socket disconnected'))
return None
self._rbuffer.write(data)
except ssl.SSLWantReadError:
return None
except ConnectionError as e:
# Extremely small chance that we have exactly 4 bytes for a
# header, but nothing to read in the body yet
if six.PY2 and e.errno == errno.EWOULDBLOCK:
return None
log.exception('%s: Error in recv', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
staged_bytes = self._rbuffer.tell()
if staged_bytes > self._next_payload_bytes:
self.close(error=Errors.KafkaError('Receive buffer has more bytes than expected?'))
if staged_bytes != self._next_payload_bytes:
return None
self._receiving = False
self._next_payload_bytes = 0
self._rbuffer.seek(0)
response = self._process_response(self._rbuffer)
self._rbuffer.seek(0)
self._rbuffer.truncate()
return response
def _recv(self):
# Not receiving is the state of reading the payload header
if not self._receiving:
try:
bytes_to_read = 4 - self._rbuffer.tell()
data = self._sock.recv(bytes_to_read)
# We expect socket.recv to raise an exception if there is not
# enough data to read the full bytes_to_read
# but if the socket is disconnected, we will get empty data
# without an exception raised
if not data:
log.error('%s: socket disconnected', self)
self.close(error=Errors.ConnectionError('socket disconnected'))
return None
self._rbuffer.write(data)
except ssl.SSLWantReadError:
return None
except ConnectionError as e:
if six.PY2 and e.errno == errno.EWOULDBLOCK:
return None
log.exception('%s: Error receiving 4-byte payload header -'
' closing socket', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
if self._rbuffer.tell() == 4:
self._rbuffer.seek(0)
self._next_payload_bytes = Int32.decode(self._rbuffer)
# reset buffer and switch state to receiving payload bytes
self._rbuffer.seek(0)
self._rbuffer.truncate()
self._receiving = True
elif self._rbuffer.tell() > 4:
raise Errors.KafkaError('this should not happen - are you threading?')
if self._receiving:
staged_bytes = self._rbuffer.tell()
try:
bytes_to_read = self._next_payload_bytes - staged_bytes
data = self._sock.recv(bytes_to_read)
# We expect socket.recv to raise an exception if there is not
# enough data to read the full bytes_to_read
# but if the socket is disconnected, we will get empty data
# without an exception raised
if bytes_to_read and not data:
log.error('%s: socket disconnected', self)
self.close(error=Errors.ConnectionError('socket disconnected'))
return None
self._rbuffer.write(data)
except ssl.SSLWantReadError:
return None
except ConnectionError as e:
# Extremely small chance that we have exactly 4 bytes for a
# header, but nothing to read in the body yet
if six.PY2 and e.errno == errno.EWOULDBLOCK:
return None
log.exception('%s: Error in recv', self)
self.close(error=Errors.ConnectionError(e))
return None
except BlockingIOError:
if six.PY3:
return None
raise
staged_bytes = self._rbuffer.tell()
if staged_bytes > self._next_payload_bytes:
self.close(error=Errors.KafkaError('Receive buffer has more bytes than expected?'))
if staged_bytes != self._next_payload_bytes:
return None
self._receiving = False
self._next_payload_bytes = 0
if self._sensors:
self._sensors.bytes_received.record(4 + self._rbuffer.tell())
self._rbuffer.seek(0)
response = self._process_response(self._rbuffer)
self._rbuffer.seek(0)
self._rbuffer.truncate()
return response
