[bindings] Add custom ref_cycle annotation

bindings/pydrake/common/BUILD.bazel

@@ -234,6 +234,20 @@ 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",
+        "@fmt",
+        "@pybind11",
+    ],
+)
+
 # N.B. Any C++ libraries that include this must include `cpp_template_py` when
 # being used in Python.
 drake_cc_library(
@@ -492,6 +506,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,100 @@
+#include "drake/bindings/pydrake/common/ref_cycle_pybind.h"
+
+#include <fmt/format.h>
+
+#include "drake/common/drake_assert.h"
+#include "drake/common/drake_throw.h"
+
+namespace drake {
+namespace pydrake {
+
+using py::handle;
+using py::detail::function_call;
+
+namespace internal {
+
+namespace {
+
+void make_ref_cycle(handle p0, handle p1) {
+  DRAKE_DEMAND(static_cast<bool>(p0));
+  DRAKE_DEMAND(static_cast<bool>(p1));
+  DRAKE_DEMAND(!p0.is_none());
+  DRAKE_DEMAND(!p1.is_none());
+  DRAKE_DEMAND(PyType_IS_GC(Py_TYPE(p0.ptr())));
+  DRAKE_DEMAND(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_internal_ref_cycle_peers";
+    if (!hasattr(a, refcycle_peers)) {
+      py::set new_set;
+      DRAKE_DEMAND(PyType_IS_GC(Py_TYPE(new_set.ptr())));
+      a.attr(refcycle_peers) = new_set;
+    }
+    handle peers = a.attr(refcycle_peers);
+    // 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);
+}
+
+}  // namespace
+
+void ref_cycle_impl(
+    size_t peer0, size_t peer1, const function_call& call, handle ret) {
+  // Returns the handle selected by the given index. Throws if the index is
+  // invalid.
+  auto get_arg = [&](size_t n) -> handle {
+    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];
+    }
+    py::pybind11_fail(fmt::format(
+        "Could not activate ref_cycle: index {} is invalid for function '{}'",
+        n, call.func.name));
+  };
+  handle p0 = get_arg(peer0);
+  handle p1 = get_arg(peer1);
+
+  // Returns false if the handle's value is None. Throws if the handle's value
+  // is not of a garbage-collectable type.
+  auto check_handle = [&](size_t n, handle p) -> bool {
+    if (p.is_none()) {
+      return false;
+    }
+    // Among the reasons the following check may fail is that one of the
+    // participating pybind11::class_ types does not declare
+    // pybind11::dynamic_attr().
+    if (!PyType_IS_GC(Py_TYPE(p.ptr()))) {
+      py::pybind11_fail(fmt::format(
+          "Could not activate ref_cycle: object type at index {} for "
+          "function '{}' is not tracked by garbage collection. Was the object "
+          "defined with `pybind11::class_<...>(... pybind11::dynamic_attr())`?",
+          n, call.func.name));
+    }
+    return true;
+  };
+  if (!check_handle(peer0, p0) || !check_handle(peer1, p1)) {
+    // At least one of the handles is None. We can't construct a ref-cycle, but
+    // neither should we complain. A None variable value could happen for any
+    // number of legitimate reasons, and does not mean that the ref_cycle call
+    // policy is defective.
+    return;
+  }
+  make_ref_cycle(p0, p1);
+}
+
+}  // namespace internal
+}  // namespace pydrake
+}  // namespace drake

bindings/pydrake/common/ref_cycle_pybind.h

@@ -0,0 +1,93 @@
+#pragma once
+
+#include "drake/bindings/pydrake/pydrake_pybind.h"
+
+namespace drake {
+namespace pydrake {
+namespace internal {
+
+/* pydrake::internal::ref_cycle is a custom call policy for pybind11.
+
+   For an overview of other call policies, See
+   https://pybind11.readthedocs.io/en/stable/advanced/functions.html#additional-call-policies
+
+  `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 be either ordinary python objects or pybind11 classes
+  that were defined with the dynamic_attr() annotation. If either of the peer
+  objects is a pybind11 class not defined with py::dynamic_attr(), the
+  binding will raise an exception at the time the function or method is
+  called. The error message will complain about objects not being tracked by
+  garbage collection. By default pybind11 python objects are not tracked; the
+  dynamic_attr() annotation makes sure the objects are tracked. Also, it
+  ensures that the implementation of `ref_cycle` can create and assign a new
+  attribute dynamically, which is how the garbage collectible bookkeeping is
+  accomplished.
+
+  `ref_cycle` causes each object to refer to the other in a cycle. It is
+  bidirectional and symmetric. The order of the template arguments does not
+  matter.
+
+  Note the consequences for object lifetimes:
+
+  * M keeps N alive, and N keeps M alive.
+  * Neither object is destroyed until:
+    * both are unreachable, and
+    * garbage collection runs.
+
+  @tparam Peer0 an argument index
+  @tparam Peer1 an argument index
+
+  The argument index starts at 1; for methods, `self` is at index 1. Index 0
+  denotes the return value.
+ */
+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 py::detail::function_call& call, py::handle ret);
+
+}  // 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) {
+    // Only generate code if this invocation doesn't need the return value.
+    if constexpr (!needs_return_value()) {
+      drake::pydrake::internal::ref_cycle_impl(Peer0, Peer1, call, handle());
+    }
+  }
+
+  // NOLINTNEXTLINE(runtime/references)
+  static void postcall(function_call& call, handle ret) {
+    // Only generate code if this invocation *does* need the return value.
+    if constexpr (needs_return_value()) {
+      drake::pydrake::internal::ref_cycle_impl(Peer0, Peer1, call, ret);
+    }
+  }
+
+ private:
+  // Returns true if either template parameter denotes the return value.
+  static constexpr bool needs_return_value() {
+    return Peer0 == 0 || Peer1 == 0;
+  }
+};
+
+}  // namespace detail
+}  // namespace pybind11