Make Ext a trait

kernel/src/net/iface/ext.rs

@@ -1,78 +1,9 @@
 // SPDX-License-Identifier: MPL-2.0
 
-use alloc::string::String;
-use core::sync::atomic::{AtomicU64, Ordering};
+use super::sched::PollScheduler;
 
-use ostd::sync::WaitQueue;
+pub struct BigtcpExt;
 
-use super::Iface;
-
-/// The iface extension.
-pub struct IfaceExt {
-    /// The name of the iface.
-    name: String,
-    /// The time when we should do the next poll.
-    /// We store the total number of milliseconds since the system booted.
-    next_poll_at_ms: AtomicU64,
-    /// The wait queue that the background polling thread will sleep on.
-    polling_wait_queue: WaitQueue,
-}
-
-impl IfaceExt {
-    pub(super) fn new(name: String) -> Self {
-        Self {
-            name,
-            next_poll_at_ms: AtomicU64::new(0),
-            polling_wait_queue: WaitQueue::new(),
-        }
-    }
-
-    pub(super) fn next_poll_at_ms(&self) -> Option<u64> {
-        let millis = self.next_poll_at_ms.load(Ordering::Relaxed);
-        if millis == 0 {
-            None
-        } else {
-            Some(millis)
-        }
-    }
-
-    pub(super) fn polling_wait_queue(&self) -> &WaitQueue {
-        &self.polling_wait_queue
-    }
-
-    fn schedule_next_poll(&self, poll_at: Option<u64>) {
-        let Some(new_instant) = poll_at else {
-            self.next_poll_at_ms.store(0, Ordering::Relaxed);
-            return;
-        };
-
-        let old_instant = self.next_poll_at_ms.load(Ordering::Relaxed);
-        self.next_poll_at_ms.store(new_instant, Ordering::Relaxed);
-
-        if old_instant == 0 || new_instant < old_instant {
-            self.polling_wait_queue.wake_all();
-        }
-    }
-}
-
-pub trait IfaceEx {
-    /// Gets the name of the iface.
-    ///
-    /// In Linux, the name is usually the driver name followed by a unit number.
-    fn name(&self) -> &str;
-
-    /// Transmits or receives packets queued in the iface, and updates socket status accordingly.
-    ///
-    /// The background polling thread is woken up to perform the next poll if necessary.
-    fn poll(&self);
-}
-
-impl IfaceEx for Iface {
-    fn name(&self) -> &str {
-        &self.ext().name
-    }
-
-    fn poll(&self) {
-        self.raw_poll(&|next_poll| self.ext().schedule_next_poll(next_poll));
-    }
+impl aster_bigtcp::ext::Ext for BigtcpExt {
+    type ScheduleNextPoll = PollScheduler;
 }

kernel/src/net/iface/init.rs

@@ -7,10 +7,7 @@ use ostd::sync::LocalIrqDisabled;
 use spin::Once;
 
 use super::{poll::poll_ifaces, Iface};
-use crate::{
-    net::iface::ext::{IfaceEx, IfaceExt},
-    prelude::*,
-};
+use crate::{net::iface::sched::PollScheduler, prelude::*};
 
 pub static IFACES: Once<Vec<Arc<Iface>>> = Once::new();
 
@@ -69,7 +66,8 @@ fn new_virtio() -> Arc<Iface> {
         EthernetAddress(ether_addr),
         Ipv4Cidr::new(VIRTIO_ADDRESS, VIRTIO_ADDRESS_PREFIX_LEN),
         VIRTIO_GATEWAY,
-        IfaceExt::new("virtio".to_owned()),
+        "virtio".to_owned(),
+        PollScheduler::new(),
     )
 }
 
@@ -100,6 +98,7 @@ fn new_loopback() -> Arc<Iface> {
     IpIface::new(
         Wrapper(Mutex::new(Loopback::new(Medium::Ip))),
         Ipv4Cidr::new(LOOPBACK_ADDRESS, LOOPBACK_ADDRESS_PREFIX_LEN),
-        IfaceExt::new("lo".to_owned()),
+        "lo".to_owned(),
+        PollScheduler::new(),
     ) as _
 }

kernel/src/net/iface/mod.rs

@@ -3,11 +3,11 @@
 mod ext;
 mod init;
 mod poll;
+mod sched;
 
-pub use ext::IfaceEx;
 pub use init::{init, IFACES};
 pub use poll::lazy_init;
 
-pub type Iface = dyn aster_bigtcp::iface::Iface<ext::IfaceExt>;
-pub type BoundTcpSocket = aster_bigtcp::socket::BoundTcpSocket<ext::IfaceExt>;
-pub type BoundUdpSocket = aster_bigtcp::socket::BoundUdpSocket<ext::IfaceExt>;
+pub type Iface = dyn aster_bigtcp::iface::Iface<ext::BigtcpExt>;
+pub type BoundTcpSocket = aster_bigtcp::socket::BoundTcpSocket<ext::BigtcpExt>;
+pub type BoundUdpSocket = aster_bigtcp::socket::BoundUdpSocket<ext::BigtcpExt>;

kernel/src/net/iface/poll.rs

@@ -6,7 +6,7 @@ use core::time::Duration;
 use log::trace;
 use ostd::timer::Jiffies;
 
-use super::{ext::IfaceEx, Iface, IFACES};
+use super::{Iface, IFACES};
 use crate::{sched::priority::Priority, thread::kernel_thread::ThreadOptions, WaitTimeout};
 
 pub fn lazy_init() {
@@ -27,14 +27,14 @@ fn spawn_background_poll_thread(iface: Arc<Iface>) {
     let task_fn = move || {
         trace!("spawn background poll thread for {}", iface.name());
 
-        let iface_ext = iface.ext();
-        let wait_queue = iface_ext.polling_wait_queue();
+        let sched_poll = iface.sched_poll();
+        let wait_queue = sched_poll.polling_wait_queue();
 
         loop {
-            let next_poll_at_ms = if let Some(next_poll_at_ms) = iface_ext.next_poll_at_ms() {
+            let next_poll_at_ms = if let Some(next_poll_at_ms) = sched_poll.next_poll_at_ms() {
                 next_poll_at_ms
             } else {
-                wait_queue.wait_until(|| iface_ext.next_poll_at_ms())
+                wait_queue.wait_until(|| sched_poll.next_poll_at_ms())
             };
 
             let now_as_ms = Jiffies::elapsed().as_duration().as_millis() as u64;
@@ -54,9 +54,9 @@ fn spawn_background_poll_thread(iface: Arc<Iface>) {
 
             let duration = Duration::from_millis(next_poll_at_ms - now_as_ms);
             let _ = wait_queue.wait_until_or_timeout(
-                // If `iface_ext.next_poll_at_ms()` changes to an earlier time, we will end the
+                // If `sched_poll.next_poll_at_ms()` changes to an earlier time, we will end the
                 // waiting.
-                || (iface_ext.next_poll_at_ms()? < next_poll_at_ms).then_some(()),
+                || (sched_poll.next_poll_at_ms()? < next_poll_at_ms).then_some(()),
                 &duration,
             );
         }

kernel/src/net/iface/sched.rs

@@ -0,0 +1,52 @@
+// SPDX-License-Identifier: MPL-2.0
+
+use core::sync::atomic::{AtomicU64, Ordering};
+
+use aster_bigtcp::iface::ScheduleNextPoll;
+use ostd::sync::WaitQueue;
+
+pub struct PollScheduler {
+    /// The time when we should do the next poll.
+    /// We store the total number of milliseconds since the system booted.
+    next_poll_at_ms: AtomicU64,
+    /// The wait queue that the background polling thread will sleep on.
+    polling_wait_queue: WaitQueue,
+}
+
+impl PollScheduler {
+    pub(super) fn new() -> Self {
+        Self {
+            next_poll_at_ms: AtomicU64::new(0),
+            polling_wait_queue: WaitQueue::new(),
+        }
+    }
+
+    pub(super) fn next_poll_at_ms(&self) -> Option<u64> {
+        let millis = self.next_poll_at_ms.load(Ordering::Relaxed);
+        if millis == 0 {
+            None
+        } else {
+            Some(millis)
+        }
+    }
+
+    pub(super) fn polling_wait_queue(&self) -> &WaitQueue {
+        &self.polling_wait_queue
+    }
+}
+
+impl ScheduleNextPoll for PollScheduler {
+    fn schedule_next_poll(&self, poll_at: Option<u64>) {
+        let Some(new_instant) = poll_at else {
+            self.next_poll_at_ms.store(0, Ordering::Relaxed);
+            return;
+        };
+
+        let old_instant = self.next_poll_at_ms.load(Ordering::Relaxed);
+        self.next_poll_at_ms.store(new_instant, Ordering::Relaxed);
+
+        if old_instant == 0 || new_instant < old_instant {
+            self.polling_wait_queue.wake_all();
+        }
+    }
+}

kernel/src/net/socket/ip/datagram/mod.rs

@@ -18,16 +18,13 @@ use crate::{
         utils::{InodeMode, Metadata, StatusFlags},
     },
     match_sock_option_mut,
-    net::{
-        iface::IfaceEx,
-        socket::{
-            options::{Error as SocketError, SocketOption},
-            util::{
-                options::SocketOptionSet, send_recv_flags::SendRecvFlags, socket_addr::SocketAddr,
-                MessageHeader,
-            },
-            Socket,
+    net::socket::{
+        options::{Error as SocketError, SocketOption},
+        util::{
+            options::SocketOptionSet, send_recv_flags::SendRecvFlags, socket_addr::SocketAddr,
+            MessageHeader,
         },
+        Socket,
     },
     prelude::*,
     process::signal::{PollHandle, Pollable, Pollee},

kernel/src/net/socket/ip/stream/mod.rs

@@ -23,16 +23,13 @@ use crate::{
         utils::{InodeMode, Metadata, StatusFlags},
     },
     match_sock_option_mut, match_sock_option_ref,
-    net::{
-        iface::IfaceEx,
-        socket::{
-            options::{Error as SocketError, SocketOption},
-            util::{
-                options::SocketOptionSet, send_recv_flags::SendRecvFlags,
-                shutdown_cmd::SockShutdownCmd, socket_addr::SocketAddr, MessageHeader,
-            },
-            Socket,
+    net::socket::{
+        options::{Error as SocketError, SocketOption},
+        util::{
+            options::SocketOptionSet, send_recv_flags::SendRecvFlags,
+            shutdown_cmd::SockShutdownCmd, socket_addr::SocketAddr, MessageHeader,
         },
+        Socket,
     },
     prelude::*,
     process::signal::{PollHandle, Pollable, Pollee},