[Mlir-commits] [mlir] [MLIR:Python] Fix race on PyOperations. (PR #139721)

[Maksim Levental]

================
@@ -1254,18 +1335,26 @@ PyOperationRef PyOperation::forOperation(PyMlirContextRef contextRef,
   nb::ft_lock_guard lock(contextRef->liveOperationsMutex);
   auto &liveOperations = contextRef->liveOperations;
   auto it = liveOperations.find(operation.ptr);
-  if (it == liveOperations.end()) {
-    // Create.
-    PyOperationRef result = createInstance(std::move(contextRef), operation,
-                                           std::move(parentKeepAlive));
-    liveOperations[operation.ptr] =
-        std::make_pair(result.getObject(), result.get());
-    return result;
+  if (it != liveOperations.end()) {
+    PyOperation *existing = it->second.second;
+    nb::handle pyRef = it->second.first;
+
+    // Try to increment the reference count of the existing entry. This can fail
+    // if the object is in the process of being destroyed by another thread.
+    if (tryIncRef(pyRef)) {
+      return PyOperationRef(existing, nb::steal<nb::object>(pyRef));
+    }
+
+    // Mark the existing entry as invalid, since we are about to replace it.
+    existing->valid = false;
----------------
makslevental wrote:

nit: `setInvalidLocked` instead of direct access

https://github.com/llvm/llvm-project/pull/139721