fix wrong udpsocket read write migration

Author: 0x_antoni

crates/network/src/service.rs

@@ -44,7 +44,9 @@ use serde::{Deserialize, Serialize};
 use crate::{
     discovery::Discovery,
     handshake::BYPASS_CRYPTOGRAPHY,
-    iolib::{IoContext, IoHandler, IoService, StreamToken, TimerToken},
+    iolib::{
+        IoContext, IoHandler, IoService, MapNonBlock, StreamToken, TimerToken,
+    },
     ip_utils::{map_external_address, select_public_address},
     node_database::NodeDatabase,
     node_table::*,
@@ -1389,12 +1391,13 @@ impl NetworkServiceInner {
         }
 
         let mut buf = [0u8; MAX_DATAGRAM_SIZE];
-        match udp_socket.recv_from(&mut buf) {
-            Ok((len, address)) => self
+        match udp_socket.recv_from(&mut buf).map_non_block() {
+            Ok(Some((len, address))) => self
                 .on_udp_packet(&buf[0..len], address)
                 .unwrap_or_else(|e| {
                     debug!("Error processing UDP packet: {:?}", e);
                 }),
+            Ok(_) => {}
             Err(e) => {
                 debug!("Error reading UDP socket: {:?}", e);
             }
@@ -1420,11 +1423,18 @@ impl NetworkServiceInner {
         let udp_socket = self.udp_socket.lock();
         let mut udp_channel = self.udp_channel.write();
         while let Some(data) = udp_channel.dequeue_send() {
-            match udp_socket.send_to(&data.payload, data.address) {
-                Ok(size) if size == data.payload.len() => {}
-                Ok(_) => {
+            match udp_socket
+                .send_to(&data.payload, data.address)
+                .map_non_block()
+            {
+                Ok(Some(size)) if size == data.payload.len() => {}
+                Ok(Some(_)) => {
                     warn!("UDP sent incomplete datagram");
                 }
+                Ok(None) => {
+                    udp_channel.requeue_send(data);
+                    return;
+                }
                 Err(e) => {
                     debug!(
                         "UDP send error: {:?}, address: {:?}",

crates/util/io/src/lib.rs

@@ -32,7 +32,7 @@ pub use crate::service_mio::{
 };
 
 use mio::{Registry, Token};
-use mio_util::NotifyError;
+pub use mio_util::{would_block, MapNonBlock, NotifyError};
 use std::{cell::Cell, env, error, fmt, io};
 
 thread_local! {

crates/util/io/src/mio_util/io.rs

@@ -0,0 +1,28 @@
+// Re-export the io::Result / Error types for convenience
+pub use std::io::{Error, ErrorKind, Read, Result, Write};
+
+// TODO: Delete this
+/// A helper trait to provide the map_non_block function on Results.
+pub trait MapNonBlock<T> {
+    /// Maps a `Result<T>` to a `Result<Option<T>>` by converting
+    /// operation-would-block errors into `Ok(None)`.
+    fn map_non_block(self) -> Result<Option<T>>;
+}
+
+impl<T> MapNonBlock<T> for Result<T> {
+    fn map_non_block(self) -> Result<Option<T>> {
+        match self {
+            Ok(value) => Ok(Some(value)),
+            Err(err) => {
+                if let ErrorKind::WouldBlock = err.kind() {
+                    Ok(None)
+                } else {
+                    Err(err)
+                }
+            }
+        }
+    }
+}
+
+/// Returns a std `WouldBlock` error without allocating
+pub fn would_block() -> Error { ErrorKind::WouldBlock.into() }

crates/util/io/src/mio_util/mod.rs

@@ -1,8 +1,10 @@
 #![allow(unused)]
 mod event_loop;
 mod handler;
+mod io;
 mod notify_error;
 
 pub use event_loop::{EventLoop, EventLoopBuilder, Sender};
 pub use handler::Handler;
+pub use io::{would_block, MapNonBlock};
 pub use notify_error::NotifyError;