Merge pull request #17421 from seberg/coercion-cache-leak-large-arrays

charris · web-flow · commit e6b8b19ee832 · 2020-10-03T08:02:22.000-06:00
BUG: Fix memory leak in array-coercion error paths
diff --git a/numpy/core/src/multiarray/array_coercion.c b/numpy/core/src/multiarray/array_coercion.c
@@ -538,7 +538,7 @@ npy_new_coercion_cache(
         cache = _coercion_cache_cache[_coercion_cache_num];
     }
     else {
-        cache = PyObject_MALLOC(sizeof(coercion_cache_obj));
+        cache = PyMem_Malloc(sizeof(coercion_cache_obj));
     }
     if (cache == NULL) {
         PyErr_NoMemory();
@@ -570,7 +570,7 @@ npy_unlink_coercion_cache(coercion_cache_obj *current)
         _coercion_cache_num++;
     }
     else {
-        PyObject_FREE(current);
+        PyMem_Free(current);
     }
     return next;
 }
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
@@ -610,6 +610,7 @@ PyArray_AssignFromCache(PyArrayObject *self, coercion_cache_obj *cache) {
         PyErr_SetString(PyExc_RuntimeError,
                 "Inconsistent object during array creation? "
                 "Content of sequences changed (cache not consumed).");
+        npy_free_coercion_cache(cache);
         return -1;
     }
     return 0;
@@ -1467,6 +1468,7 @@ PyArray_FromAny(PyObject *op, PyArray_Descr *newtype, int min_depth,
         PyErr_SetString(PyExc_TypeError,
                         "WRITEBACKIFCOPY used for non-array input.");
         Py_DECREF(dtype);
+        npy_free_coercion_cache(cache);
         return NULL;
     }
 
@@ -1475,6 +1477,7 @@ PyArray_FromAny(PyObject *op, PyArray_Descr *newtype, int min_depth,
             &PyArray_Type, dtype, ndim, dims, NULL, NULL,
             flags&NPY_ARRAY_F_CONTIGUOUS, NULL);
     if (ret == NULL) {
+        npy_free_coercion_cache(cache);
         return NULL;
     }
     if (cache == NULL) {
diff --git a/numpy/core/tests/test_array_coercion.py b/numpy/core/tests/test_array_coercion.py
@@ -635,3 +635,18 @@ def __array__(self):
         assert arr[()] is ArrayLike
         arr = np.array([ArrayLike])
         assert arr[0] is ArrayLike
+
+    @pytest.mark.skipif(
+            np.dtype(np.intp).itemsize < 8, reason="Needs 64bit platform")
+    def test_too_large_array_error_paths(self):
+        """Test the error paths, including for memory leaks"""
+        arr = np.array(0, dtype="uint8")
+        # Guarantees that a contiguous copy won't work:
+        arr = np.broadcast_to(arr, 2**62)
+
+        for i in range(5):
+            # repeat, to ensure caching cannot have an effect:
+            with pytest.raises(MemoryError):
+                np.array(arr)
+            with pytest.raises(MemoryError):
+                np.array([arr])

Original file line number	Diff line number	Diff line change
`@@ -538,7 +538,7 @@ npy_new_coercion_cache(`
`538`	`538`	`cache = _coercion_cache_cache[_coercion_cache_num];`
`539`	`539`	`}`
`540`	`540`	`else {`
`541`		`- cache = PyObject_MALLOC(sizeof(coercion_cache_obj));`
	`541`	`+ cache = PyMem_Malloc(sizeof(coercion_cache_obj));`
`542`	`542`	`}`
`543`	`543`	`if (cache == NULL) {`
`544`	`544`	`PyErr_NoMemory();`
`@@ -570,7 +570,7 @@ npy_unlink_coercion_cache(coercion_cache_obj *current)`
`570`	`570`	`_coercion_cache_num++;`
`571`	`571`	`}`
`572`	`572`	`else {`
`573`		`- PyObject_FREE(current);`
	`573`	`+ PyMem_Free(current);`
`574`	`574`	`}`
`575`	`575`	`return next;`
`576`	`576`	`}`