From 5e77efb446898337062ba3272b0f70e918a8a804 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 17 Aug 2022 13:03:58 -0700 Subject: [PATCH 1/2] Add `ip_address` as argument to listener Allows binding the listener to a specific IP address, instead of always listening on `0.0.0.0`. --- ucp/_libs/ucx_listener.pyx | 14 ++++++++++++-- ucp/core.py | 10 +++++++++- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/ucp/_libs/ucx_listener.pyx b/ucp/_libs/ucx_listener.pyx index 6436a9791..de197bf2a 100644 --- a/ucp/_libs/ucx_listener.pyx +++ b/ucp/_libs/ucx_listener.pyx @@ -48,6 +48,9 @@ cdef class UCXListener(UCXObject): Extra arguments to the call-back function cb_kwargs: dict, optional Extra keyword arguments to the call-back function + ip_address: str, optional + IP address to bind the listener to. Binds to `0.0.0.0` if not + specified. Returns ------- @@ -69,7 +72,8 @@ cdef class UCXListener(UCXObject): uint16_t port, cb_func, tuple cb_args=None, - dict cb_kwargs=None + dict cb_kwargs=None, + str ip_address=None, ): if cb_args is None: cb_args = () @@ -90,7 +94,13 @@ cdef class UCXListener(UCXObject): ) params.conn_handler.cb = _listener_cb params.conn_handler.arg = self.cb_data - if c_util_set_sockaddr(¶ms.sockaddr, NULL, port): + + cdef alloc_sockaddr_ret = ( + c_util_set_sockaddr(¶ms.sockaddr, NULL, port) + if ip_address is None else + c_util_set_sockaddr(¶ms.sockaddr, ip_address.encode(), port) + ) + if alloc_sockaddr_ret: raise MemoryError("Failed allocation of sockaddr") cdef ucs_status_t status = ucp_listener_create( diff --git a/ucp/core.py b/ucp/core.py index 6f5ddf3c0..6e484e6aa 100644 --- a/ucp/core.py +++ b/ucp/core.py @@ -234,6 +234,7 @@ def create_listener( callback_func, port=0, endpoint_error_handling=True, + ip_address=None, ): """Create and start a listener to accept incoming connections @@ -257,6 +258,9 @@ def create_listener( but prevents a process from terminating unexpectedly that may happen when disabled. If `False` endpoint endpoint error handling is disabled. + ip_address: str, optional + IP address to bind the listener to. Binds to `0.0.0.0` if not + specified. Returns ------- @@ -274,6 +278,7 @@ def create_listener( port=port, cb_func=_listener_handler, cb_args=(callback_func, self, endpoint_error_handling), + ip_address=ip_address, ) ) return ret @@ -992,11 +997,14 @@ def register_am_allocator(allocator, allocator_type): return _get_ctx().register_am_allocator(allocator, allocator_type) -def create_listener(callback_func, port=None, endpoint_error_handling=True): +def create_listener( + callback_func, port=None, endpoint_error_handling=True, ip_address=None +): return _get_ctx().create_listener( callback_func, port, endpoint_error_handling=endpoint_error_handling, + ip_address=ip_address, ) From 60fb15eed22fa4ab950c98a05e0966d172784e68 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 17 Aug 2022 13:16:56 -0700 Subject: [PATCH 2/2] Add tests to verify listener's IP/port binding --- tests/test_listener.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 tests/test_listener.py diff --git a/tests/test_listener.py b/tests/test_listener.py new file mode 100644 index 000000000..2a7c9b2bb --- /dev/null +++ b/tests/test_listener.py @@ -0,0 +1,38 @@ +import pytest + +import ucp + + +@pytest.mark.asyncio +async def test_bind_ip(): + listener = ucp.create_listener(lambda: None, ip_address="127.0.0.1") + + assert isinstance(listener.port, int) + assert listener.port >= 1024 + + assert isinstance(listener.ip, str) + assert listener.ip == "127.0.0.1" + + +@pytest.mark.asyncio +async def test_bind_port(): + listener = None + port = None + + for i in range(10000, 20000): + try: + listener = ucp.create_listener( + lambda: None, + i, + ) + except ucp.UCXError: + # Port already in use, try another + continue + else: + port = i + break + + assert isinstance(listener.port, int) + assert listener.port == port + + assert isinstance(listener.ip, str)