Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Flush unfinished function as finished when the program is exiting #481

Merged
merged 3 commits into from
Oct 14, 2024
Merged

Flush unfinished function as finished when the program is exiting #481

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c1c9011
Flush unfinished function as finished when the program is exiting
gaogaotiantian Oct 11, 2024
2f150b1
Use the current ts instead of the prev one
gaogaotiantian Oct 11, 2024
c2b5c0c
Make main use flush as finished as well
gaogaotiantian Oct 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/viztracer/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -361,7 +361,7 @@ def term_handler(signalnum, frame):
exec(code, global_dict)

if not options.log_exit:
tracer.stop()
tracer.stop(stop_option="flush_as_finish")

# The user code may forked, check it because Finalize won't execute
# if the pid is not the same
Expand Down
69 changes: 51 additions & 18 deletions src/viztracer/modules/snaptrace.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ int snaptrace_tracefunc(PyObject* obj, PyFrameObject* frame, int what, PyObject*
int snaptrace_tracefuncdisabled(PyObject* obj, PyFrameObject* frame, int what, PyObject* arg);
static PyObject* snaptrace_threadtracefunc(PyObject* obj, PyObject* args);
static PyObject* snaptrace_start(TracerObject* self, PyObject* Py_UNUSED(unused));
static PyObject* snaptrace_stop(TracerObject* self, PyObject* Py_UNUSED(unused));
static PyObject* snaptrace_stop(TracerObject* self, PyObject* stop_option);
static PyObject* snaptrace_pause(TracerObject* self, PyObject* Py_UNUSED(unused));
static PyObject* snaptrace_resume(TracerObject* self, PyObject* Py_UNUSED(unused));
static PyObject* snaptrace_load(TracerObject* self, PyObject* Py_UNUSED(unused));
Expand All @@ -45,7 +45,7 @@ static PyObject* snaptrace_getfunctionarg(TracerObject* self, PyObject* args);
static PyObject* snaptrace_getts(TracerObject* self, PyObject* Py_UNUSED(unused));
static PyObject* snaptrace_setcurrstack(TracerObject* self, PyObject* args);
static PyObject* snaptrace_setignorestackcounter(TracerObject* self, PyObject* args);
static void snaptrace_flush_unfinished(TracerObject* self);
static void snaptrace_flush_unfinished(TracerObject* self, int flush_as_finish);
static void snaptrace_threaddestructor(void* key);
static struct ThreadInfo* snaptrace_createthreadinfo(TracerObject* self);
static void log_func_args(struct FunctionNode* node, PyFrameObject* frame);
Expand Down Expand Up @@ -204,7 +204,7 @@ void clear_stack(struct FunctionNode** stack_top) {
static PyMethodDef Tracer_methods[] = {
{"threadtracefunc", (PyCFunction)snaptrace_threadtracefunc, METH_VARARGS, "trace function"},
{"start", (PyCFunction)snaptrace_start, METH_NOARGS, "start profiling"},
{"stop", (PyCFunction)snaptrace_stop, METH_NOARGS, "stop profiling"},
{"stop", (PyCFunction)snaptrace_stop, METH_O, "stop profiling"},
{"load", (PyCFunction)snaptrace_load, METH_NOARGS, "load buffer"},
{"dump", (PyCFunction)snaptrace_dump, METH_VARARGS, "dump buffer to file"},
{"clear", (PyCFunction)snaptrace_clear, METH_NOARGS, "clear buffer"},
Expand Down Expand Up @@ -508,14 +508,14 @@ snaptrace_tracefunc(PyObject* obj, PyFrameObject* frame, int what, PyObject* arg

if (CHECK_FLAG(self->check_flags, SNAPTRACE_MAX_STACK_DEPTH)) {
if (is_call) {
info->curr_stack_depth += 1;
if (info->curr_stack_depth > self->max_stack_depth) {
if (info->curr_stack_depth >= self->max_stack_depth) {
info->curr_stack_depth += 1;
return 0;
}
} else {
if (info->curr_stack_depth > 0) {
info->curr_stack_depth -= 1;
if (info->curr_stack_depth + 1 > self->max_stack_depth) {
if (info->curr_stack_depth > self->max_stack_depth) {
info->curr_stack_depth -= 1;
return 0;
}
}
Expand All @@ -524,17 +524,31 @@ snaptrace_tracefunc(PyObject* obj, PyFrameObject* frame, int what, PyObject* arg

switch (what) {
case PyTrace_CALL:
return snaptrace_pycall_callback(self, frame, info);
info->curr_stack_depth += 1;
snaptrace_pycall_callback(self, frame, info);
break;
case PyTrace_C_CALL:
return snaptrace_ccall_callback(self, frame, info, arg);
info->curr_stack_depth += 1;
snaptrace_ccall_callback(self, frame, info, arg);
break;
case PyTrace_RETURN:
return snaptrace_pyreturn_callback(self, frame, info, arg);
snaptrace_pyreturn_callback(self, frame, info, arg);
if (info->curr_stack_depth > 0) {
info->curr_stack_depth -= 1;
}
break;
case PyTrace_C_RETURN:
case PyTrace_C_EXCEPTION:
return snaptrace_creturn_callback(self, frame, info, arg);
snaptrace_creturn_callback(self, frame, info, arg);
if (info->curr_stack_depth > 0) {
info->curr_stack_depth -= 1;
}
break;
default:
return 0;
}

return 0;
}

static PyObject* snaptrace_threadtracefunc(PyObject* obj, PyObject* args)
Expand Down Expand Up @@ -584,16 +598,21 @@ snaptrace_start(TracerObject* self, PyObject* Py_UNUSED(unused))
}

static PyObject*
snaptrace_stop(TracerObject* self, PyObject* Py_UNUSED(unused))
snaptrace_stop(TracerObject* self, PyObject* stop_option)
{
if (self) {
struct ThreadInfo* info = get_thread_info(self);
self->collecting = 0;

if (PyUnicode_CheckExact(stop_option) &&
strcmp(PyUnicode_AsUTF8(stop_option), "flush_as_finish") == 0) {
snaptrace_flush_unfinished(self, 1);
} else {
snaptrace_flush_unfinished(self, 0);
}
info->curr_stack_depth = 0;
info->ignore_stack_depth = 0;
info->paused = 0;
snaptrace_flush_unfinished(self);
}

curr_tracer = NULL;
Expand Down Expand Up @@ -649,7 +668,7 @@ snaptrace_resume(TracerObject* self, PyObject* Py_UNUSED(unused))
}

static void
snaptrace_flush_unfinished(TracerObject* self)
snaptrace_flush_unfinished(TracerObject* self, int flush_as_finish)
{
SNAPTRACE_THREAD_PROTECT_START(self);

Expand All @@ -664,23 +683,36 @@ snaptrace_flush_unfinished(TracerObject* self)

struct FunctionNode* func_node = info->stack_top;

while (func_node->prev) {
while (func_node->prev && info->curr_stack_depth > 0) {
// Fake a FEE node to get the name
struct EventNode* fee_node = get_next_node(self);

fee_node->ntype = FEE_NODE;
fee_node->ts = func_node->ts;
fee_node->tid = meta_node->tid;
fee_node->data.fee.dur = 0;

if (flush_as_finish) {
fee_node->data.fee.dur = get_ts(info) - func_node->ts;
} else {
fee_node->data.fee.dur = 0;
}

if (PyCode_Check(func_node->func)) {
PyCodeObject* code = (PyCodeObject*) func_node->func;
fee_node->data.fee.type = PyTrace_CALL;
if (flush_as_finish) {
fee_node->data.fee.type = PyTrace_RETURN;
} else {
fee_node->data.fee.type = PyTrace_CALL;
}
fee_node->data.fee.code = code;
Py_INCREF(code);
} else if (PyCFunction_Check(func_node->func)) {
PyCFunctionObject* cfunc = (PyCFunctionObject*) func_node->func;
fee_node->data.fee.type = PyTrace_C_CALL;
if (flush_as_finish) {
fee_node->data.fee.type = PyTrace_C_RETURN;
} else {
fee_node->data.fee.type = PyTrace_C_CALL;
}
fee_node->data.fee.ml_name = cfunc->m_ml->ml_name;
if (cfunc->m_module) {
// The function belongs to a module
Expand All @@ -704,6 +736,7 @@ snaptrace_flush_unfinished(TracerObject* self)
Py_CLEAR(func_node->func);

func_node = func_node->prev;
info->curr_stack_depth -= 1;
}
info->stack_top = func_node;
meta_node = meta_node->next;
Expand Down
4 changes: 2 additions & 2 deletions src/viztracer/tracer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -254,11 +254,11 @@ def start(self) -> None:
self.config()
self._tracer.start()

def stop(self) -> None:
def stop(self, stop_option: Optional[str] = None) -> None:
self.enable = False
if self.log_print:
self.restore_print()
self._tracer.stop()
self._tracer.stop(stop_option)

def pause(self) -> None:
if self.enable:
Expand Down
6 changes: 3 additions & 3 deletions src/viztracer/viztracer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,9 +184,9 @@ def start(self) -> None:
self._plugin_manager.event("pre-start")
_VizTracer.start(self)

def stop(self) -> None:
def stop(self, stop_option: Optional[str] = None) -> None:
if self.enable:
_VizTracer.stop(self)
_VizTracer.stop(self, stop_option)
self._plugin_manager.event("post-stop")

def run(self, command: str, output_file: Optional[str] = None) -> None:
Expand Down Expand Up @@ -301,7 +301,7 @@ def term_handler(sig, frame):
def exit_routine(self) -> None:
# We need to avoid SIGTERM terminate our process when we dump data
signal.signal(signal.SIGTERM, lambda sig, frame: 0)
self.stop()
self.stop(stop_option="flush_as_finish")
if not self._exiting:
self._exiting = True
os.chdir(self.cwd)
Expand Down
Loading