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:43:26 2016 +0100 @@ -16,6 +16,22 @@ 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 memory block is already tracked, update the existing 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:43:26 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:43:26 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:43:26 2016 +0100 @@ -432,44 +432,64 @@ traceback_new(void) return traceback; } +static void +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; + } +} + static int tracemalloc_add_trace(void *ptr, size_t size) { traceback_t *traceback; - trace_t trace; int res; + _Py_hashtable_entry_t *entry; + assert(tracemalloc_config.tracing); #ifdef WITH_THREAD assert(PyGILState_Check()); #endif traceback = traceback_new(); - if (traceback == NULL) + if (traceback == NULL) { + /* on error, remove the old trace */ + tracemalloc_remove_trace(ptr); return -1; - - trace.size = size; - trace.traceback = traceback; - - res = _Py_HASHTABLE_SET(tracemalloc_traces, ptr, trace); - if (res == 0) { - assert(tracemalloc_traced_memory <= PY_SIZE_MAX - size); - tracemalloc_traced_memory += size; - if (tracemalloc_traced_memory > tracemalloc_peak_traced_memory) - tracemalloc_peak_traced_memory = tracemalloc_traced_memory; } - return res; -} + entry = _Py_hashtable_get_entry(tracemalloc_traces, ptr); + if (entry != NULL) { + /* a trace already exists: update the trace */ + trace_t *ptrace = (trace_t *)_Py_HASHTABLE_ENTRY_DATA(entry); -static void -tracemalloc_remove_trace(void *ptr) -{ - trace_t trace; + assert(tracemalloc_traced_memory >= ptrace->size); + tracemalloc_traced_memory -= ptrace->size; - if (_Py_hashtable_pop(tracemalloc_traces, ptr, &trace, sizeof(trace))) { - assert(tracemalloc_traced_memory >= trace.size); - tracemalloc_traced_memory -= trace.size; + ptrace->size = size; + ptrace->traceback = traceback; } + else { + trace_t trace; + trace.size = size; + trace.traceback = traceback; + + res = _Py_HASHTABLE_SET(tracemalloc_traces, ptr, trace); + if (res < 0) + return -1; + } + + assert(tracemalloc_traced_memory <= PY_SIZE_MAX - size); + tracemalloc_traced_memory += size; + if (tracemalloc_traced_memory > tracemalloc_peak_traced_memory) + tracemalloc_peak_traced_memory = tracemalloc_traced_memory; + return 0; } static void* @@ -512,7 +532,9 @@ tracemalloc_realloc(void *ctx, void *ptr /* an existing memory block has been resized */ TABLES_LOCK(); - tracemalloc_remove_trace(ptr); + + if (ptr != ptr2) + tracemalloc_remove_trace(ptr); if (tracemalloc_add_trace(ptr2, new_size) < 0) { /* Memory allocation failed. The error cannot be reported to @@ -1395,6 +1417,46 @@ parse_sys_xoptions(PyObject *value) } int +_PyTraceMalloc_Track(void *ptr, size_t size) +{ + int res; +#ifdef WITH_THREAD + PyGILState_STATE gil_state; +#endif + + if (!tracemalloc_config.tracing) { + /* tracemalloc is not tracing: do nothing */ + return 0; + } + +#ifdef WITH_THREAD + gil_state = PyGILState_Ensure(); +#endif + + TABLES_LOCK(); + 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;