import collections import logging import struct import kafka.errors as Errors from kafka.protocol.metadata import FindCoordinatorResponse from kafka.protocol.frame import KafkaBytes from kafka.protocol.schemas.fields.codecs import Int32 from kafka.version import __version__ log = logging.getLogger(__name__) class KafkaProtocol: """Manage the kafka network protocol Use an instance of KafkaProtocol to manage bytes send/recv'd from a network socket to a broker. Arguments: client_id (str): identifier string to be included in each request ident (str): Optional log-prefix identifier. receive_message_max_bytes (int): Maximum allowed message frame size. Default: 100000000 (100MB). """ def __init__(self, **config): self._ident = config.get('ident', '') self._client_id = config.get('client_id', self._gen_client_id()) self._max_frame_size = config.get('receive_message_max_bytes', 100000000) self._correlation_id = 0 self._header = KafkaBytes(4) self._rbuffer = None self._receiving = False self.in_flight_requests = collections.deque() self.bytes_to_send = [] def _next_correlation_id(self): self._correlation_id = (self._correlation_id + 1) % 2**31 return self._correlation_id def _gen_client_id(self): return 'kafka-python' + __version__ def send_request(self, request, correlation_id=None): """Encode and queue a kafka api request for sending. Arguments: request : 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 """ if correlation_id is None: correlation_id = self._next_correlation_id() log.debug('%s Sending request %d %s', self._ident, correlation_id, request) request.with_header(correlation_id=correlation_id, client_id=self._client_id) data = request.encode(framed=True, header=True) self.bytes_to_send.append(data) if request.expect_response(): self.in_flight_requests.append(request.header) return correlation_id def send_bytes(self): """Retrieve all pending bytes to send on the network""" # Short-circuit the common single-request case to avoid an extra # full-request copy through b''.join. n = len(self.bytes_to_send) if n == 0: return b'' if n == 1: data = self.bytes_to_send[0] self.bytes_to_send = [] else: data = b''.join(self.bytes_to_send) self.bytes_to_send = [] log.debug('%s Send: %r', self._ident, data) return data 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 = [] if data: log.debug('%s Recv: %r', self._ident, data) 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) self._validate_frame_size(nbytes) # 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 _validate_frame_size(self, nbytes): if nbytes < 0 or nbytes > self._max_frame_size: raise Errors.InvalidReceiveError('Invalid frame length: %d' % nbytes) def _process_response(self, read_buffer): if not self.in_flight_requests: raise Errors.CorrelationIdError('No in-flight-request found for server response') header = self.in_flight_requests.popleft() correlation_id = header.correlation_id response_type = header.get_response_class() if response_type is None: log.error('Unable to find ResponseType for api=%d version=%d', header.api_key, header.api_version) raise Errors.KafkaProtocolError('Unable to find response type for api %d v%d' % (header.api_key, header.api_version)) response_header = response_type.parse_header(read_buffer) recv_correlation_id = response_header.correlation_id # ignore correlation id mismatch for 0.8.2 quirk if (recv_correlation_id == 0 and correlation_id != 0 and response_type is FindCoordinatorResponse and header.api_version == 0): log.warning('Kafka 0.8.2 quirk -- FindCoordinatorResponse' ' 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 try: response = response_type.decode(read_buffer) except (ValueError, struct.error): read_buffer.seek(0) buf = read_buffer.read() log.error('Response %d [ResponseType: %s RequestHeader: %s]:' ' Unable to decode %d-byte buffer: %r', correlation_id, response_type, header, len(buf), buf) raise Errors.KafkaProtocolError('Unable to decode response') log.debug('%s Received response %d %s', self._ident, correlation_id, response) return (correlation_id, response) def _reset_buffer(self): self._receiving = False self._header.seek(0) self._rbuffer = None