From 224dff4e15f1195c1cdbd85d9080e13b14a151e2 Mon Sep 17 00:00:00 2001
From: DrMeepster <19316085+DrMeepster@users.noreply.github.com>
Date: Sun, 30 Oct 2022 21:39:16 -0700
Subject: [PATCH] add acquire when init once is already complete

---
 src/tools/miri/src/concurrency/init_once.rs   | 35 +++++++++--------
 src/tools/miri/src/shims/windows/sync.rs      |  6 ++-
 .../pass/concurrency/windows_init_once.rs     | 38 +++++++++++++++++++
 3 files changed, 59 insertions(+), 20 deletions(-)

diff --git a/src/tools/miri/src/concurrency/init_once.rs b/src/tools/miri/src/concurrency/init_once.rs
index 791931901e2a..b1443662e2b7 100644
--- a/src/tools/miri/src/concurrency/init_once.rs
+++ b/src/tools/miri/src/concurrency/init_once.rs
@@ -141,18 +141,11 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         // Wake up everyone.
         // need to take the queue to avoid having `this` be borrowed multiple times
         for waiter in std::mem::take(&mut init_once.waiters) {
-            // End of the wait happens-before woken-up 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,
-                );
-            }
-
             this.unblock_thread(waiter.thread);
 
             // Call callback, with the woken-up thread as `current`.
             this.set_active_thread(waiter.thread);
+            this.init_once_acquire(id);
             waiter.callback.call(this)?;
             this.set_active_thread(current_thread);
         }
@@ -172,26 +165,17 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         );
 
         // Each complete happens-before the end of the wait
-        // FIXME: should this really induce synchronization? If we think of it as a lock, then yes,
-        // but the docs don't talk about such details.
         if let Some(data_race) = &this.machine.data_race {
             data_race.validate_lock_release(&mut init_once.data_race, current_thread);
         }
 
         // Wake up one waiting thread, so they can go ahead and try to init this.
         if let Some(waiter) = init_once.waiters.pop_front() {
-            // End of the wait happens-before woken-up 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,
-                );
-            }
-
             this.unblock_thread(waiter.thread);
 
             // Call callback, with the woken-up thread as `current`.
             this.set_active_thread(waiter.thread);
+            this.init_once_acquire(id);
             waiter.callback.call(this)?;
             this.set_active_thread(current_thread);
         } else {
@@ -201,4 +185,19 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
 
         Ok(())
     }
+
+    /// Synchronize with the previous completion or failure of an InitOnce.
+    /// This is required to prevent data races.
+    #[inline]
+    fn init_once_acquire(&mut self, id: InitOnceId) {
+        let this = self.eval_context_mut();
+        let current_thread = this.get_active_thread();
+
+        if let Some(data_race) = &this.machine.data_race {
+            data_race.validate_lock_acquire(
+                &this.machine.threads.sync.init_onces[id].data_race,
+                current_thread,
+            );
+        }
+    }
 }
diff --git a/src/tools/miri/src/shims/windows/sync.rs b/src/tools/miri/src/shims/windows/sync.rs
index 8156ae8af1ef..f8980e188b45 100644
--- a/src/tools/miri/src/shims/windows/sync.rs
+++ b/src/tools/miri/src/shims/windows/sync.rs
@@ -177,8 +177,10 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                     Box::new(Callback { init_once_id: id, pending_place }),
                 )
             }
-            InitOnceStatus::Complete =>
-                this.write_scalar(this.eval_windows("c", "FALSE")?, &pending_place)?,
+            InitOnceStatus::Complete => {
+                this.init_once_acquire(id);
+                this.write_scalar(this.eval_windows("c", "FALSE")?, &pending_place)?;
+            }
         }
 
         // This always succeeds (even if the thread is blocked, we will succeed if we ever unblock).
diff --git a/src/tools/miri/tests/pass/concurrency/windows_init_once.rs b/src/tools/miri/tests/pass/concurrency/windows_init_once.rs
index d3c72c3d028c..6e5129acaf85 100644
--- a/src/tools/miri/tests/pass/concurrency/windows_init_once.rs
+++ b/src/tools/miri/tests/pass/concurrency/windows_init_once.rs
@@ -131,8 +131,46 @@ fn retry_on_fail() {
     waiter2.join().unwrap();
 }
 
+fn no_data_race_after_complete() {
+    let mut init_once = null_mut();
+    let mut pending = 0;
+
+    unsafe {
+        assert_eq!(InitOnceBeginInitialize(&mut init_once, 0, &mut pending, null_mut()), TRUE);
+        assert_eq!(pending, TRUE);
+    }
+
+    let init_once_ptr = SendPtr(&mut init_once);
+
+    let mut place = 0;
+    let place_ptr = SendPtr(&mut place);
+
+    let reader = thread::spawn(move || unsafe {
+        let mut pending = 0;
+
+        assert_eq!(InitOnceBeginInitialize(init_once_ptr.0, 0, &mut pending, null_mut()), TRUE);
+        assert_eq!(pending, FALSE);
+        // this should not data race
+        place_ptr.0.read()
+    });
+
+    unsafe {
+        // this should not data race
+        place_ptr.0.write(1);
+    }
+
+    unsafe {
+        assert_eq!(InitOnceComplete(init_once_ptr.0, 0, null_mut()), TRUE);
+    }
+    //println!("complete");
+
+    // run reader
+    assert_eq!(reader.join().unwrap(), 1);
+}
+
 fn main() {
     single_thread();
     block_until_complete();
     retry_on_fail();
+    no_data_race_after_complete();
 }