Add UDP support (#46)

* Implement UDP socket support and related transport management in the event loop

* Refactor UDP socket connection and streamline datagram endpoint tests

* Refactor UDP impl

* Code lint

* Add UDP write buffering

---------

Co-authored-by: Giovanni Barillari <giovanni.barillari@sentry.io>

Author: DVNghiem

Makefile

@@ -27,6 +27,7 @@ lint-rust:
 		-D warnings \
 		-W clippy::pedantic \
 		-W clippy::dbg_macro \
+		-A clippy::blocks_in_conditions \
 		-A clippy::cast-possible-truncation \
 		-A clippy::cast-sign-loss \
 		-A clippy::declare-interior-mutable-const \

rloop/loop.py

@@ -655,12 +655,79 @@ async def create_datagram_endpoint(
         family=0,
         proto=0,
         flags=0,
-        reuse_address=None,
+        #: not in stdlib
+        # reuse_address=None,
         reuse_port=None,
         allow_broadcast=None,
         sock=None,
     ):
-        raise NotImplementedError
+        if sock is not None:
+            if getattr(sock, 'type', None) != socket.SOCK_DGRAM:
+                raise ValueError(f'A datagram socket was expected, got {sock!r}')
+            if any((local_addr, remote_addr, family, proto, flags, reuse_port, allow_broadcast)):
+                raise ValueError('socket modifier keyword arguments can not be used when sock is specified.')
+            sock.setblocking(False)
+            r_addr = None
+        else:
+            if not (local_addr or remote_addr):
+                if family == 0:
+                    raise ValueError('unexpected address family')
+                addr_info = (family, proto, None, None)
+            elif hasattr(socket, 'AF_UNIX') and family == socket.AF_UNIX:
+                for addr in (local_addr, remote_addr):
+                    if addr is not None and not isinstance(addr, str):
+                        raise TypeError('string is expected')
+                addr_info = (family, proto, local_addr, remote_addr)
+            else:
+                addr_info, infos = None, None
+                for addr in (local_addr, remote_addr):
+                    if addr is None:
+                        continue
+                    if not (isinstance(addr, tuple) and len(addr) == 2):
+                        raise TypeError('2-tuple is expected')
+                    infos = await self._ensure_resolved(
+                        addr, family=family, type=socket.SOCK_DGRAM, proto=proto, flags=flags
+                    )
+                    break
+
+                if not infos:
+                    raise OSError('getaddrinfo() returned empty list')
+                if local_addr is not None:
+                    addr_info = (infos[0][0], infos[0][2], infos[0][4], None)
+                if remote_addr is not None:
+                    addr_info = (infos[0][0], infos[0][2], None, infos[0][4])
+                if not addr_info:
+                    raise ValueError('can not get address information')
+
+            sock = None
+            r_addr = None
+            sfam, spro, sladdr, sraddr = addr_info
+            try:
+                sock = socket.socket(family=sfam, type=socket.SOCK_DGRAM, proto=spro)
+                #: not in stdlib
+                # if reuse_address:
+                #     sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+                if reuse_port:
+                    _set_reuseport(sock)
+                if allow_broadcast:
+                    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
+                sock.setblocking(False)
+                if sladdr:
+                    sock.bind(sladdr)
+                if sraddr:
+                    if not allow_broadcast:
+                        await self.sock_connect(sock, sraddr)
+                    r_addr = sraddr
+            except OSError:
+                if sock is not None:
+                    sock.close()
+                raise
+
+        # Create the transport
+        transport, protocol = self._udp_conn((sock.fileno(), sock.family), protocol_factory, r_addr)
+        # sock is now owned by the transport, prevent close
+        sock.detach()
+        return transport, protocol
 
     #: pipes and subprocesses methods
     async def connect_read_pipe(self, protocol_factory, pipe):

src/event_loop.rs

@@ -19,13 +19,15 @@ use crate::{
     server::Server,
     tcp::{TCPReadHandle, TCPServer, TCPServerRef, TCPTransport, TCPWriteHandle},
     time::Timer,
+    udp::{UDPReadHandle, UDPTransport, UDPWriteHandle},
 };
 
 enum IOHandle {
     Py(PyHandleData),
     Signals,
     TCPListener(TCPListenerHandleData),
     TCPStream(Interest),
+    UDPSocket(Interest),
 }
 
 struct PyHandleData {
@@ -73,6 +75,7 @@ pub struct EventLoop {
     task_factory: RwLock<PyObject>,
     tcp_lstreams: papaya::HashMap<usize, papaya::HashSet<usize>>,
     tcp_transports: papaya::HashMap<usize, Py<TCPTransport>>,
+    udp_transports: papaya::HashMap<usize, Py<UDPTransport>>,
     thread_id: atomic::AtomicI64,
     watcher_child: RwLock<PyObject>,
     #[pyo3(get)]
@@ -149,6 +152,7 @@ impl EventLoop {
                         IOHandle::Py(handle) => self.handle_io_py(py, event, handle, &mut cb_handles),
                         IOHandle::TCPListener(handle) => self.handle_io_tcpl(py, handle, &io_handles, &mut cb_handles),
                         IOHandle::TCPStream(_) => self.handle_io_tcps(event, &mut cb_handles),
+                        IOHandle::UDPSocket(_) => self.handle_io_udp(event, &mut cb_handles),
                         IOHandle::Signals => self.handle_io_signals(py, &mut state.buf, &mut cb_handles),
                     }
                 }
@@ -231,7 +235,7 @@ impl EventLoop {
                 let fd = stream.as_raw_fd() as usize;
                 let token = Token(fd);
                 #[allow(clippy::cast_possible_wrap)]
-                let mut source = Source::TCPStream(fd as i32);
+                let mut source = Source::FD(fd as i32);
                 let (pytransport, stream_handle) = handle.server.new_stream(py, stream);
                 transports.insert(fd, pytransport);
                 lstreams.insert(fd);
@@ -254,6 +258,16 @@ impl EventLoop {
         }
     }
 
+    #[inline]
+    fn handle_io_udp(&self, event: &event::Event, handles_ready: &mut VecDeque<BoxedHandle>) {
+        let fd = event.token().0;
+        if event.is_readable() {
+            handles_ready.push_back(Box::new(UDPReadHandle { fd }));
+        } else if event.is_writable() {
+            handles_ready.push_back(Box::new(UDPWriteHandle { fd }));
+        }
+    }
+
     #[inline]
     fn handle_io_signals(&self, py: Python, buf: &mut [u8], handles_ready: &mut VecDeque<BoxedHandle>) {
         let mut sock_guard = self.ssock.write().unwrap();
@@ -338,7 +352,7 @@ impl EventLoop {
             },
             || {
                 #[allow(clippy::cast_possible_wrap)]
-                let mut source = Source::TCPStream(fd as i32);
+                let mut source = Source::FD(fd as i32);
                 {
                     let guard_poll = self.io.lock().unwrap();
                     _ = guard_poll.registry().register(&mut source, token, interest);
@@ -403,6 +417,83 @@ impl EventLoop {
         }
     }
 
+    #[inline]
+    pub(crate) fn udp_socket_add(&self, fd: usize, interest: Interest) {
+        let token = Token(fd);
+        self.handles_io.pin().update_or_insert_with(
+            token,
+            |io_handle| {
+                if let IOHandle::UDPSocket(interest_prev) = io_handle {
+                    if *interest_prev == interest {
+                        return IOHandle::UDPSocket(interest);
+                    }
+
+                    let interests = *interest_prev | interest;
+                    {
+                        #[allow(clippy::cast_possible_wrap)]
+                        let mut source = Source::FD(fd as i32);
+                        let guard_poll = self.io.lock().unwrap();
+                        _ = guard_poll.registry().reregister(&mut source, token, interests);
+                    }
+                    return IOHandle::UDPSocket(interests);
+                }
+                unreachable!()
+            },
+            || {
+                #[allow(clippy::cast_possible_wrap)]
+                let mut source = Source::FD(fd as i32);
+                {
+                    let guard_poll = self.io.lock().unwrap();
+                    _ = guard_poll.registry().register(&mut source, token, interest);
+                }
+                IOHandle::UDPSocket(interest)
+            },
+        );
+    }
+
+    #[inline]
+    pub(crate) fn udp_socket_rem(&self, fd: usize, interest: Interest) {
+        let token = Token(fd);
+
+        match self.handles_io.pin().remove_if(&token, |_, io_handle| {
+            if let IOHandle::UDPSocket(interest_ex) = io_handle {
+                return *interest_ex == interest;
+            }
+            false
+        }) {
+            Ok(None) => {}
+            Ok(_) => {
+                #[allow(clippy::cast_possible_wrap)]
+                let mut source = Source::FD(fd as i32);
+                let guard_poll = self.io.lock().unwrap();
+                _ = guard_poll.registry().deregister(&mut source);
+            }
+            _ => {
+                self.handles_io.pin().update(token, |io_handle| {
+                    if let IOHandle::UDPSocket(interest_ex) = io_handle {
+                        let interest_new = interest_ex.remove(interest).unwrap();
+                        #[allow(clippy::cast_possible_wrap)]
+                        let mut source = Source::FD(fd as i32);
+                        let guard_poll = self.io.lock().unwrap();
+                        _ = guard_poll.registry().reregister(&mut source, token, interest_new);
+                        return IOHandle::UDPSocket(interest_new);
+                    }
+                    unreachable!()
+                });
+            }
+        }
+    }
+
+    #[inline]
+    pub(crate) fn udp_socket_close(&self, fd: usize) {
+        self.udp_transports.pin().remove(&fd);
+    }
+
+    #[inline(always)]
+    pub(crate) fn get_udp_transport(&self, fd: usize, py: Python) -> Py<UDPTransport> {
+        self.udp_transports.pin().get(&fd).unwrap().clone_ref(py)
+    }
+
     pub(crate) fn log_exception(&self, py: Python, ctx: LogExc) -> PyResult<PyObject> {
         let handler = self.exc_handler.read().unwrap();
         handler.call1(
@@ -631,6 +722,7 @@ impl EventLoop {
             task_factory: RwLock::new(py.None()),
             tcp_lstreams: papaya::HashMap::with_capacity(32),
             tcp_transports: papaya::HashMap::with_capacity(1024),
+            udp_transports: papaya::HashMap::with_capacity(1024),
             thread_id: atomic::AtomicI64::new(0),
             watcher_child: RwLock::new(py.None()),
             _asyncgens: weakset(py)?.unbind(),
@@ -1110,6 +1202,23 @@ impl EventLoop {
         self.tcp_transports.pin().contains_key(&fd)
     }
 
+    fn _udp_conn(
+        pyself: Py<Self>,
+        py: Python,
+        sock: (i32, i32),
+        protocol_factory: PyObject,
+        remote_addr: Option<(String, u16)>,
+    ) -> PyResult<(Py<UDPTransport>, PyObject)> {
+        let rself = pyself.get();
+        let transport = UDPTransport::from_py(py, &pyself, sock, protocol_factory, remote_addr);
+        let fd = transport.fd;
+        let pytransport = Py::new(py, transport)?;
+        let proto = UDPTransport::attach(&pytransport, py)?;
+        rself.udp_transports.pin().insert(fd, pytransport.clone_ref(py));
+        rself.udp_socket_add(fd, Interest::READABLE);
+        Ok((pytransport, proto))
+    }
+
     fn _sig_add(&self, py: Python, sig: u8, callback: PyObject, args: PyObject, context: PyObject) {
         let handle = Py::new(py, CBHandle::new(callback, args, context)).unwrap();
         self.sig_handlers.pin().insert(sig, handle);

src/io.rs

@@ -8,15 +8,19 @@ use std::os::windows::io::RawSocket;
 use mio::{Interest, Registry, Token, event::Source as MioSource, net::TcpListener};
 
 pub(crate) enum Source {
-    TCPListener(TcpListener),
-    #[cfg(unix)]
-    TCPStream(RawFd),
-    #[cfg(windows)]
-    TCPStream(RawSocket),
     #[cfg(unix)]
     FD(RawFd),
     #[cfg(windows)]
     FD(RawSocket),
+    TCPListener(TcpListener),
+    // #[cfg(unix)]
+    // TCPStream(RawFd),
+    // #[cfg(windows)]
+    // TCPStream(RawSocket),
+    // #[cfg(unix)]
+    // UDPSocket(RawFd),
+    // #[cfg(windows)]
+    // UDPSocket(RawSocket),
 }
 
 #[cfg(windows)]
@@ -43,36 +47,33 @@ impl MioSource for Source {
     #[inline]
     fn register(&mut self, registry: &Registry, token: Token, interests: Interest) -> std::io::Result<()> {
         match self {
-            Self::TCPListener(inner) => inner.register(registry, token, interests),
-            Self::TCPStream(inner) => SourceFd(inner).register(registry, token, interests),
             #[cfg(unix)]
             Self::FD(inner) => SourceFd(inner).register(registry, token, interests),
             #[cfg(windows)]
             Self::FD(inner) => SourceRawSocket(inner).register(registry, token, interests),
+            Self::TCPListener(inner) => inner.register(registry, token, interests),
         }
     }
 
     #[inline]
     fn reregister(&mut self, registry: &Registry, token: Token, interests: Interest) -> std::io::Result<()> {
         match self {
-            Self::TCPListener(inner) => inner.reregister(registry, token, interests),
-            Self::TCPStream(inner) => SourceFd(inner).reregister(registry, token, interests),
             #[cfg(unix)]
             Self::FD(inner) => SourceFd(inner).reregister(registry, token, interests),
             #[cfg(windows)]
             Self::FD(inner) => SourceRawSocket(inner).register(registry, token, interests),
+            Self::TCPListener(inner) => inner.reregister(registry, token, interests),
         }
     }
 
     #[inline]
     fn deregister(&mut self, registry: &Registry) -> std::io::Result<()> {
         match self {
-            Self::TCPListener(inner) => inner.deregister(registry),
-            Self::TCPStream(inner) => SourceFd(inner).deregister(registry),
             #[cfg(unix)]
             Self::FD(inner) => SourceFd(inner).deregister(registry),
             #[cfg(windows)]
             Self::FD(inner) => SourceRawSocket(inner).register(registry, token, interests),
+            Self::TCPListener(inner) => inner.deregister(registry),
         }
     }
 }

src/lib.rs

@@ -10,6 +10,7 @@ mod server;
 mod sock;
 mod tcp;
 mod time;
+mod udp;
 mod utils;
 
 pub(crate) fn get_lib_version() -> &'static str {