Skip to content

Commit

Permalink
Wrap hook functions with GIL, add example.
Browse files Browse the repository at this point in the history
Wraps the SetHook and calls to the hook with the GIL, to prevent races.
Adds an example of using the interface for callbacks into python code.
  • Loading branch information
thouis committed Jun 18, 2012
1 parent 3237093 commit 5b5a0f4
Show file tree
Hide file tree
Showing 5 changed files with 701 additions and 7 deletions.
33 changes: 26 additions & 7 deletions numpy/core/src/multiarray/multiarraymodule.c
Expand Up @@ -3672,22 +3672,29 @@ NPY_NO_EXPORT void *_PyDataMem_eventhook_user_data;
* Returns a pointer to the previous hook or NULL. If old_data is
* non-NULL, the previous user_data pointer will be copied to it.
*
*
* If not NULL, hook will be called at the end of each PyDataMem_NEW/FREE/RENEW:
* result = PyDataMem_NEW(size) -> (*hook)(NULL, result, size, user_data)
* PyDataMem_FREE(ptr) -> (*hook)(ptr, NULL, 0, user_data)
* result = PyDataMem_RENEW(ptr, size) -> (*hook)(ptr, result, size, user_data)
*
* When the hook is called, the GIL will be held by the calling
* thread. The hook should be written to be reentrant, if it performs
* operations that might cause new allocation events (such as the
* creation/descruction numpy objects, or creating/destroying Python
* objects which might cause a gc)
*/
NPY_NO_EXPORT PyDataMem_EventHookFunc *
PyDataMem_SetEventHook(PyDataMem_EventHookFunc *newhook,
void *user_data, void **old_data)
{
PyGILState_STATE gilstate = PyGILState_Ensure();
PyDataMem_EventHookFunc *temp = _PyDataMem_eventhook;
_PyDataMem_eventhook = newhook;
if (old_data != NULL) {
*old_data = _PyDataMem_eventhook_user_data;
}
_PyDataMem_eventhook_user_data = user_data;
PyGILState_Release(gilstate);
return temp;
}

Expand All @@ -3701,8 +3708,12 @@ PyDataMem_NEW(size_t size)

result = malloc(size);
if (_PyDataMem_eventhook != NULL) {
(*_PyDataMem_eventhook)(NULL, result, size,
_PyDataMem_eventhook_user_data);
PyGILState_STATE gilstate = PyGILState_Ensure();
if (_PyDataMem_eventhook != NULL) {
(*_PyDataMem_eventhook)(NULL, result, size,
_PyDataMem_eventhook_user_data);
}
PyGILState_Release(gilstate);
}
return (char *)result;
}
Expand All @@ -3715,8 +3726,12 @@ PyDataMem_FREE(void *ptr)
{
free(ptr);
if (_PyDataMem_eventhook != NULL) {
(*_PyDataMem_eventhook)(ptr, NULL, 0,
_PyDataMem_eventhook_user_data);
PyGILState_STATE gilstate = PyGILState_Ensure();
if (_PyDataMem_eventhook != NULL) {
(*_PyDataMem_eventhook)(ptr, NULL, 0,
_PyDataMem_eventhook_user_data);
}
PyGILState_Release(gilstate);
}
}

Expand All @@ -3730,8 +3745,12 @@ PyDataMem_RENEW(void *ptr, size_t size)

result = realloc(ptr, size);
if (_PyDataMem_eventhook != NULL) {
(*_PyDataMem_eventhook)(ptr, result, size,
_PyDataMem_eventhook_user_data);
PyGILState_STATE gilstate = PyGILState_Ensure();
if (_PyDataMem_eventhook != NULL) {
(*_PyDataMem_eventhook)(ptr, result, size,
_PyDataMem_eventhook_user_data);
}
PyGILState_Release(gilstate);
}
return (char *)result;
}
Expand Down
42 changes: 42 additions & 0 deletions tools/allocation_tracking/alloc_hook.pyx
@@ -0,0 +1,42 @@
# A cython wrapper for using python functions as callbacks for
# PyDataMem_SetEventHook.

cimport numpy as np

cdef extern from "Python.h":
object PyLong_FromVoidPtr(void *)
void *PyLong_AsVoidPtr(object)

ctypedef void PyDataMem_EventHookFunc(void *inp, void *outp, size_t size,
void *user_data)
cdef extern from "numpy/arrayobject.h":
PyDataMem_EventHookFunc * \
PyDataMem_SetEventHook(PyDataMem_EventHookFunc *newhook,
void *user_data, void **old_data)

np.import_array()

cdef void pyhook(void *old, void *new, size_t size, void *user_data):
cdef object pyfunc = <object> user_data
pyfunc(PyLong_FromVoidPtr(old),
PyLong_FromVoidPtr(new),
size)

class NumpyAllocHook(object):
def __init__(self, callback):
self.callback = callback

def __enter__(self):
cdef void *old_hook, *old_data
old_hook = <void *> \
PyDataMem_SetEventHook(<PyDataMem_EventHookFunc *> pyhook,
<void *> self.callback,
<void **> &old_data)
self.old_hook = PyLong_FromVoidPtr(old_hook)
self.old_data = PyLong_FromVoidPtr(old_data)

def __exit__(self):
PyDataMem_SetEventHook(<PyDataMem_EventHookFunc *> \
PyLong_AsVoidPtr(self.old_hook),
<void *> PyLong_AsVoidPtr(self.old_data),
<void **> 0)
9 changes: 9 additions & 0 deletions tools/allocation_tracking/setup.py
@@ -0,0 +1,9 @@
from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext
import numpy

setup(
cmdclass = {'build_ext': build_ext},
ext_modules = [Extension("alloc_hook", ["alloc_hook.pyx"],
include_dirs=[numpy.get_include()])])

0 comments on commit 5b5a0f4

Please sign in to comment.