From 3aaf94f42c206de95b497d388be35bb509434ff0 Mon Sep 17 00:00:00 2001
From: Saran Tunyasuvunakool <stunya@google.com>
Date: Mon, 17 Oct 2022 00:05:02 -0700
Subject: [PATCH] Fix Python GIL-not-held issues in `python/callbacks.cc`.

The GIL bug was introduced in ec6ea6a69e43b40ab3ef30981408f9baccaea047 and is resolved by decrementing the refcount before returning from MjWrapperLookup.

PiperOrigin-RevId: 481558798
Change-Id: I2744287d48635a55ae6d2ec46c9e43fe881019d8
---
 python/mujoco/callbacks.cc | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/python/mujoco/callbacks.cc b/python/mujoco/callbacks.cc
index 929e3458..c6a99351 100644
--- a/python/mujoco/callbacks.cc
+++ b/python/mujoco/callbacks.cc
@@ -18,6 +18,7 @@
 #include <limits>
 #include <sstream>
 #include <type_traits>
+#include <utility>
 
 #include <mujoco/mujoco.h>
 #include "errors.h"
@@ -46,7 +47,7 @@ using enable_if_not_const_t =
 // table that associates raw MuJoCo struct pointers back to the pointers to
 // their corresponding wrappers.
 template <typename Raw>
-static enable_if_not_const_t<Raw, py::object> MjWrapperLookup(Raw* ptr) {
+static enable_if_not_const_t<Raw, py::handle> MjWrapperLookup(Raw* ptr) {
   using LookupFnType = MjWrapper<Raw>* (Raw*);
   static LookupFnType* const lookup = []() -> LookupFnType* {
     py::gil_scoped_acquire gil;
@@ -102,7 +103,7 @@ static enable_if_not_const_t<Raw, py::object> MjWrapperLookup(Raw* ptr) {
               "cannot find the Python instance of the MjWrapper");
         }
       } else {
-        return instance;
+        return std::move(instance);
       }
     } else {
       if (!PyErr_Occurred()) {
@@ -117,7 +118,7 @@ static enable_if_not_const_t<Raw, py::object> MjWrapperLookup(Raw* ptr) {
 }
 
 template <typename Raw>
-static const py::object MjWrapperLookup(const Raw* ptr) {
+static const py::handle MjWrapperLookup(const Raw* ptr) {
   return MjWrapperLookup(const_cast<Raw*>(ptr));
 }