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
