Keep an atexit callback alive while it is being compared

atexit.unregister() releases the callback list around each __eq__ call and identified the entry it had compared by the address of its Box. __eq__ can call atexit._clear(), which drops that Box, and atexit.register(), whose new Box lands on the freed allocation; the identity search then matched the freshly registered callback and removed it. atexit.register(a); atexit.register(b); atexit.register(c) # __eq__ runs _clear() then register(d), returns True atexit.unregister(probe) left no callbacks registered where CPython leaves d. Entries are Arc-shared now, so unregister() holds the one it is comparing and matches it with Arc::ptr_eq: an address cannot be reused while the comparison that named it is still running. Assisted-by: Claude
youknowone · Aug 16, 2026 · Aug 14, 2026 · Aug 14, 2026 · Aug 14, 2026 · Aug 14, 2026
commit fab2b3b8ff21db3aa5d05105606c12270ec253a4
diff --git a/crates/vm/src/stdlib/atexit.rs b/crates/vm/src/stdlib/atexit.rs
@@ -4,14 +4,15 @@ pub(crate) use atexit::module_def;
 #[pymodule]
 mod atexit {
     use crate::{AsObject, PyObjectRef, PyResult, VirtualMachine, function::FuncArgs};
+    use alloc::sync::Arc;
 
     #[pyfunction]
     fn register(func: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyObjectRef {
         // Callbacks go in LIFO order (insert at front)
         vm.state
             .atexit_funcs
             .lock()
-            .insert(0, Box::new((func.clone(), args)));
+            .insert(0, Arc::new((func.clone(), args)));
         func
     }
 
@@ -29,24 +30,26 @@ mod atexit {
             funcs.len() as isize - 1
         };
         while i >= 0 {
-            let (cb, entry_ptr) = {
+            let entry = {
                 let funcs = vm.state.atexit_funcs.lock();
                 if i as usize >= funcs.len() {
                     i = funcs.len() as isize;
                     i -= 1;
                     continue;
                 }
-                let entry = &funcs[i as usize];
-                (entry.0.clone(), &**entry as *const (PyObjectRef, FuncArgs))
+                // Keep the entry alive for as long as it is being compared, so
+                // it cannot be dropped and have its address handed to a
+                // callback registered from within __eq__.
+                funcs[i as usize].clone()
             };
             // Lock released: __eq__ can safely call atexit functions
-            let eq = vm.bool_eq(&func, &cb)?;
+            let eq = vm.bool_eq(&func, &entry.0)?;
             if eq {
                 // The entry may have moved during __eq__. Search backward by identity.
                 let mut funcs = vm.state.atexit_funcs.lock();
                 let mut j = (funcs.len() as isize - 1).min(i);
                 while j >= 0 {
-                    if core::ptr::eq(&**funcs.get(j as usize).unwrap(), entry_ptr) {
+                    if Arc::ptr_eq(funcs.get(j as usize).unwrap(), &entry) {
                         funcs.remove(j as usize);
                         i = j;
                         break;
@@ -70,7 +73,7 @@ mod atexit {
         let funcs: Vec<_> = core::mem::take(&mut *vm.state.atexit_funcs.lock());
         // Callbacks stored in LIFO order, iterate forward
         for entry in funcs {
-            let (func, args) = *entry;
+            let (func, args) = Arc::try_unwrap(entry).unwrap_or_else(|e| (*e).clone());
             if let Err(e) = func.call(args, vm) {
                 let exit = e.fast_isinstance(vm.ctx.exceptions.system_exit);
                 let msg = func

diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs
@@ -44,7 +44,7 @@ use crate::{
     stdlib,
     warn::WarningsState,
 };
-use alloc::{borrow::Cow, collections::BTreeMap};
+use alloc::{borrow::Cow, collections::BTreeMap, sync::Arc};
 #[cfg(all(not(unix), feature = "threading"))]
 use core::ptr::NonNull;
 #[cfg(feature = "threading")]
@@ -759,7 +759,10 @@ pub struct PyGlobalState {
     pub stacksize: AtomicCell<usize>,
     pub thread_count: AtomicCell<usize>,
     pub hash_secret: HashSecret,
-    pub atexit_funcs: PyMutex<Vec<Box<(PyObjectRef, FuncArgs)>>>,
+    /// Registered `atexit` callbacks, newest first. Shared ownership so
+    /// `atexit.unregister` can keep the entry it is comparing alive while the
+    /// list is unlocked, and still recognize it afterwards by identity.
+    pub atexit_funcs: PyMutex<Vec<Arc<(PyObjectRef, FuncArgs)>>>,
     pub codec_registry: CodecsRegistry,
     pub finalizing: AtomicBool,
     pub warnings: WarningsState,

diff --git a/extra_tests/snippets/stdlib_atexit.py b/extra_tests/snippets/stdlib_atexit.py
@@ -0,0 +1,101 @@
+"""atexit.unregister() compares callbacks with arbitrary Python code.
+
+The comparison runs with the callback list unlocked, so __eq__ may clear it
+and register something new. unregister() then has to tell whether the entry
+it compared is still there, and must not mistake a later registration that
+happens to occupy the same storage for that entry.
+"""
+
+import atexit
+
+
+def make(name):
+    def f():
+        ran.append(name)
+
+    f.tag = name
+    return f
+
+
+ran = []
+a, b, c, d = (make(n) for n in "abcd")
+
+
+class Probe:
+    def __init__(self, action=None, result=True):
+        self.action = action
+        self.result = result
+        self.seen = []
+
+    def __eq__(self, other):
+        self.seen.append(getattr(other, "tag", "?"))
+        if self.action is not None:
+            self.action()
+        return self.result
+
+
+def remaining():
+    del ran[:]
+    atexit._run_exitfuncs()
+    return list(ran)
+
+
+# A callback the probe does not match is left alone.
+atexit._clear()
+atexit.register(a)
+atexit.register(b)
+probe = Probe(result=False)
+atexit.unregister(probe)
+assert probe.seen == ["a", "b"], probe.seen
+assert remaining() == ["b", "a"], ran
+
+# Matching callbacks are dropped, oldest compared first.
+atexit._clear()
+atexit.register(a)
+atexit.register(b)
+atexit.register(c)
+probe = Probe(result=True)
+atexit.unregister(probe)
+assert probe.seen == ["a", "b", "c"], probe.seen
+assert atexit._ncallbacks() == 0
+assert remaining() == [], ran
+
+# __eq__ empties the list: there is nothing left to drop.
+atexit._clear()
+atexit.register(a)
+atexit.register(b)
+probe = Probe(action=atexit._clear, result=True)
+atexit.unregister(probe)
+assert probe.seen == ["a"], probe.seen
+assert remaining() == [], ran
+
+# __eq__ empties the list and registers a replacement. The replacement is a
+# different callback, so it survives however its storage was reused.
+atexit._clear()
+atexit.register(a)
+atexit.register(b)
+atexit.register(c)
+
+
+def replace():
+    atexit._clear()
+    atexit.register(d)
+
+
+probe = Probe(action=replace, result=True)
+atexit.unregister(probe)
+assert probe.seen == ["a", "d"], probe.seen
+assert remaining() == ["d"], ran
+
+# __eq__ registers without clearing: every entry the walk had already passed
+# stays, and so does each newly registered one.
+atexit._clear()
+atexit.register(a)
+atexit.register(b)
+probe = Probe(action=lambda: atexit.register(c), result=True)
+atexit.unregister(probe)
+assert probe.seen == ["a", "c"], probe.seen
+assert remaining() == ["c", "c", "b", "a"], ran
+
+atexit._clear()
+print("ok")