From 5a65cd32cf46d731937a9944e6beff8d6b17edb1 Mon Sep 17 00:00:00 2001 From: Rick Poyner Date: Wed, 9 Oct 2024 18:02:48 -0400 Subject: [PATCH] [bindings] Add custom ref_cycle annotation Introduce the new annotation internal::ref_cycle(). It will eventually replace existing cyclic py::keep_alive<>() annotations (which do their job, but leak their participants forever). The participants of ref_cycle<>() will be garbage collectible, so that applications can run loops that use various drake components without exhausting memory. This patch just adds the implementation and its unit test. --- bindings/pydrake/common/BUILD.bazel | 30 ++++ bindings/pydrake/common/ref_cycle_pybind.cc | 71 ++++++++ bindings/pydrake/common/ref_cycle_pybind.h | 62 +++++++ .../pydrake/common/test/ref_cycle_test.py | 153 ++++++++++++++++++ .../common/test/ref_cycle_test_util_py.cc | 75 +++++++++ 5 files changed, 391 insertions(+) create mode 100644 bindings/pydrake/common/ref_cycle_pybind.cc create mode 100644 bindings/pydrake/common/ref_cycle_pybind.h create mode 100644 bindings/pydrake/common/test/ref_cycle_test.py create mode 100644 bindings/pydrake/common/test/ref_cycle_test_util_py.cc diff --git a/bindings/pydrake/common/BUILD.bazel b/bindings/pydrake/common/BUILD.bazel index a284c291e35d..b9bc0c0bc19d 100644 --- a/bindings/pydrake/common/BUILD.bazel +++ b/bindings/pydrake/common/BUILD.bazel @@ -234,6 +234,19 @@ drake_cc_library( ], ) +drake_cc_library( + name = "ref_cycle_pybind", + srcs = ["ref_cycle_pybind.cc"], + hdrs = ["ref_cycle_pybind.h"], + declare_installed_headers = 0, + visibility = ["//visibility:public"], + deps = [ + "//bindings/pydrake:pydrake_pybind", + "//common:essential", + "@pybind11", + ], +) + # N.B. Any C++ libraries that include this must include `cpp_template_py` when # being used in Python. drake_cc_library( @@ -492,6 +505,23 @@ drake_py_unittest( ], ) +drake_pybind_library( + name = "ref_cycle_test_util_py", + testonly = True, + add_install = False, + cc_deps = [":ref_cycle_pybind"], + cc_srcs = ["test/ref_cycle_test_util_py.cc"], + package_info = PACKAGE_INFO, +) + +drake_py_unittest( + name = "ref_cycle_test", + deps = [ + ":common", + ":ref_cycle_test_util_py", + ], +) + drake_py_unittest( name = "schema_test", deps = [ diff --git a/bindings/pydrake/common/ref_cycle_pybind.cc b/bindings/pydrake/common/ref_cycle_pybind.cc new file mode 100644 index 000000000000..3dd8fc9c7db5 --- /dev/null +++ b/bindings/pydrake/common/ref_cycle_pybind.cc @@ -0,0 +1,71 @@ +#include "drake/bindings/pydrake/common/ref_cycle_pybind.h" + +#include "drake/common/drake_assert.h" +#include "drake/common/drake_throw.h" + +using pybind11::handle; +using pybind11::detail::function_call; + +namespace drake { +namespace pydrake { +namespace internal { + +void do_ref_cycle_impl(handle p0, handle p1) { + if (!p0 || !p1) { + pybind11::pybind11_fail("Could not activate ref_cycle!"); + } + + if (p1.is_none() || p0.is_none()) { + return; // Nothing to do; at least one peer is missing. + } + // Among the reasons the following checks may fail is that one of the + // participating pybind11::class_ types does not declare + // pybind11::dynamic_attr(). + DRAKE_THROW_UNLESS(PyType_IS_GC(Py_TYPE(p0.ptr()))); + DRAKE_THROW_UNLESS(PyType_IS_GC(Py_TYPE(p1.ptr()))); + + // Each peer will have a new/updated attribute, containing a set of + // handles. Insert each into the other's handle set. Create the set first + // if it is not yet existing. + auto make_link = [](handle a, handle b) { + static const char refcycle_peers[] = "_pydrake_ref_cycle_peers"; + handle peers; + if (hasattr(a, refcycle_peers)) { + peers = a.attr(refcycle_peers); + } else { + peers = PySet_New(nullptr); + DRAKE_DEMAND(PyType_IS_GC(Py_TYPE(peers.ptr()))); + a.attr(refcycle_peers) = peers; + Py_DECREF(peers.ptr()); + } + // Ensure the proper ref count on the `peers` set. If it is > 1, the + // objects will live forever. If it is < 1, the cycle will just be deleted + // immediately. + DRAKE_DEMAND(Py_REFCNT(peers.ptr()) == 1); + PySet_Add(peers.ptr(), b.ptr()); + }; + make_link(p0, p1); + make_link(p1, p0); +} + +void ref_cycle_impl( + size_t Peer0, size_t Peer1, const function_call& call, handle ret) { + auto get_arg = [&](size_t n) { + if (n == 0) { + return ret; + } + if (n == 1 && call.init_self) { + return call.init_self; + } + if (n <= call.args.size()) { + return call.args[n - 1]; + } + return handle(); + }; + + do_ref_cycle_impl(get_arg(Peer0), get_arg(Peer1)); +} + +} // namespace internal +} // namespace pydrake +} // namespace drake diff --git a/bindings/pydrake/common/ref_cycle_pybind.h b/bindings/pydrake/common/ref_cycle_pybind.h new file mode 100644 index 000000000000..cea0ff5e3485 --- /dev/null +++ b/bindings/pydrake/common/ref_cycle_pybind.h @@ -0,0 +1,62 @@ +#pragma once + +#include "drake/bindings/pydrake/pydrake_pybind.h" + +namespace drake { +namespace pydrake { +namespace internal { + +/* pydrake::internal::ref_cycle creates a reference count cycle that Python's + cyclic garbage collection can find and collect, once the cycle's objects are + no longer reachable. Both peer objects must by pybind11 classes that were + defined with the dynamic_attr() annotation. + + Wherever pybind11::keep_alive() can be used, ref_cycle() can be + used similarly. + */ +template +struct ref_cycle {}; + +/* This function is used in the template below to select peers by call/return + index. */ +void ref_cycle_impl(size_t Peer0, size_t Peer1, + const pybind11::detail::function_call& call, pybind11::handle ret); + +/* This function creates the cycle, given peers as pybind11::handles. */ +void do_ref_cycle_impl(pybind11::handle p0, pybind11::handle p1); + +} // namespace internal +} // namespace pydrake +} // namespace drake + +namespace pybind11 { +namespace detail { + +// Provide a specialization of the pybind11 internal process_attribute +// template; this allows writing an annotation that works seamlessly in +// bindings definitions. +template +class process_attribute> + : public process_attribute_default< + drake::pydrake::internal::ref_cycle> { + public: + // NOLINTNEXTLINE(runtime/references) + static void precall(function_call& call) { + if constexpr (!needs_result()) { + drake::pydrake::internal::ref_cycle_impl(Peer0, Peer1, call, handle()); + } + } + + // NOLINTNEXTLINE(runtime/references) + static void postcall(function_call& call, handle ret) { + if constexpr (needs_result()) { + drake::pydrake::internal::ref_cycle_impl(Peer0, Peer1, call, ret); + } + } + + private: + static constexpr bool needs_result() { return Peer0 == 0 || Peer1 == 0; } +}; + +} // namespace detail +} // namespace pybind11 diff --git a/bindings/pydrake/common/test/ref_cycle_test.py b/bindings/pydrake/common/test/ref_cycle_test.py new file mode 100644 index 000000000000..a5179428e6e9 --- /dev/null +++ b/bindings/pydrake/common/test/ref_cycle_test.py @@ -0,0 +1,153 @@ +"""Unit test for ref_cycle<>() annotation. + +See also ref_cycle_test_util_py.cc for the bindings used in the tests. +""" +import gc +import sys +import unittest +import weakref + +from pydrake.common.ref_cycle_test_util import ( + NotDynamic, IsDynamic, invalid_arg_index, free_function, ouroboros) + + +class TestRefCycle(unittest.TestCase): + def check_is_collectable_cycle(self, p0, p1): + # The edges of the cycle are: + # p0 -> p0.__dict__ -> p0._pydrake_ref_cycle_peers \ + # -> p1 -> p1.__dict__ -> p1._pydrake_ref_cycle_peers -> p0 + # where the object at each _pydrake_ref_cycle_peers is a set. + # + # It is impractical to check the counts of p0 and p1 here because + # callers may hold an arbitrary number of references. + + for x in [p0, p1]: + # Check the counts of the internal parts of the cycle. + # Recall that sys.getrefcount() adds 1 to the actual count. + self.assertEqual(sys.getrefcount(x.__dict__), 2) + self.assertEqual(sys.getrefcount(x._pydrake_ref_cycle_peers), 2) + + # Check that all parts are tracked by gc. + self.assertTrue(gc.is_tracked(x)) + self.assertTrue(gc.is_tracked(x.__dict__)) + self.assertTrue(gc.is_tracked(x._pydrake_ref_cycle_peers)) + + # Check that the peers refer to each other. + self.assertTrue(p1 in p0._pydrake_ref_cycle_peers) + self.assertTrue(p0 in p1._pydrake_ref_cycle_peers) + + def check_no_cycle(self, p0, p1): + for x in [p0, p1]: + self.assertFalse(hasattr(x, '_pydrake_ref_cycle_peers')) + + def test_invalid_index(self): + with self.assertRaisesRegex(RuntimeError, + "Could not activate ref_cycle.*"): + invalid_arg_index() + + def test_ouroboros(self): + # The self-cycle edges are: + # dut -> dut.__dict__ -> dut._pydrake_ref_cycle_peers -> dut + # + # This still passes check_is_collectable_cycle() -- the function just + # does redundant work. + dut = IsDynamic() + returned = ouroboros(dut) + self.assertEqual(returned, dut) + self.assertEqual(len(dut._pydrake_ref_cycle_peers), 1) + self.check_is_collectable_cycle(returned, dut) + + def test_free_function(self): + p0 = IsDynamic() + p1 = IsDynamic() + free_function(p0, p1) + self.check_is_collectable_cycle(p0, p1) + + def test_not_dynamic_add(self): + dut = NotDynamic() + peer = IsDynamic() + # Un-annotated call is fine. + dut.AddIs(peer) + self.check_no_cycle(dut, peer) + # Annotated call dies because dut is not py::dynamic_attr(). + with self.assertRaisesRegex(SystemExit, ".*PyType_IS_GC.*"): + dut.AddIsCycle(peer) + + def test_not_dynamic_return(self): + dut = NotDynamic() + # Un-annotated call is fine. + returned = dut.ReturnIs() + self.check_no_cycle(dut, returned) + # Annotated call dies because dut is not py::dynamic_attr(). + with self.assertRaisesRegex(SystemExit, ".*PyType_IS_GC.*"): + dut.ReturnIsCycle() + + def test_not_dynamic_null(self): + dut = NotDynamic() + # Un-annotated call is fine. + self.assertIsNone(dut.ReturnNullIs()) + # Annotated call does not die because one peer is missing. + self.assertIsNone(dut.ReturnNullIsCycle()) + + def test_is_dynamic_add_not(self): + dut = IsDynamic() + notpeer = NotDynamic() + dut.AddNot(notpeer) + self.check_no_cycle(dut, notpeer) + # Annotated call dies because notpeer is not py::dynamic_attr(). + with self.assertRaisesRegex(SystemExit, ".*PyType_IS_GC.*"): + dut.AddNotCycle(notpeer) + + def test_is_dynamic_return_not(self): + dut = IsDynamic() + # Un-annotated call is fine. + returned = dut.ReturnNot() + self.check_no_cycle(dut, returned) + # Annotated call dies because return is not py::dynamic_attr(). + with self.assertRaisesRegex(SystemExit, ".*PyType_IS_GC.*"): + dut.ReturnNotCycle() + + def test_is_dynamic_return_null(self): + dut = IsDynamic() + # Un-annotated call is fine. + self.assertIsNone(dut.ReturnNullNot()) + self.assertIsNone(dut.ReturnNullIs()) + # Annotated call does not die because one peer is missing. + self.assertIsNone(dut.ReturnNullNotCycle()) + self.assertIsNone(dut.ReturnNullIsCycle()) + + def test_is_dynamic_add_is(self): + dut = IsDynamic() + peer = IsDynamic() + # Un-annotated call does not implement a cycle. + dut.AddIs(peer) + self.check_no_cycle(dut, peer) + # Annotated call produces a collectable cycle. + dut.AddIsCycle(peer) + self.check_is_collectable_cycle(dut, peer) + + def test_is_dynamic_return_is(self): + dut = IsDynamic() + # Un-annotated call does not implement a cycle. + returned = dut.ReturnIs() + self.check_no_cycle(dut, returned) + # Annotated call produces a collectable cycle. + returned = dut.ReturnIsCycle() + self.check_is_collectable_cycle(dut, returned) + + def test_actual_collection(self): + + def make_a_cycle(): + dut = IsDynamic() + return dut.ReturnIsCycle() + + cycle = make_a_cycle() + finalizer = weakref.finalize(cycle, lambda: None) + # Cycle is alive while we refer to it. + self.assertTrue(finalizer.alive) + del cycle + # Cycle is alive because of the ref_cycle. + self.assertTrue(finalizer.alive) + gc.collect() + # Cycle does not survive garbage collection. + self.assertFalse(finalizer.alive) diff --git a/bindings/pydrake/common/test/ref_cycle_test_util_py.cc b/bindings/pydrake/common/test/ref_cycle_test_util_py.cc new file mode 100644 index 000000000000..d12c7e8acbc9 --- /dev/null +++ b/bindings/pydrake/common/test/ref_cycle_test_util_py.cc @@ -0,0 +1,75 @@ +// Bindings that help test the ref_cycle<>() annotation. +// See also ref_cycle_test.py. + +#include "drake/bindings/pydrake/common/ref_cycle_pybind.h" +#include "drake/bindings/pydrake/pydrake_pybind.h" + +namespace drake { +namespace pydrake { + +namespace { + +template +class TestDummyBase { + public: + using IsDynamic = TestDummyBase; + using NotDynamic = TestDummyBase; + TestDummyBase() = default; + TestDummyBase(const TestDummyBase&) = default; + ~TestDummyBase() = default; + + void AddNot(NotDynamic*) {} + NotDynamic* ReturnNot() { return new NotDynamic(); } + NotDynamic* ReturnNullNot() { return nullptr; } + void AddIs(IsDynamic*) {} + IsDynamic* ReturnIs() { return new IsDynamic(); } + IsDynamic* ReturnNullIs() { return nullptr; } +}; +using IsDynamic = TestDummyBase; +using NotDynamic = TestDummyBase; + +} // namespace + +PYBIND11_MODULE(ref_cycle_test_util, m) { + using internal::ref_cycle; + { + using Class = NotDynamic; + py::class_(m, "NotDynamic") + .def(py::init<>()) + .def("AddIs", &Class::AddIs) + .def("AddIsCycle", &Class::AddIs, ref_cycle<1, 2>()) + .def("ReturnIs", &Class::ReturnIs) + .def("ReturnIsCycle", &Class::ReturnIs, ref_cycle<0, 1>()) + .def("ReturnNullIs", &Class::ReturnNullIs) + .def("ReturnNullIsCycle", &Class::ReturnNullIs, ref_cycle<0, 1>()); + } + + { + using Class = IsDynamic; + py::class_(m, "IsDynamic", py::dynamic_attr()) + .def(py::init<>()) + .def("AddNot", &Class::AddNot) + .def("AddNotCycle", &Class::AddNot, ref_cycle<1, 2>()) + .def("ReturnNot", &Class::ReturnNot) + .def("ReturnNotCycle", &Class::ReturnNot, ref_cycle<0, 1>()) + .def("ReturnNullNot", &Class::ReturnNullNot) + .def("ReturnNullNotCycle", &Class::ReturnNullNot, ref_cycle<0, 1>()) + .def("AddIs", &Class::AddIs) + .def("AddIsCycle", &Class::AddIs, ref_cycle<1, 2>()) + .def("ReturnIs", &Class::ReturnIs) + .def("ReturnIsCycle", &Class::ReturnIs, ref_cycle<0, 1>()) + .def("ReturnNullIs", &Class::ReturnNullIs) + .def("ReturnNullIsCycle", &Class::ReturnNullIs, ref_cycle<0, 1>()); + } + + m.def( + "free_function", [](IsDynamic*, IsDynamic*) {}, ref_cycle<1, 2>()); + m.def( + "invalid_arg_index", [] {}, ref_cycle<0, 1>()); + // Returns its argument and creates a self-cycle. + m.def( + "ouroboros", [](IsDynamic* x) { return x; }, ref_cycle<0, 1>()); +} + +} // namespace pydrake +} // namespace drake