optimize wave2

Author: 0x_antoni

Cargo.toml

@@ -322,7 +322,7 @@ hyper-util = "0.1"
 async-trait = "0.1"
 tower = "0.4"
 pin-project = "1.0"
-mio = "1.0.4"
+mio = { version = "1.0.4", features = ["os-ext", "net", "os-poll"] }
 mio-timer = { git = "https://github.com/Pana/mio-timer", rev = "251170b" }
 mio-misc = { git = "https://github.com/Pana/mio-misc", rev = "27ad80f" }
 reqwest = "0.12"

crates/network/Cargo.toml

@@ -9,7 +9,7 @@ cfx-addr = { workspace = true }
 cfx-bytes = { workspace = true }
 cfx-types = { workspace = true }
 cfxkey = { workspace = true }
-mio = { workspace = true }
+mio = { workspace = true, features = ["os-ext", "net", "os-poll"] }
 parking_lot = { workspace = true }
 log = { workspace = true }
 slab = { workspace = true }

crates/network/src/connection.rs

@@ -13,7 +13,7 @@ use log::{debug, trace};
 use metrics::{
     register_meter_with_group, Gauge, GaugeUsize, Histogram, Meter, Sample,
 };
-use mio::{net::TcpStream, Interest, Poll, Token};
+use mio::{net::TcpStream, Interest, Registry, Token};
 use priority_send_queue::{PrioritySendQueue, SendQueuePriority};
 use serde::Deserialize;
 use serde_derive::Serialize;
@@ -398,7 +398,7 @@ impl Connection {
 
     /// Register this connection with the IO event loop.
     pub fn register_socket(
-        &mut self, reg: Token, event_loop: &Poll,
+        &mut self, reg: Token, poll_registry: &Registry,
     ) -> io::Result<()> {
         if self.registered.load(AtomicOrdering::SeqCst) {
             return Ok(());
@@ -409,9 +409,7 @@ impl Connection {
             reg
         );
         if let Err(e) =
-            event_loop
-                .registry()
-                .register(&mut self.socket, reg, self.interest)
+            poll_registry.register(&mut self.socket, reg, self.interest)
         {
             trace!(
                 "Failed to register socket, token = {}, reg = {:?}, err = {:?}",
@@ -427,17 +425,17 @@ impl Connection {
     /// Update connection registration. Should be called at the end of the IO
     /// handler.
     pub fn update_socket(
-        &mut self, reg: Token, event_loop: &Poll,
+        &mut self, reg: Token, poll_registry: &Registry,
     ) -> io::Result<()> {
         trace!(
             "Connection reregister, token = {}, reg = {:?}",
             self.token,
             reg
         );
         if !self.registered.load(AtomicOrdering::SeqCst) {
-            self.register_socket(reg, event_loop)
+            self.register_socket(reg, poll_registry)
         } else {
-            event_loop.registry()
+            poll_registry
                 .reregister(&mut self.socket, reg, self.interest)
                 .unwrap_or_else(|e| {
                     trace!("Failed to reregister socket, token = {}, reg = {:?}, err = {:?}", self.token, reg, e);
@@ -448,9 +446,11 @@ impl Connection {
 
     /// Delete connection registration. Should be called at the end of the IO
     /// handler.
-    pub fn deregister_socket(&mut self, event_loop: &Poll) -> io::Result<()> {
+    pub fn deregister_socket(
+        &mut self, poll_registry: &Registry,
+    ) -> io::Result<()> {
         trace!("Connection deregister, token = {}", self.token);
-        event_loop.registry().deregister(&mut self.socket).ok();
+        poll_registry.deregister(&mut self.socket).ok();
         Ok(())
     }
 

crates/network/src/service.rs

@@ -18,7 +18,7 @@ use std::{
 use keccak_hash::keccak;
 use mio::{
     net::{TcpListener, TcpStream, UdpSocket},
-    Interest, Poll, Token,
+    Interest, Registry, Token,
 };
 use parity_path::restrict_permissions_owner;
 use parking_lot::{Mutex, RwLock};
@@ -171,7 +171,6 @@ impl<'a> UdpIoContext<'a> {
 pub struct NetworkService {
     pub io_service: Option<IoService<NetworkIoMessage>>,
     pub inner: Option<Arc<NetworkServiceInner>>,
-    network_poll: Arc<Mutex<Poll>>,
     config: NetworkConfiguration,
 }
 
@@ -180,7 +179,6 @@ impl NetworkService {
         NetworkService {
             io_service: None,
             inner: None,
-            network_poll: Arc::new(Mutex::new(Poll::new().unwrap())),
             config,
         }
     }
@@ -200,7 +198,7 @@ impl NetworkService {
         &mut self, pos_pub_keys: (ConsensusPublicKey, ConsensusVRFPublicKey),
     ) -> Result<(), Error> {
         let raw_io_service =
-            IoService::<NetworkIoMessage>::start(self.network_poll.clone())?;
+            IoService::<NetworkIoMessage>::start(STOP_NET_POLL)?;
         self.io_service = Some(raw_io_service);
 
         if self.inner.is_none() {
@@ -226,18 +224,10 @@ impl NetworkService {
 
     pub fn start(&self) {
         let handler = self.inner.as_ref().unwrap().clone();
-        let main_event_loop_channel =
-            self.io_service.as_ref().unwrap().channel();
         self.io_service
             .as_ref()
             .expect("Already set")
-            .start_network_poll(
-                self.network_poll.clone(),
-                handler,
-                main_event_loop_channel,
-                MAX_SESSIONS,
-                STOP_NET_POLL,
-            );
+            .start_network_poll(handler, MAX_SESSIONS);
     }
 
     /// Add a P2P peer to the client as a trusted node
@@ -1756,19 +1746,19 @@ impl IoHandler<NetworkIoMessage> for NetworkServiceInner {
     }
 
     fn register_stream(
-        &self, stream: StreamToken, reg: Token, event_loop: &Poll,
+        &self, stream: StreamToken, reg: Token, poll_registry: &Registry,
     ) {
         match stream {
             FIRST_SESSION..=LAST_SESSION => {
                 if let Some(session) = self.sessions.get(stream) {
                     session
                         .write()
-                        .register_socket(reg, event_loop)
+                        .register_socket(reg, poll_registry)
                         .expect("Error registering socket");
                 }
             }
             TCP_ACCEPT => {
-                event_loop.registry()
+                poll_registry
                     .register(
                         &mut *self.tcp_listener.lock(),
                         Token(TCP_ACCEPT),
@@ -1777,7 +1767,7 @@ impl IoHandler<NetworkIoMessage> for NetworkServiceInner {
                     .expect("Error registering stream");
             }
             UDP_MESSAGE => {
-                event_loop.registry()
+                poll_registry
                     .register(
                         &mut *self.udp_socket.lock(),
                         reg,
@@ -1789,13 +1779,13 @@ impl IoHandler<NetworkIoMessage> for NetworkServiceInner {
         }
     }
 
-    fn deregister_stream(&self, stream: StreamToken, event_loop: &Poll) {
+    fn deregister_stream(&self, stream: StreamToken, poll_registry: &Registry) {
         match stream {
             FIRST_SESSION..=LAST_SESSION => {
                 if let Some(session) = self.sessions.get(stream) {
                     let mut sess = session.write();
                     if sess.expired() {
-                        sess.deregister_socket(event_loop)
+                        sess.deregister_socket(poll_registry)
                             .expect("Error deregistering socket");
                         if let Some(node_id) = sess.id() {
                             self.node_db.write().note_failure(
@@ -1813,18 +1803,18 @@ impl IoHandler<NetworkIoMessage> for NetworkServiceInner {
     }
 
     fn update_stream(
-        &self, stream: StreamToken, reg: Token, event_loop: &Poll,
+        &self, stream: StreamToken, reg: Token, poll_registry: &Registry,
     ) {
         match stream {
             FIRST_SESSION..=LAST_SESSION => {
                 if let Some(session) = self.sessions.get(stream) {
                     session
                         .write()
-                        .update_socket(reg, event_loop)
+                        .update_socket(reg, poll_registry)
                         .expect("Error updating socket");
                 }
             }
-            TCP_ACCEPT => event_loop.registry()
+            TCP_ACCEPT => poll_registry
                 .reregister(
                     &mut *self.tcp_listener.lock(),
                     Token(TCP_ACCEPT),
@@ -1840,7 +1830,7 @@ impl IoHandler<NetworkIoMessage> for NetworkServiceInner {
                 } else {
                     Interest::READABLE
                 };
-                event_loop.registry()
+                poll_registry
                     .reregister(
                         &mut *udp_socket,
                         reg,

crates/network/src/session.rs

@@ -17,7 +17,7 @@ use diem_crypto::{bls::BLS_PUBLIC_KEY_LENGTH, ValidCryptoMaterial};
 use diem_types::validator_config::{ConsensusPublicKey, ConsensusVRFPublicKey};
 use io::{IoContext, StreamToken};
 use log::{debug, trace};
-use mio::{net::TcpStream, Poll, Token};
+use mio::{net::TcpStream, Registry, Token};
 use priority_send_queue::SendQueuePriority;
 use rlp::{Rlp, RlpStream};
 use serde::Deserialize;
@@ -177,28 +177,28 @@ impl Session {
     /// Register event loop for the underlying connection.
     /// If session expired, no effect taken.
     pub fn register_socket(
-        &mut self, reg: Token, event_loop: &Poll,
+        &mut self, reg: Token, poll_registry: &Registry,
     ) -> Result<(), Error> {
         if !self.expired() {
-            self.connection_mut().register_socket(reg, event_loop)?;
+            self.connection_mut().register_socket(reg, poll_registry)?;
         }
 
         Ok(())
     }
 
     /// Update the event loop for the underlying connection.
     pub fn update_socket(
-        &mut self, reg: Token, event_loop: &Poll,
+        &mut self, reg: Token, poll_registry: &Registry,
     ) -> Result<(), Error> {
-        self.connection_mut().update_socket(reg, event_loop)?;
+        self.connection_mut().update_socket(reg, poll_registry)?;
         Ok(())
     }
 
     /// Deregister the event loop for the underlying connection.
     pub fn deregister_socket(
-        &mut self, event_loop: &Poll,
+        &mut self, poll_registry: &Registry,
     ) -> Result<(), Error> {
-        self.connection_mut().deregister_socket(event_loop)?;
+        self.connection_mut().deregister_socket(poll_registry)?;
         Ok(())
     }
 

crates/util/io/Cargo.toml

@@ -7,7 +7,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-mio = { workspace = true }
+mio = { workspace = true, features = ["os-ext", "net", "os-poll"] }
 crossbeam-deque = { workspace = true }
 crossbeam-channel = { workspace = true }
 parking_lot = { workspace = true }

crates/util/io/src/lib.rs

@@ -31,7 +31,7 @@ pub use crate::service_mio::{
     TimerToken, TOKENS_PER_HANDLER,
 };
 
-use mio::{Poll, Token};
+use mio::{Registry, Token};
 use mio_util::NotifyError;
 use std::{cell::Cell, env, error, fmt, io};
 
@@ -99,22 +99,21 @@ where Message: Send + Sync + 'static
     fn stream_writable(&self, _io: &IoContext<Message>, _stream: StreamToken) {}
     /// Register a new stream with the event loop
     fn register_stream(
-        &self, _stream: StreamToken, _reg: Token, _event_loop: &Poll,
+        &self, _stream: StreamToken, _reg: Token, _registry: &Registry,
     ) {
     }
     /// Re-register a stream with the event loop
     fn update_stream(
-        &self, _stream: StreamToken, _reg: Token, _event_loop: &Poll,
+        &self, _stream: StreamToken, _reg: Token, _registry: &Registry,
     ) {
     }
     /// Deregister a stream. Called when stream is removed from event loop
-    fn deregister_stream(&self, _stream: StreamToken, _event_loop: &Poll) {}
+    fn deregister_stream(&self, _stream: StreamToken, _registry: &Registry) {}
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use parking_lot::Mutex;
     use std::{
         sync::{atomic, Arc},
         thread,
@@ -143,9 +142,8 @@ mod tests {
         }
 
         let handler = Arc::new(MyHandler(atomic::AtomicBool::new(false)));
-        let poll = Arc::new(Mutex::new(Poll::new().unwrap()));
 
-        let service = IoService::<MyMessage>::start(poll)
+        let service = IoService::<MyMessage>::start(123)
             .expect("Error creating network service");
         service.register_handler(handler.clone()).unwrap();
 
@@ -178,9 +176,8 @@ mod tests {
         }
 
         let handler = Arc::new(MyHandler(atomic::AtomicBool::new(false)));
-        let poll = Arc::new(Mutex::new(Poll::new().unwrap()));
 
-        let service = IoService::<MyMessage>::start(poll)
+        let service = IoService::<MyMessage>::start(123)
             .expect("Error creating network service");
         service.register_handler(handler.clone()).unwrap();
 
@@ -210,9 +207,8 @@ mod tests {
         }
 
         let handler = Arc::new(MyHandler(atomic::AtomicUsize::new(0)));
-        let poll = Arc::new(Mutex::new(Poll::new().unwrap()));
 
-        let service = IoService::<MyMessage>::start(poll)
+        let service = IoService::<MyMessage>::start(123)
             .expect("Error creating network service");
         service.register_handler(handler.clone()).unwrap();