bpo-40630: add tracemalloc.reset_peak

Doc/library/tracemalloc.rst

@@ -289,6 +289,26 @@ Functions
    :mod:`tracemalloc` module as a tuple: ``(current: int, peak: int)``.
 
 
+.. function:: reset_peak()
+
+   Set the peak size of memory blocks traced by the :mod:`tracemalloc` module
+   to the current size.
+
+   The peak memory usage of a specific section of code can be determined as
+   follows::
+
+     tracemalloc.reset_peak()
+
+     # some code
+
+     current, peak = tracemalloc.get_traced_memory()
+
+   ``peak`` will then be the peak size of traced memory blocks after the
+   :func:`reset_peak` call, even if there was a higher peak before that.
+
+   See also :func:`get_traced_memory`.
+
+
 .. function:: get_tracemalloc_memory()
 
    Get the memory usage in bytes of the :mod:`tracemalloc` module used to store

Lib/test/test_tracemalloc.py

@@ -246,6 +246,35 @@ def test_clear_traces(self):
         traceback2 = tracemalloc.get_object_traceback(obj)
         self.assertIsNone(traceback2)
 
+    def test_reset_peak(self):
+        max_error = 2048
+        def valid(size):
+            return range(size, size + max_error)
+
+        obj_size = 1024 * 1024
+
+        tracemalloc.clear_traces()
+
+        obj, obj_traceback = allocate_bytes(obj_size)
+        # add to the peak
+        allocate_bytes(obj_size)
+
+        size, peak_size = tracemalloc.get_traced_memory()
+        self.assertIn(peak_size, valid(2 * obj_size))
+
+        tracemalloc.reset_peak()
+
+        size2, peak_size2 = tracemalloc.get_traced_memory()
+        self.assertIn(size2, valid(size))
+        # the peak should include 'obj', but not the other allocation
+        self.assertIn(peak_size2, valid(size2))
+
+        # the peak should continue to be updated
+        allocate_bytes(obj_size)
+        size3, peak_size3 = tracemalloc.get_traced_memory()
+        self.assertIn(size3, valid(obj_size))
+        self.assertIn(peak_size3, valid(2 * obj_size))
+
     def test_is_tracing(self):
         tracemalloc.stop()
         self.assertFalse(tracemalloc.is_tracing())

Misc/ACKS

@@ -1860,6 +1860,7 @@ Alex Willmer
 David Wilson
 Geoff Wilson
 Greg V. Wilson
+Huon Wilson
 J Derek Wilson
 Paul Winkler
 Jody Winston

Misc/NEWS.d/next/Library/2020-05-15-13-40-15.bpo-40630.YXEX_M.rst

@@ -0,0 +1,2 @@
+Added :func:`tracemalloc.reset_peak` to set the peak size of traced memory
+blocks to the current size, to measure the peak of specific pieces of code.

Modules/_tracemalloc.c

@@ -1643,6 +1643,28 @@ _tracemalloc_get_traced_memory_impl(PyObject *module)
     return Py_BuildValue("nn", size, peak_size);
 }
 
+/*[clinic input]
+_tracemalloc.reset_peak
+
+Set the peak size of memory blocks traced by tracemalloc to the current size.
+
+[clinic start generated code]*/
+
+static PyObject *
+_tracemalloc_reset_peak_impl(PyObject *module)
+/*[clinic end generated code: output=140c2870f691dbb2 input=5a7e962ebebf9836]*/
+{
+    if (!_Py_tracemalloc_config.tracing) {
+        Py_RETURN_NONE;
+    }
+
+    TABLES_LOCK();
+    tracemalloc_peak_traced_memory = tracemalloc_traced_memory;
+    TABLES_UNLOCK();
+
+    Py_RETURN_NONE;
+}
+
 
 static PyMethodDef module_methods[] = {
     _TRACEMALLOC_IS_TRACING_METHODDEF
@@ -1654,6 +1676,7 @@ static PyMethodDef module_methods[] = {
     _TRACEMALLOC_GET_TRACEBACK_LIMIT_METHODDEF
     _TRACEMALLOC_GET_TRACEMALLOC_MEMORY_METHODDEF
     _TRACEMALLOC_GET_TRACED_MEMORY_METHODDEF
+    _TRACEMALLOC_RESET_PEAK_METHODDEF
     /* sentinel */
     {NULL, NULL}
 };