codypiersall / pynng

Python bindings for Nanomsg Next Generation.
https://pynng.readthedocs.io
MIT License
260 stars 58 forks source link

Support for `recv_into` #95

Open sk1p opened 2 years ago

sk1p commented 2 years ago

My use case was sending numpy arrays via nng, where it really makes sense to reduce the number of allocations and copies that need to be performed.

Reading the code, I saw that there is already the general idea to support a recv_into operation on sockets. I built a prototype that looks like this, based on the normal read method:

def recv_into(s, buffer):
    """Receive data on the socket s into buffer
    If the request times out the exception
    :class:`pynng.Timeout` is raised.  If the socket cannot perform that
    operation (e.g., a :class:`Pub0`, which can only
    :meth:`~Socket.send`), the exception :class:`pynng.NotSupported`
    is raised.
    """
    from pynng import ffi, lib
    from pynng.nng import check_err
    # we DON'T set lib.NNG_FLAG_ALLOC here, because we want to give
    # our own recv buffer
    flags = 0
    size_t = ffi.new('size_t *')
    size_t[0] = len(buffer)
    ret = lib.nng_recv(s.socket, buffer, size_t, flags)
    check_err(ret)
    assert size_t[0] == len(buffer)

Used like this:

recv_mem = np.zeros(arr_shape, dtype=np.float32)
recv_buf = ffi.from_buffer(
    recv_mem,
    require_writable=True,
)
recv_into(your_socket, recv_buf)
# the received data is now available in `recv_mem`

If there is interest, I could convert this into a PR including test cases etc. and add proper support for returning the part of the buffer that was written to etc. (this prototype was only used for evaluation purposes / benchmarking). Let me know if anyone is interested!

aqc-carlodri commented 2 years ago

Thank you, yes I would definitely be interested in this!