Merge pull request rust-lang#4892 from WhySoBad/network-socket-io-polling

Handle blocking I/O using Mio, add `accept` and `connect` network socket shims

Author: RalfJung

src/tools/miri/Cargo.lock

@@ -779,9 +779,9 @@ checksum = "db13adb97ab515a3691f56e4dbab09283d0b86cb45abd991d8634a9d6f501760"
 
 [[package]]
 name = "libc"
-version = "0.2.177"
+version = "0.2.182"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976"
+checksum = "6800badb6cb2082ffd7b6a67e6125bb39f18782f793520caee8cb8846be06112"
 
 [[package]]
 name = "libffi"
@@ -923,11 +923,12 @@ dependencies = [
 
 [[package]]
 name = "mio"
-version = "1.1.0"
+version = "1.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "69d83b0086dc8ecf3ce9ae2874b2d1290252e2a30720bea58a5c6639b0092873"
+checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc"
 dependencies = [
  "libc",
+ "log",
  "wasi",
  "windows-sys 0.61.2",
 ]
@@ -950,6 +951,7 @@ dependencies = [
  "libffi",
  "libloading",
  "measureme",
+ "mio",
  "nix",
  "rand",
  "regex",

src/tools/miri/Cargo.toml

@@ -28,6 +28,7 @@ chrono-tz = "0.10"
 directories = "6"
 bitflags = "2.6"
 serde_json = { version = "1.0", optional = true }
+mio = { version = "1.1.1", features = ["os-poll", "net"] }
 
 [target.'cfg(unix)'.dependencies]
 libc = "0.2"

src/tools/miri/src/concurrency/blocking_io.rs

@@ -0,0 +1,148 @@
+use std::io;
+use std::time::Duration;
+
+use mio::event::Source;
+use mio::{Events, Interest, Poll, Token};
+use rustc_data_structures::fx::FxHashMap;
+
+use crate::*;
+
+/// Capacity of the event queue which can be polled at a time.
+/// Since we don't expect many simultaneous blocking I/O events
+/// this value can be set rather low.
+const IO_EVENT_CAPACITY: usize = 16;
+
+/// Trait for values that contain a mio [`Source`].
+pub trait WithSource {
+    /// Invoke `f` on the source inside `self`.
+    fn with_source(&self, f: &mut dyn FnMut(&mut dyn Source) -> io::Result<()>) -> io::Result<()>;
+}
+
+/// Manager for managing blocking host I/O in a non-blocking manner.
+/// We use [`Poll`] to poll for new I/O events from the OS for sources
+/// registered using this manager.
+///
+/// Since blocking host I/O is inherently non-deterministic, no method on this
+/// manager should be called when isolation is enabled. The only exception is
+/// the [`BlockingIoManager::new`] function to create the manager. Everywhere else,
+/// we assert that isolation is disabled!
+pub struct BlockingIoManager {
+    /// Poll instance to monitor I/O events from the OS.
+    /// This is only [`None`] when Miri is run with isolation enabled.
+    poll: Option<Poll>,
+    /// Buffer used to store the ready I/O events when calling [`Poll::poll`].
+    /// This is not part of the state and only stored to avoid allocating a
+    /// new buffer for every poll.
+    events: Events,
+    /// Map between threads which are currently blocked and the
+    /// underlying I/O source.
+    sources: FxHashMap<ThreadId, Box<dyn WithSource>>,
+}
+
+impl BlockingIoManager {
+    /// Create a new blocking I/O manager instance based on the availability
+    /// of communication with the host.
+    pub fn new(communicate: bool) -> Result<Self, io::Error> {
+        let manager = Self {
+            poll: communicate.then_some(Poll::new()?),
+            events: Events::with_capacity(IO_EVENT_CAPACITY),
+            sources: FxHashMap::default(),
+        };
+        Ok(manager)
+    }
+
+    /// Poll for new I/O events from the OS or wait until the timeout expired.
+    ///
+    /// - If the timeout is [`Some`] and contains [`Duration::ZERO`], the poll doesn't block and just
+    ///   reads all events since the last poll.
+    /// - If the timeout is [`Some`] and contains a non-zero duration, it blocks at most for the
+    ///   specified duration.
+    /// - If the timeout is [`None`] the poll blocks indefinitely until an event occurs.
+    ///
+    /// Returns all threads that are ready because they received an I/O event.
+    pub fn poll(&mut self, timeout: Option<Duration>) -> Result<Vec<ThreadId>, io::Error> {
+        let poll =
+            self.poll.as_mut().expect("Blocking I/O should not be called with isolation enabled");
+
+        // Poll for new I/O events from OS and store them in the events buffer.
+        poll.poll(&mut self.events, timeout)?;
+
+        let ready = self
+            .events
+            .iter()
+            .map(|event| {
+                let token = event.token();
+                ThreadId::new_unchecked(token.0.try_into().unwrap())
+            })
+            .collect::<Vec<_>>();
+
+        // Deregister all ready sources as we only want to receive one event per thread.
+        ready.iter().for_each(|thread_id| self.deregister(*thread_id));
+
+        Ok(ready)
+    }
+
+    /// Register a blocking I/O source for a thread together with it's poll interests.
+    ///
+    /// The source will be deregistered automatically once an event for it is received.
+    ///
+    /// As the OS can always produce spurious wake-ups, it's the callers responsibility to
+    /// verify the requested I/O interests are really ready and to register again if they're not.
+    pub fn register(&mut self, source: Box<dyn WithSource>, thread: ThreadId, interests: Interest) {
+        let poll =
+            self.poll.as_ref().expect("Blocking I/O should not be called with isolation enabled");
+
+        let token = Token(thread.to_u32().to_usize());
+
+        // Treat errors from registering as fatal. On UNIX hosts this can only
+        // fail due to system resource errors (e.g. ENOMEM or ENOSPC).
+        source
+            .with_source(&mut |source| source.register(poll.registry(), token, interests))
+            .unwrap();
+        self.sources
+            .try_insert(thread, source)
+            .unwrap_or_else(|_| panic!("A thread cannot be registered twice at the same time"));
+    }
+
+    /// Deregister the event source for a thread. Returns the kind of I/O the thread was
+    /// blocked on.
+    fn deregister(&mut self, thread: ThreadId) {
+        let poll =
+            self.poll.as_ref().expect("Blocking I/O should not be called with isolation enabled");
+
+        let Some(source) = self.sources.remove(&thread) else {
+            panic!("Attempt to deregister a token which isn't registered")
+        };
+
+        // Treat errors from deregistering as fatal. On UNIX hosts this can only
+        // fail due to system resource errors (e.g. ENOMEM or ENOSPC).
+        source.with_source(&mut |source| source.deregister(poll.registry())).unwrap();
+    }
+}
+
+impl<'tcx> EvalContextExt<'tcx> for MiriInterpCx<'tcx> {}
+pub trait EvalContextExt<'tcx>: MiriInterpCxExt<'tcx> {
+    /// Block the current thread until some interests on an I/O source
+    /// are fulfilled or the optional timeout exceeded.
+    /// The callback will be invoked when the thread gets unblocked.
+    ///
+    /// There can be spurious wake-ups by the OS and thus it's the callers
+    /// responsibility to verify that the requested I/O interests are
+    /// really ready and to block again if they're not.
+    #[inline]
+    fn block_thread_for_io(
+        &mut self,
+        source: impl WithSource + 'static,
+        interests: Interest,
+        timeout: Option<(TimeoutClock, TimeoutAnchor, Duration)>,
+        callback: DynUnblockCallback<'tcx>,
+    ) {
+        let this = self.eval_context_mut();
+        this.machine.blocking_io.register(
+            Box::new(source),
+            this.machine.threads.active_thread(),
+            interests,
+        );
+        this.block_thread(BlockReason::IO, timeout, callback);
+    }
+}

src/tools/miri/src/concurrency/mod.rs

@@ -1,3 +1,4 @@
+pub mod blocking_io;
 pub mod cpu_affinity;
 pub mod data_race;
 mod data_race_handler;

src/tools/miri/src/concurrency/thread.rs

@@ -1,9 +1,9 @@
 //! Implements threads.
 
-use std::mem;
 use std::sync::atomic::Ordering::Relaxed;
 use std::task::Poll;
 use std::time::{Duration, SystemTime};
+use std::{io, mem};
 
 use rand::seq::IteratorRandom;
 use rustc_abi::ExternAbi;
@@ -25,10 +25,11 @@ use crate::*;
 enum SchedulingAction {
     /// Execute step on the active thread.
     ExecuteStep,
-    /// Execute a timeout callback.
-    ExecuteTimeoutCallback,
-    /// Wait for a bit, until there is a timeout to be called.
-    Sleep(Duration),
+    /// Wait for a bit, but at most as long as the duration specified.
+    /// We wake up early if an I/O event happened.
+    /// If the duration is [`None`], we sleep indefinitely. This is
+    /// only allowed when isolation is disabled and there are threads waiting for I/O!
+    SleepAndWaitForIo(Option<Duration>),
 }
 
 /// What to do with TLS allocations from terminated threads
@@ -111,6 +112,8 @@ pub enum BlockReason {
     Eventfd,
     /// Blocked on unnamed_socket.
     UnnamedSocket,
+    /// Blocked on an IO operation.
+    IO,
     /// Blocked for any reason related to GenMC, such as `assume` statements (GenMC mode only).
     /// Will be implicitly unblocked when GenMC schedules this thread again.
     Genmc,
@@ -765,26 +768,45 @@ trait EvalContextPrivExt<'tcx>: MiriInterpCxExt<'tcx> {
         }
 
         // We are not in GenMC mode, so we control the scheduling.
-        let thread_manager = &mut this.machine.threads;
-        let clock = &this.machine.monotonic_clock;
-        let rng = this.machine.rng.get_mut();
+        let thread_manager = &this.machine.threads;
         // This thread and the program can keep going.
         if thread_manager.threads[thread_manager.active_thread].state.is_enabled()
             && !thread_manager.yield_active_thread
         {
             // The currently active thread is still enabled, just continue with it.
             return interp_ok(SchedulingAction::ExecuteStep);
         }
-        // The active thread yielded or got terminated. Let's see if there are any timeouts to take
-        // care of. We do this *before* running any other thread, to ensure that timeouts "in the
-        // past" fire before any other thread can take an action. This ensures that for
+
+        // The active thread yielded or got terminated. Let's see if there are any I/O events
+        // or timeouts to take care of.
+
+        if this.machine.communicate() {
+            // When isolation is disabled we need to check for events for
+            // threads which are blocked on host I/O.
+            // We do this before running any other threads such that the threads
+            // which received events are available for scheduling afterwards.
+
+            // Perform a non-blocking poll for newly available I/O events from the OS.
+            this.poll_and_unblock(Some(Duration::ZERO))?;
+        }
+
+        let thread_manager = &this.machine.threads;
+        let clock = &this.machine.monotonic_clock;
+
+        // We also check timeouts before running any other thread, to ensure that timeouts
+        // "in the past" fire before any other thread can take an action. This ensures that for
         // `pthread_cond_timedwait`, "an error is returned if [...] the absolute time specified by
         // abstime has already been passed at the time of the call".
         // <https://pubs.opengroup.org/onlinepubs/9699919799/functions/pthread_cond_timedwait.html>
         let potential_sleep_time = thread_manager.next_callback_wait_time(clock);
         if potential_sleep_time == Some(Duration::ZERO) {
-            return interp_ok(SchedulingAction::ExecuteTimeoutCallback);
+            // The timeout exceeded for some thread so we unblock the thread and execute its timeout callback.
+            this.run_timeout_callback()?;
         }
+
+        let thread_manager = &mut this.machine.threads;
+        let rng = this.machine.rng.get_mut();
+
         // No callbacks immediately scheduled, pick a regular thread to execute.
         // The active thread blocked or yielded. So we go search for another enabled thread.
         // We build the list of threads by starting with the threads after the current one, followed by
@@ -832,7 +854,16 @@ trait EvalContextPrivExt<'tcx>: MiriInterpCxExt<'tcx> {
             // All threads are currently blocked, but we have unexecuted
             // timeout_callbacks, which may unblock some of the threads. Hence,
             // sleep until the first callback.
-            interp_ok(SchedulingAction::Sleep(sleep_time))
+            interp_ok(SchedulingAction::SleepAndWaitForIo(Some(sleep_time)))
+        } else if thread_manager
+            .threads
+            .iter()
+            .any(|thread| thread.state.is_blocked_on(BlockReason::IO))
+        {
+            // At least one thread is blocked on host I/O but doesn't
+            // have a timeout set. Hence, we sleep indefinitely in the
+            // hope that eventually an I/O event for this thread happens.
+            interp_ok(SchedulingAction::SleepAndWaitForIo(None))
         } else {
             throw_machine_stop!(TerminationInfo::GlobalDeadlock);
         }
@@ -1300,13 +1331,38 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                         }
                     }
                 }
-                SchedulingAction::ExecuteTimeoutCallback => {
-                    this.run_timeout_callback()?;
-                }
-                SchedulingAction::Sleep(duration) => {
-                    this.machine.monotonic_clock.sleep(duration);
+                SchedulingAction::SleepAndWaitForIo(duration) => {
+                    if this.machine.communicate() {
+                        // When we're running with isolation disabled, instead of
+                        // strictly sleeping the duration we allow waking up
+                        // early for I/O events from the OS.
+
+                        this.poll_and_unblock(duration)?;
+                    } else {
+                        let duration = duration.expect(
+                            "Infinite sleep should not be triggered when isolation is enabled",
+                        );
+                        this.machine.monotonic_clock.sleep(duration);
+                    }
                 }
             }
         }
     }
+
+    /// Poll for I/O events until either an I/O event happened or the timeout expired.
+    /// The different timeout values are described in [`BlockingIoManager::poll`].
+    fn poll_and_unblock(&mut self, timeout: Option<Duration>) -> InterpResult<'tcx> {
+        let this = self.eval_context_mut();
+
+        let ready = match this.machine.blocking_io.poll(timeout) {
+            Ok(ready) => ready,
+            // We can ignore errors originating from interrupts; that's just a spurious wakeup.
+            Err(e) if e.kind() == io::ErrorKind::Interrupted => return interp_ok(()),
+            // For other errors we panic. On Linux and BSD hosts this should only be
+            // reachable when a system resource error (e.g. ENOMEM or ENOSPC) occurred.
+            Err(e) => panic!("{e}"),
+        };
+
+        ready.into_iter().try_for_each(|thread_id| this.unblock_thread(thread_id, BlockReason::IO))
+    }
 }

src/tools/miri/src/diagnostics.rs

@@ -141,13 +141,6 @@ pub enum NonHaltingDiagnostic {
         effective_failure_ordering: AtomicReadOrd,
     },
     FileInProcOpened,
-    SocketListenUnsupportedBacklog {
-        details: bool,
-        /// Unsupported backlog value provided the by the program.
-        provided: i32,
-        /// Supported backlog value by Miri.
-        supported: i32,
-    },
 }
 
 /// Level of Miri specific diagnostics
@@ -650,8 +643,6 @@ impl<'tcx> MiriMachine<'tcx> {
             | WeakMemoryOutdatedLoad { .. } =>
                 ("tracking was triggered here".to_string(), DiagLevel::Note),
             FileInProcOpened => ("open a file in `/proc`".to_string(), DiagLevel::Warning),
-            SocketListenUnsupportedBacklog { .. } =>
-                ("call to `listen` with unsupported backlog value".to_string(), DiagLevel::Warning),
         };
 
         let title = match &e {
@@ -700,17 +691,12 @@ impl<'tcx> MiriMachine<'tcx> {
                 format!("GenMC currently does not model the failure ordering for `compare_exchange`. {was_upgraded_msg}. Miri with GenMC might miss bugs related to this memory access.")
             }
             FileInProcOpened => format!("files in `/proc` can bypass the Abstract Machine and might not work properly in Miri"),
-            SocketListenUnsupportedBacklog { provided, supported, .. } => format!("called `listen` on socket with backlog value of {provided} but only {supported} is supported"),
         };
 
         let notes = match &e {
             ProgressReport { block_count } => {
                 vec![note!("so far, {block_count} basic blocks have been executed")]
             }
-            SocketListenUnsupportedBacklog { details: true, supported, .. } =>
-                vec![note!(
-                    "the given value will be ignored and a backlog of {supported} will be used instead"
-                )],
             _ => vec![],
         };