Merge pull request #4950 from WhySoBad/fix-blocking-io-manager

Fix panic when multiple threads block on same fd

Author: RalfJung

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

@@ -1,10 +1,12 @@
+use std::collections::BTreeMap;
 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::shims::{FdId, FileDescriptionRef};
 use crate::*;
 
 /// Capacity of the event queue which can be polled at a time.
@@ -18,6 +20,14 @@ pub trait WithSource {
     fn with_source(&self, f: &mut dyn FnMut(&mut dyn Source) -> io::Result<()>) -> io::Result<()>;
 }
 
+/// An interest receiver defines the action that should be taken when
+/// the associated [`Interest`] is fulfilled.
+#[derive(Debug, Hash, PartialEq, Clone, Copy, Eq, PartialOrd, Ord)]
+pub enum InterestReceiver {
+    /// The specified thread should be unblocked.
+    UnblockThread(ThreadId),
+}
+
 /// 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.
@@ -34,9 +44,10 @@ pub struct BlockingIoManager {
     /// 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>>,
+    /// Map from source ids to the actual sources and their registered receivers
+    /// together with their associated interests.
+    sources:
+        BTreeMap<FdId, (FileDescriptionRef<dyn WithSource>, FxHashMap<InterestReceiver, Interest>)>,
 }
 
 impl BlockingIoManager {
@@ -46,7 +57,7 @@ impl BlockingIoManager {
         let manager = Self {
             poll: communicate.then_some(Poll::new()?),
             events: Events::with_capacity(IO_EVENT_CAPACITY),
-            sources: FxHashMap::default(),
+            sources: BTreeMap::default(),
         };
         Ok(manager)
     }
@@ -59,8 +70,12 @@ impl BlockingIoManager {
     ///   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> {
+    /// Returns the interest receivers for all file descriptions which received an I/O event together
+    /// with the file description they were registered for.
+    pub fn poll(
+        &mut self,
+        timeout: Option<Duration>,
+    ) -> Result<Vec<(InterestReceiver, FileDescriptionRef<dyn WithSource>)>, io::Error> {
         let poll =
             self.poll.as_mut().expect("Blocking I/O should not be called with isolation enabled");
 
@@ -70,56 +85,120 @@ impl BlockingIoManager {
         let ready = self
             .events
             .iter()
-            .map(|event| {
+            .flat_map(|event| {
                 let token = event.token();
-                ThreadId::new_unchecked(token.0.try_into().unwrap())
+                // We know all tokens are valid `FdId`.
+                let fd_id = FdId::new_unchecked(token.0);
+                let (source, interests) =
+                    self.sources.get(&fd_id).expect("Source should be registered");
+                assert_eq!(source.id(), fd_id);
+                // Because we allow spurious wake-ups, we mark all interests as ready even
+                // though some may not have been fulfilled.
+                interests.keys().map(move |receiver| (*receiver, source.clone()))
             })
             .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));
+        // Deregister all ready sources as we only want to receive one event per receiver.
+        ready.iter().for_each(|(receiver, source)| self.deregister(source.id(), *receiver));
 
         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.
+    /// Register an interest for a blocking I/O source.
     ///
     /// 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) {
+    ///
+    /// It's assumed that no interest is already registered for this source with the same reason!
+    pub fn register(
+        &mut self,
+        source_fd: FileDescriptionRef<dyn WithSource>,
+        receiver: InterestReceiver,
+        interest: 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());
+        let id = source_fd.id();
+        let token = Token(id.to_usize());
+
+        let Some((_, current_interests)) = self.sources.get_mut(&id) else {
+            // The source is not yet registered.
+
+            // Treat errors from registering as fatal. On UNIX hosts this can only
+            // fail due to system resource errors (e.g. ENOMEM or ENOSPC).
+            source_fd
+                .with_source(&mut |source| poll.registry().register(source, token, interest))
+                .unwrap();
+
+            self.sources.insert(id, (source_fd, FxHashMap::from_iter([(receiver, interest)])));
+            return;
+        };
+
+        // The source is already registered. We need to check whether we need to
+        // reregister because the provided interest contains new interests for the source.
 
-        // Treat errors from registering as fatal. On UNIX hosts this can only
+        let old_interest =
+            interest_union(current_interests).expect("Source should contain at least one interest");
+
+        current_interests
+            .try_insert(receiver, interest)
+            .unwrap_or_else(|_| panic!("Receiver should be unique"));
+
+        let new_interest = old_interest.add(interest);
+
+        // Reregister the source since the overall interests might have changed.
+
+        // Treat errors from reregistering 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))
+        source_fd
+            .with_source(&mut |source| poll.registry().reregister(source, token, new_interest))
             .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) {
+    /// Deregister an interest from a blocking I/O source.
+    ///
+    /// The receiver is assumed to be registered for the provided source!
+    pub fn deregister(&mut self, source_id: FdId, receiver: InterestReceiver) {
         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")
+        let token = Token(source_id.to_usize());
+        let (fd, current_interests) =
+            self.sources.get_mut(&source_id).expect("Source should be registered");
+
+        current_interests
+            .remove(&receiver)
+            .unwrap_or_else(|| panic!("Receiver should be registered for source"));
+
+        let Some(new_interest) = interest_union(current_interests) else {
+            // There are no longer any interests in this source.
+            // We can thus deregister the source from the poll.
+
+            // Treat errors from deregistering as fatal. On UNIX hosts this can only
+            // fail due to system resource errors (e.g. ENOMEM or ENOSPC).
+            fd.with_source(&mut |source| poll.registry().deregister(source)).unwrap();
+            self.sources.remove(&source_id);
+            return;
         };
 
-        // Treat errors from deregistering as fatal. On UNIX hosts this can only
+        // Reregister the source since the overall interests might have changed.
+
+        // Treat errors from reregistering 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();
+        fd.with_source(&mut |source| poll.registry().reregister(source, token, new_interest))
+            .unwrap();
     }
 }
 
+/// Get the union of all interests for a source. Returns `None` if the map is empty.
+fn interest_union(interests: &FxHashMap<InterestReceiver, Interest>) -> Option<Interest> {
+    interests
+        .values()
+        .copied()
+        .fold(None, |acc, interest| acc.map(|acc: Interest| acc.add(interest)).or(Some(interest)))
+}
+
 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
@@ -132,15 +211,15 @@ pub trait EvalContextExt<'tcx>: MiriInterpCxExt<'tcx> {
     #[inline]
     fn block_thread_for_io(
         &mut self,
-        source: impl WithSource + 'static,
+        source_fd: FileDescriptionRef<dyn WithSource>,
         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(),
+            source_fd,
+            InterestReceiver::UnblockThread(this.machine.threads.active_thread()),
             interests,
         );
         this.block_thread(BlockReason::IO, timeout, callback);

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

@@ -18,6 +18,7 @@ use rustc_span::{DUMMY_SP, Span};
 use rustc_target::spec::Os;
 
 use crate::concurrency::GlobalDataRaceHandler;
+use crate::concurrency::blocking_io::InterestReceiver;
 use crate::shims::tls;
 use crate::*;
 
@@ -822,7 +823,12 @@ trait EvalContextPrivExt<'tcx>: MiriInterpCxExt<'tcx> {
             Err(e) => panic!("unexpected error while polling: {e}"),
         };
 
-        ready.into_iter().try_for_each(|thread_id| this.unblock_thread(thread_id, BlockReason::IO))
+        ready.into_iter().try_for_each(|(receiver, _source)| {
+            match receiver {
+                InterestReceiver::UnblockThread(thread_id) =>
+                    this.unblock_thread(thread_id, BlockReason::IO),
+            }
+        })
     }
 
     /// Find all threads with expired timeouts, unblock them and execute their timeout callbacks.

src/tools/miri/src/shims/files.rs

@@ -19,6 +19,17 @@ use crate::*;
 #[derive(Debug, Copy, Clone, Default, Eq, PartialEq, Ord, PartialOrd)]
 pub struct FdId(usize);
 
+impl FdId {
+    pub fn to_usize(self) -> usize {
+        self.0
+    }
+
+    /// Create a new fd id from a `usize` without checking if this fd exists.
+    pub fn new_unchecked(id: usize) -> Self {
+        Self(id)
+    }
+}
+
 #[derive(Debug, Clone)]
 struct FdIdWith<T: ?Sized> {
     id: FdId,

src/tools/miri/src/shims/mod.rs

@@ -23,7 +23,7 @@ pub mod time;
 pub mod tls;
 pub mod unwind;
 
-pub use self::files::FdTable;
+pub use self::files::{FdId, FdTable, FileDescriptionRef};
 #[cfg(all(feature = "native-lib", unix))]
 pub use self::native_lib::trace::{init_sv, register_retcode_sv};
 pub use self::unix::{DirTable, EpollInterestTable};

src/tools/miri/src/shims/unix/fd.rs

@@ -15,7 +15,7 @@ use crate::shims::unix::*;
 use crate::*;
 
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
-pub(crate) enum FlockOp {
+pub enum FlockOp {
     SharedLock { nonblocking: bool },
     ExclusiveLock { nonblocking: bool },
     Unlock,

src/tools/miri/src/shims/unix/socket.rs

@@ -1388,7 +1388,7 @@ impl VisitProvenance for FileDescriptionRef<Socket> {
     fn visit_provenance(&self, _visit: &mut VisitWith<'_>) {}
 }
 
-impl WithSource for FileDescriptionRef<Socket> {
+impl WithSource for Socket {
     fn with_source(&self, f: &mut dyn FnMut(&mut dyn Source) -> io::Result<()>) -> io::Result<()> {
         let mut state = self.state.borrow_mut();
         match &mut *state {

src/tools/miri/tests/pass-dep/libc/libc-blocking-io-same-fd.rs

@@ -0,0 +1,56 @@
+//@ignore-target: windows # No libc socket on Windows
+//@compile-flags: -Zmiri-disable-isolation -Zmiri-fixed-schedule
+
+#[path = "../../utils/libc.rs"]
+mod libc_utils;
+use std::thread;
+
+use libc_utils::*;
+
+// This tests that the blocking I/O implementation works when multiple threads block on the
+// same fd at the same time.
+
+fn main() {
+    let (server_sockfd, addr) = net::make_listener_ipv4(0).unwrap();
+    let client_sockfd =
+        unsafe { errno_result(libc::socket(libc::AF_INET, libc::SOCK_STREAM, 0)).unwrap() };
+
+    // Spawn the server thread.
+    let server_thread = thread::spawn(move || {
+        let (peerfd, _) = net::accept_ipv4(server_sockfd).unwrap();
+
+        // Yield back to reader threads to ensure that we have
+        // two threads being blocked on the same fd at the same time.
+        thread::yield_now();
+
+        let mut buffer = [22u8; 128];
+        let bytes_written = unsafe {
+            errno_result(net::send_all(peerfd, buffer.as_mut_ptr().cast(), buffer.len(), 0))
+                .unwrap()
+        };
+        assert_eq!(bytes_written as usize, 128);
+    });
+
+    net::connect_ipv4(client_sockfd, addr);
+
+    let reader_thread = thread::spawn(move || {
+        let mut buffer = [0u8; 8];
+        let bytes_read = unsafe {
+            errno_result(net::recv_all(client_sockfd, buffer.as_mut_ptr().cast(), buffer.len(), 0))
+                .unwrap()
+        };
+        assert_eq!(bytes_read, 8);
+        assert_eq!(&buffer, &[22u8; 8]);
+    });
+
+    let mut buffer = [0u8; 8];
+    let bytes_read = unsafe {
+        errno_result(net::recv_all(client_sockfd, buffer.as_mut_ptr().cast(), buffer.len(), 0))
+            .unwrap()
+    };
+    assert_eq!(bytes_read, 8);
+    assert_eq!(&buffer, &[22u8; 8]);
+
+    reader_thread.join().unwrap();
+    server_thread.join().unwrap();
+}