diff --git a/crates/vm/src/codecs.rs b/crates/vm/src/codecs.rs index 32724984e..b136a68bc 100644 --- a/crates/vm/src/codecs.rs +++ b/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 { @@ -277,13 +280,17 @@ impl CodecsRegistry { let res = func.call((encoding.clone(),), vm)?; let res: Option = 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 { diff --git a/extra_tests/snippets/stdlib_codecs_wide.py b/extra_tests/snippets/stdlib_codecs_wide.py index bd8307ffc..00f8a780b 100644 --- a/extra_tests/snippets/stdlib_codecs_wide.py +++ b/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()