Add context/thread-locality tests, allow testing custom policies

numpy · seberg · Oct 25, 2021 · Oct 11, 2020 · Oct 16, 2020 · Oct 18, 2020
commit 144acc6d1a9ee9f8f08e9aca9fbe24ff0de6f0d6
diff --git a/doc/source/reference/c-api/data_memory.rst b/doc/source/reference/c-api/data_memory.rst
@@ -88,7 +88,7 @@ reallocate or free the data memory of the instance.
 .. c:function:: const PyDataMem_Handler * PyDataMem_GetHandler(PyArrayObject *obj)
 
    Return the `PyDataMem_Handler` used by the ``PyArrayObject``. If ``NULL``,
-   return the current global policy that ill be used to allocate data for the
+   return the current global policy that will be used to allocate data for the
    next ``PyArrayObject``. On failure, set an exception and return ``NULL``.
 
 For an example of setting up and using the PyDataMem_Handler, see the test in

diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c
@@ -514,7 +514,7 @@ PyDataMem_SetHandler(PyDataMem_Handler *handler)
 
 /*NUMPY_API
  * Return the PyDataMem_Handler used by the PyArrayObject. If NULL, return
- * the current global policy that ill be used to allocate data
+ * the current global policy that will be used to allocate data
  * for the next PyArrayObject. On failure, set an exception and return NULL.
  */
 NPY_NO_EXPORT const PyDataMem_Handler *

diff --git a/numpy/core/tests/test_mem_policy.py b/numpy/core/tests/test_mem_policy.py
@@ -1,5 +1,7 @@
+import asyncio
 import pytest
 import numpy as np
+import threading
 from numpy.testing import extbuild
 
 
@@ -11,48 +13,51 @@ def get_module(tmp_path):
     free/calloc go via the functions here.
     """
     functions = [(
-        "test_prefix", "METH_O",
-        """
-            if (!PyArray_Check(args)) {
-                PyErr_SetString(PyExc_ValueError,
-                        "must be called with a numpy scalar or ndarray");
-            }
-            return PyUnicode_FromString(
-                            PyDataMem_GetHandler((PyArrayObject*)args)->name);
-        """
-        ),
-        ("set_new_policy", "METH_NOARGS",
+        "set_secret_data_policy", "METH_NOARGS",
          """
-             const PyDataMem_Handler *old = PyDataMem_SetHandler(&new_handler);
-             return PyUnicode_FromString(old->name);
+             PyDataMem_Handler *old = (PyDataMem_Handler *) PyDataMem_SetHandler(&secret_data_handler);
+             return PyCapsule_New(old, NULL, NULL);
          """),
-        ("set_old_policy", "METH_NOARGS",
+        ("set_old_policy", "METH_O",
          """
-             const PyDataMem_Handler *old = PyDataMem_SetHandler(NULL);
-             return PyUnicode_FromString(old->name);
+             PyDataMem_Handler *old = NULL;
+             if (args != NULL && PyCapsule_CheckExact(args)) {
+                 old = (PyDataMem_Handler *) PyCapsule_GetPointer(args, NULL);
+             }
+             PyDataMem_SetHandler(old);
+             Py_RETURN_NONE;
          """),
         ]
     prologue = '''
         #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
         #include <numpy/arrayobject.h>
+        /*
+         * This struct allows the dynamic configuration of the allocator funcs
+         * of the `secret_data_allocator`. It is provided here for
+         * demonstration purposes, as a valid `ctx` use-case scenario.
+         */
         typedef struct {
             void *(*malloc)(size_t);
             void *(*calloc)(size_t, size_t);
             void *(*realloc)(void *, size_t);
             void (*free)(void *);
-        } Allocator;
+        } SecretDataAllocatorFuncs;
         NPY_NO_EXPORT void *
-        shift_alloc(Allocator *ctx, size_t sz) {
-            char *real = (char *)ctx->malloc(sz + 64);
+        shift_alloc(void *ctx, size_t sz) {
+            SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *) ctx;
+
+            char *real = (char *)funcs->malloc(sz + 64);
             if (real == NULL) {
                 return NULL;
             }
             snprintf(real, 64, "originally allocated %ld", (unsigned long)sz);
             return (void *)(real + 64);
         }
         NPY_NO_EXPORT void *
-        shift_zero(Allocator *ctx, size_t sz, size_t cnt) {
-            char *real = (char *)ctx->calloc(sz + 64, cnt);
+        shift_zero(void *ctx, size_t sz, size_t cnt) {
+            SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *) ctx;
+
+            char *real = (char *)funcs->calloc(sz + 64, cnt);
             if (real == NULL) {
                 return NULL;
             }
@@ -61,7 +66,9 @@ def get_module(tmp_path):
             return (void *)(real + 64);
         }
         NPY_NO_EXPORT void
-        shift_free(Allocator *ctx, void * p, npy_uintp sz) {
+        shift_free(void *ctx, void * p, npy_uintp sz) {
+            SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *) ctx;
+
             if (p == NULL) {
                 return ;
             }
@@ -70,34 +77,36 @@ def get_module(tmp_path):
                 fprintf(stdout, "uh-oh, unmatched shift_free, "
                         "no appropriate prefix\\n");
                 /* Make C runtime crash by calling free on the wrong address */
-                ctx->free((char *)p + 10);
-                /* ctx->free(real); */
+                funcs->free((char *)p + 10);
+                /* funcs->free(real); */
             }
             else {
                 npy_uintp i = (npy_uintp)atoi(real +20);
                 if (i != sz) {
                     fprintf(stderr, "uh-oh, unmatched shift_free"
                             "(ptr, %ld) but allocated %ld\\n", sz, i);
                     /* This happens in some places, only print */
-                    ctx->free(real);
+                    funcs->free(real);
                 }
                 else {
-                    ctx->free(real);
+                    funcs->free(real);
                 }
             }
         }
         NPY_NO_EXPORT void *
-        shift_realloc(Allocator *ctx, void * p, npy_uintp sz) {
+        shift_realloc(void *ctx, void * p, npy_uintp sz) {
+            SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *) ctx;
+
             if (p != NULL) {
                 char *real = (char *)p - 64;
                 if (strncmp(real, "originally allocated", 20) != 0) {
                     fprintf(stdout, "uh-oh, unmatched shift_realloc\\n");
                     return realloc(p, sz);
                 }
-                return (void *)((char *)ctx->realloc(real, sz + 64) + 64);
+                return (void *)((char *)funcs->realloc(real, sz + 64) + 64);
             }
             else {
-                char *real = (char *)ctx->realloc(p, sz + 64);
+                char *real = (char *)funcs->realloc(p, sz + 64);
                 if (real == NULL) {
                     return NULL;
                 }
@@ -106,20 +115,21 @@ def get_module(tmp_path):
                 return (void *)(real + 64);
             }
         }
-        static Allocator new_handler_ctx = {
+        /* As an example, we use the standard {m|c|re}alloc/free funcs. */
+        static SecretDataAllocatorFuncs secret_data_handler_ctx = {
             malloc,
             calloc,
             realloc,
             free
         };
-        static PyDataMem_Handler new_handler = {
+        static PyDataMem_Handler secret_data_handler = {
             "secret_data_allocator",
             {
-                &new_handler_ctx,
-                shift_alloc,      /* malloc */
-                shift_zero, /* calloc */
-                shift_realloc,      /* realloc */
-                shift_free       /* free */
+                &secret_data_handler_ctx, /* ctx */
+                shift_alloc,              /* malloc */
+                shift_zero,               /* calloc */
+                shift_realloc,            /* realloc */
+                shift_free                /* free */
             }
         };
         '''
@@ -138,36 +148,124 @@ def get_module(tmp_path):
 
 
 def test_set_policy(get_module):
-    a = np.arange(10)
-    orig_policy = get_module.test_prefix(a)
-    assert orig_policy == np.core.multiarray.get_handler_name()
-    assert orig_policy == np.core.multiarray.get_handler_name(a)
-    assert get_module.set_new_policy() == orig_policy
-    if orig_policy == 'default_allocator':
-        get_module.set_old_policy()
+    orig_policy_name = np.core.multiarray.get_handler_name()
+
+    a = np.arange(10).reshape((2, 5)) # a doesn't own its own data
+    assert np.core.multiarray.get_handler_name(a) == orig_policy_name
+
+    orig_policy = get_module.set_secret_data_policy()
+
+    b = np.arange(10).reshape((2, 5)) # b doesn't own its own data
+    assert np.core.multiarray.get_handler_name(b) == 'secret_data_allocator'
+
+    if orig_policy_name == 'default_allocator':
+        get_module.set_old_policy(None)
+
+        assert np.core.multiarray.get_handler_name() == 'default_allocator'
+    else:
+        get_module.set_old_policy(orig_policy)
+
+        assert np.core.multiarray.get_handler_name() == orig_policy_name
+
+
+async def concurrent_context1(get_module, event):
+    get_module.set_secret_data_policy()
+
+    assert np.core.multiarray.get_handler_name() == 'secret_data_allocator'
+
+    event.set()
+
+
+async def concurrent_context2(get_module, orig_policy_name, event):
+    await event.wait()
+
+    assert np.core.multiarray.get_handler_name() == orig_policy_name
+
+
+async def secret_data_context(get_module):
+    assert np.core.multiarray.get_handler_name() == 'secret_data_allocator'
+
+    get_module.set_old_policy(None)
+
+
+async def async_test_context_locality(get_module):
+    orig_policy_name = np.core.multiarray.get_handler_name()
+
+    event = asyncio.Event()
+    concurrent_task1 = asyncio.create_task(concurrent_context1(get_module, event))
+    concurrent_task2 = asyncio.create_task(concurrent_context2(get_module, orig_policy_name, event))
+    await concurrent_task1
+    await concurrent_task2
+
+    assert np.core.multiarray.get_handler_name() == orig_policy_name
+
+    orig_policy = get_module.set_secret_data_policy()
+
+    await asyncio.create_task(secret_data_context(get_module))
+
+    assert np.core.multiarray.get_handler_name() == 'secret_data_allocator'
+
+    get_module.set_old_policy(orig_policy)
+
+
+def test_context_locality(get_module):
+D889
    asyncio.run(async_test_context_locality(get_module))
+
+
+def concurrent_thread1(get_module, event):
+    assert np.core.multiarray.get_handler_name() == 'default_allocator'
+
+    get_module.set_secret_data_policy()
+
+    assert np.core.multiarray.get_handler_name() == 'secret_data_allocator'
+
+    event.set()
+
+
+def concurrent_thread2(get_module, event):
+    event.wait()
+
+    assert np.core.multiarray.get_handler_name() == 'default_allocator'
+
+
+def test_thread_locality(get_module):
+    orig_policy_name = np.core.multiarray.get_handler_name()
+
+    event = threading.Event()
+    concurrent_task1 = threading.Thread(target=concurrent_thread1, args=(get_module, event))
+    concurrent_task2 = threading.Thread(target=concurrent_thread2, args=(get_module, event))
+    concurrent_task1.start()
+    concurrent_task2.start()
+    concurrent_task1.join()
+    concurrent_task2.join()
+
+    assert np.core.multiarray.get_handler_name() == orig_policy_name
 
 
 @pytest.mark.slow
 def test_new_policy(get_module):
     a = np.arange(10)
-    orig_policy = get_module.test_prefix(a)
-    assert get_module.set_new_policy() == orig_policy
-    b = np.arange(10).reshape((2, 5))
-    assert get_module.test_prefix(b) == 'secret_data_allocator'
+    orig_policy_name = np.core.multiarray.get_handler_name(a)
+
+    orig_policy = get_module.set_secret_data_policy()
+
+    b = np.arange(10)
+    assert np.core.multiarray.get_handler_name(b) == 'secret_data_allocator'
 
     # test array manipulation. This is slow
-    if orig_policy == 'default_allocator':
+    if orig_policy_name == 'default_allocator':
         # when the np.core.test tests recurse into this test, the
         # policy will be set so this "if" will be false, preventing
         # infinite recursion
         #
         # if needed, debug this by
         # - running tests with -- -s (to not capture stdout/stderr
         # - setting extra_argv=['-vv'] here
-        np.core.test(verbose=2, extra_argv=['-vv'])
+        assert np.core.test('full', verbose=2, extra_argv=['-vv'])
         # also try the ma tests, the pickling test is quite tricky
-        np.ma.test(verbose=2, extra_argv=['-vv'])
-    get_module.set_old_policy()
-    assert get_module.test_prefix(a) == orig_policy
+        assert np.ma.test('full', verbose=2, extra_argv=['-vv'])
+
+    get_module.set_old_policy(orig_policy)
+
     c = np.arange(10)
-    assert get_module.test_prefix(c) == 'default_allocator'
+    assert np.core.multiarray.get_handler_name(c) == orig_policy_name