File: C:/Users/fred/anaconda3/Lib/site-packages/distributed/comm/tcp.py
from __future__ import annotations
import asyncio
import ctypes
import errno
import functools
import logging
import socket
import ssl
import struct
import sys
import weakref
from ssl import SSLCertVerificationError, SSLError
from typing import Any, ClassVar
from tlz import sliding_window
from tornado import gen, netutil
from tornado.iostream import IOStream, StreamClosedError
from tornado.tcpclient import TCPClient
from tornado.tcpserver import TCPServer
import dask
from dask.utils import parse_timedelta
from distributed.comm.addressing import parse_host_port, unparse_host_port
from distributed.comm.core import (
BaseListener,
Comm,
CommClosedError,
Connector,
FatalCommClosedError,
)
from distributed.comm.registry import Backend
from distributed.comm.utils import (
ensure_concrete_host,
from_frames,
get_tcp_server_address,
to_frames,
)
from distributed.protocol.utils import host_array, pack_frames_prelude, unpack_frames
from distributed.system import MEMORY_LIMIT
from distributed.utils import ensure_ip, ensure_memoryview, get_ip, get_ipv6, nbytes
logger = logging.getLogger(__name__)
# We must not load more than this into a buffer at a time
# It's currently unclear why that is
# see
# - https://github.com/dask/distributed/pull/5854
# - https://bugs.python.org/issue42853
# - https://github.com/dask/distributed/pull/8507
C_INT_MAX = 256 ** ctypes.sizeof(ctypes.c_int) // 2 - 1
MAX_BUFFER_SIZE = MEMORY_LIMIT / 2
def set_tcp_timeout(comm):
"""
Set kernel-level TCP timeout on the stream.
"""
if comm.closed():
return
timeout = dask.config.get("distributed.comm.timeouts.tcp")
timeout = int(parse_timedelta(timeout, default="seconds"))
sock = comm.socket
# Default (unsettable) value on Windows
# https://msdn.microsoft.com/en-us/library/windows/desktop/dd877220(v=vs.85).aspx
nprobes = 10
assert timeout >= nprobes + 1, "Timeout too low"
idle = max(2, timeout // 4)
interval = max(1, (timeout - idle) // nprobes)
idle = timeout - interval * nprobes
assert idle > 0
try:
if sys.platform.startswith("win"):
logger.debug("Setting TCP keepalive: idle=%d, interval=%d", idle, interval)
sock.ioctl(socket.SIO_KEEPALIVE_VALS, (1, idle * 1000, interval * 1000))
else:
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
try:
TCP_KEEPIDLE = socket.TCP_KEEPIDLE
TCP_KEEPINTVL = socket.TCP_KEEPINTVL
TCP_KEEPCNT = socket.TCP_KEEPCNT
except AttributeError:
if sys.platform == "darwin":
TCP_KEEPIDLE = 0x10 # (named "TCP_KEEPALIVE" in C)
TCP_KEEPINTVL = 0x101
TCP_KEEPCNT = 0x102
else:
TCP_KEEPIDLE = None
if TCP_KEEPIDLE is not None:
logger.debug(
"Setting TCP keepalive: nprobes=%d, idle=%d, interval=%d",
nprobes,
idle,
interval,
)
sock.setsockopt(socket.SOL_TCP, TCP_KEEPCNT, nprobes)
sock.setsockopt(socket.SOL_TCP, TCP_KEEPIDLE, idle)
sock.setsockopt(socket.SOL_TCP, TCP_KEEPINTVL, interval)
if sys.platform.startswith("linux"):
logger.debug("Setting TCP user timeout: %d ms", timeout * 1000)
TCP_USER_TIMEOUT = 18 # since Linux 2.6.37
sock.setsockopt(socket.SOL_TCP, TCP_USER_TIMEOUT, timeout * 1000)
except OSError:
logger.exception("Could not set timeout on TCP stream.")
def get_stream_address(comm):
"""
Get a stream's local address.
"""
# raise OSError in case the comm is closed, s.t.
# retry code can handle it appropriately; see also
# https://github.com/dask/distributed/issues/7953
if comm.closed():
raise CommClosedError()
return unparse_host_port(*comm.socket.getsockname()[:2])
def convert_stream_closed_error(obj, exc):
"""
Re-raise StreamClosedError as CommClosedError.
"""
if exc.real_error is not None:
# The stream was closed because of an underlying OS error
exc = exc.real_error
if isinstance(exc, ssl.SSLError):
if exc.reason and "UNKNOWN_CA" in exc.reason:
raise FatalCommClosedError(f"in {obj}: {exc.__class__.__name__}: {exc}")
raise CommClosedError(f"in {obj}: {exc.__class__.__name__}: {exc}") from exc
else:
raise CommClosedError(f"in {obj}: {exc}") from exc
def _close_comm(ref):
"""Callback to close Dask Comm when Tornado Stream closes
Parameters
----------
ref: weak reference to a Dask comm
"""
comm = ref()
if comm:
comm._closed = True
class TCP(Comm):
"""
An established communication based on an underlying Tornado IOStream.
"""
max_shard_size: ClassVar[int] = dask.utils.parse_bytes(
dask.config.get("distributed.comm.shard")
)
stream: IOStream | None
def __init__(
self,
stream: IOStream,
local_addr: str,
peer_addr: str,
deserialize: bool = True,
):
self._closed = False
super().__init__(deserialize=deserialize)
self._local_addr = local_addr
self._peer_addr = peer_addr
self.stream = stream
self._finalizer = weakref.finalize(self, self._get_finalizer())
self._finalizer.atexit = False
self._extra: dict = {}
ref = weakref.ref(self)
stream.set_close_callback(functools.partial(_close_comm, ref))
stream.set_nodelay(True)
set_tcp_timeout(stream)
self._read_extra()
def _read_extra(self):
pass
def _get_finalizer(self):
r = repr(self)
def finalize(stream=self.stream, r=r):
# stream is None if a StreamClosedError is raised during interpreter
# shutdown
if stream is not None and not stream.closed():
logger.warning(f"Closing dangling stream in {r}")
stream.close()
return finalize
@property
def local_address(self) -> str:
return self._local_addr
@property
def peer_address(self) -> str:
return self._peer_addr
async def read(self, deserializers=None):
stream = self.stream
if stream is None:
raise CommClosedError()
fmt = "Q"
fmt_size = struct.calcsize(fmt)
try:
# Don't store multiple numpy or parquet buffers into the same buffer, or
# none will be released until all are released.
frames_nosplit_nbytes_bin = await stream.read_bytes(fmt_size)
(frames_nosplit_nbytes,) = struct.unpack(fmt, frames_nosplit_nbytes_bin)
frames_nosplit = await read_bytes_rw(stream, frames_nosplit_nbytes)
frames, buffers_nbytes = unpack_frames(frames_nosplit, partial=True)
for buffer_nbytes in buffers_nbytes:
buffer = await read_bytes_rw(stream, buffer_nbytes)
frames.append(buffer)
except StreamClosedError as e:
self.stream = None
self._closed = True
convert_stream_closed_error(self, e)
except BaseException:
# Some OSError, CancelledError or another "low-level" exception.
# We do not really know what was already read from 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, #6548).
self.abort()
raise
else:
try:
msg = await from_frames(
frames,
deserialize=self.deserialize,
deserializers=deserializers,
allow_offload=self.allow_offload,
)
except EOFError:
# Frames possibly garbled or truncated by communication error
self.abort()
raise CommClosedError("aborted stream on truncated data")
return msg
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,
},
frame_split_size=self.max_shard_size,
)
frames, frames_nbytes, frames_nbytes_total = _add_frames_header(frames)
try:
# trick to enqueue all frames for writing beforehand
for each_frame_nbytes, each_frame in zip(frames_nbytes, frames):
if each_frame_nbytes:
# Make sure that `len(data) == data.nbytes`
# See <https://github.com/tornadoweb/tornado/pull/2996>
each_frame = ensure_memoryview(each_frame)
for i, j in sliding_window(
2,
range(
0,
each_frame_nbytes + C_INT_MAX,
C_INT_MAX,
),
):
chunk = each_frame[i:j]
chunk_nbytes = chunk.nbytes
if stream._write_buffer is None:
raise StreamClosedError()
stream._write_buffer.append(chunk)
stream._total_write_index += chunk_nbytes
# start writing frames
stream.write(b"")
except StreamClosedError as e:
self.stream = None
self._closed = True
convert_stream_closed_error(self, e)
except BaseException:
# 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).
# In case of, for instance, KeyboardInterrupts or other
# BaseExceptions that could be handled further upstream, we equally
# want to discard this comm
self.abort()
raise
return frames_nbytes_total
@gen.coroutine
def close(self):
# We use gen.coroutine here rather than async def to avoid errors like
# Task was destroyed but it is pending!
# Triggered by distributed.deploy.tests.test_local::test_silent_startup
stream, self.stream = self.stream, None
self._closed = True
if stream is not None and not stream.closed():
try:
# Flush the stream's write buffer by waiting for a last write.
if stream.writing():
yield stream.write(b"")
stream.socket.shutdown(socket.SHUT_RDWR)
except OSError:
pass
finally:
self._finalizer.detach()
stream.close()
def abort(self) -> None:
stream, self.stream = self.stream, None
self._closed = True
if stream is not None and not stream.closed():
self._finalizer.detach()
stream.close()
def closed(self) -> bool:
return self._closed
@property
def extra_info(self):
return self._extra
async def read_bytes_rw(stream: IOStream, n: int) -> memoryview:
"""Read n bytes from stream. Unlike stream.read_bytes, allow for
very large messages and return a writeable buffer.
"""
buf = host_array(n)
for i, j in sliding_window(
2,
range(0, n + C_INT_MAX, C_INT_MAX),
):
chunk = buf[i:j]
actual = await stream.read_into(chunk) # type: ignore[arg-type]
assert actual == chunk.nbytes
return buf
def _add_frames_header(
frames: list[bytes | memoryview],
) -> tuple[list[bytes | memoryview], list[int], int]:
""" """
frames_nbytes = [nbytes(f) for f in frames]
frames_nbytes_total = sum(frames_nbytes)
# Calculate the number of bytes that are inclusive of:
# - prelude
# - msgpack header
# - simple pickle bytes
# - compressed buffers
# - first uncompressed buffer (possibly sharded), IFF the pickle bytes are
# negligible in size
#
# All these can be fetched by read() into a single buffer with a single call to
# Tornado, because they will be dereferenced soon after they are deserialized.
# Read uncompressed numpy/parquet buffers, which will survive indefinitely past
# the end of read(), into their own host arrays so that their memory can be
# released independently.
frames_nbytes_nosplit = 0
first_uncompressed_buffer: object = None
for frame, nb in zip(frames, frames_nbytes):
buffer = frame.obj if isinstance(frame, memoryview) else frame
if not isinstance(buffer, bytes):
# Uncompressed buffer; it will be referenced by the unpickled object
if first_uncompressed_buffer is None:
if frames_nbytes_nosplit > max(2048, nb * 0.05):
# Don't extend the lifespan of non-trivial amounts of pickled bytes
# to that of the buffers
break
first_uncompressed_buffer = buffer
elif first_uncompressed_buffer is not buffer: # don't split sharded frame
# Always store 2+ separate numpy/parquet objects onto separate
# buffers
break
frames_nbytes_nosplit += nb
header = pack_frames_prelude(frames)
header = struct.pack("Q", nbytes(header) + frames_nbytes_nosplit) + header
header_nbytes = nbytes(header)
frames = [header, *frames]
frames_nbytes = [header_nbytes, *frames_nbytes]
frames_nbytes_total += header_nbytes
if frames_nbytes_total < 2**17 or ( # 128 kiB total
frames_nbytes_total < 2**25 # 32 MiB total
and frames_nbytes_total // len(frames) < 2**15 # 32 kiB mean
):
# very small or very fragmented; send in one go
frames = [b"".join(frames)]
frames_nbytes = [frames_nbytes_total]
return frames, frames_nbytes, frames_nbytes_total
class TLS(TCP):
"""
A TLS-specific version of TCP.
"""
# Workaround for OpenSSL 1.0.2 (can drop with OpenSSL 1.1.1)
max_shard_size = min(C_INT_MAX, TCP.max_shard_size)
def _read_extra(self):
TCP._read_extra(self)
sock = self.stream.socket
if sock is not None:
self._extra.update(peercert=sock.getpeercert(), cipher=sock.cipher())
cipher, proto, bits = self._extra["cipher"]
logger.debug(
"TLS connection with %r: protocol=%s, cipher=%s, bits=%d",
self._peer_addr,
proto,
cipher,
bits,
)
def _expect_tls_context(connection_args):
ctx = connection_args.get("ssl_context")
if not isinstance(ctx, ssl.SSLContext):
raise TypeError(
"TLS expects a `ssl_context` argument of type "
"ssl.SSLContext (perhaps check your TLS configuration?)"
f" Instead got {ctx!r}"
)
return ctx
class RequireEncryptionMixin:
def _check_encryption(self, address, connection_args):
if not self.encrypted and connection_args.get("require_encryption"):
# XXX Should we have a dedicated SecurityError class?
raise RuntimeError(
"encryption required by Dask configuration, "
"refusing communication from/to %r" % (self.prefix + address,)
)
_NUMERIC_ONLY = socket.AI_NUMERICHOST | socket.AI_NUMERICSERV
async def _getaddrinfo(host, port, *, family, type=socket.SOCK_STREAM):
# If host and port are numeric, then getaddrinfo doesn't block and we
# can skip get_running_loop().getaddrinfo which is implemented by
# running in a ThreadPoolExecutor. So we try first with the
# _NUMERIC_ONLY flags set, and then only use the threadpool if that
# fails with EAI_NONAME:
try:
return socket.getaddrinfo(
host,
port,
family=family,
type=type,
flags=_NUMERIC_ONLY,
)
except socket.gaierror as e:
if e.errno != socket.EAI_NONAME:
raise
# That failed; it's a real hostname. We better use a thread.
return await asyncio.get_running_loop().getaddrinfo(
host, port, family=family, type=socket.SOCK_STREAM
)
class _DefaultLoopResolver(netutil.Resolver):
"""
Resolver implementation using `asyncio.loop.getaddrinfo`.
backport from Tornado 6.2+
https://github.com/tornadoweb/tornado/blob/3de78b7a15ba7134917a18b0755ea24d7f8fde94/tornado/netutil.py#L416-L432
With an additional optimization based on
https://github.com/python-trio/trio/blob/4edfd41bd5519a2e626e87f6c6ca9fb32b90a6f4/trio/_socket.py#L125-L192
(Copyright Contributors to the Trio project.)
And proposed to cpython in https://github.com/python/cpython/pull/31497/
"""
async def resolve(
self, host: str, port: int, family: socket.AddressFamily = socket.AF_UNSPEC
) -> list[tuple[int, Any]]:
# On Solaris, getaddrinfo fails if the given port is not found
# in /etc/services and no socket type is given, so we must pass
# one here. The socket type used here doesn't seem to actually
# matter (we discard the one we get back in the results),
# so the addresses we return should still be usable with SOCK_DGRAM.
return [
(fam, address)
for fam, _, _, _, address in await _getaddrinfo(
host, port, family=family, type=socket.SOCK_STREAM
)
]
class BaseTCPConnector(Connector, RequireEncryptionMixin):
client: ClassVar[TCPClient] = TCPClient(resolver=_DefaultLoopResolver())
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:
# server_hostname option (for SNI) only works with tornado.iostream.IOStream
if "server_hostname" in kwargs:
stream = await self.client.connect(
ip, port, max_buffer_size=MAX_BUFFER_SIZE
)
stream = await stream.start_tls(False, **kwargs)
else:
stream = await self.client.connect(
ip, port, max_buffer_size=MAX_BUFFER_SIZE, **kwargs
)
# Under certain circumstances tornado will have a closed connection 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
class TCPConnector(BaseTCPConnector):
prefix = "tcp://"
comm_class = TCP
encrypted = False
def _get_connect_args(self, **connection_args):
return {}
class TLSConnector(BaseTCPConnector):
prefix = "tls://"
comm_class = TLS
encrypted = True
def _get_connect_args(self, **connection_args):
tls_args = {"ssl_options": _expect_tls_context(connection_args)}
if connection_args.get("server_hostname"):
tls_args["server_hostname"] = connection_args["server_hostname"]
return tls_args
class BaseTCPListener(BaseListener, RequireEncryptionMixin):
def __init__(
self,
address,
comm_handler,
deserialize=True,
allow_offload=True,
default_host=None,
default_port=0,
**connection_args,
):
super().__init__()
self._check_encryption(address, connection_args)
self.ip, self.port = parse_host_port(address, default_port)
self.default_host = default_host
self.comm_handler = comm_handler
self.deserialize = deserialize
self.allow_offload = allow_offload
self.server_args = self._get_server_args(**connection_args)
self.tcp_server = None
self.bound_address = None
async def start(self):
self.tcp_server = TCPServer(max_buffer_size=MAX_BUFFER_SIZE, **self.server_args)
self.tcp_server.handle_stream = self._handle_stream
backlog = int(dask.config.get("distributed.comm.socket-backlog"))
for _ in range(5):
try:
# When shuffling data between workers, there can
# really be O(cluster size) connection requests
# on a single worker socket, make sure the backlog
# is large enough not to lose any.
sockets = netutil.bind_sockets(
self.port, address=self.ip, backlog=backlog
)
except OSError as e:
# EADDRINUSE can happen sporadically when trying to bind
# to an ephemeral port
if self.port != 0 or e.errno != errno.EADDRINUSE:
raise
exc = e
else:
self.tcp_server.add_sockets(sockets)
break
else:
raise exc
self.get_host_port() # trigger assignment to self.bound_address
def stop(self):
tcp_server, self.tcp_server = self.tcp_server, None
if tcp_server is not None:
tcp_server.stop()
def _check_started(self):
if self.tcp_server is None:
raise ValueError("invalid operation on non-started TCPListener")
async def _handle_stream(self, stream, address):
address = self.prefix + unparse_host_port(*address[:2])
stream = await self._prepare_stream(stream, address)
if stream is None:
# Preparation failed
return
logger.debug("Incoming connection from %r to %r", address, self.contact_address)
local_address = self.prefix + get_stream_address(stream)
comm = self.comm_class(stream, local_address, address, self.deserialize)
comm.allow_offload = self.allow_offload
try:
await self.on_connection(comm)
except CommClosedError:
logger.info("Connection from %s closed before handshake completed", address)
return
await self.comm_handler(comm)
def get_host_port(self):
"""
The listening address as a (host, port) tuple.
"""
self._check_started()
if self.bound_address is None:
self.bound_address = get_tcp_server_address(self.tcp_server)
# IPv6 getsockname() can return more a 4-len tuple
return self.bound_address[:2]
@property
def listen_address(self):
"""
The listening address as a string.
"""
return self.prefix + unparse_host_port(*self.get_host_port())
@property
def contact_address(self):
"""
The contact address as a string.
"""
host, port = self.get_host_port()
host = ensure_concrete_host(host, default_host=self.default_host)
return self.prefix + unparse_host_port(host, port)
class TCPListener(BaseTCPListener):
prefix = "tcp://"
comm_class = TCP
encrypted = False
def _get_server_args(self, **connection_args):
return {}
async def _prepare_stream(self, stream, address):
return stream
class TLSListener(BaseTCPListener):
prefix = "tls://"
comm_class = TLS
encrypted = True
def _get_server_args(self, **connection_args):
ctx = _expect_tls_context(connection_args)
return {"ssl_options": ctx}
async def _prepare_stream(self, stream, address):
try:
await stream.wait_for_handshake()
except OSError as e:
# The handshake went wrong, log and ignore
logger.warning(
"Listener on %r: TLS handshake failed with remote %r: %s",
self.listen_address,
address,
getattr(e, "real_error", None) or e,
)
else:
return stream
class BaseTCPBackend(Backend):
# I/O
def get_connector(self):
return self._connector_class()
def get_listener(self, loc, handle_comm, deserialize, **connection_args):
return self._listener_class(loc, handle_comm, deserialize, **connection_args)
# Address handling
def get_address_host(self, loc):
return parse_host_port(loc)[0]
def get_address_host_port(self, loc):
return parse_host_port(loc)
def resolve_address(self, loc):
host, port = parse_host_port(loc)
return unparse_host_port(ensure_ip(host), port)
def get_local_address_for(self, loc):
host, port = parse_host_port(loc)
host = ensure_ip(host)
if ":" in host:
local_host = get_ipv6(host)
else:
local_host = get_ip(host)
return unparse_host_port(local_host, None)
class TCPBackend(BaseTCPBackend):
_connector_class = TCPConnector
_listener_class = TCPListener
class TLSBackend(BaseTCPBackend):
_connector_class = TLSConnector
_listener_class = TLSListener