ENH: extend the address normalization to allow domain sockets

Author: mrakitin

src/bluesky/callbacks/zmq.py

@@ -16,10 +16,41 @@
 import copy
 import pickle
 import warnings
+from typing import Union
 
 from ..run_engine import Dispatcher, DocumentNames
 
 
+def _normalize_address(inp: Union[str, tuple, int]):
+    if isinstance(inp, str):
+        if "://" in inp:
+            protocol, _, rest_str = inp.partition("://")
+        else:
+            protocol = "tcp"
+            rest_str = inp
+    elif isinstance(inp, tuple):
+        if inp[0] in ["tcp", "ipc"]:
+            protocol, *rest = inp
+        else:
+            protocol = "tcp"
+            rest = list(inp)
+        if protocol == "tcp":
+            if len(rest) == 2:
+                rest_str = ":".join(str(r) for r in rest)
+            else:
+                (rest_str,) = rest
+        else:
+            (rest_str,) = rest
+    elif isinstance(inp, int):
+        protocol = "tcp"
+        rest_str = f"localhost:{inp}"
+
+    else:
+        raise TypeError(f"Input expected to be str or tuple, not {type(inp)}")
+
+    return f"{protocol}://{rest_str}"
+
+
 class Bluesky0MQDecodeError(Exception):
     """Custom exception class for things that go wrong reading message from wire."""
 
@@ -73,20 +104,20 @@ def __init__(self, address, *, prefix=b"", RE=None, zmq=None, serializer=pickle.
             raise ValueError(f"prefix {prefix!r} may not contain b' '")
         if zmq is None:
             import zmq
-        if isinstance(address, str):
-            address = address.split(":", maxsplit=1)
-        self.address = (address[0], int(address[1]))
+
+        self.address = _normalize_address(address)
         self.RE = RE
-        url = "tcp://%s:%d" % self.address
+
         self._prefix = bytes(prefix)
         self._context = zmq.Context()
         self._socket = self._context.socket(zmq.PUB)
-        self._socket.connect(url)
+        self._socket.connect(self.address)
         if RE:
             self._subscription_token = RE.subscribe(self)
         self._serializer = serializer
 
     def __call__(self, name, doc):
+        print(f"{name = }\n{doc}\n")
         doc = copy.deepcopy(doc)
         message = b" ".join([self._prefix, name.encode(), self._serializer(doc)])
         self._socket.send(message)
@@ -102,23 +133,40 @@ class Proxy:
     """
     Start a 0MQ proxy on the local host.
 
+    The addresses can be specified flexibly.  It is best to use
+    a domain_socket (available on unix):
+
+     - ``'icp:///tmp/domain_socket'``
+     - ``('ipc', '/tmp/domain_socket')``
+
+    tcp sockets are also supported:
+
+     - ``'tcp://localhost:6557'``
+     - ``6657``  (implicitly binds to ``'tcp://localhost:6657'``
+     - ``('tcp', 'localhost', 6657)``
+     - ``('localhost', 6657)``
+
     Parameters
     ----------
-    in_port : int, optional
-        Port that RunEngines should broadcast to. If None, a random port is
-        used.
-    out_port : int, optional
-        Port that subscribers should subscribe to. If None, a random port is
-        used.
+    in_address : str or tuple or int, optional
+        Address that RunEngines should broadcast to.
+
+        If None, a random tcp port on all interfaces is used.
+
+    out_address : str or tuple or int, optional
+        Address that subscribers should subscribe to.
+
+        If None, a random tcp port on all interfaces is used.
+
     zmq : object, optional
         By default, the 'zmq' module is imported and used. Anything else
         mocking its interface is accepted.
 
     Attributes
     ----------
-    in_port : int
+    in_port : int or str
         Port that RunEngines should broadcast to.
-    out_port : int
+    out_port : int or str
         Port that subscribers should subscribe to.
     closed : boolean
         True if the Proxy has already been started and subsequently
@@ -146,7 +194,7 @@ class Proxy:
     >>> proxy.start()  # runs until interrupted
     """
 
-    def __init__(self, in_port=None, out_port=None, *, zmq=None):
+    def __init__(self, in_address=None, out_address=None, *, zmq=None):
         if zmq is None:
             import zmq
         self.zmq = zmq
@@ -155,19 +203,22 @@ def __init__(self, in_port=None, out_port=None, *, zmq=None):
             context = zmq.Context(1)
             # Socket facing clients
             frontend = context.socket(zmq.SUB)
-            if in_port is None:
+            if in_address is None:
                 in_port = frontend.bind_to_random_port("tcp://*")
             else:
-                frontend.bind("tcp://*:%d" % in_port)
+                in_address = _normalize_address(in_address)
+                in_port = frontend.bind(in_address)
 
             frontend.setsockopt_string(zmq.SUBSCRIBE, "")
 
             # Socket facing services
             backend = context.socket(zmq.PUB)
-            if out_port is None:
+            if out_address is None:
                 out_port = backend.bind_to_random_port("tcp://*")
             else:
-                backend.bind("tcp://*:%d" % out_port)
+                out_address = _normalize_address(out_address)
+                out_port = backend.bind(out_address)
+
         except BaseException:
             # Clean up whichever components we have defined so far.
             try:
@@ -257,10 +308,8 @@ def __init__(
             import zmq
         if zmq_asyncio is None:
             import zmq.asyncio as zmq_asyncio
-        if isinstance(address, str):
-            address = address.split(":", maxsplit=1)
         self._deserializer = deserializer
-        self.address = (address[0], int(address[1]))
+        self.address = _normalize_address(address)
 
         if loop is None:
             loop = asyncio.new_event_loop()
@@ -274,8 +323,7 @@ def __finish_setup():
             self._context = zmq_asyncio.Context()
             self._socket = self._context.socket(zmq.SUB)
 
-            url = "tcp://%s:%d" % self.address
-            self._socket.connect(url)
+            self._socket.connect(self.address)
             self._socket.setsockopt_string(zmq.SUBSCRIBE, "")
 
         self.__factory = __finish_setup
@@ -332,6 +380,7 @@ async def _poll(self):
                             f"\n\n{e}"
                         )
                         continue
+                print(f"{name = }\n{doc}")
                 self.loop.call_soon(self.process, DocumentNames[name], doc)
 
     def start(self):

src/bluesky/commandline/zmq_proxy.py

@@ -7,15 +7,15 @@
 logger = logging.getLogger("bluesky")
 
 
-def start_dispatcher(host, port, logfile=None):
+def start_dispatcher(out_address, logfile=None):
     """The dispatcher function
     Parameters
     ----------
     logfile : string
         string come from user command. ex --logfile=temp.log
         logfile will be "temp.log". logfile could be empty.
     """
-    dispatcher = RemoteDispatcher((host, port))
+    dispatcher = RemoteDispatcher(out_address)
     if logfile is not None:
         raise ValueError(
             "Parameter 'logfile' is deprecated and will be removed in future releases. "
@@ -41,8 +41,8 @@ def log_writer(name, doc):
 def main():
     DESC = "Start a 0MQ proxy for publishing bluesky documents over a network."
     parser = argparse.ArgumentParser(description=DESC)
-    parser.add_argument("in_port", type=int, nargs=1, help="port that RunEngines should broadcast to")
-    parser.add_argument("out_port", type=int, nargs=1, help="port that subscribers should subscribe to")
+    parser.add_argument("--in-address", help="port that RunEngines should broadcast to")
+    parser.add_argument("--out-address", help="port that subscribers should subscribe to")
     parser.add_argument(
         "--verbose",
         "-v",
@@ -51,8 +51,6 @@ def main():
     )
     parser.add_argument("--logfile", type=str, help="Redirect logging output to a file on disk.")
     args = parser.parse_args()
-    in_port = args.in_port[0]
-    out_port = args.out_port[0]
 
     if args.verbose:
         from bluesky.log import config_bluesky_logging
@@ -64,11 +62,11 @@ def main():
         else:
             config_bluesky_logging(level=level)
         # Set daemon to kill all threads upon IPython exit
-        threading.Thread(target=start_dispatcher, args=("localhost", out_port), daemon=True).start()
+        threading.Thread(target=start_dispatcher, args=(args.out_address), daemon=True).start()
 
     print("Connecting...")
-    proxy = Proxy(in_port, out_port)
-    print("Receiving on port %d; publishing to port %d." % (in_port, out_port))
+    proxy = Proxy(args.in_address, args.out_address)
+    print("Receiving on address %s; publishing to address %s." % (args.in_address, args.out_address))
     print("Use Ctrl+C to exit.")
     try:
         proxy.start()

src/bluesky/tests/test_zmq.py

@@ -11,7 +11,7 @@
 from event_model import sanitize_doc
 
 from bluesky import Msg
-from bluesky.callbacks.zmq import Proxy, Publisher, RemoteDispatcher
+from bluesky.callbacks.zmq import Proxy, Publisher, RemoteDispatcher, _normalize_address
 from bluesky.plans import count
 from bluesky.tests import uses_os_kill_sigint
 
@@ -126,21 +126,6 @@ def delayed_sigint(delay):
     gc.collect()
 
 
-@pytest.mark.parametrize("host", ["localhost:5555", ("localhost", 5555)])
-def test_zmq_RD_ports_spec(host):
-    # test that two ways of specifying address are equivalent
-    d = RemoteDispatcher(host)
-    assert d.address == ("localhost", 5555)
-    assert d._socket is None
-    assert d._context is None
-    assert not d.closed
-    d.stop()
-    assert d._socket is None
-    assert d._context is None
-    assert d.closed
-    del d
-
-
 def test_zmq_no_RE(RE):
     # COMPONENT 1
     # Run a 0MQ proxy on a separate process.
@@ -345,3 +330,46 @@ def local_cb(name, doc):
     ra = sanitize_doc(remote_accumulator)
     la = sanitize_doc(local_accumulator)
     assert ra == la
+
+
+@pytest.mark.parametrize(
+    "host",
+    ["localhost:5555", ("localhost", 5555)],
+)
+def test_zmq_RD_ports_spec(host):
+    # test that two ways of specifying address are equivalent
+    d = RemoteDispatcher(host)
+    assert d.address == "tcp://localhost:5555"
+    assert d._socket is None
+    assert d._context is None
+    assert not d.closed
+    d.stop()
+    assert d._socket is None
+    assert d._context is None
+    assert d.closed
+    del d
+
+
+@pytest.mark.parametrize(
+    "address",
+    [
+        ("localhost", "tcp://localhost"),
+        ("localhost:9", "tcp://localhost:9"),
+        ("remote.host", "tcp://remote.host"),
+        ("remote.host:9", "tcp://remote.host:9"),
+        ("tcp://remote.host", "tcp://remote.host"),
+        ("tcp://localhost", "tcp://localhost"),
+        ("tcp://localhost:9", "tcp://localhost:9"),
+        ("tcp://remote.host:9", "tcp://remote.host:9"),
+        ("ipc:///tmp/path", "ipc:///tmp/path"),
+        (("localhost",), "tcp://localhost"),
+        (("localhost", 9), "tcp://localhost:9"),
+        (("ipc", "/tmp/path"), "ipc:///tmp/path"),
+        (("tcp", "localhost"), "tcp://localhost"),
+        (("tcp", "localhost", 9), "tcp://localhost:9"),
+        (("tcp", "localhost", "9"), "tcp://localhost:9"),
+    ],
+)
+def test_address_normaliaztion(address):
+    inp, outp = address
+    assert _normalize_address(inp) == outp