def send(self, request, expect_response=True): if self._writer is None: raise Errors.ConnectionError( "No connection to broker at {0}:{1}".format( self._host, self._port)) correlation_id = self._next_correlation_id() header = RequestHeader(request, correlation_id=correlation_id, client_id=self._client_id) message = header.encode() + request.encode() size = struct.pack(">i", len(message)) try: self._writer.write(size + message) except OSError as err: self.close(reason=CloseReason.CONNECTION_BROKEN) raise Errors.ConnectionError( "Connection at {0}:{1} broken: {2}".format( self._host, self._port, err)) if not expect_response: return self._writer.drain() fut = create_future(loop=self._loop) self._requests.append((correlation_id, request.RESPONSE_TYPE, fut)) return asyncio.wait_for(fut, self._request_timeout, loop=self._loop)
def close(self, reason=None, exc=None): self.log.debug("Closing connection at %s:%s", self._host, self._port) if self._reader is not None: self._writer.close() self._writer = self._reader = None if not self._read_task.done(): self._read_task.cancel() self._read_task = None for _, _, fut in self._requests: if not fut.done(): error = Errors.ConnectionError( "Connection at {0}:{1} closed".format( self._host, self._port)) if exc is not None: error.__cause__ = exc error.__context__ = exc fut.set_exception(error) self._requests = [] if self._on_close_cb is not None: self._on_close_cb(self, reason) self._on_close_cb = None if self._idle_handle is not None: self._idle_handle.cancel() # transport.close() will close socket, but not right ahead. Return # a future in case we need to wait on it. return self._closed_fut
def _send_sasl_token(self, payload): if self._writer is None: raise Errors.ConnectionError( "No connection to broker at {0}:{1}".format( self._host, self._port)) size = struct.pack(">i", len(payload)) try: self._writer.write(size + payload) except OSError as err: self.close(reason=CloseReason.CONNECTION_BROKEN) raise Errors.ConnectionError( "Connection at {0}:{1} broken: {2}".format( self._host, self._port, err)) fut = create_future(loop=self._loop) self._requests.append((None, None, fut)) return asyncio.wait_for(fut, self._request_timeout, loop=self._loop)
def _on_read_task_error(self, read_task): try: read_task.result() except Exception as exc: conn_exc = Errors.ConnectionError( "Connection at {0}:{1} broken".format(self._host, self._port)) conn_exc.__cause__ = exc conn_exc.__context__ = exc for _, _, fut in self._requests: fut.set_exception(conn_exc) self.close(reason=CloseReason.CONNECTION_BROKEN)
def _read(self): try: while True: resp = yield from self._reader.readexactly(4) size, = self.HEADER.unpack(resp) resp = yield from self._reader.readexactly(size) recv_correlation_id, = self.HEADER.unpack(resp[:4]) correlation_id, resp_type, fut = self._requests.pop(0) if (self._api_version == (0, 8, 2) and resp_type is GroupCoordinatorResponse and correlation_id != 0 and recv_correlation_id == 0): self.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 correlation_id != recv_correlation_id: error = Errors.CorrelationIdError( 'Correlation ids do not match: sent {}, recv {}' .format(correlation_id, recv_correlation_id)) if not fut.done(): fut.set_exception(error) self.close(reason=CloseReason.OUT_OF_SYNC) break if not fut.done(): response = resp_type.decode(resp[4:]) self.log.debug('%s Response %d: %s', self, correlation_id, response) fut.set_result(response) # Update idle timer. self._last_action = self._loop.time() except (OSError, EOFError, ConnectionError) as exc: for _, _, fut in self._requests: conn_exc = Errors.ConnectionError( "Connection at {0}:{1} broken" .format(self._host, self._port)) conn_exc.__cause__ = exc conn_exc.__context__ = exc fut.set_exception(conn_exc) self.close(reason=CloseReason.CONNECTION_BROKEN) except asyncio.CancelledError: pass
def close(self): if self._reader is not None: self._writer.close() self._writer = self._reader = None self._read_task.cancel() self._read_task = None error = Errors.ConnectionError( "Connection at {0}:{1} closed".format(self._host, self._port)) for _, _, fut in self._requests: if not fut.done(): fut.set_exception(error) self._requests = [] if self._idle_handle is not None: self._idle_handle.cancel() # transport.close() will close socket, but not right ahead. Return # a future in case we need to wait on it. return self._closed_fut