Fix type/instruction cache races with atomic ops

youknowone · youknowone · commit 2adc355fe43c · 2026-03-05T14:26:46.000+09:00
- Type cache: use try_to_owned (CAS-based safe_inc) instead of
  ManuallyDrop+clone to prevent use-after-free when a concurrent
  writer drops the cached value between pointer load and refcount
  increment.

- Instruction cache: replace split 4×u16 read/write (torn
  read/write under concurrent specialization) with single
  AtomicUsize load/store via pointer_cache side array.

- Instruction cache reads: use try_to_owned with deoptimize
  fallback instead of blind to_owned on cached descriptor pointers.

- Call modified() BEFORE attribute modification in set_attr and
  set___type_params__ so cached descriptor pointers remain alive
  when type_cache_clear_version drops cache entries.
diff --git a/crates/compiler-core/src/bytecode.rs b/crates/compiler-core/src/bytecode.rs
@@ -12,7 +12,7 @@ use core::{
     cell::UnsafeCell,
     hash, mem,
     ops::Deref,
-    sync::atomic::{AtomicU8, AtomicU16, Ordering},
+    sync::atomic::{AtomicU8, AtomicU16, AtomicUsize, Ordering},
 };
 use itertools::Itertools;
 use malachite_bigint::BigInt;
@@ -411,6 +411,10 @@ impl TryFrom<&[u8]> for CodeUnit {
 pub struct CodeUnits {
     units: UnsafeCell<Box<[CodeUnit]>>,
     adaptive_counters: Box<[AtomicU16]>,
+    /// Pointer-sized cache entries for descriptor pointers.
+    /// Single atomic load/store prevents torn reads when multiple threads
+    /// specialize the same instruction concurrently.
+    pointer_cache: Box<[AtomicUsize]>,
 }
 
 // SAFETY: All cache operations use atomic read/write instructions.
@@ -432,9 +436,15 @@ impl Clone for CodeUnits {
             .iter()
             .map(|c| AtomicU16::new(c.load(Ordering::Relaxed)))
             .collect();
+        let pointer_cache = self
+            .pointer_cache
+            .iter()
+            .map(|c| AtomicUsize::new(c.load(Ordering::Relaxed)))
+            .collect();
         Self {
             units: UnsafeCell::new(units),
             adaptive_counters,
+            pointer_cache,
         }
     }
 }
@@ -472,13 +482,19 @@ impl<const N: usize> From<[CodeUnit; N]> for CodeUnits {
 impl From<Vec<CodeUnit>> for CodeUnits {
     fn from(value: Vec<CodeUnit>) -> Self {
         let units = value.into_boxed_slice();
-        let adaptive_counters = (0..units.len())
+        let len = units.len();
+        let adaptive_counters = (0..len)
             .map(|_| AtomicU16::new(0))
             .collect::<Vec<_>>()
             .into_boxed_slice();
+        let pointer_cache = (0..len)
+            .map(|_| AtomicUsize::new(0))
+            .collect::<Vec<_>>()
+            .into_boxed_slice();
         Self {
             units: UnsafeCell::new(units),
             adaptive_counters,
+            pointer_cache,
         }
     }
 }
@@ -600,25 +616,22 @@ impl CodeUnits {
         lo | (hi << 16)
     }
 
-    /// Write a u64 value across four consecutive CACHE code units starting at `index`.
+    /// Store a pointer-sized value atomically in the pointer cache at `index`.
+    ///
+    /// Uses a single `AtomicUsize` store to prevent torn writes when
+    /// multiple threads specialize the same instruction concurrently.
     ///
     /// # Safety
-    /// Same requirements as `write_cache_u16`.
-    pub unsafe fn write_cache_u64(&self, index: usize, value: u64) {
-        unsafe {
-            self.write_cache_u32(index, value as u32);
-            self.write_cache_u32(index + 2, (value >> 32) as u32);
-        }
+    /// - `index` must be in bounds.
+    pub unsafe fn write_cache_ptr(&self, index: usize, value: usize) {
+        self.pointer_cache[index].store(value, Ordering::Relaxed);
     }
 
-    /// Read a u64 value from four consecutive CACHE code units starting at `index`.
+    /// Load a pointer-sized value atomically from the pointer cache at `index`.
     ///
-    /// # Panics
-    /// Panics if `index + 3` is out of bounds.
-    pub fn read_cache_u64(&self, index: usize) -> u64 {
-        let lo = self.read_cache_u32(index) as u64;
-        let hi = self.read_cache_u32(index + 2) as u64;
-        lo | (hi << 32)
+    /// Uses a single `AtomicUsize` load to prevent torn reads.
+    pub fn read_cache_ptr(&self, index: usize) -> usize {
+        self.pointer_cache[index].load(Ordering::Relaxed)
     }
 
     /// Read adaptive counter bits for instruction at `index`.
diff --git a/crates/vm/src/builtins/type.rs b/crates/vm/src/builtins/type.rs
@@ -701,8 +701,11 @@ impl PyType {
     }
 
     pub fn set_attr(&self, attr_name: &'static PyStrInterned, value: PyObjectRef) {
-        self.attributes.write().insert(attr_name, value);
+        // Invalidate caches BEFORE modifying attributes so that cached
+        // descriptor pointers are still alive when type_cache_clear_version
+        // drops the cache's strong references.
         self.modified();
+        self.attributes.write().insert(attr_name, value);
     }
 
     /// Internal get_attr implementation for fast lookup on a class.
@@ -735,23 +738,20 @@ impl PyType {
             {
                 let ptr = entry.value.load(Ordering::Acquire);
                 if !ptr.is_null() {
-                    // SAFETY: The value pointer was stored via PyObjectRef::into_raw
-                    // and is valid as long as the version hasn't changed. We create
-                    // a temporary reference (ManuallyDrop prevents decrement), clone
-                    // it to get our own strong reference, then re-check the version
-                    // to confirm the entry wasn't invalidated during our read.
-                    let cloned = unsafe {
-                        let tmp = core::mem::ManuallyDrop::new(PyObjectRef::from_raw(
-                            NonNull::new_unchecked(ptr),
-                        ));
-                        (*tmp).clone()
-                    };
-                    // SeqLock validation: if version changed, discard our clone
-                    let v2 = entry.version.load(Ordering::Acquire);
-                    if v2 == v1 {
-                        return Some(cloned);
+                    // Use try_to_owned (CAS-based safe_inc) to prevent a
+                    // use-after-free race: between loading the pointer and
+                    // incrementing the refcount, a concurrent writer may
+                    // invalidate the entry and drop the old value. safe_inc
+                    // fails atomically if the refcount has already reached 0.
+                    let obj: &PyObject = unsafe { &*ptr };
+                    if let Some(cloned) = obj.try_to_owned() {
+                        // SeqLock validation: if version changed, discard
+                        let v2 = entry.version.load(Ordering::Acquire);
+                        if v2 == v1 {
+                            return Some(cloned);
+                        }
+                        drop(cloned);
                     }
-                    drop(cloned);
                 }
             }
         }
@@ -1498,8 +1498,8 @@ impl PyType {
             PySetterValue::Assign(ref val) => {
                 let key = identifier!(vm, __type_params__);
                 self.check_set_special_type_attr(key, vm)?;
-                self.attributes.write().insert(key, val.clone().into());
                 self.modified();
+                self.attributes.write().insert(key, val.clone().into());
             }
             PySetterValue::Delete => {
                 // For delete, we still need to check if the type is immutable
@@ -1510,8 +1510,8 @@ impl PyType {
                     )));
                 }
                 let key = identifier!(vm, __type_params__);
-                self.attributes.write().shift_remove(&key);
                 self.modified();
+                self.attributes.write().shift_remove(&key);
             }
         }
         Ok(())
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
