bpo-40630: Add tracemalloc.reset_peak (GH-20102)

The reset_peak function sets the peak memory size to the current size,
representing a resetting of that metric. This allows for recording the
peak of specific sections of code, ignoring other code that may have
had a higher peak (since the most recent `tracemalloc.start()` or
tracemalloc.clear_traces()` call).

2 files changed, 46 insertions, 1 deletions

diff --git a/Modules/_tracemalloc.c b/Modules/_tracemalloc.c
index 4522d1a..5675716 100644
--- a/Modules/_tracemalloc.c
+++ b/Modules/_tracemalloc.c
@@ -1643,6 +1643,30 @@ _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.
+
+Do nothing if the tracemalloc module is not tracing memory allocations.
+
+[clinic start generated code]*/
+
+static PyObject *
+_tracemalloc_reset_peak_impl(PyObject *module)
+/*[clinic end generated code: output=140c2870f691dbb2 input=18afd0635066e9ce]*/
+{
+    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 +1678,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}
 };
diff --git a/Modules/clinic/_tracemalloc.c.h b/Modules/clinic/_tracemalloc.c.h
index 68fafdc..049cacd 100644
--- a/Modules/clinic/_tracemalloc.c.h
+++ b/Modules/clinic/_tracemalloc.c.h
@@ -197,4 +197,24 @@ _tracemalloc_get_traced_memory(PyObject *module, PyObject *Py_UNUSED(ignored))
 {
     return _tracemalloc_get_traced_memory_impl(module);
 }
-/*[clinic end generated code: output=1bc96dc569706afa input=a9049054013a1b77]*/
+
+PyDoc_STRVAR(_tracemalloc_reset_peak__doc__,
+"reset_peak($module, /)\n"
+"--\n"
+"\n"
+"Set the peak size of memory blocks traced by tracemalloc to the current size.\n"
+"\n"
+"Do nothing if the tracemalloc module is not tracing memory allocations.");
+
+#define _TRACEMALLOC_RESET_PEAK_METHODDEF    \
+    {"reset_peak", (PyCFunction)_tracemalloc_reset_peak, METH_NOARGS, _tracemalloc_reset_peak__doc__},
+
+static PyObject *
+_tracemalloc_reset_peak_impl(PyObject *module);
+
+static PyObject *
+_tracemalloc_reset_peak(PyObject *module, PyObject *Py_UNUSED(ignored))
+{
+    return _tracemalloc_reset_peak_impl(module);
+}
+/*[clinic end generated code: output=a130117b1af821da input=a9049054013a1b77]*/