Make Ext a trait

kernel/libs/aster-bigtcp/src/ext.rs

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: MPL-2.0
+
+use crate::iface::ScheduleNextPoll;
+
+/// Extension to be implemented by users of this crate.
+///
+/// This should be implemented on an empty type that carries no data, since the type will never
+/// actually be instantiated.
+///
+/// The purpose of having this trait is to allow users of this crate to inject multiple types
+/// without the hassle of writing multiple trait bounds, which can be achieved by using the types
+/// associated with this trait.
+pub trait Ext {
+    /// The type for ifaces to schedule the next poll.
+    type ScheduleNextPoll: ScheduleNextPoll;
+}

kernel/libs/aster-bigtcp/src/iface/common.rs

@@ -6,6 +6,7 @@ use alloc::{
         btree_map::{BTreeMap, Entry},
         btree_set::BTreeSet,
     },
+    string::String,
     sync::Arc,
 };
 
@@ -25,41 +26,52 @@ use super::{
 };
 use crate::{
     errors::BindError,
+    ext::Ext,
     socket::{
         BoundTcpSocket, BoundTcpSocketInner, BoundUdpSocket, BoundUdpSocketInner, UnboundTcpSocket,
         UnboundUdpSocket,
     },
 };
 
-pub struct IfaceCommon<E> {
+pub struct IfaceCommon<E: Ext> {
+    name: String,
     interface: SpinLock<smoltcp::iface::Interface, LocalIrqDisabled>,
     used_ports: SpinLock<BTreeMap<u16, usize>, PreemptDisabled>,
     tcp_sockets: SpinLock<BTreeSet<KeyableArc<BoundTcpSocketInner<E>>>, LocalIrqDisabled>,
     udp_sockets: SpinLock<BTreeSet<KeyableArc<BoundUdpSocketInner<E>>>, LocalIrqDisabled>,
-    ext: E,
+    sched_poll: E::ScheduleNextPoll,
 }
 
-impl<E> IfaceCommon<E> {
-    pub(super) fn new(interface: smoltcp::iface::Interface, ext: E) -> Self {
+impl<E: Ext> IfaceCommon<E> {
+    pub(super) fn new(
+        name: String,
+        interface: smoltcp::iface::Interface,
+        sched_poll: E::ScheduleNextPoll,
+    ) -> Self {
         Self {
+            name,
             interface: SpinLock::new(interface),
             used_ports: SpinLock::new(BTreeMap::new()),
             tcp_sockets: SpinLock::new(BTreeSet::new()),
             udp_sockets: SpinLock::new(BTreeSet::new()),
-            ext,
+            sched_poll,
         }
     }
 
+    pub(super) fn name(&self) -> &str {
+        &self.name
+    }
+
     pub(super) fn ipv4_addr(&self) -> Option<Ipv4Address> {
         self.interface.lock().ipv4_addr()
     }
 
-    pub(super) fn ext(&self) -> &E {
-        &self.ext
+    pub(super) fn sched_poll(&self) -> &E::ScheduleNextPoll {
+        &self.sched_poll
     }
 }
 
-impl<E> IfaceCommon<E> {
+impl<E: Ext> IfaceCommon<E> {
     /// Acquires the lock to the interface.
     pub(crate) fn interface(&self) -> SpinLockGuard<smoltcp::iface::Interface, LocalIrqDisabled> {
         self.interface.lock()
@@ -69,7 +81,7 @@ impl<E> IfaceCommon<E> {
 const IP_LOCAL_PORT_START: u16 = 32768;
 const IP_LOCAL_PORT_END: u16 = 60999;
 
-impl<E> IfaceCommon<E> {
+impl<E: Ext> IfaceCommon<E> {
     pub(super) fn bind_tcp(
         &self,
         iface: Arc<dyn Iface<E>>,
@@ -159,7 +171,7 @@ impl<E> IfaceCommon<E> {
     }
 }
 
-impl<E> IfaceCommon<E> {
+impl<E: Ext> IfaceCommon<E> {
     #[allow(clippy::mutable_key_type)]
     fn remove_dead_tcp_sockets(&self, sockets: &mut BTreeSet<KeyableArc<BoundTcpSocketInner<E>>>) {
         sockets.retain(|socket| {
@@ -192,7 +204,7 @@ impl<E> IfaceCommon<E> {
     }
 }
 
-impl<E> IfaceCommon<E> {
+impl<E: Ext> IfaceCommon<E> {
     pub(super) fn poll<D, P, Q>(
         &self,
         device: &mut D,

kernel/libs/aster-bigtcp/src/iface/iface.rs

@@ -7,6 +7,7 @@ use smoltcp::wire::Ipv4Address;
 use super::port::BindPortConfig;
 use crate::{
     errors::BindError,
+    ext::Ext,
     socket::{BoundTcpSocket, BoundUdpSocket, UnboundTcpSocket, UnboundUdpSocket},
 };
 
@@ -18,19 +19,10 @@ use crate::{
 /// private network (VPN) connections.
 pub trait Iface<E>: internal::IfaceInternal<E> + Send + Sync {
     /// Transmits or receives packets queued in the iface, and updates socket status accordingly.
-    ///
-    /// The `schedule_next_poll` callback is invoked with the time at which the next poll should be
-    /// performed, or `None` if no next poll is required. It's up to the caller to determine the
-    /// mechanism to ensure that the next poll happens at the right time (e.g. by setting a timer).
-    fn raw_poll(&self, schedule_next_poll: &dyn Fn(Option<u64>));
+    fn poll(&self);
 }
 
-impl<E> dyn Iface<E> {
-    /// Gets the extension of the iface.
-    pub fn ext(&self) -> &E {
-        self.common().ext()
-    }
-
+impl<E: Ext> dyn Iface<E> {
     /// Binds a socket to the iface.
     ///
     /// After binding the socket to the iface, the iface will handle all packets to and from the
@@ -60,19 +52,33 @@ impl<E> dyn Iface<E> {
         common.bind_udp(self.clone(), socket, config)
     }
 
+    /// Gets the name of the iface.
+    ///
+    /// In Linux, the name is usually the driver name followed by a unit number.
+    pub fn name(&self) -> &str {
+        self.common().name()
+    }
+
     /// Gets the IPv4 address of the iface, if any.
     ///
     /// FIXME: One iface may have multiple IPv4 addresses.
     pub fn ipv4_addr(&self) -> Option<Ipv4Address> {
         self.common().ipv4_addr()
     }
+
+    /// Returns a reference to the associated [`ScheduleNextPoll`].
+    pub fn sched_poll(&self) -> &E::ScheduleNextPoll {
+        self.common().sched_poll()
+    }
 }
 
 pub(super) mod internal {
-    use crate::iface::common::IfaceCommon;
+    use crate::{ext::Ext, iface::common::IfaceCommon};
 
     /// An internal trait that abstracts the common part of different ifaces.
     pub trait IfaceInternal<E> {
-        fn common(&self) -> &IfaceCommon<E>;
+        fn common(&self) -> &IfaceCommon<E>
+        where
+            E: Ext;
     }
 }

kernel/libs/aster-bigtcp/src/iface/mod.rs

@@ -6,8 +6,10 @@ mod iface;
 mod phy;
 mod poll;
 mod port;
+mod sched;
 mod time;
 
 pub use iface::Iface;
 pub use phy::{EtherIface, IpIface};
 pub use port::BindPortConfig;
+pub use sched::ScheduleNextPoll;

kernel/libs/aster-bigtcp/src/iface/phy/ether.rs

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: MPL-2.0
 
-use alloc::{collections::btree_map::BTreeMap, sync::Arc};
+use alloc::{collections::btree_map::BTreeMap, string::String, sync::Arc};
 
 use ostd::sync::{LocalIrqDisabled, SpinLock};
 use smoltcp::{
@@ -14,25 +14,28 @@ use smoltcp::{
 
 use crate::{
     device::WithDevice,
+    ext::Ext,
     iface::{
         common::IfaceCommon, iface::internal::IfaceInternal, time::get_network_timestamp, Iface,
+        ScheduleNextPoll,
     },
 };
 
-pub struct EtherIface<D, E> {
+pub struct EtherIface<D, E: Ext> {
     driver: D,
     common: IfaceCommon<E>,
     ether_addr: EthernetAddress,
     arp_table: SpinLock<BTreeMap<Ipv4Address, EthernetAddress>, LocalIrqDisabled>,
 }
 
-impl<D: WithDevice, E> EtherIface<D, E> {
+impl<D: WithDevice, E: Ext> EtherIface<D, E> {
     pub fn new(
         driver: D,
         ether_addr: EthernetAddress,
         ip_cidr: Ipv4Cidr,
         gateway: Ipv4Address,
-        ext: E,
+        name: String,
+        sched_poll: E::ScheduleNextPoll,
     ) -> Arc<Self> {
         let interface = driver.with(|device| {
             let config = Config::new(wire::HardwareAddress::Ethernet(ether_addr));
@@ -50,7 +53,7 @@ impl<D: WithDevice, E> EtherIface<D, E> {
             interface
         });
 
-        let common = IfaceCommon::new(interface, ext);
+        let common = IfaceCommon::new(name, interface, sched_poll);
 
         Arc::new(Self {
             driver,
@@ -61,26 +64,26 @@ impl<D: WithDevice, E> EtherIface<D, E> {
     }
 }
 
-impl<D, E> IfaceInternal<E> for EtherIface<D, E> {
+impl<D, E: Ext> IfaceInternal<E> for EtherIface<D, E> {
     fn common(&self) -> &IfaceCommon<E> {
         &self.common
     }
 }
 
-impl<D: WithDevice + 'static, E: Send + Sync> Iface<E> for EtherIface<D, E> {
-    fn raw_poll(&self, schedule_next_poll: &dyn Fn(Option<u64>)) {
+impl<D: WithDevice + 'static, E: Ext> Iface<E> for EtherIface<D, E> {
+    fn poll(&self) {
         self.driver.with(|device| {
             let next_poll = self.common.poll(
                 &mut *device,
                 |data, iface_cx, tx_token| self.process(data, iface_cx, tx_token),
                 |pkt, iface_cx, tx_token| self.dispatch(pkt, iface_cx, tx_token),
             );
-            schedule_next_poll(next_poll);
+            self.common.sched_poll().schedule_next_poll(next_poll);
         });
     }
 }
 
-impl<D, E> EtherIface<D, E> {
+impl<D, E: Ext> EtherIface<D, E> {
     fn process<'pkt, T: TxToken>(
         &self,
         data: &'pkt [u8],

kernel/libs/aster-bigtcp/src/iface/phy/ip.rs

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: MPL-2.0
 
-use alloc::sync::Arc;
+use alloc::{string::String, sync::Arc};
 
 use smoltcp::{
     iface::Config,
@@ -10,18 +10,25 @@ use smoltcp::{
 
 use crate::{
     device::WithDevice,
+    ext::Ext,
     iface::{
         common::IfaceCommon, iface::internal::IfaceInternal, time::get_network_timestamp, Iface,
+        ScheduleNextPoll,
     },
 };
 
-pub struct IpIface<D, E> {
+pub struct IpIface<D, E: Ext> {
     driver: D,
     common: IfaceCommon<E>,
 }
 
-impl<D: WithDevice, E> IpIface<D, E> {
-    pub fn new(driver: D, ip_cidr: Ipv4Cidr, ext: E) -> Arc<Self> {
+impl<D: WithDevice, E: Ext> IpIface<D, E> {
+    pub fn new(
+        driver: D,
+        ip_cidr: Ipv4Cidr,
+        name: String,
+        sched_poll: E::ScheduleNextPoll,
+    ) -> Arc<Self> {
         let interface = driver.with(|device| {
             let config = Config::new(smoltcp::wire::HardwareAddress::Ip);
             let now = get_network_timestamp();
@@ -34,20 +41,20 @@ impl<D: WithDevice, E> IpIface<D, E> {
             interface
         });
 
-        let common = IfaceCommon::new(interface, ext);
+        let common = IfaceCommon::new(name, interface, sched_poll);
 
         Arc::new(Self { driver, common })
     }
 }
 
-impl<D, E> IfaceInternal<E> for IpIface<D, E> {
+impl<D, E: Ext> IfaceInternal<E> for IpIface<D, E> {
     fn common(&self) -> &IfaceCommon<E> {
         &self.common
     }
 }
 
-impl<D: WithDevice + 'static, E: Send + Sync> Iface<E> for IpIface<D, E> {
-    fn raw_poll(&self, schedule_next_poll: &dyn Fn(Option<u64>)) {
+impl<D: WithDevice + 'static, E: Ext> Iface<E> for IpIface<D, E> {
+    fn poll(&self) {
         self.driver.with(|device| {
             let next_poll = self.common.poll(
                 device,
@@ -64,7 +71,7 @@ impl<D: WithDevice + 'static, E: Send + Sync> Iface<E> for IpIface<D, E> {
                     });
                 },
             );
-            schedule_next_poll(next_poll);
+            self.common.sched_poll().schedule_next_poll(next_poll);
         });
     }
 }

kernel/libs/aster-bigtcp/src/iface/sched.rs

@@ -0,0 +1,11 @@
+// SPDX-License-Identifier: MPL-2.0
+
+/// A trait to provide the `schedule_next_poll` method for ifaces.
+pub trait ScheduleNextPoll: Send + Sync {
+    /// Schedules the next poll at the specific time.
+    ///
+    /// This is invoked with the time at which the next poll should be performed, or `None` if no
+    /// next poll is required. It's up to the caller to determine the mechanism to ensure that the
+    /// next poll happens at the right time (e.g. by setting a timer).
+    fn schedule_next_poll(&self, ms: Option<u64>);
+}

kernel/libs/aster-bigtcp/src/lib.rs

@@ -16,6 +16,7 @@
 
 pub mod device;
 pub mod errors;
+pub mod ext;
 pub mod iface;
 pub mod socket;
 pub mod time;

kernel/libs/aster-bigtcp/src/socket/bound.rs

@@ -21,9 +21,9 @@ use super::{
     event::{SocketEventObserver, SocketEvents},
     RawTcpSocket, RawUdpSocket, TcpStateCheck,
 };
-use crate::iface::Iface;
+use crate::{ext::Ext, iface::Iface};
 
-pub struct BoundSocket<T: AnySocket, E>(Arc<BoundSocketInner<T, E>>);
+pub struct BoundSocket<T: AnySocket, E: Ext>(Arc<BoundSocketInner<T, E>>);
 
 /// [`TcpSocket`] or [`UdpSocket`].
 pub trait AnySocket {
@@ -33,7 +33,7 @@ pub trait AnySocket {
     fn new(socket: Box<Self::RawSocket>) -> Self;
 
     /// Called by [`BoundSocket::drop`].
-    fn on_drop<E>(this: &Arc<BoundSocketInner<Self, E>>)
+    fn on_drop<E: Ext>(this: &Arc<BoundSocketInner<Self, E>>)
     where
         Self: Sized;
 }
@@ -176,7 +176,7 @@ impl AnySocket for UdpSocket {
         Self::new(socket)
     }
 
-    fn on_drop<E>(this: &Arc<BoundSocketInner<Self, E>>) {
+    fn on_drop<E: Ext>(this: &Arc<BoundSocketInner<Self, E>>) {
         this.socket.lock().close();
 
         // A UDP socket can be removed immediately.
@@ -184,7 +184,7 @@ impl AnySocket for UdpSocket {
     }
 }
 
-impl<T: AnySocket, E> Drop for BoundSocket<T, E> {
+impl<T: AnySocket, E: Ext> Drop for BoundSocket<T, E> {
     fn drop(&mut self) {
         T::on_drop(&self.0);
     }
@@ -193,7 +193,7 @@ impl<T: AnySocket, E> Drop for BoundSocket<T, E> {
 pub(crate) type BoundTcpSocketInner<E> = BoundSocketInner<TcpSocket, E>;
 pub(crate) type BoundUdpSocketInner<E> = BoundSocketInner<UdpSocket, E>;
 
-impl<T: AnySocket, E> BoundSocket<T, E> {
+impl<T: AnySocket, E: Ext> BoundSocket<T, E> {
     pub(crate) fn new(
         iface: Arc<dyn Iface<E>>,
         port: u16,
@@ -215,7 +215,7 @@ impl<T: AnySocket, E> BoundSocket<T, E> {
     }
 }
 
-impl<T: AnySocket, E> BoundSocket<T, E> {
+impl<T: AnySocket, E: Ext> BoundSocket<T, E> {
     /// Sets the observer whose `on_events` will be called when certain iface events happen. After
     /// setting, the new observer will fire once immediately to avoid missing any events.
     ///
@@ -269,7 +269,7 @@ impl Deref for NeedIfacePoll {
     }
 }
 
-impl<E> BoundTcpSocket<E> {
+impl<E: Ext> BoundTcpSocket<E> {
     /// Connects to a remote endpoint.
     ///
     /// Polling the iface is _always_ required after this method succeeds.
@@ -378,7 +378,7 @@ impl<E> BoundTcpSocket<E> {
     }
 }
 
-impl<E> BoundUdpSocket<E> {
+impl<E: Ext> BoundUdpSocket<E> {
     /// Binds to a specified endpoint.
     ///
     /// Polling the iface is _not_ required after this method succeeds.

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},