Fix max_stack_depth when combined with other filters (#521)

gaogaotiantian · Nov 18, 2024 · 9ece82f · 9ece82f
1 parent 1275073
commit 9ece82f
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 25 deletions.
diff --git a/src/viztracer/modules/snaptrace.c b/src/viztracer/modules/snaptrace.c
@@ -338,25 +338,17 @@ prepare_before_trace(TracerObject* self, int is_call, struct ThreadInfo** info_o
     }
 
     if (info->ignore_stack_depth > 0) {
-        if (is_call) {
-            info->ignore_stack_depth += 1;
-            return 0;
-        } else {
-            info->ignore_stack_depth -= 1;
-            return 0;
-        }
+        return 0;
     }
 
     if (CHECK_FLAG(self->check_flags, SNAPTRACE_MAX_STACK_DEPTH)) {
         if (is_call) {
             if (info->curr_stack_depth >= self->max_stack_depth) {
-                info->curr_stack_depth += 1;
                 return 0;
             }
         } else {
             if (info->curr_stack_depth > 0) {
                 if (info->curr_stack_depth > self->max_stack_depth) {
-                    info->curr_stack_depth -= 1;
                     return 0;
                 }
             }
@@ -374,16 +366,15 @@ tracer_pycall_callback(TracerObject* self, PyCodeObject* code)
     if (prepare_before_trace(self, 1, &info) <= 0) {
         // For now we think -1 and 0 should both return because we should not
         // have the -1 case.
-        return 0;
+        goto cleanup_ignore;
     }
 
     PyObject* co_filename = code->co_filename;
 
     if (!CHECK_FLAG(self->check_flags, SNAPTRACE_TRACE_SELF)) {
         if (self->lib_file_path && co_filename && PyUnicode_Check(co_filename) &&
                 startswith(PyUnicode_AsUTF8(co_filename), self->lib_file_path)) {
-            info->ignore_stack_depth += 1;
-            goto cleanup;
+            goto cleanup_ignore;
         }
     }
 
@@ -408,18 +399,16 @@ tracer_pycall_callback(TracerObject* self, PyCodeObject* code)
                 }
             }
             if (record == 0) {
-                info->ignore_stack_depth += 1;
-                goto cleanup;
+                goto cleanup_ignore;
             }
         } else {
-            goto cleanup;
+            goto cleanup_ignore;
         }
     }
 
     if (CHECK_FLAG(self->check_flags, SNAPTRACE_IGNORE_FROZEN)) {
         if (startswith(PyUnicode_AsUTF8(co_filename), "<frozen")) {
-            info->ignore_stack_depth += 1;
-            goto cleanup;
+            goto cleanup_ignore;
         }
     }
 
@@ -456,10 +445,17 @@ tracer_pycall_callback(TracerObject* self, PyCodeObject* code)
         log_func_args(info->stack_top, PyEval_GetFrame(), self->log_func_repr);
     }
 
-cleanup:
-
     info->curr_stack_depth += 1;
 
+    return 0;
+
+cleanup_ignore:
+
+    if (info) {
+        info->ignore_stack_depth += 1;
+        info->curr_stack_depth += 1;
+    }
+
     return 0;
 }
 
@@ -471,15 +467,13 @@ tracer_ccall_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
     if (prepare_before_trace(self, 1, &info) <= 0) {
         // For now we think -1 and 0 should both return because we should not
         // have the -1 case.
-        return 0;
+        goto cleanup_ignore;
     }
 
     PyCFunctionObject* cfunc = (PyCFunctionObject*) arg;
 
     if (cfunc->m_self == (PyObject*)self) {
-        info->ignore_stack_depth += 1;
-        info->curr_stack_depth += 1;
-        return 0;
+        goto cleanup_ignore;
     }
 
     // If it's a call, we need a new node, and we need to update the stack
@@ -496,6 +490,15 @@ tracer_ccall_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
 
     info->curr_stack_depth += 1;
 
+    return 0;
+
+cleanup_ignore:
+
+    if (info) {
+        info->ignore_stack_depth += 1;
+        info->curr_stack_depth += 1;
+    }
+
     return 0;
 }
 
@@ -507,7 +510,7 @@ tracer_pyreturn_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
     if (prepare_before_trace(self, 0, &info) <= 0) {
         // For now we think -1 and 0 should both return because we should not
         // have the -1 case.
-        return 0;
+        goto cleanup_ignore;
     }
 
     struct FunctionNode* stack_top = info->stack_top;
@@ -571,6 +574,20 @@ tracer_pyreturn_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
         info->curr_stack_depth -= 1;
     } 
 
+    return 0;
+
+cleanup_ignore:
+
+    if (info) {
+        if (info->curr_stack_depth > 0) {
+            info->curr_stack_depth -= 1;
+        }
+
+        if (info->ignore_stack_depth > 0) {
+            info->ignore_stack_depth -= 1;
+        }
+    }
+
     return 0;
 }
 
@@ -582,7 +599,7 @@ tracer_creturn_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
     if (prepare_before_trace(self, 0, &info) <= 0) {
         // For now we think -1 and 0 should both return because we should not
         // have the -1 case.
-        return 0;
+        goto cleanup_ignore;
     }
 
     struct FunctionNode* stack_top = info->stack_top;
@@ -639,6 +656,20 @@ tracer_creturn_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
         info->curr_stack_depth -= 1;
     }
 
+    return 0;
+
+cleanup_ignore:
+
+    if (info) {
+        if (info->curr_stack_depth > 0) {
+            info->curr_stack_depth -= 1;
+        }
+
+        if (info->ignore_stack_depth > 0) {
+            info->ignore_stack_depth -= 1;
+        }
+    }
+
     return 0;
 }
 

diff --git a/tests/test_regression.py b/tests/test_regression.py
@@ -308,6 +308,30 @@ def test_issue162_os_popen(self):
                       script=issue162_code_os_popen, expected_stdout=r".*test_issue162.*")
 
 
+class TestIssue508(CmdlineTmpl):
+    def test_issue508(self):
+        script = """
+            import inspect
+            import os
+            import viztracer
+
+            exclude = os.path.dirname(inspect.__file__)
+
+            def call_self(n):
+                if n == 0:
+                    return
+                call_self(n - 1)
+
+            with viztracer.VizTracer(exclude_files=[exclude], max_stack_depth=6):
+                inspect.getsource(call_self)
+                call_self(10)
+        """
+
+        self.template(["python", "cmdline_test.py"], script=script,
+                      expected_output_file="result.json",
+                      expected_entries=6)
+
+
 file_timestamp_disorder = """
 def g():
     pass