from __future__ import annotations import asyncio import itertools import logging import os import sys import threading import weakref from collections import deque, namedtuple from tornado.concurrent import Future from tornado.ioloop import IOLoop from distributed.comm.core import BaseListener, Comm, CommClosedError, Connector from distributed.comm.registry import Backend, backends from distributed.protocol.serialize import _nested_deserialize from distributed.utils import get_ip logger = logging.getLogger(__name__) ConnectionRequest = namedtuple( "ConnectionRequest", ("c2s_q", "s2c_q", "c_loop", "c_addr", "conn_event") ) class Manager: """ An object coordinating listeners and their addresses. """ def __init__(self): self.listeners = weakref.WeakValueDictionary() self.addr_suffixes = itertools.count(1) self._ip = None self.lock = threading.Lock() @property def ip(self): if not self._ip: try: self._ip = get_ip() except OSError: self._ip = "127.0.0.1" return self._ip def add_listener(self, addr, listener): with self.lock: if addr in self.listeners: raise RuntimeError(f"already listening on {addr!r}") self.listeners[addr] = listener def remove_listener(self, addr): with self.lock: try: del self.listeners[addr] except KeyError: pass def get_listener_for(self, addr): with self.lock: self.validate_address(addr) return self.listeners.get(addr) def new_address(self): return "%s/%d/%s" % (self.ip, os.getpid(), next(self.addr_suffixes)) def validate_address(self, addr): """ Validate the address' IP and pid. """ ip, pid, suffix = addr.split("/") if ip != self.ip or int(pid) != os.getpid(): raise ValueError( "inproc address %r does not match host (%r) or pid (%r)" % (addr, self.ip, os.getpid()) ) global_manager = Manager() def new_address(): """ Generate a new address. """ return "inproc://" + global_manager.new_address() class QueueEmpty(Exception): pass def _set_result_unless_cancelled(fut, result): """Helper setting the result only if the future was not cancelled.""" if fut.cancelled(): return fut.set_result(result) class Queue: """ A single-reader, single-writer, non-threadsafe, peekable queue. """ def __init__(self): self._q = deque() self._read_future = None def get_nowait(self): q = self._q if not q: raise QueueEmpty return q.popleft() def get(self): assert not self._read_future, "Only one reader allowed" fut = Future() q = self._q if q: fut.set_result(q.popleft()) else: self._read_future = fut return fut def put_nowait(self, value): q = self._q fut = self._read_future if fut is not None: assert len(q) == 0 self._read_future = None _set_result_unless_cancelled(fut, value) else: q.append(value) put = put_nowait _omitted = object() def peek(self, default=_omitted): """ Get the next object in the queue without removing it from the queue. """ q = self._q if q: return q[0] elif default is not self._omitted: return default else: raise QueueEmpty _EOF = object() class InProc(Comm): """ An established communication based on a pair of in-process queues. Reminder: a Comm must always be used from a single thread. Its peer Comm can be running in any thread. """ _initialized = False def __init__( # type: ignore[no-untyped-def] self, local_addr: str, peer_addr: str, read_q, write_q, write_loop, deserialize: bool = True, ): super().__init__(deserialize=deserialize) self._local_addr = local_addr self._peer_addr = peer_addr self._read_q = read_q self._write_q = write_q self._write_loop = write_loop self._closed = False self._finalizer = weakref.finalize(self, self._get_finalizer()) self._finalizer.atexit = False self._initialized = True def _get_finalizer(self): r = repr(self) def finalize( read_q=self._read_q, write_q=self._write_q, write_loop=self._write_loop, is_finalizing=sys.is_finalizing, r=r, ): if read_q.peek(None) is _EOF or is_finalizing(): return logger.warning(f"Closing dangling queue in {r}") write_loop.add_callback(write_q.put_nowait, _EOF) return finalize @property def local_address(self) -> str: return self._local_addr @property def peer_address(self) -> str: return self._peer_addr @property def same_host(self) -> bool: return True async def read(self, deserializers="ignored"): if self._closed: raise CommClosedError() msg = await self._read_q.get() if msg is _EOF: self._closed = True self._finalizer.detach() raise CommClosedError() msg = _nested_deserialize(msg, self.deserialize) return msg async def write(self, msg, serializers=None, on_error=None): if self.closed(): raise CommClosedError() # Ensure we feed the queue in the same thread it is read from. self._write_loop.add_callback(self._write_q.put_nowait, msg) return 1 async def close(self): self.abort() def abort(self): if not self.closed(): # Putting EOF is cheap enough that we do it on abort() too self._write_loop.add_callback(self._write_q.put_nowait, _EOF) self._read_q.put_nowait(_EOF) self._write_q = self._read_q = None self._closed = True self._finalizer.detach() def closed(self): """ Whether this comm is closed. An InProc comm is closed if: 1) close() or abort() was called on this comm 2) close() or abort() was called on the other end and the read queue is empty """ if self._closed: return True # NOTE: repr() is called by finalize() during __init__()... if self._initialized and self._read_q.peek(None) is _EOF: self._closed = True self._finalizer.detach() return True else: return False class InProcListener(BaseListener): prefix = "inproc" def __init__(self, address, comm_handler, deserialize=True): super().__init__() self.manager = global_manager self.address = address or self.manager.new_address() self.comm_handler = comm_handler self.deserialize = deserialize self.listen_q = Queue() async def _handle_stream(self, comm): try: await self.on_connection(comm) except CommClosedError: logger.debug("Connection closed before handshake completed") return await self.comm_handler(comm) async def _listen(self): while True: conn_req = await self.listen_q.get() if conn_req is None: break comm = InProc( local_addr="inproc://" + self.address, peer_addr="inproc://" + conn_req.c_addr, read_q=conn_req.c2s_q, write_q=conn_req.s2c_q, write_loop=conn_req.c_loop, deserialize=self.deserialize, ) # Notify connector conn_req.c_loop.add_callback(conn_req.conn_event.set) IOLoop.current().add_callback(self._handle_stream, comm) def connect_threadsafe(self, conn_req): self.loop.add_callback(self.listen_q.put_nowait, conn_req) async def start(self): self.loop = IOLoop.current() self._listen_future = asyncio.ensure_future(self._listen()) self.manager.add_listener(self.address, self) def stop(self): self.listen_q.put_nowait(None) self.manager.remove_listener(self.address) @property def listen_address(self): return "inproc://" + self.address @property def contact_address(self): return "inproc://" + self.address class InProcConnector(Connector): def __init__(self, manager): self.manager = manager async def connect(self, address, deserialize=True, **connection_args): listener = self.manager.get_listener_for(address) if listener is None: raise OSError(f"no endpoint for inproc address {address!r}") conn_req = ConnectionRequest( c2s_q=Queue(), s2c_q=Queue(), c_loop=IOLoop.current(), c_addr=self.manager.new_address(), conn_event=asyncio.Event(), ) listener.connect_threadsafe(conn_req) # Wait for connection acknowledgement # (do not pretend we're connected if the other comm never gets # created, for example if the listener was stopped in the meantime) await conn_req.conn_event.wait() comm = InProc( local_addr="inproc://" + conn_req.c_addr, peer_addr="inproc://" + address, read_q=conn_req.s2c_q, write_q=conn_req.c2s_q, write_loop=listener.loop, deserialize=deserialize, ) return comm class InProcBackend(Backend): manager = global_manager # I/O def get_connector(self): return InProcConnector(self.manager) def get_listener(self, loc, handle_comm, deserialize, **connection_args): return InProcListener(loc, handle_comm, deserialize) # Address handling def get_address_host(self, loc): self.manager.validate_address(loc) return self.manager.ip def resolve_address(self, loc): return loc def get_local_address_for(self, loc): self.manager.validate_address(loc) return self.manager.new_address() backends["inproc"] = InProcBackend()