gh-150494: Sampling mode for tracemalloc

Author: danielsn

Doc/c-api/init_config.rst

@@ -493,6 +493,10 @@ Configuration Options
      - :c:member:`tracemalloc <PyConfig.tracemalloc>`
      - ``int``
      - Read-only
+   * - ``"tracemalloc_sample_interval"``
+     - :c:member:`tracemalloc_sample_interval <PyConfig.tracemalloc_sample_interval>`
+     - ``int``
+     - Read-only
    * - ``"use_environment"``
      - :c:member:`use_environment <PyConfig.use_environment>`
      - ``bool``
@@ -1891,13 +1895,31 @@ PyConfig
 
       Enable tracemalloc?
 
-      If non-zero, call :func:`tracemalloc.start` at startup.
+      If non-zero, call :func:`tracemalloc.start` at startup with
+      :c:member:`tracemalloc` as the traceback limit and
+      :c:member:`tracemalloc_sample_interval` as the sampling interval.
 
-      Set by :option:`-X tracemalloc=N <-X>` command line option and by the
-      :envvar:`PYTHONTRACEMALLOC` environment variable.
+      Set by :option:`-X tracemalloc=NFRAME[:INTERVAL] <-X>` command line
+      option and by the :envvar:`PYTHONTRACEMALLOC` environment variable.
 
       Default: ``-1`` in Python mode, ``0`` in isolated mode.
 
+   .. c:member:: int tracemalloc_sample_interval
+
+      Set the :mod:`tracemalloc` sampling interval in bytes at startup.
+
+      If ``0``, every allocation is traced.  If greater than ``0``,
+      allocations are sampled using a Poisson process with a mean
+      inter-arrival of :c:member:`tracemalloc_sample_interval` bytes.
+
+      Only used when :c:member:`tracemalloc` is non-zero.
+
+      Set by the ``INTERVAL`` part of
+      :option:`-X tracemalloc=NFRAME:INTERVAL <-X>` command line option and
+      by the :envvar:`PYTHONTRACEMALLOC` environment variable.
+
+      Default: ``0``.
+
    .. c:member:: int perf_profiling
 
       Enable the Linux ``perf`` profiler support?

Doc/library/tracemalloc.rst

@@ -363,13 +363,29 @@ Functions
     See also :func:`start` and :func:`stop` functions.
 
 
-.. function:: start(nframe: int=1)
+.. function:: start(nframe: int=1, sample_interval: 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 *sample_interval* is ``0`` (the default), every allocation is traced.
+   If *sample_interval* is greater than zero, allocations are sampled using
+   a `Poisson process <https://en.wikipedia.org/wiki/Poisson_point_process>`_
+   with a mean inter-arrival of *sample_interval* bytes.  Sampling can
+   significantly reduce overhead while producing useful aggregate estimates.
+   In sampled mode, :attr:`Trace.size` is an upscaled estimate of the bytes
+   represented by the sample, while :attr:`Trace.real_size` is the actual
+   allocation size.
+
+   Both *nframe* and *sample_interval* can be passed by keyword:
+
+   .. code-block:: python
+
+      tracemalloc.start(sample_interval=512 * 1024)
+      tracemalloc.start(nframe=25, sample_interval=512 * 1024)
+
    You can still read the original number of total frames that composed the
    traceback by looking at the :attr:`Traceback.total_nframe` attribute.
 
@@ -382,12 +398,17 @@ Functions
    to measure how much memory is used by the :mod:`!tracemalloc` module.
 
    The :envvar:`PYTHONTRACEMALLOC` environment variable
-   (``PYTHONTRACEMALLOC=NFRAME``) and the :option:`-X` ``tracemalloc=NFRAME``
+   (``PYTHONTRACEMALLOC=NFRAME`` or ``PYTHONTRACEMALLOC=NFRAME:INTERVAL``)
+   and the :option:`-X` ``tracemalloc=NFRAME[:INTERVAL]``
    command line option can be used to start tracing at startup.
 
    See also :func:`stop`, :func:`is_tracing` and :func:`get_traceback_limit`
    functions.
 
+   .. versionchanged:: 3.16
+      Added the *sample_interval* parameter.  Both arguments can now be
+      passed by keyword.
+
 
 .. function:: stop()
 
@@ -697,7 +718,17 @@ Trace
 
    .. attribute:: size
 
-      Size of the memory block in bytes (``int``).
+      Size of the memory block in bytes (``int``).  When sampling is enabled,
+      this is an upscaled estimate of the total bytes the trace represents.
+      Use this value for aggregation.
+
+   .. attribute:: real_size
+
+      Actual allocation size in bytes (``int``).  In exact mode
+      (``sample_interval=0``), this equals :attr:`size`.  In sampled mode,
+      this is the real number of bytes requested by the allocation call.
+
+      .. versionadded:: 3.16
 
    .. attribute:: traceback
 

Doc/using/cmdline.rst

@@ -542,12 +542,17 @@ Miscellaneous options
    * ``-X tracemalloc`` to start tracing Python memory allocations using the
      :mod:`tracemalloc` module. By default, only the most recent frame is
      stored in a traceback of a trace. Use ``-X tracemalloc=NFRAME`` to start
-     tracing with a traceback limit of *NFRAME* frames.
+     tracing with a traceback limit of *NFRAME* frames.  Use
+     ``-X tracemalloc=NFRAME:INTERVAL`` to enable Poisson sampling with a
+     mean interval of *INTERVAL* bytes, which can reduce tracing overhead.
      See :func:`tracemalloc.start` and :envvar:`PYTHONTRACEMALLOC`
      for more information.
 
      .. versionadded:: 3.4
 
+     .. versionchanged:: 3.16
+        Added sampling support via ``NFRAME:INTERVAL`` syntax.
+
    * ``-X int_max_str_digits`` configures the :ref:`integer string conversion
      length limitation <int_max_str_digits>`.  See also
      :envvar:`PYTHONINTMAXSTRDIGITS`.
@@ -1033,12 +1038,17 @@ conflict.
    Python memory allocations using the :mod:`tracemalloc` module. The value of
    the variable is the maximum number of frames stored in a traceback of a
    trace. For example, ``PYTHONTRACEMALLOC=1`` stores only the most recent
-   frame.
+   frame.  Use ``PYTHONTRACEMALLOC=NFRAME:INTERVAL`` to enable Poisson
+   sampling with a mean interval of *INTERVAL* bytes (e.g.
+   ``PYTHONTRACEMALLOC=1:524288`` for 512 KB sampling).
    See the :func:`tracemalloc.start` function for more information.
    This is equivalent to setting the :option:`-X` ``tracemalloc`` option.
 
    .. versionadded:: 3.4
 
+   .. versionchanged:: 3.16
+      Added ``NFRAME:INTERVAL`` syntax for sampling support.
+
 
 .. envvar:: PYTHONPROFILEIMPORTTIME
 

Include/cpython/initconfig.h

@@ -141,6 +141,7 @@ typedef struct PyConfig {
     unsigned long hash_seed;
     int faulthandler;
     int tracemalloc;
+    int tracemalloc_sample_interval;
     int perf_profiling;
     int remote_debug;
     int import_time;

Include/internal/pycore_global_objects_fini_generated.h

@@ -666,6 +666,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(newline)
         STRUCT_FOR_ID(newlines)
         STRUCT_FOR_ID(next)
+        STRUCT_FOR_ID(nframe)
         STRUCT_FOR_ID(nlocals)
         STRUCT_FOR_ID(node_depth)
         STRUCT_FOR_ID(node_offset)
@@ -767,6 +768,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(reversed)
         STRUCT_FOR_ID(rounding)
         STRUCT_FOR_ID(salt)
+        STRUCT_FOR_ID(sample_interval)
         STRUCT_FOR_ID(sample_interval_us)
         STRUCT_FOR_ID(sched_priority)
         STRUCT_FOR_ID(scheduler)

Include/internal/pycore_global_strings.h

@@ -30,6 +30,10 @@ struct _PyTraceMalloc_Config {
     /* limit of the number of frames in a traceback, 1 by default.
        Variable protected by the GIL. */
     int max_nframe;
+
+    /* Poisson sampling interval in bytes. 0 means trace every allocation.
+       Variable protected by the GIL. */
+    size_t sample_interval;
 };
 
 
@@ -113,6 +117,7 @@ struct _tracemalloc_runtime_state {
             .initialized = TRACEMALLOC_NOT_INITIALIZED, \
             .tracing = 0, \
             .max_nframe = 1, \
+            .sample_interval = 0, \
         }, \
         .reentrant_key = Py_tss_NEEDS_INIT, \
     }
@@ -148,7 +153,7 @@ extern PyObject* _PyTraceMalloc_GetObjectTraceback(PyObject *obj);
 extern PyStatus _PyTraceMalloc_Init(void);
 
 /* Start tracemalloc */
-extern int _PyTraceMalloc_Start(int max_nframe);
+extern int _PyTraceMalloc_Start(int max_nframe, size_t sample_interval);
 
 /* Stop tracemalloc */
 extern void _PyTraceMalloc_Stop(void);

Include/internal/pycore_runtime_init_generated.h

@@ -23,6 +23,14 @@ struct _gc_thread_state {
 };
 #endif
 
+/* Per-thread tracemalloc Poisson sampling state.
+   Zero-initialized; prng_state == 0 means "needs initialization". */
+struct _tracemalloc_sampling_state {
+    size_t bytes_since_last_sample;
+    size_t threshold;
+    uint64_t prng_state;
+};
+
 
 // Every PyThreadState is actually allocated as a _PyThreadStateImpl. The
 // PyThreadState fields are exposed as part of the C API, although most fields
@@ -103,6 +111,8 @@ typedef struct _PyThreadStateImpl {
     struct _PyJitTracerState *jit_tracer_state;
 #endif
 
+    struct _tracemalloc_sampling_state tracemalloc_sampling;
+
 #ifdef Py_GIL_DISABLED
     // gh-144438: Add padding to ensure that the fields above don't share a
     // cache line with other allocations.