Avoid codec registry deadlocks during object cleanup (#9159)

Assisted-by: Codex:gpt-6

crates/vm/src/codecs.rs

@@ -241,23 +241,26 @@ impl CodecsRegistry {
             return;
         }
 
-        for (i, item) in inner.search_path.iter().enumerate() {
-            if item.get_id() == search_function.get_id() {
-                if !inner.search_cache.is_empty() {
-                    inner.search_cache.clear();
-                }
-                inner.search_path.remove(i);
-                return;
-            }
+        if let Some(i) = inner
+            .search_path
+            .iter()
+            .position(|item| item.get_id() == search_function.get_id())
+        {
+            let cache = core::mem::take(&mut inner.search_cache);
+            let function = inner.search_path.remove(i);
+            drop(inner);
+            drop((cache, function));
         }
     }
 
     pub(crate) fn register_manual(&self, name: &str, codec: PyCodec) {
         let name = normalize_registry_name(name);
-        self.inner
+        let previous = self
+            .inner
             .write()
             .search_cache
             .insert(name.into_owned(), codec);
+        drop(previous);
     }
 
     pub fn lookup(&self, encoding: &str, vm: &VirtualMachine) -> PyResult<PyCodec> {
@@ -277,13 +280,17 @@ impl CodecsRegistry {
             let res = func.call((encoding.clone(),), vm)?;
             let res: Option<PyCodec> = res.try_into_value(vm)?;
             if let Some(codec) = res {
-                let mut inner = self.inner.write();
-                // someone might have raced us to this, so use theirs
-                let codec = inner
-                    .search_cache
-                    .entry(encoding.as_str().to_owned())
-                    .or_insert(codec);
-                return Ok(codec.clone());
+                let cached = {
+                    let mut inner = self.inner.write();
+                    // someone might have raced us to this, so use theirs
+                    // Keep our tuple alive until the registry lock is released.
+                    inner
+                        .search_cache
+                        .entry(encoding.as_str().to_owned())
+                        .or_insert_with(|| codec.clone())
+                        .clone()
+                };
+                return Ok(cached);
             }
         }
 
@@ -537,7 +544,8 @@ impl CodecsRegistry {
                 "cannot un-register built-in error handler '{name}'"
             )));
         }
-        Ok(self.inner.write().errors.remove(name).is_some())
+        let handler = self.inner.write().errors.remove(name);
+        Ok(handler.is_some())
     }
 
     pub fn lookup_error_opt(&self, name: &str) -> Option<PyObjectRef> {

extra_tests/snippets/stdlib_codecs_wide.py

@@ -140,3 +140,63 @@ for args, kwargs in (
         pass
     else:
         assert False, "invalid normalization arguments were accepted"
+
+
+# Registry removal and competing lookup results may run Python destructors.
+def test_registry_reentrant_cleanup():
+    released = []
+
+    class ErrorHandler:
+        def __call__(self, error):
+            raise AssertionError("handler should not be called")
+
+        def __del__(self):
+            codecs.lookup_error("strict")
+            released.append("handler")
+
+    codecs.register_error("test_registry_cleanup", ErrorHandler())
+    assert _codecs._unregister_error("test_registry_cleanup")
+    assert released == ["handler"]
+    released.clear()
+
+    class Codec(tuple):
+        def __new__(cls, label):
+            result = super().__new__(cls, (None, None, None, None))
+            result.label = label
+            return result
+
+        def __del__(self):
+            codecs.lookup_error("strict")
+            released.append(self.label)
+
+    def search(name):
+        if name == "test_registry_cleanup":
+            return Codec("cached")
+
+    codecs.register(search)
+    codecs.lookup("test_registry_cleanup")
+    codecs.unregister(search)
+    assert released == ["cached"]
+    released.clear()
+
+    entered = False
+
+    def reentrant_search(name):
+        nonlocal entered
+        if name == "test_registry_reentry":
+            if not entered:
+                entered = True
+                codecs.lookup(name)
+                return Codec("outer")
+            return Codec("inner")
+
+    codecs.register(reentrant_search)
+    try:
+        codecs.lookup("test_registry_reentry")
+        assert len(released) == 1
+    finally:
+        codecs.unregister(reentrant_search)
+    assert sorted(released) == ["inner", "outer"]
+
+
+test_registry_reentrant_cleanup()