154 lines
6 KiB
Rust
154 lines
6 KiB
Rust
// Per-thread module state (MEMRECLAIM.REVIEW F4): shadow stacks, wrapper result slots and fuel
|
|
// belong to one thread's entry, never shared.
|
|
// 1. Re-entry on one thread, A -> host -> B -> host -> A: the inner A call continues below the
|
|
// outer A's live shadow frames (it used to restart at the top of A's shadow stack, over them).
|
|
// 2. Two threads in one module at once: each keeps its own shadow frames and result slots.
|
|
// 3. Fuel is per thread and per call: another thread entering the module does not refill it.
|
|
// Run with MAKEPAD_RUN_MODULE_FUEL=2000000.
|
|
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
|
|
use std::sync::Arc;
|
|
|
|
pub static TO_B: AtomicUsize = AtomicUsize::new(0);
|
|
pub static TO_A: AtomicUsize = AtomicUsize::new(0);
|
|
pub static ENTRIES: AtomicU64 = AtomicU64::new(0);
|
|
pub static WANT: AtomicBool = AtomicBool::new(false);
|
|
|
|
/// Host API: A calls into B.
|
|
pub fn host_to_b(x: u64) -> u64 {
|
|
let f: fn(u64) -> u64 = unsafe { std::mem::transmute(TO_B.load(Ordering::SeqCst)) };
|
|
f(x)
|
|
}
|
|
/// Host API: B calls back into A.
|
|
pub fn host_to_a(x: u64) -> u64 {
|
|
let f: fn(u64) -> u64 = unsafe { std::mem::transmute(TO_A.load(Ordering::SeqCst)) };
|
|
f(x)
|
|
}
|
|
/// Host API: lets the other thread run (both threads' frames are live meanwhile).
|
|
pub fn host_yield(_x: u64) -> u64 {
|
|
for _ in 0..20 { std::thread::yield_now(); }
|
|
0
|
|
}
|
|
/// Host API: waits until another thread has entered the module since this call.
|
|
pub fn host_tick(_x: u64) -> u64 {
|
|
let n = ENTRIES.load(Ordering::SeqCst);
|
|
WANT.store(true, Ordering::SeqCst);
|
|
let t = std::time::Instant::now();
|
|
while ENTRIES.load(Ordering::SeqCst) == n && t.elapsed() < std::time::Duration::from_secs(2) {
|
|
std::hint::spin_loop();
|
|
}
|
|
0
|
|
}
|
|
|
|
const A: &str = r#"
|
|
pub fn outer(x: u64) -> u64 {
|
|
// live shadow frame across the host call
|
|
let mut keep = [0u64; 256];
|
|
for i in 0..256 { keep[i] = x.wrapping_mul(i as u64 + 1); }
|
|
let r = test_crate::host_to_b(x);
|
|
let mut ok = true;
|
|
for i in 0..256 { ok &= keep[i] == x.wrapping_mul(i as u64 + 1); }
|
|
if ok { r } else { u64::MAX }
|
|
}
|
|
pub fn inner(x: u64) -> u64 {
|
|
// a big frame: overwrites whatever is at the top of A's shadow stack
|
|
let mut scratch = [0u64; 1024];
|
|
for i in 0..1024 { scratch[i] = !(x + i as u64); }
|
|
let mut s = 0u64;
|
|
for i in 0..1024 { s = s.wrapping_add(scratch[i]); }
|
|
s
|
|
}
|
|
"#;
|
|
const B: &str = "pub fn mid(x: u64) -> u64 { let mut v = [x; 64]; v[3] += 1; test_crate::host_to_a(v[3]) + v[0] }\n";
|
|
const T: &str = r#"
|
|
pub fn work(seed: u64) -> u64 {
|
|
let mut a = [0u64; 512];
|
|
for i in 0..512 { a[i] = seed.wrapping_mul(2654435761).wrapping_add(i as u64); }
|
|
test_crate::host_yield(seed);
|
|
let mut bad = 0u64;
|
|
for i in 0..512 { if a[i] != seed.wrapping_mul(2654435761).wrapping_add(i as u64) { bad += 1; } }
|
|
bad
|
|
}
|
|
pub fn pair(seed: u64) -> (u64, [u64; 6]) {
|
|
let v = [seed, seed + 1, seed + 2, seed + 3, seed + 4, seed + 5];
|
|
test_crate::host_yield(seed);
|
|
(seed, v)
|
|
}
|
|
pub fn quick(x: u64) -> u64 { x + 1 }
|
|
pub fn spin(n: u64) -> u64 {
|
|
let mut x = 1u64;
|
|
let mut i = 0;
|
|
while i < n {
|
|
x = x.wrapping_mul(6364136223846793005).wrapping_add(i);
|
|
if i % 20000 == 0 { test_crate::host_tick(i); }
|
|
i += 1;
|
|
}
|
|
x
|
|
}
|
|
"#;
|
|
|
|
fn inner_want(x: u64) -> u64 {
|
|
let mut s = 0u64;
|
|
for i in 0..1024u64 { s = s.wrapping_add(!(x + i)); }
|
|
s
|
|
}
|
|
|
|
#[test]
|
|
fn shadow_threads() {
|
|
use std::relax::Module;
|
|
// 1. A -> host -> B -> host -> A
|
|
let a = Module::load_source("sh_a", "sh_a.rs", A, &["test_crate"]).unwrap();
|
|
let b = Module::load_source("sh_b", "sh_b.rs", B, &["test_crate"]).unwrap();
|
|
let outer: fn(u64) -> u64 = a.get("outer").unwrap();
|
|
let inner: fn(u64) -> u64 = a.get("inner").unwrap();
|
|
let mid: fn(u64) -> u64 = b.get("mid").unwrap();
|
|
TO_B.store(mid as usize, Ordering::SeqCst);
|
|
TO_A.store(inner as usize, Ordering::SeqCst);
|
|
let x = 7;
|
|
let r = outer(x);
|
|
println!("out re-entry keeps the outer frames: {}", r == inner_want(x + 1) + x);
|
|
|
|
// 2. two threads in one module
|
|
let t = Module::load_source("sh_t", "sh_t.rs", T, &["test_crate"]).unwrap();
|
|
let work: fn(u64) -> u64 = t.get("work").unwrap();
|
|
let pair: fn(u64) -> (u64, [u64; 6]) = t.get("pair").unwrap();
|
|
let bad = Arc::new(AtomicU64::new(0));
|
|
let mut hs = Vec::new();
|
|
for k in 0..4u64 {
|
|
let bad = bad.clone();
|
|
hs.push(std::thread::spawn(move || {
|
|
for j in 0..200u64 {
|
|
let seed = k * 1_000_000 + j;
|
|
bad.fetch_add(work(seed), Ordering::Relaxed);
|
|
let (s, v) = pair(seed);
|
|
if s != seed || v != [seed, seed + 1, seed + 2, seed + 3, seed + 4, seed + 5] {
|
|
bad.fetch_add(1, Ordering::Relaxed);
|
|
}
|
|
}
|
|
}));
|
|
}
|
|
for h in hs { h.join().unwrap(); }
|
|
eprintln!("shadow_threads: {} corrupted words/results across threads", bad.load(Ordering::SeqCst));
|
|
println!("out threads keep their own frames: {}", bad.load(Ordering::SeqCst) == 0);
|
|
|
|
// 3. fuel: a long call is stopped although another thread enters the module meanwhile
|
|
let quick: fn(u64) -> u64 = t.get("quick").unwrap();
|
|
let spin: fn(u64) -> u64 = t.get("spin").unwrap();
|
|
let stop = Arc::new(AtomicBool::new(false));
|
|
let s2 = stop.clone();
|
|
let hammer = std::thread::spawn(move || {
|
|
while !s2.load(Ordering::SeqCst) {
|
|
if WANT.swap(false, Ordering::SeqCst) {
|
|
quick(1);
|
|
ENTRIES.fetch_add(1, Ordering::SeqCst);
|
|
}
|
|
std::hint::spin_loop();
|
|
}
|
|
});
|
|
let long = std::relax::guard(0, || spin(10_000_000));
|
|
stop.store(true, Ordering::SeqCst);
|
|
hammer.join().unwrap();
|
|
eprintln!("shadow_threads: entries by the other thread during the long call: {}", ENTRIES.load(Ordering::SeqCst));
|
|
println!("out fuel is per thread: {}", long.is_err() && ENTRIES.load(Ordering::SeqCst) > 0);
|
|
let short = std::relax::guard(0, || spin(10_000));
|
|
println!("out short call still runs: {}", short.is_ok());
|
|
}
|