First draft of PubSub

distributed/__init__.py

@@ -10,6 +10,7 @@
                      Future, futures_of)
 from .lock import Lock
 from .nanny import Nanny
+from .pubsub import Pub, Sub
 from .queues import Queue
 from .scheduler import Scheduler
 from .threadpoolexecutor import rejoin

distributed/client.py

@@ -52,6 +52,7 @@
 from .protocol import to_serialize
 from .protocol.pickle import dumps, loads
 from .publish import Datasets
+from .pubsub import PubSubClientExtension
 from .security import Security
 from .sizeof import sizeof
 from .threadpoolexecutor import rejoin
@@ -69,6 +70,11 @@
 _global_client_index = [0]
 
 
+DEFAULT_EXTENSIONS = [
+        PubSubClientExtension,
+]
+
+
 def _get_global_client():
     for k in sorted(_global_clients, reverse=True):
         c = _global_clients[k]
@@ -506,7 +512,9 @@ def __init__(self, address=None, loop=None, timeout=no_default,
                  set_as_default=True, scheduler_file=None,
                  security=None, asynchronous=False,
                  name=None, heartbeat_interval=None,
-                 serializers=None, deserializers=None, **kwargs):
+                 serializers=None, deserializers=None,
+                 extensions=DEFAULT_EXTENSIONS,
+                 **kwargs):
         if timeout == no_default:
             timeout = dask.config.get('distributed.comm.timeouts.connect')
         if timeout is not None:
@@ -581,7 +589,7 @@ def __init__(self, address=None, loop=None, timeout=no_default,
             self._previous_shuffle = dask.config.get('shuffle', None)
             dask.config.set(shuffle='tasks')
 
-        self._handlers = {
+        self._stream_handlers = {
             'key-in-memory': self._handle_key_in_memory,
             'lost-data': self._handle_lost_data,
             'cancelled-key': self._handle_cancelled_key,
@@ -601,6 +609,9 @@ def __init__(self, address=None, loop=None, timeout=no_default,
                                      serializers=serializers,
                                      deserializers=deserializers)
 
+        for ext in extensions:
+            ext(self)
+
         self.start(timeout=timeout)
 
         from distributed.recreate_exceptions import ReplayExceptionClient
@@ -958,7 +969,7 @@ def _handle_report(self):
                             break
 
                         try:
-                            handler = self._handlers[op]
+                            handler = self._stream_handlers[op]
                             handler(**msg)
                         except Exception as e:
                             logger.exception(e)

distributed/core.py

@@ -17,7 +17,7 @@
 from tornado.ioloop import IOLoop
 from tornado.locks import Event
 
-from .compatibility import PY3
+from .compatibility import PY3, get_thread_identity
 from .comm import (connect, listen, CommClosedError,
                    normalize_address,
                    unparse_host_port, get_address_host_port)
@@ -114,7 +114,7 @@ def __init__(self, handlers, stream_handlers=None, connection_limit=512,
 
         self.listener = None
         self.io_loop = io_loop or IOLoop.current()
-        self.loop = io_loop
+        self.loop = self.io_loop
 
         # Statistics counters for various events
         with ignoring(ImportError):
@@ -139,6 +139,14 @@ def __init__(self, handlers, stream_handlers=None, connection_limit=512,
         )
         self.periodic_callbacks['tick'] = pc
 
+        self.thread_id = 0
+
+        @gen.coroutine
+        def set_thread_ident():
+            self.thread_id = get_thread_identity()
+
+        self.io_loop.add_callback(set_thread_ident)
+
         self.__stopped = False
 
     def start_periodic_callbacks(self):
@@ -153,7 +161,7 @@ def start_pcs():
             for pc in self.periodic_callbacks.values():
                 if not pc.is_running():
                     pc.start()
-        self.loop.add_callback(start_pcs)
+        self.io_loop.add_callback(start_pcs)
 
     def stop(self):
         if not self.__stopped: