diff -r adb6b029b102 Include/pymem.h --- a/Include/pymem.h Wed Mar 09 15:02:31 2016 +0100 +++ b/Include/pymem.h Thu Mar 10 17:05:33 2016 +0100 @@ -16,6 +16,23 @@ PyAPI_FUNC(void *) PyMem_RawMalloc(size_ PyAPI_FUNC(void *) PyMem_RawCalloc(size_t nelem, size_t elsize); PyAPI_FUNC(void *) PyMem_RawRealloc(void *ptr, size_t new_size); PyAPI_FUNC(void) PyMem_RawFree(void *ptr); + +/* Track an allocated memory block in the tracemalloc module. + Return 0 on success, return -1 on error (failed to allocate memory to store + the trace). + + Do nothing (return 0) if tracemalloc is not tracing Python memory + allocations. + + If the pointer was already tracked, remove the old trace and add a new + trace. */ +PyAPI_FUNC(int) _PyTraceMalloc_Track(void *ptr, size_t size); + +/* Untrack an allocated memory block in the tracemalloc module. + Do nothing if the block was not tracked. + + Do nothing if tracemalloc is not tracing Python memory allocations. */ +PyAPI_FUNC(void) _PyTraceMalloc_Untrack(void *ptr); #endif diff -r adb6b029b102 Lib/test/test_tracemalloc.py --- a/Lib/test/test_tracemalloc.py Wed Mar 09 15:02:31 2016 +0100 +++ b/Lib/test/test_tracemalloc.py Thu Mar 10 17:05:33 2016 +0100 @@ -11,9 +11,15 @@ try: import threading except ImportError: threading = None +try: + import _testcapi +except ImportError: + _testcapi = None + EMPTY_STRING_SIZE = sys.getsizeof(b'') + def get_frames(nframe, lineno_delta): frames = [] frame = sys._getframe(1) @@ -816,12 +822,103 @@ class TestCommandLine(unittest.TestCase) assert_python_ok('-X', 'tracemalloc', '-c', code) +@unittest.skipIf(_testcapi is None, 'need _testcapi') +class TestCAPI(unittest.TestCase): + def setUp(self): + if tracemalloc.is_tracing(): + self.skipTest("tracemalloc must be stopped before the test") + + # allocate an object with tracemalloc disabled + self.obj = object() + + # for the type "object", id(obj) is the address of its memory block. + # This type is not tracked by the garbage collector + self.ptr = id(self.obj) + + # use a fake size + self.size = 123456 + + def tearDown(self): + tracemalloc.stop() + + def check_track(self, release_gil): + nframe = 5 + + # Call tracemalloc_track() just after start() and just before + # get_traced_memory() to only track our pointer and be able to check + # the size parameter + frames = get_frames(nframe, 2) + tracemalloc.start(nframe) + _testcapi.tracemalloc_track(self.ptr, self.size, release_gil) + size = tracemalloc.get_traced_memory()[0] + + self.assertEqual(size, self.size) + + self.assertEqual(tracemalloc.get_object_traceback(self.obj), + tracemalloc.Traceback(frames)) + + def test_track(self): + self.check_track(False) + + def test_track_without_gil(self): + # check that calling _PyTraceMalloc_Track() without holding the GIL + # works too + self.check_track(True) + + def test_track_already_tracked(self): + nframe = 5 + tracemalloc.start(nframe) + + # track a first time + _testcapi.tracemalloc_track(self.ptr, self.size) + + # calling _PyTraceMalloc_Track() must remove the old trace and add + # a new trace with the new traceback + frames = get_frames(nframe, 1) + _testcapi.tracemalloc_track(self.ptr, self.size) + self.assertEqual(tracemalloc.get_object_traceback(self.obj), + tracemalloc.Traceback(frames)) + + def test_untrack(self): + tracemalloc.start() + + _testcapi.tracemalloc_track(self.ptr, self.size) + self.assertIsNotNone(tracemalloc.get_object_traceback(self.obj)) + + # untrack must remove the trace + _testcapi.tracemalloc_untrack(self.ptr) + self.assertIsNone(tracemalloc.get_object_traceback(self.obj)) + + # calling _PyTraceMalloc_Untrack() multiple times must not crash + _testcapi.tracemalloc_untrack(self.ptr) + _testcapi.tracemalloc_untrack(self.ptr) + + def test_stop_track(self): + tracemalloc.start() + tracemalloc.stop() + + # Calling _PyTraceMalloc_Track() with tracemalloc disabled must do + # nothing (and not crash) + _testcapi.tracemalloc_track(self.ptr, self.size) + self.assertIsNone(tracemalloc.get_object_traceback(self.obj)) + + def test_stop_untrack(self): + tracemalloc.start() + _testcapi.tracemalloc_track(self.ptr, self.size) + + tracemalloc.stop() + # Calling _PyTraceMalloc_Untrack() with tracemalloc disabled must do + # nothing (and not crash) + _testcapi.tracemalloc_untrack(self.ptr) + + def test_main(): support.run_unittest( TestTracemallocEnabled, TestSnapshot, TestFilters, TestCommandLine, + TestCAPI, ) if __name__ == "__main__": diff -r adb6b029b102 Modules/_testcapimodule.c --- a/Modules/_testcapimodule.c Wed Mar 09 15:02:31 2016 +0100 +++ b/Modules/_testcapimodule.c Thu Mar 10 17:05:33 2016 +0100 @@ -3616,6 +3616,49 @@ get_recursion_depth(PyObject *self, PyOb return PyLong_FromLong(tstate->recursion_depth - 1); } +static PyObject * +tracemalloc_track(PyObject *self, PyObject *args) +{ + PyObject *ptr_obj; + void *ptr; + Py_ssize_t size; + int release_gil = 0; + int res; + + if (!PyArg_ParseTuple(args, "On|i", &ptr_obj, &size, &release_gil)) + return NULL; + ptr = PyLong_AsVoidPtr(ptr_obj); + if (PyErr_Occurred()) + return NULL; + + if (release_gil) { + Py_BEGIN_ALLOW_THREADS + res = _PyTraceMalloc_Track(ptr, size); + Py_END_ALLOW_THREADS + } + else { + res = _PyTraceMalloc_Track(ptr, size); + } + return PyLong_FromLong(res); +} + +static PyObject * +tracemalloc_untrack(PyObject *self, PyObject *args) +{ + PyObject *ptr_obj; + void *ptr; + + if (!PyArg_ParseTuple(args, "O", &ptr_obj)) + return NULL; + ptr = PyLong_AsVoidPtr(ptr_obj); + if (PyErr_Occurred()) + return NULL; + + _PyTraceMalloc_Untrack(ptr); + + Py_RETURN_NONE; +} + static PyMethodDef TestMethods[] = { {"raise_exception", raise_exception, METH_VARARGS}, @@ -3798,6 +3841,8 @@ static PyMethodDef TestMethods[] = { {"PyTime_AsMilliseconds", test_PyTime_AsMilliseconds, METH_VARARGS}, {"PyTime_AsMicroseconds", test_PyTime_AsMicroseconds, METH_VARARGS}, {"get_recursion_depth", get_recursion_depth, METH_NOARGS}, + {"tracemalloc_track", tracemalloc_track, METH_VARARGS}, + {"tracemalloc_untrack", tracemalloc_untrack, METH_VARARGS}, {NULL, NULL} /* sentinel */ }; diff -r adb6b029b102 Modules/_tracemalloc.c --- a/Modules/_tracemalloc.c Wed Mar 09 15:02:31 2016 +0100 +++ b/Modules/_tracemalloc.c Thu Mar 10 17:05:33 2016 +0100 @@ -439,6 +439,7 @@ tracemalloc_add_trace(void *ptr, size_t trace_t trace; int res; + assert(tracemalloc_config.tracing); #ifdef WITH_THREAD assert(PyGILState_Check()); #endif @@ -466,6 +467,8 @@ tracemalloc_remove_trace(void *ptr) { trace_t trace; + assert(tracemalloc_config.tracing); + if (_Py_hashtable_pop(tracemalloc_traces, ptr, &trace, sizeof(trace))) { assert(tracemalloc_traced_memory >= trace.size); tracemalloc_traced_memory -= trace.size; @@ -1395,6 +1398,51 @@ parse_sys_xoptions(PyObject *value) } int +_PyTraceMalloc_Track(void *ptr, size_t size) +{ + int res; +#ifdef WITH_THREAD + PyGILState_STATE gil_state; + + if (!tracemalloc_config.tracing) { + /* tracemalloc is not tracing: do nothing */ + return 0; + } + + gil_state = PyGILState_Ensure(); +#endif + + TABLES_LOCK(); + + /* don't trust the caller: remove the existing trace + if it already exists. tracemalloc_add_trace() requires that the ptr + is not already tracked. */ + tracemalloc_remove_trace(ptr); + + res = tracemalloc_add_trace(ptr, size); + + TABLES_UNLOCK(); + +#ifdef WITH_THREAD + PyGILState_Release(gil_state); +#endif + return res; +} + +void +_PyTraceMalloc_Untrack(void *ptr) +{ + if (!tracemalloc_config.tracing) { + /* tracemalloc is not tracing: do nothing */ + return; + } + + TABLES_LOCK(); + tracemalloc_remove_trace(ptr); + TABLES_UNLOCK(); +} + +int _PyTraceMalloc_Init(void) { char *p;