BUG: fix race initializing legacy dtype casts

numpy · charris · Feb 11, 2025 · Feb 6, 2025 · Feb 6, 2025 · Feb 6, 2025
commit 6f7f21e80778c80c36fca774430d1707d3852f4e
diff --git a/numpy/_core/src/multiarray/convert_datatype.c b/numpy/_core/src/multiarray/convert_datatype.c
@@ -63,46 +63,24 @@ static PyObject *
 PyArray_GetObjectToGenericCastingImpl(void);
 
 
-/**
- * Fetch the casting implementation from one DType to another.
- *
- * @param from The implementation to cast from
- * @param to The implementation to cast to
- *
- * @returns A castingimpl (PyArrayDTypeMethod *), None or NULL with an
- *          error set.
- */
-NPY_NO_EXPORT PyObject *
-PyArray_GetCastingImpl(PyArray_DTypeMeta *from, PyArray_DTypeMeta *to)
+static PyObject *
+create_casting_impl(PyArray_DTypeMeta *from, PyArray_DTypeMeta *to)
 {
-    PyObject *res;
-    if (from == to) {
-        res = (PyObject *)NPY_DT_SLOTS(from)->within_dtype_castingimpl;
-    }
-    else {
-        res = PyDict_GetItemWithError(NPY_DT_SLOTS(from)->castingimpls, (PyObject *)to);
-    }
-    if (res != NULL || PyErr_Occurred()) {
-        Py_XINCREF(res);
-        return res;
-    }
     /*
-     * The following code looks up CastingImpl based on the fact that anything
+     * Look up CastingImpl based on the fact that anything
      * can be cast to and from objects or structured (void) dtypes.
-     *
-     * The last part adds casts dynamically based on legacy definition
      */
     if (from->type_num == NPY_OBJECT) {
-        res = PyArray_GetObjectToGenericCastingImpl();
+        return PyArray_GetObjectToGenericCastingImpl();
     }
     else if (to->type_num == NPY_OBJECT) {
-        res = PyArray_GetGenericToObjectCastingImpl();
+        return PyArray_GetGenericToObjectCastingImpl();
     }
     else if (from->type_num == NPY_VOID) {
-        res = PyArray_GetVoidToGenericCastingImpl();
+        return PyArray_GetVoidToGenericCastingImpl();
     }
     else if (to->type_num == NPY_VOID) {
-        res = PyArray_GetGenericToVoidCastingImpl();
+        return PyArray_GetGenericToVoidCastingImpl();
     }
     /*
      * Reject non-legacy dtypes. They need to use the new API to add casts and
@@ -125,43 +103,104 @@ PyArray_GetCastingImpl(PyArray_DTypeMeta *from, PyArray_DTypeMeta *to)
             PyArray_VectorUnaryFunc *castfunc = PyArray_GetCastFunc(
                     from->singleton, to->type_num);
             if (castfunc == NULL) {
-                PyErr_Clear();
-                /* Remember that this cast is not possible */
-                if (PyDict_SetItem(NPY_DT_SLOTS(from)->castingimpls,
-                            (PyObject *) to, Py_None) < 0) {
-                    return NULL;
-                }
-                Py_RETURN_NONE;
+                    PyErr_Clear();
+                    Py_RETURN_NONE;
             }
         }
-
-        /* PyArray_AddLegacyWrapping_CastingImpl find the correct casting level: */
-        /*
-         * TODO: Possibly move this to the cast registration time. But if we do
-         *       that, we have to also update the cast when the casting safety
-         *       is registered.
+        /* Create a cast using the state of the legacy casting setup defined
+         * during the setup of the DType.
+         *
+         * Ideally we would do this when we create the DType, but legacy user
+         * DTypes don't have a way to signal that a DType is done setting up
+         * casts. Without such a mechanism, the safest way to know that a
+         * DType is done setting up is to register the cast lazily the first
+         * time a user does the cast.
+         *
+         * We *could* register the casts when we create the wrapping
+         * DTypeMeta, but that means the internals of the legacy user DType
+         * system would need to update the state of the casting safety flags
+         * in the cast implementations stored on the DTypeMeta. That's an
+         * inversion of abstractions and would be tricky to do without
+         * creating circular dependencies inside NumPy.
          */
         if (PyArray_AddLegacyWrapping_CastingImpl(from, to, -1) < 0) {
             return NULL;
         }
+        /* castingimpls is unconditionally filled by
+         * AddLegacyWrapping_CastingImpl, so this won't create a recursive
+         * critical section
+         */
         return PyArray_GetCastingImpl(from, to);
     }
+}
 
-    if (res == NULL) {
+static PyObject *
+ensure_castingimpl_exists(PyArray_DTypeMeta *from, PyArray_DTypeMeta *to)
+{
+    int return_error = 0;
+    PyObject *res = NULL;
+
+    /* Need to create the cast. This might happen at runtime so we enter a
+       critical section to avoid races */
+
+    Py_BEGIN_CRITICAL_SECTION(NPY_DT_SLOTS(from)->castingimpls);
+
+    /* check if another thread filled it while this thread was blocked on
+       acquiring the critical section */
+    if (PyDict_GetItemRef(NPY_DT_SLOTS(from)->castingimpls, (PyObject *)to,
+                          &res) < 0) {
+        return_error = 1;
+    }
+
+    if (!return_error) {
+        res = create_casting_impl(from, to);
+        if (res == NULL) {
+            return_error = 1;
+        }
+    }
+    if (!return_error &&
+        PyDict_SetItem(NPY_DT_SLOTS(from)->castingimpls,
+                       (PyObject *)to, res) < 0) {
+        return_error = 1;
+    }
+    Py_END_CRITICAL_SECTION();
+    if (return_error) {
+        Py_XDECREF(res);
         return NULL;
     }
-    if (from == to) {
+    if (from == to && res == Py_None) {
         PyErr_Format(PyExc_RuntimeError,
                 "Internal NumPy error, within-DType cast missing for %S!", from);
-        Py_DECREF(res);
         return NULL;
     }
-    if (PyDict_SetItem(NPY_DT_SLOTS(from)->castingimpls,
-                (PyObject *)to, res) < 0) {
-        Py_DECREF(res);
+    return res;
+}
+
+/**
+ * Fetch the casting implementation from one DType to another.
+ *
+ * @param from The implementation to cast from
+ * @param to The implementation to cast to
+ *
+ * @returns A castingimpl (PyArrayDTypeMethod *), None or NULL with an
+ *          error set.
+ */
+NPY_NO_EXPORT PyObject *
+PyArray_GetCastingImpl(PyArray_DTypeMeta *from, PyArray_DTypeMeta *to)
+{
+    PyObject *res;
+    if (from == to) {
+        res = Py_XNewRef((PyObject *)NPY_DT_SLOTS(from)->within_dtype_castingimpl);
+    }
+    else if (PyDict_GetItemRef(NPY_DT_SLOTS(from)->castingimpls,
+                               (PyObject *)to, &res) < 0) {
         return NULL;
     }
-    return res;
+    if (res != NULL) {
+        return res;
+    }
+
+    return ensure_castingimpl_exists(from, to);
 }
 
 
@@ -410,7 +449,7 @@ _get_cast_safety_from_castingimpl(PyArrayMethodObject *castingimpl,
  * implementations fully to have them available for doing the actual cast
  * later.
  *
- * @param from The descriptor to cast from 
+ * @param from The descriptor to cast from
  * @param to The descriptor to cast to (may be NULL)
  * @param to_dtype If `to` is NULL, must pass the to_dtype (otherwise this
  *        is ignored).
@@ -2021,6 +2060,11 @@ PyArray_AddCastingImplementation(PyBoundArrayMethodObject *meth)
 /**
  * Add a new casting implementation using a PyArrayMethod_Spec.
  *
+ * Using this function outside of module initialization without holding a
+ * critical section on the castingimpls dict may lead to a race to fill the
+ * dict. Use PyArray_GetGastingImpl to lazily register casts at runtime
+ * safely.
+ *
  * @param spec The specification to use as a source
  * @param private If private, allow slots not publicly exposed.
  * @return 0 on success -1 on failure

diff --git a/numpy/_core/tests/test_multithreading.py b/numpy/_core/tests/test_multithreading.py
@@ -7,6 +7,7 @@
 
 from numpy.testing import IS_WASM
 from numpy.testing._private.utils import run_threaded
+from numpy._core import _rational_tests
 
 if IS_WASM:
     pytest.skip(allow_module_level=True, reason="no threading support in wasm")
@@ -254,3 +255,11 @@ def func(arr):
 
         for f in futures:
             f.result()
+
+
+def test_legacy_usertype_cast_init_thread_safety():
+    def closure(b):
+        b.wait()
+        np.full((10, 10), 1, _rational_tests.rational)
+
+    run_threaded(closure, 1000, pass_barrier=True)