[experimental] support socket.FD on draft thread-safe sockets (#2103)

gated by warning, may be removed if it causes problems

will have no effect if/when getsockopt(ZMQ_FD) is implemented for threadsafe sockets

---------

Co-authored-by: Matthew Hawn <[email protected]>

Author: minrk

tests/test_asyncio.py

@@ -405,6 +405,33 @@ async def test_poll_leak():
         assert len(s._recv_futures) == 0
 
 
+async def test_draft_asyncio():
+    if not zmq.has("draft"):
+        pytest.skip("draft API")
+    if zmq.zmq_version_info() < (4, 3, 2):
+        pytest.skip("requires libzmq 4.3.2 for zmq_poller_fd")
+    with zmq.asyncio.Context() as ctx, ctx.socket(zmq.CLIENT) as client, ctx.socket(
+        zmq.SERVER
+    ) as server:
+        server.bind_to_random_port("tcp://127.0.0.1")
+        client.connect(server.last_endpoint)
+        server.rcvtimeo = client.rcvtimeo = 100
+        with pytest.raises(zmq.Again):
+            await server.recv()
+        with pytest.raises(zmq.Again):
+            await client.recv()
+        server.rcvtimeo = client.rcvtimeo = server.sndtimeo = client.sndtimeo = 3000
+        recv_future = asyncio.ensure_future(server.recv(copy=False))
+        assert not recv_future.done()
+        await client.send(b'request')
+        msg = await recv_future
+        recv_future = asyncio.ensure_future(client.recv())
+        assert not recv_future.done()
+        await server.send(msg)
+        response = await recv_future
+        assert response == b'request'
+
+
 class ProcessForTeardownTest(Process):
     def run(self):
         """Leave context, socket and event loop upon implicit disposal"""

tests/test_draft.py

@@ -6,15 +6,13 @@
 import pytest
 
 import zmq
-from zmq_test_utils import BaseZMQTestCase
+from zmq_test_utils import BaseZMQTestCase, skip_pypy
 
+pytestmark = pytest.mark.skipif(not zmq.DRAFT_API, reason="draft api unavailable")
 
-class TestDraftSockets(BaseZMQTestCase):
-    def setUp(self):
-        if not zmq.DRAFT_API:
-            pytest.skip("draft api unavailable")
-        super().setUp()
 
+class TestDraftSockets(BaseZMQTestCase):
+    @skip_pypy
     def test_client_server(self):
         client, server = self.create_bound_pair(zmq.CLIENT, zmq.SERVER)
         client.send(b'request')
@@ -24,6 +22,15 @@ def test_client_server(self):
         reply = self.recv(client)
         assert reply == b'reply'
 
+    def test_client_server_frame(self):
+        client, server = self.create_bound_pair(zmq.CLIENT, zmq.SERVER)
+        client.send(b'request')
+        msg = self.recv(server, copy=False)
+        server.send(msg)
+        reply = self.recv(client)
+        assert reply == b'request'
+
+    @skip_pypy
     def test_radio_dish(self):
         dish, radio = self.create_bound_pair(zmq.DISH, zmq.RADIO)
         dish.rcvtimeo = 250
@@ -45,3 +52,13 @@ def test_radio_dish(self):
                 received_count += 1
         # assert that we got *something*
         assert len(received.intersection(sent)) >= 5
+
+
+def test_draft_fd():
+    if zmq.zmq_version_info() < (4, 3, 2):
+        pytest.skip("requires libzmq 4.3.2 for zmq_poller_fd")
+    with zmq.Context() as ctx, ctx.socket(zmq.SERVER) as s:
+        fd = s.FD
+        assert isinstance(fd, int)
+        fd_2 = s.get(zmq.FD)
+        assert fd_2 == fd

zmq/backend/cffi/_cdefs.h

@@ -68,6 +68,12 @@ typedef struct
 
 int zmq_poll(zmq_pollitem_t *items, int nitems, long timeout);
 
+// draft poller
+void *zmq_poller_new ();
+int zmq_poller_destroy (void **poller_p_);
+int zmq_poller_add (void *poller_, void *socket_, void *user_data_, short events_);
+int zmq_poller_fd (void *poller_, ZMQ_FD_T *fd_);
+
 // miscellany
 void * memcpy(void *restrict s1, const void *restrict s2, size_t n);
 void * malloc(size_t sz);

zmq/backend/cffi/socket.py

@@ -4,10 +4,11 @@
 # Distributed under the terms of the Modified BSD License.
 
 import errno as errno_mod
+import warnings
 
 import zmq
 from zmq.constants import SocketOption, _OptType
-from zmq.error import ZMQError, _check_rc
+from zmq.error import ZMQError, _check_rc, _check_version
 
 from ._cffi import ffi
 from ._cffi import lib as C
@@ -49,7 +50,8 @@ def value_binary_data(val, length):
     return ffi.new(f'char[{length + 1:d}]', val), ffi.sizeof('char') * length
 
 
-ZMQ_FD_64BIT = ffi.sizeof('ZMQ_FD_T') == 8
+_fd_size = ffi.sizeof('ZMQ_FD_T')
+ZMQ_FD_64BIT = _fd_size == 8
 
 IPC_PATH_MAX_LEN = C.get_ipc_path_max_len()
 
@@ -100,6 +102,8 @@ class Socket:
     _closed = None
     _ref = None
     _shadow = False
+    _draft_poller = None
+    _draft_poller_ptr = None
     copy_threshold = 0
 
     def __init__(self, context=None, socket_type=None, shadow=0, copy_threshold=None):
@@ -108,6 +112,7 @@ def __init__(self, context=None, socket_type=None, shadow=0, copy_threshold=None
         self.copy_threshold = copy_threshold
 
         self.context = context
+        self._draft_poller = self._draft_poller_ptr = None
         if shadow:
             self._zmq_socket = ffi.cast("void *", shadow)
             self._shadow = True
@@ -152,6 +157,10 @@ def closed(self):
     def close(self, linger=None):
         rc = 0
         if not self._closed and hasattr(self, '_zmq_socket'):
+            if self._draft_poller_ptr is not None:
+                rc = C.zmq_poller_destroy(self._draft_poller_ptr)
+                self._draft_poller = self._draft_poller_ptr = None
+
             if self._zmq_socket is not None:
                 if linger is not None:
                     self.set(zmq.LINGER, linger)
@@ -242,11 +251,55 @@ def get(self, option):
         else:
             opt_type = option._opt_type
 
+        if option == zmq.FD and self._draft_poller is not None:
+            c_value_pointer, _ = new_pointer_from_opt(option)
+            C.zmq_poller_fd(self._draft_poller, ffi.cast('void*', c_value_pointer))
+            return int(c_value_pointer[0])
+
         c_value_pointer, c_sizet_pointer = new_pointer_from_opt(option, length=255)
 
-        _retry_sys_call(
-            C.zmq_getsockopt, self._zmq_socket, option, c_value_pointer, c_sizet_pointer
-        )
+        try:
+            _retry_sys_call(
+                C.zmq_getsockopt,
+                self._zmq_socket,
+                option,
+                c_value_pointer,
+                c_sizet_pointer,
+            )
+        except ZMQError as e:
+            if (
+                option == SocketOption.FD
+                and e.errno == zmq.Errno.EINVAL
+                and self.get(SocketOption.THREAD_SAFE)
+            ):
+                _check_version((4, 3, 2), "draft socket FD support via zmq_poller_fd")
+                if not zmq.has('draft'):
+                    raise RuntimeError("libzmq must be built with draft support")
+                warnings.warn(zmq.error.DraftFDWarning(), stacklevel=2)
+
+                # create a poller and retrieve its fd
+                self._draft_poller_ptr = ffi.new("void*[1]")
+                self._draft_poller_ptr[0] = self._draft_poller = C.zmq_poller_new()
+                if self._draft_poller == ffi.NULL:
+                    # failed (why?), raise original error
+                    self._draft_poller_ptr = self._draft_poller = None
+                    raise
+                # register self with poller
+                rc = C.zmq_poller_add(
+                    self._draft_poller,
+                    self._zmq_socket,
+                    ffi.NULL,
+                    zmq.POLLIN | zmq.POLLOUT,
+                )
+                _check_rc(rc)
+                # use poller fd as proxy for ours
+                rc = C.zmq_poller_fd(
+                    self._draft_poller, ffi.cast('void *', c_value_pointer)
+                )
+                _check_rc(rc)
+                return int(c_value_pointer[0])
+            else:
+                raise
 
         sz = c_sizet_pointer[0]
         v = value_from_opt_pointer(option, c_value_pointer, sz)

zmq/backend/cython/_zmq.pxd

@@ -44,6 +44,7 @@ cdef class Socket:
     cdef public bint _closed    # bool property for a closed socket.
     cdef public int copy_threshold # threshold below which pyzmq will always copy messages
     cdef int _pid               # the pid of the process which created me (for fork safety)
+    cdef void *_draft_poller  # The C handle for the zmq poller for draft socket zmq.FD support
 
     # cpdef methods for direct-cython access:
     cpdef object send(self, data, int flags=*, bint copy=*, bint track=*)

zmq/backend/cython/_zmq.py

@@ -88,13 +88,16 @@
     ZMQ_ENOTSOCK,
     ZMQ_ETERM,
     ZMQ_EVENT_ALL,
+    ZMQ_FD,
     ZMQ_IDENTITY,
     ZMQ_IO_THREADS,
     ZMQ_LINGER,
     ZMQ_POLLIN,
+    ZMQ_POLLOUT,
     ZMQ_RCVMORE,
     ZMQ_ROUTER,
     ZMQ_SNDMORE,
+    ZMQ_THREAD_SAFE,
     ZMQ_TYPE,
     _zmq_version,
     fd_t,
@@ -131,6 +134,10 @@
     zmq_msg_set_routing_id,
     zmq_msg_size,
     zmq_msg_t,
+    zmq_poller_add,
+    zmq_poller_destroy,
+    zmq_poller_fd,
+    zmq_poller_new,
     zmq_pollitem_t,
     zmq_proxy,
     zmq_proxy_steerable,
@@ -700,6 +707,7 @@ def __init__(
     ):
         # pre-init
         self.handle = NULL
+        self._draft_poller = NULL
         self._pid = 0
         self._shadow = False
         self.context = None
@@ -756,6 +764,12 @@ def close(self, linger: int | None = None):
         if self.handle != NULL and not self._closed and getpid() == self._pid:
             if setlinger:
                 zmq_setsockopt(self.handle, ZMQ_LINGER, address(linger_c), sizeof(int))
+
+            # teardown draft poller
+            if self._draft_poller != NULL:
+                zmq_poller_destroy(address(self._draft_poller))
+                self._draft_poller = NULL
+
             rc = zmq_close(self.handle)
             if rc < 0 and _zmq_errno() != ENOTSOCK:
                 # ignore ENOTSOCK (closed by Context)
@@ -834,6 +848,10 @@ def get(self, option: C.int):
 
         See the 0MQ API documentation for details on specific options.
 
+        .. versionchanged:: 27
+            Added experimental support for ZMQ_FD for draft sockets via `zmq_poller_fd`.
+            Requires libzmq >=4.3.2 built with draft support.
+
         Parameters
         ----------
         option : int
@@ -882,11 +900,49 @@ def get(self, option: C.int):
                 self.handle, option, cast(p_void, address(optval_int64_c)), address(sz)
             )
             result = optval_int64_c
+        elif option == ZMQ_FD and self._draft_poller != NULL:
+            # draft sockets use FD of a draft zmq_poller as proxy
+            rc = zmq_poller_fd(self._draft_poller, address(optval_fd_c))
+            _check_rc(rc)
+            result = optval_fd_c
         elif opt_type == _OptType.fd:
             sz = sizeof(fd_t)
-            _getsockopt(
-                self.handle, option, cast(p_void, address(optval_fd_c)), address(sz)
-            )
+            try:
+                _getsockopt(
+                    self.handle, option, cast(p_void, address(optval_fd_c)), address(sz)
+                )
+            except ZMQError as e:
+                # threadsafe sockets don't support ZMQ_FD (yet!)
+                # fallback on zmq_poller_fd as proxy with the same behavior
+                # until libzmq fixes this.
+                # if upstream fixes it, this branch will never be taken
+                if (
+                    option == ZMQ_FD
+                    and e.errno == zmq.Errno.EINVAL
+                    and self.get(ZMQ_THREAD_SAFE)
+                ):
+                    _check_version(
+                        (4, 3, 2), "draft socket FD support via zmq_poller_fd"
+                    )
+                    if not zmq.has('draft'):
+                        raise RuntimeError("libzmq must be built with draft support")
+                    warnings.warn(zmq.error.DraftFDWarning(), stacklevel=2)
+
+                    # create a poller and retrieve its fd
+                    self._draft_poller = zmq_poller_new()
+                    if self._draft_poller == NULL:
+                        # failed (why?), raise original error
+                        raise
+                    # register self with poller
+                    rc = zmq_poller_add(
+                        self._draft_poller, self.handle, NULL, ZMQ_POLLIN | ZMQ_POLLOUT
+                    )
+                    _check_rc(rc)
+                    # use poller fd as proxy for ours
+                    rc = zmq_poller_fd(self._draft_poller, address(optval_fd_c))
+                    _check_rc(rc)
+                else:
+                    raise
             result = optval_fd_c
         else:
             # default is to assume int, which is what most new sockopts will be

zmq/backend/cython/libzmq.pxd

@@ -119,3 +119,10 @@ cdef extern from "zmq.h" nogil:
     uint32_t zmq_msg_routing_id(zmq_msg_t *msg)
     int zmq_msg_set_group(zmq_msg_t *msg, const char *group)
     const char *zmq_msg_group(zmq_msg_t *msg)
+
+    void *zmq_poller_new ()
+    int zmq_poller_destroy (void **poller_p_)
+    int zmq_poller_add (void *poller_, void *socket_, void *user_data_, short events_)
+    int zmq_poller_modify (void *poller_, void *socket_, short events_)
+    int zmq_poller_remove (void *poller_, void *socket_)
+    int zmq_poller_fd (void *poller_, fd_t *fd_)

zmq/error.py

@@ -7,6 +7,22 @@
 from errno import EINTR
 
 
+class DraftFDWarning(RuntimeWarning):
+    """Warning for using experimental FD on draft sockets.
+
+    .. versionadded:: 27
+    """
+
+    def __init__(self, msg=""):
+        if not msg:
+            msg = (
+                "pyzmq's back-fill socket.FD support on thread-safe sockets is experimental, and may be removed."
+                " This warning will go away automatically if/when libzmq implements socket.FD on thread-safe sockets."
+                " You can suppress this warning with `warnings.simplefilter('ignore', zmq.error.DraftFDWarning)"
+            )
+        super().__init__(msg)
+
+
 class ZMQBaseError(Exception):
     """Base exception class for 0MQ errors in Python."""
 
@@ -201,6 +217,7 @@ def _check_version(
 
 
 __all__ = [
+    "DraftFDWarning",
     "ZMQBaseError",
     "ZMQBindError",
     "ZMQError",

zmq/utils/zmq_compat.h

@@ -60,6 +60,9 @@
     #if ZMQ_VERSION >= 40200
         #define PYZMQ_DRAFT_42
     #endif
+    #if ZMQ_VERSION >= 40302
+        #define PYZMQ_DRAFT_432
+    #endif
 #endif
 #ifndef PYZMQ_DRAFT_42
     #define zmq_join(s, group) _missing
@@ -68,6 +71,14 @@
     #define zmq_msg_routing_id(msg) 0
     #define zmq_msg_set_group(msg, group) _missing
     #define zmq_msg_group(msg) NULL
+    #define zmq_poller_new() NULL
+    #define zmq_poller_destroy(poller_p) _missing
+    #define zmq_poller_add(poller, socket, userdata, events) _missing
+    #define zmq_poller_modify(poller, socket, events) _missing
+    #define zmq_poller_remove(poller, socket) _missing
+#endif
+#ifndef PYZMQ_DRAFT_432
+    #define zmq_poller_fd(poller, fd) _missing
 #endif
 
 #if ZMQ_VERSION >= 40100