Fix Python plugin reference leak

include/mitsuba/core/object.h

@@ -118,6 +118,11 @@ class MI_EXPORT_LIB Object {
      */
     virtual std::string to_string() const;
 
+    /// Set a cleanup callback for Python-defined plugins.
+    void set_py_cleanup(const std::function<void()> &callback) {
+        m_py_cleanup_callback = callback;
+    }
+
 protected:
     /** \brief Virtual protected deconstructor.
      * (Will only be called by \ref ref)
@@ -128,6 +133,9 @@ class MI_EXPORT_LIB Object {
     mutable std::atomic<int> m_ref_count { 0 };
 
     static Class *m_class;
+
+    // Cleanup function for Python-defined objects.
+    std::function<void()> m_py_cleanup_callback;
 };
 
 /**

include/mitsuba/python/python.h

@@ -78,8 +78,12 @@ PYBIND11_DECLARE_HOLDER_TYPE(T, mitsuba::ref<T>, true);
                         o = constructor(&p);                                   \
                     }                                                          \
                                                                                \
-                    ref<Name> o2 = o.cast<Name*>();                            \
-                    o.release();                                               \
+                    Name *o2 = o.cast<Name*>();                                \
+                    auto handle = o.release(); /* Manage ref count manually. */\
+                    o2->set_py_cleanup([handle]() {                            \
+                        py::gil_scoped_acquire gil;                            \
+                        handle.dec_ref();                                      \
+                    });                                                        \
                     return o2;                                                 \
                 },                                                             \
                 nullptr);                                                      \

src/core/object.cpp

@@ -7,6 +7,16 @@ NAMESPACE_BEGIN(mitsuba)
 
 void Object::dec_ref(bool dealloc) const noexcept {
     uint32_t ref_count = m_ref_count.fetch_sub(1);
+
+    // If the Object has a Python object attached to it, we need to clean up
+    // when the reference counter is decreased below 2. This accounts for the
+    // fact that there is an internal reference by the Python object itself. The
+    // cleanup callback then releases the Python object, which recursively
+    // de-allocates the Mitsuba Object instance.
+    if (ref_count == 2 && m_py_cleanup_callback) {
+        m_py_cleanup_callback();
+        return;
+    }
     if (ref_count <= 0) {
         fprintf(stderr, "Internal error: Object reference count < 0!\n");
         abort();