[bindings] Add custom ref_cycle annotation

Introduce the new annotation internal::ref_cycle<M, N>(). 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

@@ -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 = [

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

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<M, N>() can be used, ref_cycle<M, N>() can be
+  used similarly.
+ */
+template <size_t Peer0, size_t Peer1>
+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 <size_t Peer0, size_t Peer1>
+class process_attribute<drake::pydrake::internal::ref_cycle<Peer0, Peer1>>
+    : public process_attribute_default<
+          drake::pydrake::internal::ref_cycle<Peer0, Peer1>> {
+ 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

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)