diff -r 3c34ab550358 Doc/library/tracemalloc.rst --- a/Doc/library/tracemalloc.rst Wed Dec 04 01:47:46 2013 +0100 +++ b/Doc/library/tracemalloc.rst Wed Dec 04 01:54:53 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 3c34ab550358 Lib/test/test_tracemalloc.py --- a/Lib/test/test_tracemalloc.py Wed Dec 04 01:47:46 2013 +0100 +++ b/Lib/test/test_tracemalloc.py Wed Dec 04 01:54:53 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 3c34ab550358 Modules/_tracemalloc.c --- a/Modules/_tracemalloc.c Wed Dec 04 01:47:46 2013 +0100 +++ b/Modules/_tracemalloc.c Wed Dec 04 01:54:53 2013 +0100 @@ -42,7 +42,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, 1}; + + /* Memory limit in bytes, 0 if disabled. + Variable protected by the GIL. */ + Py_ssize_t memory_limit; +} tracemalloc_config = {TRACEMALLOC_NOT_INITIALIZED, 0, 1, 0}; #if defined(TRACE_RAW_MALLOC) && defined(WITH_THREAD) /* This lock is needed because tracemalloc_free() is called without @@ -439,6 +443,16 @@ traceback_new(void) } 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 int tracemalloc_add_trace(void *ptr, size_t size) { traceback_t *traceback; @@ -462,6 +476,8 @@ tracemalloc_add_trace(void *ptr, size_t tracemalloc_traced_memory += size; if (tracemalloc_traced_memory> tracemalloc_peak_traced_memory) tracemalloc_peak_traced_memory = tracemalloc_traced_memory; + assert((tracemalloc_config.memory_limit < 1) + || (tracemalloc_traced_memory <= (size_t)tracemalloc_config.memory_limit)); } return res; @@ -484,6 +500,10 @@ tracemalloc_malloc(void *ctx, size_t siz PyMemAllocator *alloc = (PyMemAllocator *)ctx; void *ptr; + if (tracemalloc_config.memory_limit> 0 + && !tracemalloc_check_limit(size)) + return NULL; + ptr = alloc->malloc(alloc->ctx, size); if (ptr == NULL) return NULL; @@ -505,6 +525,22 @@ tracemalloc_realloc(void *ctx, void *ptr PyMemAllocator *alloc = (PyMemAllocator *)ctx; void *ptr2; + if (tracemalloc_config.memory_limit> 0) { + Py_ssize_t 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)) + return NULL; + } + ptr2 = alloc->realloc(alloc->ctx, ptr, new_size); if (ptr2 == NULL) return NULL; @@ -832,7 +868,7 @@ tracemalloc_deinit(void) } static int -tracemalloc_start(int max_nframe) +tracemalloc_start(int max_nframe, Py_ssize_t memory_limit) { PyMemAllocator alloc; size_t size; @@ -847,6 +883,7 @@ tracemalloc_start(int max_nframe) assert(1 <= max_nframe && max_nframe <= MAX_NFRAME); tracemalloc_config.max_nframe = max_nframe; + tracemalloc_config.memory_limit = memory_limit; /* allocate a buffer to store a new traceback */ size = TRACEBACK_SIZE(max_nframe); @@ -1181,18 +1218,21 @@ py_tracemalloc_get_object_traceback(PyOb } 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."); static PyObject* -py_tracemalloc_start(PyObject *self, PyObject *args) +py_tracemalloc_start(PyObject *self, PyObject *args, PyObject *kwargs) { + static char *kwlist[] = {"nframe", "memory_limit", NULL}; Py_ssize_t nframe = 1; int nframe_int; + Py_ssize_t memory_limit; - 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) { @@ -1203,7 +1243,7 @@ py_tracemalloc_start(PyObject *self, PyO } nframe_int = Py_SAFE_DOWNCAST(nframe, Py_ssize_t, int); - if (tracemalloc_start(nframe_int) < 0) + if (tracemalloc_start(nframe_int, memory_limit) < 0) return NULL; Py_RETURN_NONE; @@ -1295,7 +1335,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, @@ -1417,7 +1457,7 @@ int } } - return tracemalloc_start(nframe); + return tracemalloc_start(nframe, 0); } void

AltStyle によって変換されたページ (->オリジナル) /