chore: little improvement to axnet

Author: AsakuraMizu

modules/axdriver/src/dummy.rs

@@ -3,9 +3,10 @@
 #![allow(unused_imports)]
 #![allow(dead_code)]
 
-use super::prelude::*;
 use cfg_if::cfg_if;
 
+use super::prelude::*;
+
 cfg_if! {
     if #[cfg(net_dev = "dummy")] {
         use axdriver_net::{EthernetAddress, NetBuf, NetBufBox, NetBufPool, NetBufPtr};
@@ -176,7 +177,7 @@ cfg_if! {
             fn abort(&mut self, _cid: VsockConnId) -> DevResult<()> {
                 Err(DevError::Unsupported)
             }
-            fn poll_event(&mut self, _buf: &mut [u8]) -> DevResult<Option<VsockDriverEvent>> {
+            fn poll_event(&mut self) -> DevResult<Option<VsockDriverEvent>> {
                 Err(DevError::Unsupported)
             }
         }

modules/axnet/Cargo.toml

@@ -2,10 +2,7 @@
 name = "axnet"
 version.workspace = true
 edition.workspace = true
-authors = [
-  "Yuekai Jia <equation618@gmail.com>",
-  "ChengXiang Qi <kuangjux@outlook.com>",
-]
+authors = ["Mivik <mivikq@gmail.com>"]
 description = "ArceOS network module"
 license.workspace = true
 homepage.workspace = true
@@ -16,29 +13,27 @@ documentation = "https://arceos-org.github.io/arceos/axnet/index.html"
 vsock = ["axdriver/vsock"]
 
 [dependencies]
+async-channel = { version = "2.5", default-features = false }
+async-trait = "0.1"
 axconfig = { workspace = true }
 axdriver = { workspace = true, features = ["net"] }
-axhal = { workspace = true }
-axsync = { workspace = true }
-axtask = { workspace = true }
-
 axerrno = { workspace = true }
 axfs = { workspace = true }
 axfs-ng-vfs = { workspace = true }
+axhal = { workspace = true }
 axio = { workspace = true }
 axpoll = { workspace = true }
+axsync = { workspace = true }
+axtask = { workspace = true }
 bitflags = "2.9.1"
 cfg-if = { workspace = true }
 enum_dispatch = { workspace = true }
+event-listener = { version = "5.4", default-features = false }
 hashbrown = "0.16"
-lazyinit = { workspace = true }
 lazy_static = { workspace = true }
 log = { workspace = true }
-spin = { workspace = true }
 ringbuf = { version = "0.4.8", default-features = false, features = ["alloc"] }
-async-channel = { version = "2.5.0", default-features = false }
-event-listener = { version = "5.4.0", default-features = false }
-async-trait = "0.1.88"
+spin = { workspace = true }
 
 [dependencies.smoltcp]
 package = "starry-smoltcp"

modules/axnet/src/device/vsock.rs

@@ -1,4 +1,4 @@
-use alloc::collections::VecDeque;
+use alloc::{collections::VecDeque, string::ToString};
 use core::{
     sync::atomic::{AtomicBool, AtomicU64, Ordering},
     time::Duration,
@@ -9,7 +9,7 @@ use axerrno::{AxError, AxResult, ax_bail};
 use axsync::Mutex;
 use axtask::future::{block_on, interruptible};
 
-use crate::{alloc::string::ToString, vsock::connection_manager::VSOCK_CONN_MANAGER};
+use crate::vsock::connection_manager::VSOCK_CONN_MANAGER;
 
 // we need a global and static only one vsock device
 static VSOCK_DEVICE: Mutex<Option<AxVsockDevice>> = Mutex::new(None);
@@ -210,7 +210,10 @@ fn handle_vsock_event(event: VsockDriverEvent, dev: &mut AxVsockDevice, buf: &mu
 
         VsockDriverEvent::CreditUpdate(conn_id) => {
             if let Err(e) = manager.on_credit_update(conn_id) {
-                warn!("Failed to handle credit update: {:?}, error={:?}", conn_id, e);
+                warn!(
+                    "Failed to handle credit update: {:?}, error={:?}",
+                    conn_id, e
+                );
             }
         }
 
@@ -242,7 +245,7 @@ pub fn vsock_connect(conn_id: VsockConnId) -> AxResult<()> {
 }
 
 pub fn vsock_send(conn_id: VsockConnId, buf: &[u8]) -> AxResult<usize> {
-    let max_retries  = 10;    // Tests have shown that no more than two retries will be notified
+    let max_retries = 10; // Tests have shown that no more than two retries will be notified
     for _ in 0..max_retries {
         let result = {
             let mut guard = VSOCK_DEVICE.lock();
@@ -258,7 +261,7 @@ pub fn vsock_send(conn_id: VsockConnId, buf: &[u8]) -> AxResult<usize> {
                     conn.lock().wait_for_tx();
                 };
             }
-            Err(e) => return Err(map_dev_err(e)), 
+            Err(e) => return Err(map_dev_err(e)),
         }
     }
     Err(map_dev_err(DevError::Again))

modules/axnet/src/general.rs

@@ -9,7 +9,7 @@ use axpoll::{IoEvents, Pollable};
 use axtask::future::{block_on, poll_io, timeout};
 
 use crate::{
-    SERVICE,
+    get_service,
     options::{Configurable, GetSocketOption, SetSocketOption},
 };
 
@@ -70,7 +70,7 @@ impl GeneralOptions {
     }
 
     pub fn register_waker(&self, waker: &Waker) {
-        SERVICE.lock().register_waker(self.device_mask(), waker);
+        get_service().register_waker(self.device_mask(), waker);
     }
 
     pub fn send_poller<P: Pollable, F: FnMut() -> AxResult<T>, T>(

modules/axnet/src/lib.rs

@@ -38,11 +38,11 @@ use alloc::{borrow::ToOwned, boxed::Box};
 
 use axdriver::{AxDeviceContainer, prelude::*};
 use axsync::Mutex;
-use lazyinit::LazyInit;
 use smoltcp::wire::{EthernetAddress, Ipv4Address, Ipv4Cidr};
-pub use socket::*;
+use spin::{Lazy, Once};
 
-use crate::{
+pub use self::socket::*;
+use self::{
     consts::{GATEWAY, IP, IP_PREFIX},
     device::{EthernetDevice, LoopbackDevice},
     listen_table::ListenTable,
@@ -51,10 +51,17 @@ use crate::{
     wrapper::SocketSetWrapper,
 };
 
-static LISTEN_TABLE: LazyInit<ListenTable> = LazyInit::new();
-static SOCKET_SET: LazyInit<SocketSetWrapper> = LazyInit::new();
+static LISTEN_TABLE: Lazy<ListenTable> = Lazy::new(ListenTable::new);
+static SOCKET_SET: Lazy<SocketSetWrapper> = Lazy::new(SocketSetWrapper::new);
 
-static SERVICE: LazyInit<Mutex<Service>> = LazyInit::new();
+static SERVICE: Once<Mutex<Service>> = Once::new();
+
+fn get_service() -> axsync::MutexGuard<'static, Service> {
+    SERVICE
+        .get()
+        .expect("Network service not initialized")
+        .lock()
+}
 
 /// Initializes the network subsystem by NIC devices.
 pub fn init_network(mut net_devs: AxDeviceContainer<AxNetDevice>) {
@@ -111,16 +118,13 @@ pub fn init_network(mut net_devs: AxDeviceContainer<AxNetDevice>) {
             ip_addrs.push(eth0_ip.into()).unwrap();
         }
     });
-    SERVICE.init_once(Mutex::new(service));
-
-    SOCKET_SET.init_once(SocketSetWrapper::new());
-    LISTEN_TABLE.init_once(ListenTable::new());
+    SERVICE.call_once(|| Mutex::new(service));
 }
 
 /// Init vsock subsystem by vsock devices.
 #[cfg(feature = "vsock")]
 pub fn init_vsock(mut vsock_devs: AxDeviceContainer<AxVsockDevice>) {
-    use crate::device::register_vsock_device;
+    use self::device::register_vsock_device;
     info!("Initialize vsock subsystem...");
     if let Some(dev) = vsock_devs.take_one() {
         info!("  use vsock 0: {:?}", dev.device_name());
@@ -133,5 +137,5 @@ pub fn init_vsock(mut vsock_devs: AxDeviceContainer<AxVsockDevice>) {
 }
 
 pub fn poll_interfaces() {
-    while SERVICE.lock().poll(&mut SOCKET_SET.inner.lock()) {}
+    while get_service().poll(&mut SOCKET_SET.inner.lock()) {}
 }

modules/axnet/src/listen_table.rs

@@ -16,12 +16,12 @@ use crate::{
 
 const PORT_NUM: usize = 65536;
 
-struct ListenTableEntry {
+struct ListenTableEntryInner {
     listen_endpoint: IpListenEndpoint,
     syn_queue: VecDeque<SocketHandle>,
 }
 
-impl ListenTableEntry {
+impl ListenTableEntryInner {
     pub fn new(listen_endpoint: IpListenEndpoint) -> Self {
         Self {
             listen_endpoint,
@@ -30,16 +30,18 @@ impl ListenTableEntry {
     }
 }
 
-impl Drop for ListenTableEntry {
+impl Drop for ListenTableEntryInner {
     fn drop(&mut self) {
         for &handle in &self.syn_queue {
             SOCKET_SET.remove(handle);
         }
     }
 }
 
+type ListenTableEntry = Arc<Mutex<Option<Box<ListenTableEntryInner>>>>;
+
 pub struct ListenTable {
-    tcp: Box<[Arc<Mutex<Option<Box<ListenTableEntry>>>>]>,
+    tcp: Box<[ListenTableEntry]>,
 }
 
 impl ListenTable {
@@ -63,7 +65,7 @@ impl ListenTable {
         assert_ne!(port, 0);
         let mut entry = self.tcp[port as usize].lock();
         if entry.is_none() {
-            *entry = Some(Box::new(ListenTableEntry::new(listen_endpoint)));
+            *entry = Some(Box::new(ListenTableEntryInner::new(listen_endpoint)));
             Ok(())
         } else {
             warn!("socket already listening on port {port}");
@@ -76,7 +78,7 @@ impl ListenTable {
         *self.tcp[port as usize].lock() = None;
     }
 
-    fn listen_entry(&self, port: u16) -> Arc<Mutex<Option<Box<ListenTableEntry>>>> {
+    fn listen_entry(&self, port: u16) -> Arc<Mutex<Option<Box<ListenTableEntryInner>>>> {
         self.tcp[port as usize].clone()
     }
 

modules/axnet/src/tcp.rs

@@ -16,11 +16,12 @@ use smoltcp::{
     wire::{IpEndpoint, IpListenEndpoint},
 };
 
-use super::{LISTEN_TABLE, SOCKET_SET};
 use crate::{
-    RecvFlags, RecvOptions, SERVICE, SendOptions, Shutdown, Socket, SocketAddrEx, SocketOps,
+    LISTEN_TABLE, RecvFlags, RecvOptions, SOCKET_SET, SendOptions, Shutdown, Socket, SocketAddrEx,
+    SocketOps,
     consts::{TCP_RX_BUF_LEN, TCP_TX_BUF_LEN},
     general::GeneralOptions,
+    get_service,
     options::{Configurable, GetSocketOption, SetSocketOption},
     poll_interfaces,
     state::*,
@@ -71,7 +72,7 @@ impl TcpSocket {
         result.with_smol_socket(|socket| {
             result
                 .general
-                .set_device_mask(SERVICE.lock().device_mask_for(&socket.get_bound_endpoint()));
+                .set_device_mask(get_service().device_mask_for(&socket.get_bound_endpoint()));
         });
         result
     }
@@ -232,7 +233,7 @@ impl SocketOps for TcpSocket {
                     };
                     socket.set_bound_endpoint(endpoint);
                     self.general
-                        .set_device_mask(SERVICE.lock().device_mask_for(&endpoint));
+                        .set_device_mask(get_service().device_mask_for(&endpoint));
                     Ok(())
                 })?;
                 debug!("TCP socket {}: binding to {}", self.handle, local_addr);
@@ -260,7 +261,7 @@ impl SocketOps for TcpSocket {
                     self.with_smol_socket(|socket| socket.get_bound_endpoint());
                 if bound_endpoint.addr.is_none() {
                     bound_endpoint.addr =
-                        Some(SERVICE.lock().get_source_address(&remote_endpoint.addr));
+                        Some(get_service().get_source_address(&remote_endpoint.addr));
                 }
                 if bound_endpoint.port == 0 {
                     bound_endpoint.port = get_ephemeral_port()?;
@@ -273,10 +274,10 @@ impl SocketOps for TcpSocket {
                 self.with_smol_socket(|socket| {
                     socket.set_bound_endpoint(bound_endpoint);
                     self.general
-                        .set_device_mask(SERVICE.lock().device_mask_for(&bound_endpoint));
+                        .set_device_mask(get_service().device_mask_for(&bound_endpoint));
                     socket
                         .connect(
-                            crate::SERVICE.lock().iface.context(),
+                            get_service().iface.context(),
                             remote_endpoint,
                             bound_endpoint,
                         )

modules/axnet/src/udp.rs

@@ -18,9 +18,10 @@ use smoltcp::{
 use spin::RwLock;
 
 use crate::{
-    RecvFlags, RecvOptions, SERVICE, SOCKET_SET, SendOptions, Shutdown, SocketAddrEx, SocketOps,
+    RecvFlags, RecvOptions, SOCKET_SET, SendOptions, Shutdown, SocketAddrEx, SocketOps,
     consts::{UDP_RX_BUF_LEN, UDP_TX_BUF_LEN},
     general::GeneralOptions,
+    get_service,
     options::{Configurable, GetSocketOption, SetSocketOption},
     poll_interfaces,
 };
@@ -141,7 +142,7 @@ impl SocketOps for UdpSocket {
             })
         })?;
         self.general
-            .set_device_mask(SERVICE.lock().device_mask_for(&endpoint));
+            .set_device_mask(get_service().device_mask_for(&endpoint));
 
         *guard = Some(local_endpoint);
         info!("UDP socket {}: bound on {}", self.handle, endpoint);
@@ -159,7 +160,7 @@ impl SocketOps for UdpSocket {
         }
 
         let remote_addr = IpEndpoint::from(remote_addr);
-        let src = SERVICE.lock().get_source_address(&remote_addr.addr);
+        let src = get_service().get_source_address(&remote_addr.addr);
         *guard = Some((remote_addr, src));
         debug!("UDP socket {}: connected to {}", self.handle, remote_addr);
         Ok(())
@@ -169,7 +170,7 @@ impl SocketOps for UdpSocket {
         let (remote_addr, source_addr) = match options.to {
             Some(addr) => {
                 let addr = IpEndpoint::from(addr.into_ip()?);
-                let src = SERVICE.lock().get_source_address(&addr.addr);
+                let src = get_service().get_source_address(&addr.addr);
                 (addr, src)
             }
             None => self.remote_endpoint()?,

modules/axnet/src/unix.rs modules/axnet/src/unix/mod.rsmodules/axnet/src/unix.rs renamed to modules/axnet/src/unix/mod.rs

@@ -7,6 +7,7 @@ use axtask::WaitQueue;
 use ringbuf::{HeapCons, HeapProd, HeapRb, traits::*};
 
 use super::{VsockAddr, VsockConnId};
+use crate::device::{start_vsock_poll, stop_vsock_poll};
 
 pub const VSOCK_RX_BUFFER_SIZE: usize = 64 * 1024; // 64KB receive buffer
 const VSOCK_ACCEPT_QUEUE_SIZE: usize = 128; // accept queue size
@@ -367,7 +368,7 @@ impl VsockConnectionManager {
         if self.connections.contains_key(&conn_id) {
             info!("Connection {:?} already exists, overwriting", conn_id);
         } else {
-            crate::device::start_vsock_poll();
+            start_vsock_poll();
         }
         self.connections.insert(conn_id, conn.clone());
         debug!(
@@ -386,7 +387,8 @@ impl VsockConnectionManager {
     pub fn remove_connection(&mut self, conn_id: VsockConnId) {
         if let Some(conn) = self.connections.remove(&conn_id) {
             let conn = conn.lock();
-            crate::device::stop_vsock_poll();
+
+            stop_vsock_poll();
             debug!(
                 "Removed connection {:?}: rx={} bytes, tx={} bytes, dropped={} bytes",
                 conn_id, conn.rx_bytes, conn.tx_bytes, conn.dropped_bytes
@@ -420,7 +422,7 @@ impl VsockConnectionManager {
 
         // 加入 accept 队列
         let mut queue_guard = queue.lock();
-        if let Err(_) = queue_guard.accept_queue.push(conn_id) {
+        if queue_guard.accept_queue.push(conn_id).is_err() {
             info!(
                 "Accept queue full for port {}, dropping connection from {:?}",
                 conn_id.local_port, conn_id.peer_addr