add windows one time initialization

beepster4096 · RalfJung · commit 3244c117f8b1 · 2022-10-20T22:19:06.000+02:00
diff --git a/src/tools/miri/src/concurrency/sync.rs b/src/tools/miri/src/concurrency/sync.rs
@@ -7,6 +7,7 @@ use log::trace;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_index::vec::{Idx, IndexVec};
 
+use super::thread::MachineCallback;
 use super::vector_clock::VClock;
 use crate::*;
 
@@ -149,13 +150,68 @@ struct FutexWaiter {
     bitset: u32,
 }
 
+declare_id!(InitOnceId);
+
+struct InitOnceWaiter<'mir, 'tcx> {
+    thread: ThreadId,
+    callback: Box<dyn MachineCallback<'mir, 'tcx> + 'tcx>,
+}
+
+impl<'mir, 'tcx> std::fmt::Debug for InitOnceWaiter<'mir, 'tcx> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("InitOnce")
+            .field("thread", &self.thread)
+            .field("callback", &"dyn MachineCallback")
+            .finish()
+    }
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+/// The current status of a one time initialization.
+pub enum InitOnceStatus {
+    Uninitialized,
+    Begun,
+    Complete,
+}
+
+impl Default for InitOnceStatus {
+    fn default() -> Self {
+        Self::Uninitialized
+    }
+}
+
+/// The one time initialization state.
+#[derive(Default, Debug)]
+struct InitOnce<'mir, 'tcx> {
+    status: InitOnceStatus,
+    waiters: VecDeque<InitOnceWaiter<'mir, 'tcx>>,
+    data_race: VClock,
+}
+
+impl<'mir, 'tcx> VisitTags for InitOnce<'mir, 'tcx> {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        for waiter in self.waiters.iter() {
+            waiter.callback.visit_tags(visit);
+        }
+    }
+}
+
 /// The state of all synchronization variables.
 #[derive(Default, Debug)]
-pub(crate) struct SynchronizationState {
+pub(crate) struct SynchronizationState<'mir, 'tcx> {
     mutexes: IndexVec<MutexId, Mutex>,
     rwlocks: IndexVec<RwLockId, RwLock>,
     condvars: IndexVec<CondvarId, Condvar>,
     futexes: FxHashMap<u64, Futex>,
+    init_onces: IndexVec<InitOnceId, InitOnce<'mir, 'tcx>>,
+}
+
+impl<'mir, 'tcx> VisitTags for SynchronizationState<'mir, 'tcx> {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        for init_once in self.init_onces.iter() {
+            init_once.visit_tags(visit);
+        }
+    }
 }
 
 // Private extension trait for local helper methods
@@ -581,4 +637,144 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             futex.waiters.retain(|waiter| waiter.thread != thread);
         }
     }
+
+    #[inline]
+    /// Create state for a new one time initialization.
+    fn init_once_create(&mut self) -> InitOnceId {
+        let this = self.eval_context_mut();
+        this.machine.threads.sync.init_onces.push(Default::default())
+    }
+
+    #[inline]
+    /// Provides the closure with the next InitOnceId. Creates that InitOnce if the closure returns None,
+    /// otherwise returns the value from the closure
+    fn init_once_get_or_create<F>(&mut self, existing: F) -> InterpResult<'tcx, InitOnceId>
+    where
+        F: FnOnce(
+            &mut MiriInterpCx<'mir, 'tcx>,
+            InitOnceId,
+        ) -> InterpResult<'tcx, Option<InitOnceId>>,
+    {
+        let this = self.eval_context_mut();
+        let next_index = this.machine.threads.sync.init_onces.next_index();
+        if let Some(old) = existing(this, next_index)? {
+            Ok(old)
+        } else {
+            let new_index = this.machine.threads.sync.init_onces.push(Default::default());
+            assert_eq!(next_index, new_index);
+            Ok(new_index)
+        }
+    }
+
+    #[inline]
+    fn init_once_status(&mut self, id: InitOnceId) -> InitOnceStatus {
+        let this = self.eval_context_ref();
+        this.machine.threads.sync.init_onces[id].status
+    }
+
+    #[inline]
+    /// Put the thread into the queue waiting for the initialization.
+    fn init_once_enqueue_and_block(
+        &mut self,
+        id: InitOnceId,
+        thread: ThreadId,
+        callback: Box<dyn MachineCallback<'mir, 'tcx> + 'tcx>,
+    ) {
+        let this = self.eval_context_mut();
+        let init_once = &mut this.machine.threads.sync.init_onces[id];
+        assert_ne!(init_once.status, InitOnceStatus::Complete, "queueing on complete init once");
+        init_once.waiters.push_back(InitOnceWaiter { thread, callback });
+        this.block_thread(thread);
+    }
+
+    #[inline]
+    fn init_once_begin(&mut self, id: InitOnceId) {
+        let this = self.eval_context_mut();
+        let init_once = &mut this.machine.threads.sync.init_onces[id];
+        assert_eq!(
+            init_once.status,
+            InitOnceStatus::Uninitialized,
+            "begining already begun or complete init once"
+        );
+        init_once.status = InitOnceStatus::Begun;
+    }
+
+    #[inline]
+    fn init_once_complete(&mut self, id: InitOnceId) -> InterpResult<'tcx> {
+        let this = self.eval_context_mut();
+        let current_thread = this.get_active_thread();
+        let init_once = &mut this.machine.threads.sync.init_onces[id];
+
+        assert_eq!(
+            init_once.status,
+            InitOnceStatus::Begun,
+            "completing already complete or uninit init once"
+        );
+
+        init_once.status = InitOnceStatus::Complete;
+
+        // Each complete happens-before the end of the wait
+        if let Some(data_race) = &this.machine.data_race {
+            data_race.validate_lock_release(&mut init_once.data_race, current_thread);
+        }
+
+        // need to take the queue to avoid having `this` be borrowed multiple times
+        for waiter in std::mem::take(&mut init_once.waiters) {
+            this.unblock_thread(waiter.thread);
+
+            this.set_active_thread(waiter.thread);
+            waiter.callback.call(this)?;
+            this.set_active_thread(current_thread);
+
+            if let Some(data_race) = &this.machine.data_race {
+                data_race.validate_lock_acquire(
+                    &this.machine.threads.sync.init_onces[id].data_race,
+                    waiter.thread,
+                );
+            }
+        }
+
+        Ok(())
+    }
+
+    #[inline]
+    fn init_once_fail(&mut self, id: InitOnceId) -> InterpResult<'tcx> {
+        let this = self.eval_context_mut();
+        let current_thread = this.get_active_thread();
+        let init_once = &mut this.machine.threads.sync.init_onces[id];
+        assert_eq!(
+            init_once.status,
+            InitOnceStatus::Begun,
+            "failing already completed or uninit init once"
+        );
+
+        // Each complete happens-before the end of the wait
+        if let Some(data_race) = &this.machine.data_race {
+            data_race.validate_lock_release(&mut init_once.data_race, current_thread);
+        }
+
+        // the behavior of failing the initialization is left vague by the docs
+        // it had to be determined experimentally
+        if let Some(waiter) = init_once.waiters.pop_front() {
+            // try initializing again on a different thread
+            init_once.status = InitOnceStatus::Begun;
+
+            this.unblock_thread(waiter.thread);
+
+            this.set_active_thread(waiter.thread);
+            waiter.callback.call(this)?;
+            this.set_active_thread(current_thread);
+
+            if let Some(data_race) = &this.machine.data_race {
+                data_race.validate_lock_acquire(
+                    &this.machine.threads.sync.init_onces[id].data_race,
+                    waiter.thread,
+                );
+            }
+        } else {
+            init_once.status = InitOnceStatus::Uninitialized;
+        }
+
+        Ok(())
+    }
 }
diff --git a/src/tools/miri/src/concurrency/thread.rs b/src/tools/miri/src/concurrency/thread.rs
@@ -30,8 +30,7 @@ pub enum SchedulingAction {
     Stop,
 }
 
-/// Timeout callbacks can be created by synchronization primitives to tell the
-/// scheduler that they should be called once some period of time passes.
+/// Trait for callbacks that can be executed when some event happens, such as after a timeout.
 pub trait MachineCallback<'mir, 'tcx>: VisitTags {
     fn call(&self, ecx: &mut InterpCx<'mir, 'tcx, MiriMachine<'mir, 'tcx>>) -> InterpResult<'tcx>;
 }
@@ -269,7 +268,7 @@ pub struct ThreadManager<'mir, 'tcx> {
     threads: IndexVec<ThreadId, Thread<'mir, 'tcx>>,
     /// This field is pub(crate) because the synchronization primitives
     /// (`crate::sync`) need a way to access it.
-    pub(crate) sync: SynchronizationState,
+    pub(crate) sync: SynchronizationState<'mir, 'tcx>,
     /// A mapping from a thread-local static to an allocation id of a thread
     /// specific allocation.
     thread_local_alloc_ids: RefCell<FxHashMap<(DefId, ThreadId), Pointer<Provenance>>>,
@@ -303,7 +302,7 @@ impl VisitTags for ThreadManager<'_, '_> {
             timeout_callbacks,
             active_thread: _,
             yield_active_thread: _,
-            sync: _,
+            sync,
         } = self;
 
         for thread in threads {
@@ -315,6 +314,7 @@ impl VisitTags for ThreadManager<'_, '_> {
         for callback in timeout_callbacks.values() {
             callback.callback.visit_tags(visit);
         }
+        sync.visit_tags(visit);
     }
 }
 
diff --git a/src/tools/miri/src/lib.rs b/src/tools/miri/src/lib.rs
@@ -87,7 +87,7 @@ pub use crate::concurrency::{
         AtomicFenceOrd, AtomicReadOrd, AtomicRwOrd, AtomicWriteOrd,
         EvalContextExt as DataRaceEvalContextExt,
     },
-    sync::{CondvarId, EvalContextExt as SyncEvalContextExt, MutexId, RwLockId},
+    sync::{CondvarId, EvalContextExt as SyncEvalContextExt, InitOnceId, MutexId, RwLockId},
     thread::{
         EvalContextExt as ThreadsEvalContextExt, SchedulingAction, ThreadId, ThreadManager,
         ThreadState, Time,
diff --git a/src/tools/miri/src/shims/windows/foreign_items.rs b/src/tools/miri/src/shims/windows/foreign_items.rs
@@ -261,6 +261,18 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 let ret = this.TryAcquireSRWLockShared(ptr)?;
                 this.write_scalar(ret, dest)?;
             }
+            "InitOnceBeginInitialize" => {
+                let [ptr, flags, pending, context] =
+                    this.check_shim(abi, Abi::System { unwind: false }, link_name, args)?;
+                let result = this.InitOnceBeginInitialize(ptr, flags, pending, context)?;
+                this.write_scalar(result, dest)?;
+            }
+            "InitOnceComplete" => {
+                let [ptr, flags, context] =
+                    this.check_shim(abi, Abi::System { unwind: false }, link_name, args)?;
+                let result = this.InitOnceComplete(ptr, flags, context)?;
+                this.write_scalar(result, dest)?;
+            }
 
             // Dynamic symbol loading
             "GetProcAddress" => {
diff --git a/src/tools/miri/src/shims/windows/sync.rs b/src/tools/miri/src/shims/windows/sync.rs
