From b85163164a5d33af30e236f7124b8a4c101a1011 Mon Sep 17 00:00:00 2001
From: Ryo Onodera <ryoqun@gmail.com>
Date: Thu, 4 Apr 2024 14:25:54 +0900
Subject: [PATCH] Test UB to illustrate limitation of TokenCell

---
 ci/test-miri.sh                    |  4 ++
 unified-scheduler-logic/src/lib.rs | 64 +++++++++++++++++++++++++++++-
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/ci/test-miri.sh b/ci/test-miri.sh
index 05cf3d752dd447..53d9457ce69065 100755
--- a/ci/test-miri.sh
+++ b/ci/test-miri.sh
@@ -7,3 +7,7 @@ source ci/rust-version.sh nightly
 # miri is very slow; so only run very few of selective tests!
 cargo "+${rust_nightly}" miri test -p solana-program -- hash:: account_info::
 cargo "+${rust_nightly}" miri test -p solana-unified-scheduler-logic
+
+# run intentionally-#[ignored] ub triggering tests for each to make sure they fail
+! cargo "+${rust_nightly}" miri test -p solana-unified-scheduler-logic -- \
+  --ignored --exact "utils::tests::test_ub_illegally_shared_token_cell"
diff --git a/unified-scheduler-logic/src/lib.rs b/unified-scheduler-logic/src/lib.rs
index 2aea54a4b41351..6a0cd6c643984d 100644
--- a/unified-scheduler-logic/src/lib.rs
+++ b/unified-scheduler-logic/src/lib.rs
@@ -286,7 +286,10 @@ mod utils {
 
     #[cfg(test)]
     mod tests {
-        use super::Token;
+        use {
+            super::{Token, TokenCell},
+            std::{sync::Arc, thread},
+        };
 
         #[test]
         #[should_panic(
@@ -299,6 +302,65 @@ mod utils {
                 let _ = Token::<usize>::assume_exclusive_mutating_thread();
             }
         }
+
+        // As documented above, it's illegal to share (= co-own) the same instance of TokenCell
+        // across threads. Unfortunately, we can't prevent this from happening with some
+        // type-safety magic to cause compile errors... So sanity-check here test fails due to a
+        // runtime error of the known UB, when run under miri.
+        #[test]
+        // cause (harmless) UB unless running under miri by conditionally #[ignore]-ing, confirming
+        // false-positive result to coversely show the merit of miri!
+        #[cfg_attr(miri, ignore)]
+        fn test_ub_illegally_shared_token_cell() {
+            #[derive(Debug)]
+            struct FakeQueue {
+                v: Vec<u8>,
+            }
+
+            let queue1 = Arc::new(TokenCell::new(FakeQueue {
+                v: Vec::with_capacity(20),
+            }));
+            let queue2 = queue1.clone();
+            #[cfg(not(miri))]
+            let queue3 = queue1.clone();
+
+            // Usually miri immediately detects the data race; but just repeat enough time to avoid
+            // being flaky
+            for _ in 0..10 {
+                let (queue1, queue2) = (queue1.clone(), queue2.clone());
+                let thread1 = thread::spawn(move || {
+                    let mut token =
+                        unsafe { Token::<FakeQueue>::assume_exclusive_mutating_thread() };
+                    queue1.with_borrow_mut(&mut token, |queue| {
+                        // this is UB
+                        queue.v.push(3);
+                    });
+                });
+                // immediately spawn next thread without joining thread1 to ensure there's a dace race
+                // definitely. Otherwise, joining here wouldn't cause UB.
+                let thread2 = thread::spawn(move || {
+                    let mut token =
+                        unsafe { Token::<FakeQueue>::assume_exclusive_mutating_thread() };
+                    queue2.with_borrow_mut(&mut token, |queue| {
+                        // this is UB
+                        queue.v.push(4);
+                    });
+                });
+
+                thread1.join().unwrap();
+                thread2.join().unwrap();
+            }
+
+            // it's in ub already, so we can't assert reliably, so dbg!(...) just for fun
+            #[cfg(not(miri))]
+            {
+                drop((queue1, queue2));
+                dbg!(Arc::into_inner(queue3).unwrap().0.into_inner());
+            }
+
+            // return successfully to indicate an unexpected outcome, because this test should
+            // abort by now
+        }
     }
 }