diff -r 9bbada125b3f Doc/library/tracemalloc.rst --- a/Doc/library/tracemalloc.rst Wed Nov 27 09:18:54 2013 -0600 +++ b/Doc/library/tracemalloc.rst Wed Nov 27 18:57:25 2013 +0100 @@ -295,13 +295,17 @@ Functions See also :func:`start` and :func:`stop` functions. -.. function:: start(nframe: int=1) +.. function:: start(nframe: int=1, memory_limit: int=0) Start tracing Python memory allocations: install hooks on Python memory allocators. Collected tracebacks of traces will be limited to *nframe* frames. By default, a trace of a memory block only stores the most recent frame: the limit is ``1``. *nframe* must be greater or equal to ``1``. + If *memory_limit* is greater than 0, a memory allocation fails with a + :exc:`MemoryError` if it would make the traced memory greater than + *memory_limit* bytes. See the :func:`get_traced_memory` function. + Storing more than ``1`` frame is only useful to compute statistics grouped by ``'traceback'`` or to compute cumulative statistics: see the :meth:`Snapshot.compare_to` and :meth:`Snapshot.statistics` methods. diff -r 9bbada125b3f Lib/test/test_tracemalloc.py --- a/Lib/test/test_tracemalloc.py Wed Nov 27 09:18:54 2013 -0600 +++ b/Lib/test/test_tracemalloc.py Wed Nov 27 18:57:25 2013 +0100 @@ -110,13 +110,15 @@ class TestTracemallocEnabled(unittest.Te self.assertRaises(ValueError, tracemalloc.start, -1) tracemalloc.stop() - tracemalloc.start(10) + # test 10 frames with keyword syntax + tracemalloc.start(nframe=10) obj2, obj2_traceback = allocate_bytes(obj_size) traceback = tracemalloc.get_object_traceback(obj2) self.assertEqual(len(traceback), 10) self.assertEqual(traceback, obj2_traceback) tracemalloc.stop() + # test 1 frame with indexed parameter tracemalloc.start(1) obj, obj_traceback = allocate_bytes(obj_size) traceback = tracemalloc.get_object_traceback(obj) @@ -288,6 +290,13 @@ class TestTracemallocEnabled(unittest.Te exitcode = os.WEXITSTATUS(status) self.assertEqual(exitcode, 0) + def test_memory_limit(self): + obj_size = 1024 * 1024 + tracemalloc.stop() + tracemalloc.start(memory_limit=obj_size // 2) + self.assertRaises(MemoryError, allocate_bytes, obj_size) + obj, obj_traceback = allocate_bytes(obj_size // 4) + class TestSnapshot(unittest.TestCase): maxDiff = 4000 diff -r 9bbada125b3f Modules/_tracemalloc.c --- a/Modules/_tracemalloc.c Wed Nov 27 09:18:54 2013 -0600 +++ b/Modules/_tracemalloc.c Wed Nov 27 18:57:25 2013 +0100 @@ -51,7 +51,11 @@ static struct { /* limit of the number of frames in a traceback, 1 by default. Variable protected by the GIL. */ int max_nframe; -} tracemalloc_config = {TRACEMALLOC_NOT_INITIALIZED, 0, 0, 1}; + + /* limit of the number of frames in a traceback, 1 by default. + Variable protected by the GIL. */ + Py_ssize_t memory_limit; +} tracemalloc_config = {TRACEMALLOC_NOT_INITIALIZED, 0, 0, 1, 0}; #if defined(TRACE_RAW_MALLOC) && defined(WITH_THREAD) /* This lock is needed because tracemalloc_free() is called without @@ -466,12 +470,24 @@ tracemalloc_log_alloc(void *ptr, size_t tracemalloc_traced_memory += size; if (tracemalloc_traced_memory > tracemalloc_max_traced_memory) tracemalloc_max_traced_memory = tracemalloc_traced_memory; + assert((tracemalloc_config.memory_limit < 1) + || (tracemalloc_traced_memory <= (size_t)tracemalloc_config.memory_limit)); } TABLES_UNLOCK(); return res; } +static int +tracemalloc_check_limit(size_t size) +{ + size_t free = (size_t)tracemalloc_config.memory_limit; + + assert(tracemalloc_traced_memory <= free); + free -= tracemalloc_traced_memory; + return (size <= (size_t)free); +} + static void tracemalloc_log_free(void *ptr) { @@ -498,6 +514,10 @@ tracemalloc_malloc(void *ctx, size_t siz return alloc->malloc(alloc->ctx, size); } + if (tracemalloc_config.memory_limit > 0 + && !tracemalloc_check_limit(size)) + return NULL; + /* Ignore reentrant call. PyObjet_Malloc() calls PyMem_Malloc() for allocations larger than 512 bytes. PyGILState_Ensure() may call PyMem_RawMalloc() indirectly which would call PyGILState_Ensure() if @@ -511,6 +531,7 @@ tracemalloc_malloc(void *ctx, size_t siz assert(gil_held); #endif #endif + ptr = alloc->malloc(alloc->ctx, size); if (ptr != NULL) { @@ -538,6 +559,7 @@ tracemalloc_realloc(void *ctx, void *ptr PyGILState_STATE gil_state; #endif void *ptr2; + size_t old_size; if (get_reentrant()) { /* Reentrant call to PyMem_Realloc() and PyMem_RawRealloc(). @@ -565,7 +587,27 @@ tracemalloc_realloc(void *ctx, void *ptr assert(gil_held); #endif #endif - ptr2 = alloc->realloc(alloc->ctx, ptr, new_size); + + if (tracemalloc_config.memory_limit > 0) { + old_size = 0; + if (ptr != NULL) { + trace_t trace; + + TABLES_LOCK(); + if (_Py_HASHTABLE_GET(tracemalloc_traces, ptr, trace)) + old_size = trace.size; + TABLES_UNLOCK(); + } + + if (new_size <= old_size + || tracemalloc_check_limit(new_size - old_size)) + ptr2 = alloc->realloc(alloc->ctx, ptr, new_size); + else + ptr2 = NULL; + } + else { + ptr2 = alloc->realloc(alloc->ctx, ptr, new_size); + } if (ptr2 != NULL) { if (ptr != NULL) { @@ -1185,17 +1227,23 @@ done: } PyDoc_STRVAR(tracemalloc_start_doc, - "start(nframe: int=1)\n" + "start(nframe: int=1, memory_limit: int=0)\n" "\n" "Start tracing Python memory allocations. Set also the maximum number \n" - "of frames stored in the traceback of a trace to nframe."); + "of frames stored in the traceback of a trace to nframe.\n" + "\n" + "If memory_limit is greater than 0, memory allocations will fail\n" + "with a MemoryError if the traced memory would become greater than\n" + "memory_limit."); static PyObject* -py_tracemalloc_start(PyObject *self, PyObject *args) +py_tracemalloc_start(PyObject *self, PyObject *args, PyObject *kwargs) { - Py_ssize_t nframe = 1; + static char *kwlist[] = {"nframe", "memory_limit", NULL}; + Py_ssize_t nframe=1, memory_limit=0; - if (!PyArg_ParseTuple(args, "|n:start", &nframe)) + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|nn:start", kwlist, + &nframe, &memory_limit)) return NULL; if (nframe < 1 || nframe > MAX_NFRAME) { @@ -1205,6 +1253,7 @@ py_tracemalloc_start(PyObject *self, PyO return NULL; } tracemalloc_config.max_nframe = Py_SAFE_DOWNCAST(nframe, Py_ssize_t, int); + tracemalloc_config.memory_limit = memory_limit; if (tracemalloc_start() < 0) return NULL; @@ -1298,7 +1347,7 @@ static PyMethodDef module_methods[] = { {"_get_object_traceback", (PyCFunction)py_tracemalloc_get_object_traceback, METH_O, tracemalloc_get_object_traceback_doc}, {"start", (PyCFunction)py_tracemalloc_start, - METH_VARARGS, tracemalloc_start_doc}, + METH_VARARGS | METH_KEYWORDS, tracemalloc_start_doc}, {"stop", (PyCFunction)py_tracemalloc_stop, METH_NOARGS, tracemalloc_stop_doc}, {"get_traceback_limit", (PyCFunction)py_tracemalloc_get_traceback_limit,