Add adaptive trace buffer growth option

Author: akx

docs/source/viztracer.rst

@@ -20,6 +20,8 @@ VizTracer
                  log_torch=False,\
                  log_audit=False,\
                  pid_suffix=False,\
+                 grow_on_overflow=False,\
+                 max_tracer_entries=None,\
                  file_info=True,\
                  register_global=True,\
                  trace_self=False,\
@@ -266,6 +268,26 @@ VizTracer
 
             viztracer --pid_suffix
 
+    .. py:attribute:: grow_on_overflow
+        :type: bool
+        :value: False
+
+        Whether grow the internal trace buffer instead of overwriting old events
+        when it becomes full.
+
+        Equivalent to
+
+        .. code-block::
+
+            viztracer --grow_on_overflow
+
+    .. py:attribute:: max_tracer_entries
+        :type: Optional[int]
+        :value: None
+
+        Maximum buffer size to grow to when ``grow_on_overflow`` is enabled.
+        ``None`` means there is no explicit cap.
+
     .. py:attribute:: file_info
         :type: bool
         :value: False

src/viztracer/main.py

@@ -75,6 +75,19 @@ def create_parser(self) -> argparse.ArgumentParser:
             default=1000000,
             help="size of circular buffer. How many entries can it store",
         )
+        parser.add_argument(
+            "--grow_on_overflow",
+            action="store_true",
+            default=False,
+            help="grow the trace buffer instead of overwriting old entries when it fills",
+        )
+        parser.add_argument(
+            "--max_tracer_entries",
+            nargs="?",
+            type=int,
+            default=None,
+            help="maximum buffer size to grow to when --grow_on_overflow is enabled",
+        )
         filename_group = parser.add_mutually_exclusive_group()
         filename_group.add_argument(
             "--output_file",
@@ -500,6 +513,12 @@ def parse(self, argv: list[str]) -> VizProcedureResult:
                 f"Can't convert {options.min_duration} to time. Format should be 0.3ms or 13us",
             )
 
+        if (
+            options.max_tracer_entries is not None
+            and options.max_tracer_entries < options.tracer_entries
+        ):
+            return False, "--max_tracer_entries can't be smaller than --tracer_entries"
+
         if options.log_torch:
             try:
                 import torch  # type: ignore  # noqa: F401
@@ -543,6 +562,8 @@ def parse(self, argv: list[str]) -> VizProcedureResult:
             "log_audit": options.log_audit,
             "log_torch": options.log_torch,
             "pid_suffix": options.pid_suffix,
+            "grow_on_overflow": options.grow_on_overflow,
+            "max_tracer_entries": options.max_tracer_entries,
             "file_info": False,
             "register_global": True,
             "report_endpoint": options.report_endpoint,

src/viztracer/modules/snaptrace.c

@@ -65,25 +65,86 @@ get_ts()
 #endif
 }
 
+static inline long
+get_buffered_entry_count(TracerObject* self)
+{
+    if (self->buffer_tail_idx >= self->buffer_head_idx) {
+        return self->buffer_tail_idx - self->buffer_head_idx;
+    }
+    return self->buffer_size - self->buffer_head_idx + self->buffer_tail_idx;
+}
+
+static int
+grow_buffer(TracerObject* self)
+{
+    long curr_entries = self->buffer_size - 1;
+    long new_entries = curr_entries * 2;
+    long buffered_entries = get_buffered_entry_count(self);
+    struct EventNode* new_buffer = NULL;
+
+    if (self->max_buffer_entries > 0 && new_entries > self->max_buffer_entries) {
+        new_entries = self->max_buffer_entries;
+    }
+
+    if (new_entries <= curr_entries) {
+        return -1;
+    }
+
+    new_buffer = (struct EventNode*) PyMem_Calloc(new_entries + 1, sizeof(struct EventNode));
+    if (!new_buffer) {
+        return -1;
+    }
+
+    for (long i = 0; i < buffered_entries; i++) {
+        long src_idx = self->buffer_head_idx + i;
+        if (src_idx >= self->buffer_size) {
+            src_idx -= self->buffer_size;
+        }
+        new_buffer[i] = self->buffer[src_idx];
+    }
+
+    PyMem_FREE(self->buffer);
+    self->buffer = new_buffer;
+    self->buffer_size = new_entries + 1;
+    self->buffer_head_idx = 0;
+    self->buffer_tail_idx = buffered_entries;
+    self->total_entries = buffered_entries;
+
+    return 0;
+}
+
 static inline struct EventNode*
 get_next_node(TracerObject* self)
 {
     struct EventNode* node = NULL;
+    long next_tail_idx = 0;
+    int overwrote = 0;
 
     SNAPTRACE_THREAD_PROTECT_START(self);
-    node = self->buffer + self->buffer_tail_idx;
     // This is actually faster than modulo
-    self->buffer_tail_idx = self->buffer_tail_idx + 1;
-    if (self->buffer_tail_idx >= self->buffer_size) {
-        self->buffer_tail_idx = 0;
-    }
-    if (self->buffer_tail_idx == self->buffer_head_idx) {
-        self->buffer_head_idx = self->buffer_head_idx + 1;
-        if (self->buffer_head_idx >= self->buffer_size) {
-            self->buffer_head_idx = 0;
+    next_tail_idx = self->buffer_tail_idx + 1;
+    if (next_tail_idx >= self->buffer_size) {
+        next_tail_idx = 0;
+    }
+    if (next_tail_idx == self->buffer_head_idx) {
+        if (!(self->grow_on_full && grow_buffer(self) == 0)) {
+            overwrote = 1;
+            self->overflowed = 1;
+            clear_node(self->buffer + next_tail_idx);
+            self->buffer_head_idx = next_tail_idx + 1;
+            if (self->buffer_head_idx >= self->buffer_size) {
+                self->buffer_head_idx = 0;
+            }
+        } else {
+            next_tail_idx = self->buffer_tail_idx + 1;
+            if (next_tail_idx >= self->buffer_size) {
+                next_tail_idx = 0;
+            }
         }
-        clear_node(self->buffer + self->buffer_tail_idx);
-    } else {
+    }
+    node = self->buffer + self->buffer_tail_idx;
+    self->buffer_tail_idx = next_tail_idx;
+    if (!overwrote) {
         self->total_entries += 1;
     }
     SNAPTRACE_THREAD_PROTECT_END(self);
@@ -1072,6 +1133,11 @@ tracer_start(TracerObject* self, PyObject* Py_UNUSED(unused))
         curr_tracer = self;
     }
 
+    if (self->buffer_head_idx == self->buffer_tail_idx) {
+        self->total_entries = 0;
+        self->overflowed = 0;
+    }
+
     self->collecting = 1;
 #if PY_VERSION_HEX >= 0x030C0000
     if (enable_monitoring(self) != 0) {
@@ -1493,7 +1559,6 @@ tracer_dump(TracerObject* self, PyObject* args, PyObject* kw)
     SNAPTRACE_THREAD_PROTECT_START(self);
     struct EventNode* curr = self->buffer + self->buffer_head_idx;
     unsigned long pid = 0;
-    uint8_t overflowed = ((self->buffer_tail_idx + 1) % self->buffer_size) == self->buffer_head_idx;
     struct MetadataNode* metadata_node = NULL;
     PyObject* task_dict = NULL;
 
@@ -1681,7 +1746,7 @@ tracer_dump(TracerObject* self, PyObject* args, PyObject* kw)
 
     self->buffer_tail_idx = self->buffer_head_idx;
     fseek(fptr, -1, SEEK_CUR);
-    fprintf(fptr, "], \"viztracer_metadata\": {\"overflow\":%s", overflowed? "true": "false");
+    fprintf(fptr, "], \"viztracer_metadata\": {\"overflow\":%s", self->overflowed? "true": "false");
 
     if (self->sync_marker > 0)
     {
@@ -1708,6 +1773,8 @@ tracer_clear(TracerObject* self, PyObject* Py_UNUSED(unused))
         }
     }
     self->buffer_tail_idx = self->buffer_head_idx;
+    self->total_entries = 0;
+    self->overflowed = 0;
 
     Py_RETURN_NONE;
 }
@@ -2062,10 +2129,13 @@ Tracer_New(PyTypeObject* type, PyObject* args, PyObject* kwargs)
         self->collecting = 0;
         self->fix_pid = 0;
         self->total_entries = 0;
+        self->overflowed = 0;
+        self->grow_on_full = 0;
         self->check_flags = 0;
         self->verbose = 0;
         self->lib_file_path = NULL;
         self->max_stack_depth = 0;
+        self->max_buffer_entries = 0;
         self->include_files = NULL;
         self->exclude_files = NULL;
         self->min_duration = 0;

src/viztracer/modules/snaptrace.h

@@ -108,10 +108,13 @@ typedef struct {
     // otherwise it uses this pid
     long fix_pid;
     unsigned long total_entries;
+    int overflowed;
+    int grow_on_full;
     unsigned int check_flags;
     int verbose;
     char* lib_file_path;
     int max_stack_depth;
+    long max_buffer_entries;
     PyObject* process_name;
     PyObject* include_files;
     PyObject* exclude_files;

src/viztracer/modules/snaptrace_member.c

@@ -431,6 +431,71 @@ Tracer_trace_self_getter(TracerObject* self, void* closure)
     }
 }
 
+static int
+Tracer_grow_on_full_setter(TracerObject* self, PyObject* value, void* closure)
+{
+    if (value == NULL) {
+        PyErr_SetString(PyExc_AttributeError, "Cannot delete the attribute");
+        return -1;
+    }
+
+    if (!PyBool_Check(value)) {
+        PyErr_SetString(PyExc_TypeError, "grow_on_full must be a boolean");
+        return -1;
+    }
+
+    self->grow_on_full = (value == Py_True);
+
+    return 0;
+}
+
+static PyObject*
+Tracer_grow_on_full_getter(TracerObject* self, void* closure)
+{
+    if (self->grow_on_full) {
+        Py_RETURN_TRUE;
+    } else {
+        Py_RETURN_FALSE;
+    }
+}
+
+static int
+Tracer_max_buffer_entries_setter(TracerObject* self, PyObject* value, void* closure)
+{
+    if (value == NULL) {
+        PyErr_SetString(PyExc_AttributeError, "Cannot delete the attribute");
+        return -1;
+    }
+
+    if (!PyLong_Check(value)) {
+        PyErr_SetString(PyExc_TypeError, "max_buffer_entries must be an integer");
+        return -1;
+    }
+
+    self->max_buffer_entries = PyLong_AsLong(value);
+    if (self->max_buffer_entries < 0) {
+        self->max_buffer_entries = 0;
+    }
+
+    return 0;
+}
+
+static PyObject*
+Tracer_max_buffer_entries_getter(TracerObject* self, void* closure)
+{
+    return PyLong_FromLong(self->max_buffer_entries);
+}
+
+static PyObject*
+Tracer_overflowed_getter(TracerObject* self, void* closure)
+{
+    if (self->overflowed) {
+        Py_RETURN_TRUE;
+    } else {
+        Py_RETURN_FALSE;
+    }
+}
+
 static int
 Tracer_log_func_repr_setter(TracerObject* self, PyObject* value, void* closure)
 {
@@ -478,6 +543,9 @@ PyGetSetDef Tracer_getsetters[] = {
     {"log_func_args", (getter)Tracer_log_func_args_getter, (setter)Tracer_log_func_args_setter, "log_func_args", NULL},
     {"log_async", (getter)Tracer_log_async_getter, (setter)Tracer_log_async_setter, "log_async", NULL},
     {"trace_self", (getter)Tracer_trace_self_getter, (setter)Tracer_trace_self_setter, "trace_self", NULL},
+    {"grow_on_full", (getter)Tracer_grow_on_full_getter, (setter)Tracer_grow_on_full_setter, "grow_on_full", NULL},
+    {"max_buffer_entries", (getter)Tracer_max_buffer_entries_getter, (setter)Tracer_max_buffer_entries_setter, "max_buffer_entries", NULL},
+    {"overflowed", (getter)Tracer_overflowed_getter, NULL, "overflowed", NULL},
     {"log_func_repr", (getter)Tracer_log_func_repr_getter, (setter)Tracer_log_func_repr_setter, "log_func_repr", NULL},
     {NULL}
 };

src/viztracer/report_builder.py

@@ -336,6 +336,13 @@ def print_messages(self):
                         "WARNING",
                         '    If you need more buffer, use "viztracer --tracer_entries <entry_number>"',
                     )
+                    color_print(
+                        "WARNING",
+                        (
+                            "    Or, enable --grow_on_overflow to expand the buffer "
+                            "automatically as needed"
+                        ),
+                    )
                     color_print(
                         "WARNING",
                         "    Or, you can try the filter options to filter out some data you don't need",

src/viztracer/snaptrace.pyi

@@ -8,6 +8,9 @@ class Tracer:
 
     include_files: list[str] | None
     exclude_files: list[str] | None
+    grow_on_full: bool
+    max_buffer_entries: int
+    overflowed: bool
 
     def __init__(self, tracer_entries: int, /) -> None: ...
     def start(self) -> None: ...