Skip to content

Commit

Permalink
Record unfinished functions in the report (#444)
Browse files Browse the repository at this point in the history
  • Loading branch information
@gaogaotiantian
gaogaotiantian authored Jun 12, 2024
1 parent ff90654 commit aaf4831
Show file tree
Hide file tree
Showing 3 changed files with 117 additions and 13 deletions.
113 changes: 101 additions & 12 deletions src/viztracer/modules/snaptrace.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ static PyObject* snaptrace_getfunctionarg(TracerObject* self, PyObject* args);
static PyObject* snaptrace_getts(TracerObject* self, PyObject* args);
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_threaddestructor(void* key);
static struct ThreadInfo* snaptrace_createthreadinfo(TracerObject* self);
static void log_func_args(struct FunctionNode* node, PyFrameObject* frame);
Expand Down Expand Up @@ -324,6 +325,8 @@ snaptrace_pycall_callback(TracerObject* self, PyFrameObject* frame, struct Threa
}
info->stack_top = info->stack_top->next;
info->stack_top->ts = get_ts(info);
info->stack_top->func = (PyObject*) code;
Py_INCREF(code);
if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_FUNCTION_ARGS)) {
log_func_args(info->stack_top, frame);
}
Expand All @@ -336,7 +339,7 @@ snaptrace_pycall_callback(TracerObject* self, PyFrameObject* frame, struct Threa
}

int
snaptrace_ccall_callback(TracerObject* self, PyFrameObject* frame, struct ThreadInfo* info)
snaptrace_ccall_callback(TracerObject* self, PyFrameObject* frame, struct ThreadInfo* info, PyObject* arg)
{
// If it's a call, we need a new node, and we need to update the stack
if (!info->stack_top->next) {
Expand All @@ -345,6 +348,8 @@ snaptrace_ccall_callback(TracerObject* self, PyFrameObject* frame, struct Thread
}
info->stack_top = info->stack_top->next;
info->stack_top->ts = get_ts(info);
info->stack_top->func = arg;
Py_INCREF(arg);
if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_FUNCTION_ARGS)) {
log_func_args(info->stack_top, frame);
}
Expand All @@ -363,7 +368,7 @@ snaptrace_pyreturn_callback(TracerObject* self, PyFrameObject* frame, struct Thr

if (log_this_entry) {
struct EventNode* node = get_next_node(self);
PyCodeObject* code = PyFrame_GetCode(frame);
PyCodeObject* code = (PyCodeObject*) stack_top->func;

node->ntype = FEE_NODE;
node->ts = info->stack_top->ts;
Expand All @@ -390,8 +395,6 @@ snaptrace_pyreturn_callback(TracerObject* self, PyFrameObject* frame, struct Thr
Py_INCREF(info->curr_task);
}
}

Py_DECREF(code);
}
// Finish return whether to log the data
info->stack_top = info->stack_top->prev;
Expand All @@ -401,6 +404,8 @@ snaptrace_pyreturn_callback(TracerObject* self, PyFrameObject* frame, struct Thr
stack_top->args = NULL;
}

Py_CLEAR(stack_top->func);

if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC) &&
info->curr_task &&
frame == info->curr_task_frame) {
Expand All @@ -425,7 +430,7 @@ snaptrace_creturn_callback(TracerObject* self, PyFrameObject* frame, struct Thre

if (log_this_entry) {
struct EventNode* node = get_next_node(self);
PyCFunctionObject* cfunc = (PyCFunctionObject*) arg;
PyCFunctionObject* cfunc = (PyCFunctionObject*) stack_top->func;
node->ntype = FEE_NODE;
node->ts = info->stack_top->ts;
node->data.fee.dur = dur;
Expand Down Expand Up @@ -462,6 +467,8 @@ snaptrace_creturn_callback(TracerObject* self, PyFrameObject* frame, struct Thre
Py_DECREF(stack_top->args);
stack_top->args = NULL;
}

Py_CLEAR(stack_top->func);
}

return 0;
Expand Down Expand Up @@ -534,7 +541,7 @@ snaptrace_tracefunc(PyObject* obj, PyFrameObject* frame, int what, PyObject* arg
case PyTrace_CALL:
return snaptrace_pycall_callback(self, frame, info);
case PyTrace_C_CALL:
return snaptrace_ccall_callback(self, frame, info);
return snaptrace_ccall_callback(self, frame, info, arg);
case PyTrace_RETURN:
return snaptrace_pyreturn_callback(self, frame, info, arg);
case PyTrace_C_RETURN:
Expand Down Expand Up @@ -601,8 +608,9 @@ snaptrace_stop(TracerObject* self, PyObject* args)
info->curr_stack_depth = 0;
info->ignore_stack_depth = 0;
info->paused = 0;
clear_stack(&info->stack_top);
snaptrace_flush_unfinished(self);
}

curr_tracer = NULL;
PyEval_SetProfile(NULL, NULL);

Expand Down Expand Up @@ -655,6 +663,73 @@ snaptrace_resume(TracerObject* self, PyObject* args)
Py_RETURN_NONE;
}

static void
snaptrace_flush_unfinished(TracerObject* self)
{
SNAPTRACE_THREAD_PROTECT_START(self);

struct MetadataNode* meta_node = self->metadata_head;
while(meta_node) {
struct ThreadInfo* info = meta_node->thread_info;

if (info == NULL) {
meta_node = meta_node->next;
continue;
}

struct FunctionNode* func_node = info->stack_top;

while (func_node->prev) {
// 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 (PyCode_Check(func_node->func)) {
PyCodeObject* code = (PyCodeObject*) func_node->func;
fee_node->data.fee.type = PyTrace_CALL;
fee_node->data.fee.co_name = code->co_name;
fee_node->data.fee.co_filename = code->co_filename;
fee_node->data.fee.co_firstlineno = code->co_firstlineno;
Py_INCREF(fee_node->data.fee.co_name);
Py_INCREF(fee_node->data.fee.co_filename);
} else if (PyCFunction_Check(func_node->func)) {
PyCFunctionObject* cfunc = (PyCFunctionObject*) func_node->func;
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
fee_node->data.fee.m_module = cfunc->m_module;
Py_INCREF(cfunc->m_module);
} else {
// The function is a class method
fee_node->data.fee.m_module = NULL;
if (cfunc->m_self) {
// It's not a static method, has __self__
fee_node->data.fee.tp_name = cfunc->m_self->ob_type->tp_name;
} else {
// It's a static method, does not have __self__
fee_node->data.fee.tp_name = NULL;
}
}
}

// Clean up the node
Py_CLEAR(func_node->args);
Py_CLEAR(func_node->func);

func_node = func_node->prev;
}
info->stack_top = func_node;
meta_node = meta_node->next;
}

SNAPTRACE_THREAD_PROTECT_END(self);
}

static PyObject*
snaptrace_load(TracerObject* self, PyObject* args)
{
Expand Down Expand Up @@ -795,9 +870,14 @@ snaptrace_load(TracerObject* self, PyObject* args)
case FEE_NODE:
name = get_name_from_fee_node(node, func_name_dict);

PyObject* dur = PyFloat_FromDouble(node->data.fee.dur / 1000);
PyDict_SetItemString(dict, "dur", dur);
Py_DECREF(dur);
if (node->data.fee.type == PyTrace_CALL || node->data.fee.type == PyTrace_C_CALL) {
PyDict_SetItemString(dict, "ph", ph_B);
} else {
PyDict_SetItemString(dict, "ph", ph_X);
PyObject* dur = PyFloat_FromDouble(node->data.fee.dur / 1000);
PyDict_SetItemString(dict, "dur", dur);
Py_DECREF(dur);
}
PyDict_SetItemString(dict, "name", name);
Py_DECREF(name);

Expand All @@ -817,7 +897,6 @@ snaptrace_load(TracerObject* self, PyObject* args)
Py_DECREF(arg_dict);
}

PyDict_SetItemString(dict, "ph", ph_X);
PyDict_SetItemString(dict, "cat", cat_fee);
break;
case INSTANT_NODE:
Expand Down Expand Up @@ -1026,7 +1105,11 @@ snaptrace_dump(TracerObject* self, PyObject* args)
case FEE_NODE:
;
long long dur_long = node->data.fee.dur;
fprintf(fptr, "\"ph\":\"X\",\"cat\":\"fee\",\"dur\":%lld.%03lld,\"name\":\"", dur_long / 1000, dur_long % 1000);
char ph = 'X';
if (node->data.fee.type == PyTrace_CALL || node->data.fee.type == PyTrace_C_CALL) {
ph = 'B';
}
fprintf(fptr, "\"ph\":\"%c\",\"cat\":\"fee\",\"dur\":%lld.%03lld,\"name\":\"", ph, dur_long / 1000, dur_long % 1000);
fprintfeename(fptr, node, sanitize_function_name);
fputc('\"', fptr);

Expand Down Expand Up @@ -1550,6 +1633,8 @@ static struct ThreadInfo* snaptrace_createthreadinfo(TracerObject* self) {
if (node->tid == info->tid) {
Py_DECREF(node->name);
node->name = thread_name;
node->thread_info = info;
info->metadata_node = node;
found_node = 1;
break;
}
Expand All @@ -1564,6 +1649,8 @@ static struct ThreadInfo* snaptrace_createthreadinfo(TracerObject* self) {
}
node->name = thread_name;
node->tid = info->tid;
node->thread_info = info;
info->metadata_node = node;
node->next = self->metadata_head;
self->metadata_head = node;
}
Expand Down Expand Up @@ -1597,13 +1684,15 @@ static void snaptrace_threaddestructor(void* key) {
Py_DECREF(tmp->args);
tmp->args = NULL;
}
Py_CLEAR(tmp->func);
info->stack_top = info->stack_top->next;
PyMem_FREE(tmp);
}
}
info->stack_top = NULL;
info->curr_task = NULL;
info->curr_task_frame = NULL;
info->metadata_node->thread_info = NULL;
PyMem_FREE(info);
PyGILState_Release(state);
}
Expand Down
6 changes: 5 additions & 1 deletion src/viztracer/modules/snaptrace.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ struct FunctionNode {
struct FunctionNode* prev;
double ts;
PyObject* args;
// PyCodeObject* for Python function, PyCFunctionObject* for C function
PyObject* func;
};

struct ThreadInfo {
Expand All @@ -43,12 +45,14 @@ struct ThreadInfo {
PyObject* curr_task;
PyFrameObject* curr_task_frame;
double prev_ts;
struct MetadataNode* metadata_node;
};

struct MetadataNode {
struct MetadataNode* next;
unsigned long tid;
PyObject* name;
struct MetadataNode* next;
struct ThreadInfo* thread_info;
};

typedef struct {
Expand Down
11 changes: 11 additions & 0 deletions tests/test_basic.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,17 @@ def test_tracer_entries(self):
tracer.parse()
self.assertEventNumber(tracer.data, 10)

def test_unfinished_function(self):
def f():
tracer.stop()
tracer.parse()

tracer = VizTracer(tracer_entries=10)
tracer.start()
f()
self.assertEventNumber(tracer.data, 1)
self.assertEqual(tracer.data["traceEvents"][-1]["ph"], "B")

def test_save(self):
tracer = VizTracer(tracer_entries=10)
tracer.start()
Expand Down

0 comments on commit aaf4831

Please sign in to comment.