def on_closed(self, exc_info=False):
if self._connect_future:
if exc_info:
self._connect_future.set_exception(
exc_info[1] if isinstance(exc_info, tuple) else exc_info)
else:
self._connect_future.set_exception(StreamClosedError(None))
self._connect_future = None
if self._connect_ssl_future:
if exc_info:
self._connect_ssl_future.set_exception(
exc_info[1] if isinstance(exc_info, tuple) else exc_info)
else:
self._connect_ssl_future.set_exception(StreamClosedError(None))
self._connect_ssl_future = None
if self._read_future:
if exc_info:
self._read_future.set_exception(
exc_info[1] if isinstance(exc_info, tuple) else exc_info)
else:
self._read_future.set_exception(StreamClosedError(None))
self._read_future = None
if self._close_callback:
close_callback, self._close_callback = self._close_callback, None
self._loop.call_soon(close_callback)
self._closed = True
def read(self, length):
'''Read from socket.
This function reads only from the socket and not the buffer.
This function uses a blocking fast path consecutively for
every 100 reads.
'''
self.stop_monitor_for_close()
if self._blocking_counter < 100:
try:
data = self._socket.recv(length)
except ssl.SSLError as error:
if error.errno != ssl.SSL_ERROR_WANT_READ:
raise
except IOError as error:
if error.errno not in (errno.EWOULDBLOCK, errno.EINPROGRESS):
raise
else:
if data:
self._blocking_counter += 1
raise tornado.gen.Return(data)
else:
self.close()
raise StreamClosedError('Stream unexpectedly closed.')
else:
self._blocking_counter = 0
while True:
events = yield self._wait_event(READ | ERROR,
timeout=self._rw_timeout)
if events & ERROR:
self._raise_socket_error()
try:
data = self._socket.recv(length)
except ssl.SSLError as error:
if error.errno == ssl.SSL_ERROR_WANT_READ:
continue
else:
raise
else:
break
if not data:
self.close()
raise StreamClosedError('Stream unexpectedly closed.')
raise tornado.gen.Return(data)
def write(self, data):
assert isinstance(data, bytes)
if self._closed:
raise StreamClosedError(real_error=self.error)
if not data:
if self._write_future:
return self._write_future
future = Future()
future.set_result(None)
return future
if self._write_buffer_size:
self._write_buffer += data
else:
self._write_buffer = bytearray(data)
self._write_buffer_size += len(data)
future = self._write_future = Future()
self._handle_write()
if self._write_buffer_size:
if not self._state & self.io_loop.WRITE:
self._state = self._state | self.io_loop.WRITE
self.io_loop.update_handler(self.fileno(), self._state)
return future
async def connect(self, address, deserialize=True, **connection_args):
self._check_encryption(address, connection_args)
ip, port = parse_host_port(address)
kwargs = self._get_connect_args(**connection_args)
try:
stream = await self.client.connect(
ip, port, max_buffer_size=MAX_BUFFER_SIZE, **kwargs
)
# Under certain circumstances tornado will have a closed connnection with an
# error and not raise a StreamClosedError.
#
# This occurs with tornado 5.x and openssl 1.1+
if stream.closed() and stream.error:
raise StreamClosedError(stream.error)
except StreamClosedError as e:
# The socket connect() call failed
convert_stream_closed_error(self, e)
except SSLCertVerificationError as err:
raise FatalCommClosedError(
"TLS certificate does not match. Check your security settings. "
"More info at https://distributed.dask.org/en/latest/tls.html"
) from err
except SSLError as err:
raise FatalCommClosedError() from err
local_address = self.prefix + get_stream_address(stream)
comm = self.comm_class(
stream, local_address, self.prefix + address, deserialize
)
return comm
async def write(self, msg, serializers=None, on_error="message"):
stream = self.stream
if stream is None:
raise CommClosedError()
frames = await to_frames(
msg,
allow_offload=self.allow_offload,
serializers=serializers,
on_error=on_error,
context={
"sender": self.local_info,
"recipient": self.remote_info,
**self.handshake_options,
},
)
frames_nbytes = sum(map(nbytes, frames))
header = pack_frames_prelude(frames)
header = struct.pack("Q", nbytes(header) + frames_nbytes) + header
frames = [header, *frames]
frames_nbytes += nbytes(header)
if frames_nbytes < 2 ** 17: # 128kiB
# small enough, send in one go
frames = [b"".join(frames)]
try:
# trick to enque all frames for writing beforehand
for each_frame in frames:
each_frame_nbytes = nbytes(each_frame)
if each_frame_nbytes:
if stream._write_buffer is None:
raise StreamClosedError()
if isinstance(each_frame, memoryview):
# Make sure that len(data) == data.nbytes`
# See <https://github.com/tornadoweb/tornado/pull/2996>
each_frame = memoryview(each_frame).cast("B")
stream._write_buffer.append(each_frame)
stream._total_write_index += each_frame_nbytes
# start writing frames
stream.write(b"")
except StreamClosedError as e:
self.stream = None
self._closed = True
if not shutting_down():
convert_stream_closed_error(self, e)
except Exception:
# Some OSError or a another "low-level" exception. We do not really know
# what was already written to the underlying socket, so it is not even safe
# to retry here using the same stream. The only safe thing to do is to
# abort. (See also GitHub #4133).
self.abort()
raise
return frames_nbytes
async def connect(self, address, deserialize=True, **connection_args):
self._check_encryption(address, connection_args)
ip, port = parse_host_port(address)
kwargs = self._get_connect_args(**connection_args)
try:
stream = await self.client.connect(
ip, port, max_buffer_size=MAX_BUFFER_SIZE, **kwargs
)
# Under certain circumstances tornado will have a closed connnection with an error and not raise
# a StreamClosedError.
#
# This occurs with tornado 5.x and openssl 1.1+
if stream.closed() and stream.error:
raise StreamClosedError(stream.error)
except StreamClosedError as e:
# The socket connect() call failed
convert_stream_closed_error(self, e)
local_address = self.prefix + get_stream_address(stream)
return self.comm_class(
stream, local_address, self.prefix + address, deserialize
)
def read_until(end):
nonlocal lineno
yield gen.sleep(0.01)
if lineno == len(stdout):
raise StreamClosedError()
lineno += 1
return '{}\n'.format(stdout[lineno - 1])
def test_handle_stream_closed(self):
self.layer.socks_conn = Mock()
self.layer.socks_conn.send = Mock(side_effect=self.collect_send_event)
socks_request = Request(REQ_COMMAND["CONNECT"], ADDR_TYPE["IPV4"],
u"1.2.3.4", self.port)
addr_not_support_status = RESP_STATUS["ADDRESS_TYPE_NOT_SUPPORTED"]
network_unreach_status = RESP_STATUS["NETWORK_UNREACHABLE"]
general_fail_status = RESP_STATUS["GENRAL_FAILURE"]
error_cases = [(errno.ENOEXEC, addr_not_support_status),
(errno.EBADF, addr_not_support_status),
(errno.ETIMEDOUT, network_unreach_status),
(errno.EADDRINUSE, general_fail_status),
(55566, general_fail_status)]
for code, expect_status in error_cases:
self.layer.create_dest_stream = Mock(
side_effect=self.create_raise_exception_function(
StreamClosedError((code, ))))
result_future = self.layer.handle_request_and_create_destination(
socks_request)
with self.assertRaises(DestNotConnectedError):
yield result_future
self.assertIsNotNone(self.event)
self.assertIsInstance(self.event, Response)
self.assertEqual(self.event.status, expect_status)
self.assertEqual(self.event.atyp, ADDR_TYPE["IPV4"])
self.assertEqual(str(self.event.addr), "1.2.3.4")
self.assertEqual(self.event.port, self.port)
def send(self, msg):
""" Send a message to the other side
This completes quickly and synchronously
"""
try:
if self.stream is None: # not yet started
self.buffer.append(msg)
return
if self.stream._closed:
raise StreamClosedError()
if self.buffer:
self.buffer.append(msg)
return
# If we're new and early,
now = default_timer()
if (now < self.last_transmission + self.interval
or not self.last_send._done):
self.buffer.append(msg)
self.loop.add_callback(self.send_next)
return
self.buffer.append(msg)
self.loop.add_callback(self.send_next, wait=False)
except Exception as e:
logger.exception(e)
raise
def write(self, data):
'''Write all data to socket.
This function uses a blocking fast path consecutively for
every 100 writes.
'''
self.stop_monitor_for_close()
total_bytes_sent = 0
while total_bytes_sent < len(data):
if self._blocking_counter < 100:
try:
bytes_sent = self._socket.send(data[total_bytes_sent:])
except ssl.SSLError as error:
if error.errno != ssl.SSL_ERROR_WANT_WRITE:
raise
except IOError as error:
if error.errno not in (errno.EWOULDBLOCK,
errno.EINPROGRESS):
raise
else:
if not bytes_sent:
self.close()
raise StreamClosedError('Stream unexpectedly closed.')
else:
total_bytes_sent += bytes_sent
continue
else:
self._blocking_counter = 0
events = yield self._wait_event(WRITE | ERROR,
timeout=self._rw_timeout)
if events & ERROR:
self._raise_socket_error()
bytes_sent = self._socket.send(data[total_bytes_sent:])
if not bytes_sent:
self.close()
raise StreamClosedError('Stream unexpectedly closed.')
total_bytes_sent += bytes_sent
def consume(self, amount):
while not self.closed and self.size <= 0:
yield self.cond.wait()
if self.closed:
raise StreamClosedError()
if self.size < amount:
amount = self.size
if self.parent is not None:
amount = yield self.parent.consume(amount)
self.size -= amount
raise gen.Return(amount)
async def connect(self, address, deserialize=True, **connection_args):
kwargs = self._get_connect_args(**connection_args)
try:
request = HTTPRequest(f"{self.prefix}{address}", **kwargs)
sock = await websocket_connect(request,
max_message_size=10_000_000_000)
if sock.stream.closed() and sock.stream.error:
raise StreamClosedError(sock.stream.error)
except StreamClosedError as e:
convert_stream_closed_error(self, e)
except SSLError as err:
raise FatalCommClosedError() from err
return self.comm_class(sock, deserialize=deserialize)
def send(self, msg):
""" Schedule a message for sending to the other side
This completes quickly and synchronously
"""
if self.stream is not None and self.stream._closed:
raise StreamClosedError()
self.message_count += 1
self.buffer.append(msg)
# Avoid spurious wakeups if possible
if self.next_deadline is None:
self.waker.set()
def read_bytes(self, num_bytes):
assert self._read_future is None, "Already reading"
if self._closed:
raise StreamClosedError(IOError('Already Closed'))
future = self._read_future = Future()
self._read_bytes = num_bytes
if self._read_buffer_size >= self._read_bytes:
future, self._read_future = self._read_future, None
self._read_buffer, data = bytearray(), self._read_buffer
self._read_buffer_size = 0
self._read_bytes = 0
future.set_result(data)
return future
def close_callback():
if not future.done():
# Note that unlike most Futures returned by IOStream,
# this one passes the underlying error through directly
# instead of wrapping everything in a StreamClosedError
# with a real_error attribute. This is because once the
# connection is established it's more helpful to raise
# the SSLError directly than to hide it behind a
# StreamClosedError (and the client is expecting SSL
# issues rather than network issues since this method is
# named start_tls).
future.set_exception(ssl_stream.error or StreamClosedError())
if orig_close_callback is not None:
orig_close_callback()
def write(self, data):
assert isinstance(data, bytes)
if self._closed:
raise StreamClosedError(real_error=self.error)
if data:
self._write_buffer.append(data)
self._write_buffer_size += len(data)
if not self._connecting:
self._handle_write()
if self._write_buffer:
if not self._state & self.io_loop.WRITE:
self._state = self._state | self.io_loop.WRITE
self.io_loop.update_handler(self.fileno(), self._state)
def read(self, num_bytes):
assert self._read_future is None, "Already reading"
if self._closed:
raise StreamClosedError(real_error=self.error)
future = self._read_future = Future()
self._read_bytes = num_bytes
self._read_partial = False
if self._read_buffer_size >= self._read_bytes:
future, self._read_future = self._read_future, None
self._read_buffer, data = bytearray(), self._read_buffer
self._read_buffer_size = 0
self._read_bytes = 0
future.set_result(data)
return future
def read(self, num_bytes):
assert self._read_future is None, "Already reading"
if self._closed:
raise StreamClosedError(real_error=self.error)
future = self._read_future = TracebackFuture()
self._read_bytes = num_bytes
self._read_partial = False
if self._read_buffer_size >= self._read_bytes:
future, self._read_future = self._read_future, None
data = b"".join(self._read_buffer)
self._read_buffer.clear()
self._read_buffer_size = 0
self._read_bytes = 0
future.set_result(data)
return future
async def connect(self, address, deserialize=True, **connection_args):
kwargs = self._get_connect_args(**connection_args)
try:
request = HTTPRequest(f"{self.prefix}{address}", **kwargs)
sock = await websocket_connect(request,
max_message_size=10_000_000_000)
if sock.stream.closed() and sock.stream.error:
raise StreamClosedError(sock.stream.error)
except StreamClosedError as e:
convert_stream_closed_error(self, e)
except SSLError as err:
raise FatalCommClosedError(
"TLS expects a `ssl_context` argument of type "
"ssl.SSLContext (perhaps check your TLS configuration?)"
) from err
return self.comm_class(sock, deserialize=deserialize)
def connect(address, deserialize = True, **connection_args):
ip, port = parse_host_port(address)
#kwargs = self._get_connect_args(**connection_args)
kwargs = {} # The method in Dask just returns {} as far as I can tell
client = TCPClient()
try:
stream = yield client.connect(ip, port, max_buffer_size = MAX_BUFFER_SIZE, **kwargs)
# Under certain circumstances tornado will have a closed connnection with an error and not raise
# a StreamClosedError.
#
# This occurs with tornado 5.x and openssl 1.1+
if stream.closed() and stream.error:
raise StreamClosedError(stream.error)
except StreamClosedError as e:
# The socket connect() call failed
convert_stream_closed_error("Lambda", e)
local_address = prefix + get_stream_address(stream)
raise gen.Return(TCP(stream, local_address, prefix + address, deserialize))
def handle_frame(self, frame):
if frame.type == constants.FrameType.SETTINGS:
self._handle_settings_frame(frame)
elif frame.type == constants.FrameType.WINDOW_UPDATE:
self._handle_window_update_frame(frame)
elif frame.type == constants.FrameType.PING:
self._handle_ping_frame(frame)
elif frame.type == constants.FrameType.GOAWAY:
self.stream.close()
# TODO: shut down all open streams.
raise StreamClosedError()
elif frame.type in (constants.FrameType.DATA,
constants.FrameType.HEADERS,
constants.FrameType.PRIORITY,
constants.FrameType.RST_STREAM,
constants.FrameType.PUSH_PROMISE,
constants.FrameType.CONTINUATION):
raise ConnectionError(
constants.ErrorCode.PROTOCOL_ERROR,
"invalid frame type %s for stream 0" % frame.type)
def read(self, deserializers=None):
# print("[ {} ] Attempting to read in TCP Comm...".format(datetime.datetime.utcnow()))
stream = self.stream
if stream is None:
raise CommClosedError
try:
n_frames = yield stream.read_bytes(8)
n_frames = struct.unpack("Q", n_frames)[0]
lengths = yield stream.read_bytes(8 * n_frames)
lengths = struct.unpack("Q" * n_frames, lengths)
frames = []
# print("[ {} ] Reading {} lengths now...".format(datetime.datetime.utcnow(), len(lengths)))
for length in lengths:
if length:
if PY3 and self._iostream_has_read_into:
frame = bytearray(length)
n = yield stream.read_into(frame)
assert n == length, (n, length)
else:
frame = yield stream.read_bytes(length)
else:
frame = b""
frames.append(frame)
except StreamClosedError as e:
self.stream = None
print("StreamClosedError...")
raise StreamClosedError("Stream closed...")
else:
try:
msg = yield serialization.from_frames(
frames,
deserialize=self.deserialize,
deserializers=deserializers)
except EOFError:
# Frames possibly garbled or truncated by communication error
self.abort()
# print("aborted stream on truncated data")
raise CommClosedError("aborted stream on truncated data")
raise gen.Return(msg)
def _invoke(self, method_name, *args, **kwargs):
trace = kwargs.pop("trace", None)
if trace:
update_dict_with_trace(kwargs, trace)
new_trace_id = kwargs.get('trace_id', self.trace_id)
if new_trace_id != self.trace_id:
self.log = get_trace_adapter(log, new_trace_id)
self.trace_id = new_trace_id
self.log.debug("`%s` Tx method `%s` call: %.300s %.300s",
self.service_name, method_name, args, kwargs)
if self._done:
raise ChokeEvent()
if self.pipe is None:
raise StreamClosedError()
for method_id, (method, tx_tree) in six.iteritems(self.tx_tree):
if method == method_name:
self.log.debug("method `%s` has been found in API map",
method_name)
headers = manage_headers(kwargs, self._header_table)
packed_data = msgpack_packb(
[self.session_id, method_id, args, headers])
self.log.info(
'send message to `%s`: channel id: %s, type: %s, length: %s bytes',
self.service_name, self.session_id, method_name,
len(packed_data))
self.pipe.write(packed_data)
if tx_tree == {}: # last transition
self.done()
elif tx_tree is not None: # not a recursive transition
self.tx_tree = tx_tree
raise Return(None)
raise AttributeError(method_name)
async def connect(self, address, deserialize=True, **connection_args):
ip, port, sni = parse_gateway_address(address)
ctx = connection_args.get("ssl_context")
if not isinstance(ctx, ssl.SSLContext):
raise TypeError("Gateway expects a `ssl_context` argument of type "
"ssl.SSLContext, instead got %s" % ctx)
try:
plain_stream = await self.client.connect(
ip, port, max_buffer_size=MAX_BUFFER_SIZE)
stream = await plain_stream.start_tls(False,
ssl_options=ctx,
server_hostname=sni)
if stream.closed() and stream.error:
raise StreamClosedError(stream.error)
except StreamClosedError as e:
# The socket connect() call failed
convert_stream_closed_error(self, e)
local_address = "tls://" + get_stream_address(stream)
peer_address = "gateway://" + address
return TLS(stream, local_address, peer_address, deserialize)
def call(self, clazz, method, request, response):
if self._stream is None:
raise StreamClosedError()
request_header = ProtobufRpcEngine_pb2.RequestHeaderProto()
request_header.methodName = method
request_header.declaringClassProtocolName = clazz
request_header.clientProtocolVersion = 1
# make call context
call_id = self._next_call(response)
# send request
yield self._write([
self._rpc_request_header,
request_header,
request,
])
# check future
future = self._callbacks[call_id][1]
yield future
raise gen.Return(future.result())
def write(self, data):
if self._closed:
raise StreamClosedError(IOError('Already Closed'))
self._transport.write(data)
def on_connection_close(self):
if not self.connect_future.done():
self.connect_future.set_exception(StreamClosedError())
self.on_message(None)
self.tcp_client.close()
super(WebSocketClientConnection, self).on_connection_close()
def on_connection_close(self):
raise StreamClosedError(
"Connection is closed without finishing the request")
def send(self, msg):
if self._broken:
raise StreamClosedError('Batch Stream is Closed')
else:
self.send_q.put_nowait(msg)
def recv(self):
result = yield self.recv_q.get()
if result == 'close':
raise StreamClosedError('Batched Stream is Closed')
else:
raise gen.Return(result)
