Implement garbage collection for Rust data holding refs to Python objects (#416)

Author: adriangb

src/serializers/mod.rs

@@ -2,6 +2,7 @@ use std::fmt::Debug;
 
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyDict};
+use pyo3::{PyTraverseError, PyVisit};
 
 use crate::build_context::BuildContext;
 use crate::validators::SelfValidator;
@@ -136,6 +137,21 @@ impl SchemaSerializer {
             self.serializer, self.slots
         )
     }
+
+    fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+        self.serializer.py_gc_traverse(&visit)?;
+        for slot in self.slots.iter() {
+            slot.py_gc_traverse(&visit)?;
+        }
+        Ok(())
+    }
+
+    fn __clear__(&mut self) {
+        self.serializer.py_gc_clear();
+        for slot in self.slots.iter_mut() {
+            slot.py_gc_clear();
+        }
+    }
 }
 
 #[allow(clippy::too_many_arguments)]

src/serializers/shared.rs

@@ -2,9 +2,9 @@ use std::borrow::Cow;
 use std::fmt::Debug;
 
 use pyo3::exceptions::PyTypeError;
-use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PySet};
+use pyo3::{intern, PyTraverseError, PyVisit};
 
 use enum_dispatch::enum_dispatch;
 use serde::Serialize;
@@ -225,6 +225,10 @@ impl BuildSerializer for CombinedSerializer {
 
 #[enum_dispatch(CombinedSerializer)]
 pub(crate) trait TypeSerializer: Send + Sync + Clone + Debug {
+    fn py_gc_traverse(&self, _visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        Ok(())
+    }
+    fn py_gc_clear(&mut self) {}
     fn to_python(
         &self,
         value: &PyAny,

src/serializers/type_serializers/model.rs

@@ -55,6 +55,12 @@ impl ModelSerializer {
 }
 
 impl TypeSerializer for ModelSerializer {
+    fn py_gc_traverse(&self, visit: &pyo3::PyVisit<'_>) -> Result<(), pyo3::PyTraverseError> {
+        visit.call(&self.class)?;
+        self.serializer.py_gc_traverse(visit)?;
+        Ok(())
+    }
+
     fn to_python(
         &self,
         value: &PyAny,

src/validators/mod.rs

@@ -2,10 +2,10 @@ use std::fmt::Debug;
 
 use enum_dispatch::enum_dispatch;
 
-use pyo3::intern;
 use pyo3::once_cell::GILOnceCell;
 use pyo3::prelude::*;
 use pyo3::types::{PyAny, PyDict};
+use pyo3::{intern, PyTraverseError, PyVisit};
 
 use crate::build_context::BuildContext;
 use crate::build_tools::{py_err, py_error_type, SchemaDict, SchemaError};
@@ -211,6 +211,22 @@ impl SchemaValidator {
             self.slots,
         )
     }
+
+    fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+        self.validator.py_gc_traverse(&visit)?;
+        visit.call(&self.schema)?;
+        for slot in self.slots.iter() {
+            slot.py_gc_traverse(&visit)?;
+        }
+        Ok(())
+    }
+
+    fn __clear__(&mut self) {
+        self.validator.py_gc_clear();
+        for slot in self.slots.iter_mut() {
+            slot.py_gc_clear();
+        }
+    }
 }
 
 impl SchemaValidator {
@@ -548,6 +564,10 @@ pub enum CombinedValidator {
 /// validators defined in `build_validator` also need `EXPECTED_TYPE` as a const, but that can't be part of the trait
 #[enum_dispatch(CombinedValidator)]
 pub trait Validator: Send + Sync + Clone + Debug {
+    fn py_gc_traverse(&self, _visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        Ok(())
+    }
+    fn py_gc_clear(&mut self) {}
     /// Do the actual validation for this schema/type
     fn validate<'s, 'data>(
         &'s self,

src/validators/model.rs

@@ -65,6 +65,11 @@ impl BuildValidator for ModelValidator {
 }
 
 impl Validator for ModelValidator {
+    fn py_gc_traverse(&self, visit: &pyo3::PyVisit<'_>) -> Result<(), pyo3::PyTraverseError> {
+        visit.call(&self.class)?;
+        self.validator.py_gc_traverse(visit)?;
+        Ok(())
+    }
     fn validate<'s, 'data>(
         &'s self,
         py: Python<'data>,

tests/test_garbage_collection.py

@@ -0,0 +1,82 @@
+import gc
+import platform
+from typing import Any
+from weakref import WeakValueDictionary
+
+import pytest
+
+from pydantic_core import SchemaSerializer, SchemaValidator, core_schema
+
+
+@pytest.mark.xfail(
+    condition=platform.python_implementation() == 'PyPy', reason='https://foss.heptapod.net/pypy/pypy/-/issues/3899'
+)
+def test_gc_schema_serializer() -> None:
+    # test for https://github.com/pydantic/pydantic/issues/5136
+    class BaseModel:
+        __schema__: SchemaSerializer
+
+        def __init_subclass__(cls) -> None:
+            cls.__schema__ = SchemaSerializer(
+                core_schema.model_schema(
+                    cls,
+                    core_schema.typed_dict_schema(
+                        {'x': core_schema.typed_dict_field(core_schema.definition_reference_schema('model'))}
+                    ),
+                    ref='model',
+                )
+            )
+
+    cache: 'WeakValueDictionary[int, Any]' = WeakValueDictionary()
+
+    for _ in range(10_000):
+
+        class MyModel(BaseModel):
+            pass
+
+        cache[id(MyModel)] = MyModel
+
+        del MyModel
+
+    gc.collect(0)
+    gc.collect(1)
+    gc.collect(2)
+
+    assert len(cache) == 0
+
+
+@pytest.mark.xfail(
+    condition=platform.python_implementation() == 'PyPy', reason='https://foss.heptapod.net/pypy/pypy/-/issues/3899'
+)
+def test_gc_schema_validator() -> None:
+    # test for https://github.com/pydantic/pydantic/issues/5136
+    class BaseModel:
+        __validator__: SchemaValidator
+
+        def __init_subclass__(cls) -> None:
+            cls.__validator__ = SchemaValidator(
+                core_schema.model_schema(
+                    cls,
+                    core_schema.typed_dict_schema(
+                        {'x': core_schema.typed_dict_field(core_schema.definition_reference_schema('model'))}
+                    ),
+                    ref='model',
+                )
+            )
+
+    cache: 'WeakValueDictionary[int, Any]' = WeakValueDictionary()
+
+    for _ in range(10_000):
+
+        class MyModel(BaseModel):
+            pass
+
+        cache[id(MyModel)] = MyModel
+
+        del MyModel
+
+    gc.collect(0)
+    gc.collect(1)
+    gc.collect(2)
+
+    assert len(cache) == 0