Source code distributed/comm/inproc.py

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
import asyncio
from collections import deque, namedtuple
import itertools
import logging
import os
import threading
import weakref
import warnings

from tornado.concurrent import Future
from tornado.ioloop import IOLoop

from ..protocol import nested_deserialize
from ..utils import get_ip

from .registry import Backend, backends
from .core import Comm, Connector, Listener, CommClosedError


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)
        with warnings.catch_warnings():
            # Avoid immediate warning for unreachable network
            # (will still warn for other get_ip() calls when actually used)
            warnings.simplefilter("ignore")
            try:
                self.ip = get_ip()
            except OSError:
                self.ip = "127.0.0.1"
        self.lock = threading.Lock()

    def add_listener(self, addr, listener):
        with self.lock:
            if addr in self.listeners:
                raise RuntimeError("already listening on %r" % (addr,))
            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


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
            fut.set_result(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__(
        self, local_addr, peer_addr, read_q, write_q, write_loop, deserialize=True
    ):
        Comm.__init__(self)
        self._local_addr = local_addr
        self._peer_addr = peer_addr
        self.deserialize = deserialize
        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):
        def finalize(write_q=self._write_q, write_loop=self._write_loop, r=repr(self)):
            logger.warning("Closing dangling queue in %s" % (r,))
            write_loop.add_callback(write_q.put_nowait, _EOF)

        return finalize

    @property
    def local_address(self):
        return self._local_addr

    @property
    def peer_address(self):
        return self._peer_addr

    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

        if self.deserialize:
            msg = nested_deserialize(msg)
        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(Listener):
    prefix = "inproc"

    def __init__(self, address, comm_handler, deserialize=True):
        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 _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)
            try:
                await self.on_connection(comm)
            except CommClosedError:
                logger.debug("Connection closed before handshake completed")
                return
            IOLoop.current().add_callback(self.comm_handler, 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 IOError("no endpoint for inproc address %r" % (address,))

        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()