viztracer 1.1.1__cp313-cp313-win32.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (109) hide show
  1. viztracer/__init__.py +19 -0
  2. viztracer/__main__.py +8 -0
  3. viztracer/attach.py +67 -0
  4. viztracer/attach_process/LICENSE +203 -0
  5. viztracer/attach_process/__init__.py +0 -0
  6. viztracer/attach_process/add_code_to_python_process.py +582 -0
  7. viztracer/attach_process/attach_x86.dll +0 -0
  8. viztracer/attach_process/inject_dll_amd64.exe +0 -0
  9. viztracer/attach_process/linux_and_mac/lldb_prepare.py +54 -0
  10. viztracer/attach_process/run_code_on_dllmain_amd64.dll +0 -0
  11. viztracer/attach_process/run_code_on_dllmain_x86.dll +0 -0
  12. viztracer/cellmagic.py +70 -0
  13. viztracer/code_monkey.py +353 -0
  14. viztracer/decorator.py +164 -0
  15. viztracer/event_base.py +81 -0
  16. viztracer/functree.py +135 -0
  17. viztracer/html/flamegraph.html +34 -0
  18. viztracer/html/trace_viewer_embedder.html +203 -0
  19. viztracer/html/trace_viewer_full.html +10207 -0
  20. viztracer/main.py +701 -0
  21. viztracer/modules/eventnode.c +188 -0
  22. viztracer/modules/eventnode.h +73 -0
  23. viztracer/modules/pythoncapi_compat.h +1726 -0
  24. viztracer/modules/quicktime.c +177 -0
  25. viztracer/modules/quicktime.h +104 -0
  26. viztracer/modules/snaptrace.c +2207 -0
  27. viztracer/modules/snaptrace.h +134 -0
  28. viztracer/modules/snaptrace_member.c +483 -0
  29. viztracer/modules/util.c +45 -0
  30. viztracer/modules/util.h +22 -0
  31. viztracer/modules/vcompressor/vc_dump.c +1131 -0
  32. viztracer/modules/vcompressor/vc_dump.h +49 -0
  33. viztracer/modules/vcompressor/vcompressor.c +396 -0
  34. viztracer/modules/vcompressor/vcompressor.h +15 -0
  35. viztracer/patch.py +317 -0
  36. viztracer/report_builder.py +311 -0
  37. viztracer/snaptrace.cp313-win32.pyd +0 -0
  38. viztracer/snaptrace.pyi +77 -0
  39. viztracer/util.py +196 -0
  40. viztracer/vcompressor.cp313-win32.pyd +0 -0
  41. viztracer/vcompressor.pyi +10 -0
  42. viztracer/viewer.py +529 -0
  43. viztracer/vizcounter.py +20 -0
  44. viztracer/vizevent.py +31 -0
  45. viztracer/vizlogging.py +20 -0
  46. viztracer/vizobject.py +28 -0
  47. viztracer/vizplugin.py +143 -0
  48. viztracer/viztracer.py +472 -0
  49. viztracer/web_dist/LICENSE +189 -0
  50. viztracer/web_dist/index.html +127 -0
  51. viztracer/web_dist/service_worker.js +279 -0
  52. viztracer/web_dist/trace_processor +300 -0
  53. viztracer/web_dist/v52.0-6b9586def/assets/MaterialSymbolsOutlined.woff2 +0 -0
  54. viztracer/web_dist/v52.0-6b9586def/assets/Roboto-100.woff2 +0 -0
  55. viztracer/web_dist/v52.0-6b9586def/assets/Roboto-300.woff2 +0 -0
  56. viztracer/web_dist/v52.0-6b9586def/assets/Roboto-400.woff2 +0 -0
  57. viztracer/web_dist/v52.0-6b9586def/assets/Roboto-500.woff2 +0 -0
  58. viztracer/web_dist/v52.0-6b9586def/assets/RobotoCondensed-Light.woff2 +0 -0
  59. viztracer/web_dist/v52.0-6b9586def/assets/RobotoCondensed-Regular.woff2 +0 -0
  60. viztracer/web_dist/v52.0-6b9586def/assets/RobotoMono-Regular.woff2 +0 -0
  61. viztracer/web_dist/v52.0-6b9586def/assets/brand.png +0 -0
  62. viztracer/web_dist/v52.0-6b9586def/assets/catapult_trace_viewer.html +3946 -0
  63. viztracer/web_dist/v52.0-6b9586def/assets/catapult_trace_viewer.js +7539 -0
  64. viztracer/web_dist/v52.0-6b9586def/assets/favicon.png +0 -0
  65. viztracer/web_dist/v52.0-6b9586def/assets/logo-128.png +0 -0
  66. viztracer/web_dist/v52.0-6b9586def/assets/logo-3d.png +0 -0
  67. viztracer/web_dist/v52.0-6b9586def/assets/rec_atrace.png +0 -0
  68. viztracer/web_dist/v52.0-6b9586def/assets/rec_battery_counters.png +0 -0
  69. viztracer/web_dist/v52.0-6b9586def/assets/rec_board_voltage.png +0 -0
  70. viztracer/web_dist/v52.0-6b9586def/assets/rec_cpu_coarse.png +0 -0
  71. viztracer/web_dist/v52.0-6b9586def/assets/rec_cpu_fine.png +0 -0
  72. viztracer/web_dist/v52.0-6b9586def/assets/rec_cpu_freq.png +0 -0
  73. viztracer/web_dist/v52.0-6b9586def/assets/rec_cpu_voltage.png +0 -0
  74. viztracer/web_dist/v52.0-6b9586def/assets/rec_frame_timeline.png +0 -0
  75. viztracer/web_dist/v52.0-6b9586def/assets/rec_ftrace.png +0 -0
  76. viztracer/web_dist/v52.0-6b9586def/assets/rec_gpu_mem_total.png +0 -0
  77. viztracer/web_dist/v52.0-6b9586def/assets/rec_java_heap_dump.png +0 -0
  78. viztracer/web_dist/v52.0-6b9586def/assets/rec_lmk.png +0 -0
  79. viztracer/web_dist/v52.0-6b9586def/assets/rec_logcat.png +0 -0
  80. viztracer/web_dist/v52.0-6b9586def/assets/rec_long_trace.png +0 -0
  81. viztracer/web_dist/v52.0-6b9586def/assets/rec_mem_hifreq.png +0 -0
  82. viztracer/web_dist/v52.0-6b9586def/assets/rec_meminfo.png +0 -0
  83. viztracer/web_dist/v52.0-6b9586def/assets/rec_native_heap_profiler.png +0 -0
  84. viztracer/web_dist/v52.0-6b9586def/assets/rec_one_shot.png +0 -0
  85. viztracer/web_dist/v52.0-6b9586def/assets/rec_profiling.png +0 -0
  86. viztracer/web_dist/v52.0-6b9586def/assets/rec_ps_stats.png +0 -0
  87. viztracer/web_dist/v52.0-6b9586def/assets/rec_ring_buf.png +0 -0
  88. viztracer/web_dist/v52.0-6b9586def/assets/rec_syscalls.png +0 -0
  89. viztracer/web_dist/v52.0-6b9586def/assets/rec_vmstat.png +0 -0
  90. viztracer/web_dist/v52.0-6b9586def/assets/scheduling_latency.png +0 -0
  91. viztracer/web_dist/v52.0-6b9586def/assets/vscode-icon.png +0 -0
  92. viztracer/web_dist/v52.0-6b9586def/engine_bundle.js +3 -0
  93. viztracer/web_dist/v52.0-6b9586def/frontend_bundle.js +5495 -0
  94. viztracer/web_dist/v52.0-6b9586def/index.html +127 -0
  95. viztracer/web_dist/v52.0-6b9586def/manifest.json +52 -0
  96. viztracer/web_dist/v52.0-6b9586def/perfetto.css +5737 -0
  97. viztracer/web_dist/v52.0-6b9586def/stdlib_docs.json +1 -0
  98. viztracer/web_dist/v52.0-6b9586def/trace_config_utils.wasm +0 -0
  99. viztracer/web_dist/v52.0-6b9586def/trace_processor.wasm +0 -0
  100. viztracer/web_dist/v52.0-6b9586def/trace_processor_memory64.wasm +0 -0
  101. viztracer/web_dist/v52.0-6b9586def/traceconv.wasm +0 -0
  102. viztracer/web_dist/v52.0-6b9586def/traceconv_bundle.js +2 -0
  103. viztracer-1.1.1.dist-info/METADATA +326 -0
  104. viztracer-1.1.1.dist-info/RECORD +109 -0
  105. viztracer-1.1.1.dist-info/WHEEL +5 -0
  106. viztracer-1.1.1.dist-info/entry_points.txt +3 -0
  107. viztracer-1.1.1.dist-info/licenses/LICENSE +222 -0
  108. viztracer-1.1.1.dist-info/licenses/NOTICE.txt +27 -0
  109. viztracer-1.1.1.dist-info/top_level.txt +1 -0
@@ -0,0 +1,2207 @@
1
+ // Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
2
+ // For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
3
+
4
+ #define PY_SSIZE_T_CLEAN
5
+ #include <Python.h>
6
+ #include <stdlib.h>
7
+ #include <frameobject.h>
8
+ #if _WIN32
9
+ #include <windows.h>
10
+ #elif defined(__APPLE__)
11
+ #include <pthread.h>
12
+ #elif defined(__FreeBSD__)
13
+ #include <pthread_np.h>
14
+ #else
15
+ #include <pthread.h>
16
+ #include <sys/syscall.h>
17
+ #endif
18
+
19
+ #include "pythoncapi_compat.h"
20
+ #include "snaptrace.h"
21
+ #include "quicktime.h"
22
+ #include "util.h"
23
+ #include "eventnode.h"
24
+
25
+
26
+ TracerObject* curr_tracer = NULL;
27
+ PyObject* threading_module = NULL;
28
+ PyObject* multiprocessing_module = NULL;
29
+ PyObject* json_module = NULL;
30
+ PyObject* asyncio_module = NULL;
31
+ PyObject* asyncio_tasks_module = NULL;
32
+ PyObject* trio_module = NULL;
33
+ PyObject* trio_lowlevel_module = NULL;
34
+ PyObject* sys_module = NULL;
35
+ PyObject* sys_monitoring_missing = NULL;
36
+
37
+ PyObject* curr_task_getters[2] = {0};
38
+
39
+ // =============================================================================
40
+ // Utility function
41
+ // =============================================================================
42
+
43
+ int64_t prev_ts = 0;
44
+
45
+ static inline int64_t
46
+ get_ts()
47
+ {
48
+ #if defined(QUICKTIME_RDTSC)
49
+ return get_system_ts();
50
+ #else
51
+ int64_t curr_ts = get_system_ts();
52
+ if (curr_ts <= prev_ts) {
53
+ // We use artificial timestamp to avoid timestamp conflict.
54
+ // 20 ns should be a safe granularity because that's normally
55
+ // how long clock_gettime() takes.
56
+ // It's possible to have three same timestamp in a row so we
57
+ // need to check if curr_ts <= prev_ts instead of ==
58
+ #if _WIN32 || defined(__APPLE__)
59
+ curr_ts = prev_ts + 1;
60
+ #else
61
+ curr_ts = prev_ts + 20;
62
+ #endif
63
+ }
64
+ prev_ts = curr_ts;
65
+ return curr_ts;
66
+ #endif
67
+ }
68
+
69
+ static inline struct EventNode*
70
+ get_next_node(TracerObject* self)
71
+ {
72
+ struct EventNode* node = NULL;
73
+
74
+ SNAPTRACE_THREAD_PROTECT_START(self);
75
+ node = self->buffer + self->buffer_tail_idx;
76
+ // This is actually faster than modulo
77
+ self->buffer_tail_idx = self->buffer_tail_idx + 1;
78
+ if (self->buffer_tail_idx >= self->buffer_size) {
79
+ self->buffer_tail_idx = 0;
80
+ }
81
+ if (self->buffer_tail_idx == self->buffer_head_idx) {
82
+ self->buffer_head_idx = self->buffer_head_idx + 1;
83
+ if (self->buffer_head_idx >= self->buffer_size) {
84
+ self->buffer_head_idx = 0;
85
+ }
86
+ clear_node(self->buffer + self->buffer_tail_idx);
87
+ } else {
88
+ self->total_entries += 1;
89
+ }
90
+ SNAPTRACE_THREAD_PROTECT_END(self);
91
+
92
+ return node;
93
+ }
94
+
95
+ static void
96
+ log_func_args(struct FunctionNode* node, PyFrameObject* frame, PyObject* log_func_repr)
97
+ {
98
+ PyObject* func_arg_dict = PyDict_New();
99
+ PyCodeObject* code = PyFrame_GetCode(frame);
100
+ PyObject* names = PyCode_GetVarnames(code);
101
+
102
+ #if PY_VERSION_HEX >= 0x030D0000
103
+ PyObject* locals = PyEval_GetFrameLocals();
104
+ #else
105
+ PyObject* locals = PyEval_GetLocals();
106
+ #endif
107
+
108
+ int idx = 0;
109
+ if (!node->args) {
110
+ node->args = PyDict_New();
111
+ }
112
+
113
+ int name_length = code->co_argcount + code->co_kwonlyargcount;
114
+ if (code->co_flags & CO_VARARGS) {
115
+ name_length ++;
116
+ }
117
+
118
+ if (code->co_flags & CO_VARKEYWORDS) {
119
+ name_length ++;
120
+ }
121
+
122
+ while (idx < name_length) {
123
+ // Borrowed
124
+ PyObject* name = PyTuple_GET_ITEM(names, idx);
125
+ PyObject* repr = NULL;
126
+ // New
127
+ if (log_func_repr) {
128
+ repr = PyObject_CallOneArg(log_func_repr, PyDict_GetItem(locals, name));
129
+ } else {
130
+ repr = PyObject_Repr(PyDict_GetItem(locals, name));
131
+ }
132
+ if (!repr) {
133
+ repr = PyUnicode_FromString("Not Displayable");
134
+ PyErr_Clear();
135
+ }
136
+ PyDict_SetItem(func_arg_dict, name, repr);
137
+ Py_DECREF(repr);
138
+ idx++;
139
+ }
140
+
141
+ #if PY_VERSION_HEX >= 0x030D0000
142
+ Py_DECREF(locals);
143
+ #endif
144
+
145
+ PyDict_SetItemString(node->args, "func_args", func_arg_dict);
146
+ Py_DECREF(func_arg_dict);
147
+
148
+ Py_XDECREF(code);
149
+ Py_XDECREF(names);
150
+ }
151
+
152
+ static void
153
+ verbose_printf(TracerObject* self, int v, const char* fmt, ...)
154
+ {
155
+ va_list args;
156
+ if (self->verbose >= v) {
157
+ va_start(args, fmt);
158
+ vprintf(fmt, args);
159
+ va_end(args);
160
+ fflush(stdout);
161
+ }
162
+ }
163
+
164
+ void
165
+ clear_stack(struct FunctionNode** stack_top) {
166
+ Py_CLEAR((*stack_top)->args);
167
+ Py_CLEAR((*stack_top)->func);
168
+ while ((*stack_top)->prev) {
169
+ (*stack_top) = (*stack_top) -> prev;
170
+ Py_CLEAR((*stack_top)->args);
171
+ Py_CLEAR((*stack_top)->func);
172
+ }
173
+ }
174
+
175
+ // =============================================================================
176
+ // Thread info related functions
177
+ // =============================================================================
178
+
179
+ static struct ThreadInfo*
180
+ snaptrace_createthreadinfo(TracerObject* self) {
181
+ struct ThreadInfo* info = PyMem_Calloc(1, sizeof(struct ThreadInfo));
182
+ info->stack_top = (struct FunctionNode*) PyMem_Calloc(1, sizeof(struct FunctionNode));
183
+
184
+ #if _WIN32
185
+ info->tid = GetCurrentThreadId();
186
+ #elif defined(__APPLE__)
187
+ __uint64_t tid = 0;
188
+ if (pthread_threadid_np(NULL, &tid)) {
189
+ info->tid = (unsigned long)pthread_self();
190
+ } else {
191
+ info->tid = tid;
192
+ }
193
+ #elif defined(__FreeBSD__)
194
+ info->tid = pthread_getthreadid_np();
195
+ #else
196
+ info->tid = syscall(SYS_gettid);
197
+ #endif
198
+
199
+ #if _WIN32
200
+ TlsSetValue(self->dwTlsIndex, info);
201
+ #else
202
+ pthread_setspecific(self->thread_key, info);
203
+ #endif
204
+
205
+ PyGILState_STATE state = PyGILState_Ensure();
206
+ SNAPTRACE_THREAD_PROTECT_START(self);
207
+
208
+ PyObject* current_thread = PyObject_CallMethod(threading_module, "current_thread", "");
209
+ if (!current_thread) {
210
+ PyErr_SetString(PyExc_RuntimeError, "Failed to get current thread");
211
+ goto cleanup;
212
+ }
213
+ PyObject* thread_name = PyObject_GetAttrString(current_thread, "name");
214
+ if (!thread_name) {
215
+ // It's okay not having a name
216
+ PyErr_Clear();
217
+ thread_name = PyUnicode_FromString("Unknown");
218
+ }
219
+
220
+ Py_DECREF(current_thread);
221
+
222
+ // Check for existing node for the same tid first
223
+ struct MetadataNode* node = self->metadata_head;
224
+ int found_node = 0;
225
+
226
+ while (node) {
227
+ if (node->tid == info->tid) {
228
+ Py_DECREF(node->name);
229
+ node->name = thread_name;
230
+ node->thread_info = info;
231
+ info->metadata_node = node;
232
+ found_node = 1;
233
+ break;
234
+ }
235
+ node = node->next;
236
+ }
237
+
238
+ if (!found_node) {
239
+ node = (struct MetadataNode*) PyMem_Calloc(1, sizeof(struct MetadataNode));
240
+ if (!node) {
241
+ PyErr_SetString(PyExc_RuntimeError, "Out of memory!");
242
+ info = NULL;
243
+ goto cleanup;
244
+ }
245
+ node->name = thread_name;
246
+ node->tid = info->tid;
247
+ node->thread_info = info;
248
+ info->metadata_node = node;
249
+ node->next = self->metadata_head;
250
+ self->metadata_head = node;
251
+ }
252
+
253
+ info->curr_task = NULL;
254
+ info->curr_task_frame = NULL;
255
+
256
+ cleanup:
257
+
258
+ SNAPTRACE_THREAD_PROTECT_END(self);
259
+ PyGILState_Release(state);
260
+
261
+ return info;
262
+ }
263
+
264
+ static struct ThreadInfo*
265
+ get_thread_info(TracerObject* self)
266
+ {
267
+ // self is non-NULL value
268
+ struct ThreadInfo* info = NULL;
269
+ #if _WIN32
270
+ info = TlsGetValue(self->dwTlsIndex);
271
+ #else
272
+ info = pthread_getspecific(self->thread_key);
273
+ #endif
274
+ if (!info) {
275
+ info = snaptrace_createthreadinfo(self);
276
+ }
277
+ return info;
278
+ }
279
+
280
+ static void
281
+ snaptrace_threaddestructor(void* key) {
282
+ struct ThreadInfo* info = key;
283
+ struct FunctionNode* tmp = NULL;
284
+ if (info) {
285
+ PyGILState_STATE state = PyGILState_Ensure();
286
+ info->paused = 0;
287
+ info->curr_stack_depth = 0;
288
+ info->ignore_stack_depth = 0;
289
+ info->tid = 0;
290
+ if (info->stack_top) {
291
+ while (info->stack_top->prev) {
292
+ info->stack_top = info->stack_top->prev;
293
+ }
294
+ while (info->stack_top) {
295
+ tmp = info->stack_top;
296
+ Py_CLEAR(tmp->args);
297
+ Py_CLEAR(tmp->func);
298
+ info->stack_top = info->stack_top->next;
299
+ PyMem_FREE(tmp);
300
+ }
301
+ }
302
+ info->stack_top = NULL;
303
+ Py_CLEAR(info->curr_task);
304
+ Py_CLEAR(info->curr_task_frame);
305
+ info->metadata_node->thread_info = NULL;
306
+ PyMem_FREE(info);
307
+ PyGILState_Release(state);
308
+ }
309
+ }
310
+
311
+ // =============================================================================
312
+ // Tracing function, triggered when FEE
313
+ // =============================================================================
314
+
315
+ // This function is called before we actually start tracing.
316
+ // * Prepare the thread info and create one if not exist
317
+ // * Check if we should trace based on all the flags
318
+ // * -1: Error
319
+ // * 0: Not trace
320
+ // * 1: Trace
321
+ int
322
+ prepare_before_trace(TracerObject* self, int is_call, struct ThreadInfo** info_out) {
323
+
324
+ if (!self->collecting) {
325
+ return 0;
326
+ }
327
+
328
+ struct ThreadInfo* info = get_thread_info(self);
329
+
330
+ if (!info) {
331
+ self->collecting = 0;
332
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to create thread info. This should not happen.");
333
+ return -1;
334
+ }
335
+
336
+ *info_out = info;
337
+
338
+ if (info->paused) {
339
+ return 0;
340
+ }
341
+
342
+ if (info->ignore_stack_depth > 0) {
343
+ return 0;
344
+ }
345
+
346
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_MAX_STACK_DEPTH)) {
347
+ if (is_call) {
348
+ if (info->curr_stack_depth >= self->max_stack_depth) {
349
+ return 0;
350
+ }
351
+ } else {
352
+ if (info->curr_stack_depth > 0) {
353
+ if (info->curr_stack_depth > self->max_stack_depth) {
354
+ return 0;
355
+ }
356
+ }
357
+ }
358
+ }
359
+
360
+ return 1;
361
+ }
362
+
363
+ int
364
+ tracer_pycall_callback(TracerObject* self, PyCodeObject* code)
365
+ {
366
+ struct ThreadInfo* info = NULL;
367
+
368
+ if (prepare_before_trace(self, 1, &info) <= 0) {
369
+ // For now we think -1 and 0 should both return because we should not
370
+ // have the -1 case.
371
+ goto cleanup_ignore;
372
+ }
373
+
374
+ PyObject* co_filename = code->co_filename;
375
+
376
+ if (!CHECK_FLAG(self->check_flags, SNAPTRACE_TRACE_SELF)) {
377
+ if (self->lib_file_path && co_filename && PyUnicode_Check(co_filename) &&
378
+ startswith(PyUnicode_AsUTF8(co_filename), self->lib_file_path)) {
379
+ goto cleanup_ignore;
380
+ }
381
+ }
382
+
383
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_INCLUDE_FILES | SNAPTRACE_EXCLUDE_FILES)) {
384
+ if (info->ignore_stack_depth == 0) {
385
+ PyObject* files = NULL;
386
+ int record = 0;
387
+ int is_include = CHECK_FLAG(self->check_flags, SNAPTRACE_INCLUDE_FILES);
388
+ if (is_include) {
389
+ files = self->include_files;
390
+ record = 0;
391
+ } else {
392
+ files = self->exclude_files;
393
+ record = 1;
394
+ }
395
+ Py_ssize_t length = PyList_GET_SIZE(files);
396
+ for (int i = 0; i < length; i++) {
397
+ PyObject* f = PyList_GET_ITEM(files, i);
398
+ if (startswith(PyUnicode_AsUTF8(co_filename), PyUnicode_AsUTF8(f))) {
399
+ record = 1 - record;
400
+ break;
401
+ }
402
+ }
403
+ if (record == 0) {
404
+ goto cleanup_ignore;
405
+ }
406
+ } else {
407
+ goto cleanup_ignore;
408
+ }
409
+ }
410
+
411
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_IGNORE_FROZEN)) {
412
+ if (startswith(PyUnicode_AsUTF8(co_filename), "<frozen")) {
413
+ goto cleanup_ignore;
414
+ }
415
+ }
416
+
417
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC) &&
418
+ info->curr_task == NULL &&
419
+ (code->co_flags & CO_COROUTINE) != 0) {
420
+ PyObject* curr_task = Py_None;
421
+ info->paused = 1;
422
+ for (size_t i = 0; i < sizeof(curr_task_getters)/sizeof(curr_task_getters[0]); i++) {
423
+ if (curr_task_getters[i] != NULL) {
424
+ curr_task = PyObject_CallNoArgs(curr_task_getters[i]);
425
+ if (!curr_task) {
426
+ PyErr_Clear(); // RuntimeError, probably
427
+ curr_task = Py_None;
428
+ } else if (curr_task != Py_None) {
429
+ break; // got a valid task
430
+ }
431
+ }
432
+ }
433
+ info->paused = 0;
434
+ info->curr_task = Py_NewRef(curr_task);
435
+ info->curr_task_frame = (PyFrameObject*)Py_NewRef(PyEval_GetFrame());
436
+ }
437
+
438
+ // If it's a call, we need a new node, and we need to update the stack
439
+ if (!info->stack_top->next) {
440
+ info->stack_top->next = (struct FunctionNode*) PyMem_Calloc(1, sizeof(struct FunctionNode));
441
+ info->stack_top->next->prev = info->stack_top;
442
+ }
443
+ info->stack_top = info->stack_top->next;
444
+ info->stack_top->ts = get_ts();
445
+ info->stack_top->func = Py_NewRef(code);
446
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_FUNCTION_ARGS)) {
447
+ log_func_args(info->stack_top, PyEval_GetFrame(), self->log_func_repr);
448
+ }
449
+
450
+ info->curr_stack_depth += 1;
451
+
452
+ return 0;
453
+
454
+ cleanup_ignore:
455
+
456
+ if (info) {
457
+ info->ignore_stack_depth += 1;
458
+ info->curr_stack_depth += 1;
459
+ }
460
+
461
+ return 0;
462
+ }
463
+
464
+ int
465
+ tracer_ccall_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
466
+ {
467
+ struct ThreadInfo* info = NULL;
468
+
469
+ if (prepare_before_trace(self, 1, &info) <= 0) {
470
+ // For now we think -1 and 0 should both return because we should not
471
+ // have the -1 case.
472
+ goto cleanup_ignore;
473
+ }
474
+
475
+ PyCFunctionObject* cfunc = (PyCFunctionObject*) arg;
476
+
477
+ if (cfunc->m_self == (PyObject*)self) {
478
+ goto cleanup_ignore;
479
+ }
480
+
481
+ // If it's a call, we need a new node, and we need to update the stack
482
+ if (!info->stack_top->next) {
483
+ info->stack_top->next = (struct FunctionNode*) PyMem_Calloc(1, sizeof(struct FunctionNode));
484
+ info->stack_top->next->prev = info->stack_top;
485
+ }
486
+ info->stack_top = info->stack_top->next;
487
+ info->stack_top->ts = get_ts();
488
+ info->stack_top->func = Py_NewRef(arg);
489
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_FUNCTION_ARGS)) {
490
+ log_func_args(info->stack_top, PyEval_GetFrame(), self->log_func_repr);
491
+ }
492
+
493
+ info->curr_stack_depth += 1;
494
+
495
+ return 0;
496
+
497
+ cleanup_ignore:
498
+
499
+ if (info) {
500
+ info->ignore_stack_depth += 1;
501
+ info->curr_stack_depth += 1;
502
+ }
503
+
504
+ return 0;
505
+ }
506
+
507
+ int
508
+ tracer_pyreturn_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
509
+ {
510
+ struct ThreadInfo* info = NULL;
511
+
512
+ if (prepare_before_trace(self, 0, &info) <= 0) {
513
+ // For now we think -1 and 0 should both return because we should not
514
+ // have the -1 case.
515
+ goto cleanup_ignore;
516
+ }
517
+
518
+ struct FunctionNode* stack_top = info->stack_top;
519
+ if (stack_top->prev) {
520
+ // if stack_top has prev, it's not the fake node so it's at least root
521
+ int64_t dur = get_ts() - info->stack_top->ts;
522
+ int log_this_entry = self->min_duration == 0 || dur_ts_to_ns(dur) >= self->min_duration;
523
+
524
+ if (log_this_entry) {
525
+ PyCodeObject* call_code = (PyCodeObject*) stack_top->func;
526
+
527
+ if (!PyCode_Check(call_code) || call_code != code) {
528
+ self->collecting = 0;
529
+ PyErr_WarnEx(PyExc_RuntimeWarning,
530
+ "VizTracer: Unexpected function return, tracing is stopped", 1);
531
+ return 0;
532
+ }
533
+
534
+ struct EventNode* node = get_next_node(self);
535
+
536
+ node->ntype = FEE_NODE;
537
+ node->ts = info->stack_top->ts;
538
+ node->data.fee.dur = dur;
539
+ node->tid = info->tid;
540
+ node->data.fee.type = PyTrace_RETURN;
541
+ node->data.fee.code = (PyCodeObject*)Py_NewRef(code);
542
+ // steal the reference when return
543
+ node->data.fee.args = Py_XNewRef(stack_top->args);
544
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_RETURN_VALUE)) {
545
+ PyObject* repr = NULL;
546
+ if (self->log_func_repr) {
547
+ repr = PyObject_CallOneArg(self->log_func_repr, arg);
548
+ } else {
549
+ repr = PyObject_Repr(arg);
550
+ }
551
+ if (!repr) {
552
+ repr = PyUnicode_FromString("Not Displayable");
553
+ PyErr_Clear();
554
+ }
555
+ node->data.fee.retval = repr;
556
+ }
557
+
558
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC)) {
559
+ node->data.fee.asyncio_task = Py_XNewRef(info->curr_task);
560
+ }
561
+ }
562
+ // Finish return whether to log the data
563
+ info->stack_top = info->stack_top->prev;
564
+
565
+ Py_CLEAR(stack_top->args);
566
+ Py_CLEAR(stack_top->func);
567
+
568
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC) &&
569
+ info->curr_task &&
570
+ PyEval_GetFrame() == info->curr_task_frame) {
571
+ Py_CLEAR(info->curr_task);
572
+ Py_CLEAR(info->curr_task_frame);
573
+ }
574
+ }
575
+
576
+ if (info->curr_stack_depth > 0) {
577
+ info->curr_stack_depth -= 1;
578
+ }
579
+
580
+ return 0;
581
+
582
+ cleanup_ignore:
583
+
584
+ if (info) {
585
+ if (info->curr_stack_depth > 0) {
586
+ info->curr_stack_depth -= 1;
587
+ }
588
+
589
+ if (info->ignore_stack_depth > 0) {
590
+ info->ignore_stack_depth -= 1;
591
+ }
592
+ }
593
+
594
+ return 0;
595
+ }
596
+
597
+ int
598
+ tracer_creturn_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
599
+ {
600
+ struct ThreadInfo* info = NULL;
601
+
602
+ if (prepare_before_trace(self, 0, &info) <= 0) {
603
+ // For now we think -1 and 0 should both return because we should not
604
+ // have the -1 case.
605
+ goto cleanup_ignore;
606
+ }
607
+
608
+ struct FunctionNode* stack_top = info->stack_top;
609
+ if (stack_top->prev) {
610
+ // if stack_top has prev, it's not the fake node so it's at least root
611
+ int64_t dur = get_ts() - info->stack_top->ts;
612
+ int log_this_entry = self->min_duration == 0 || dur_ts_to_ns(dur) >= self->min_duration;
613
+
614
+ if (log_this_entry) {
615
+ PyCFunctionObject* cfunc = (PyCFunctionObject*) stack_top->func;
616
+
617
+ if (!PyCFunction_Check(cfunc)) {
618
+ self->collecting = 0;
619
+ PyErr_WarnEx(PyExc_RuntimeWarning,
620
+ "VizTracer: Unexpected function return, tracing is stopped", 1);
621
+ return 0;
622
+ }
623
+
624
+ struct EventNode* node = get_next_node(self);
625
+
626
+ node->ntype = FEE_NODE;
627
+ node->ts = info->stack_top->ts;
628
+ node->data.fee.dur = dur;
629
+ node->tid = info->tid;
630
+ node->data.fee.type = PyTrace_C_RETURN;
631
+ node->data.fee.ml_name = cfunc->m_ml->ml_name;
632
+ if (cfunc->m_module) {
633
+ // The function belongs to a module
634
+ node->data.fee.m_module = Py_NewRef(cfunc->m_module);
635
+ } else {
636
+ // The function is a class method
637
+ node->data.fee.m_module = NULL;
638
+ if (cfunc->m_self) {
639
+ // It's not a static method, has __self__
640
+ node->data.fee.tp_name = cfunc->m_self->ob_type->tp_name;
641
+ } else {
642
+ // It's a static method, does not have __self__
643
+ node->data.fee.tp_name = NULL;
644
+ }
645
+ }
646
+
647
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC)) {
648
+ node->data.fee.asyncio_task = Py_XNewRef(info->curr_task);
649
+ }
650
+ }
651
+ // Finish return whether to log the data
652
+ info->stack_top = info->stack_top->prev;
653
+
654
+ Py_CLEAR(stack_top->args);
655
+ Py_CLEAR(stack_top->func);
656
+ }
657
+
658
+
659
+ if (info->curr_stack_depth > 0) {
660
+ info->curr_stack_depth -= 1;
661
+ }
662
+
663
+ return 0;
664
+
665
+ cleanup_ignore:
666
+
667
+ if (info) {
668
+ if (info->curr_stack_depth > 0) {
669
+ info->curr_stack_depth -= 1;
670
+ }
671
+
672
+ if (info->ignore_stack_depth > 0) {
673
+ info->ignore_stack_depth -= 1;
674
+ }
675
+ }
676
+
677
+ return 0;
678
+ }
679
+
680
+ // sys.setprofile mechanism
681
+
682
+ int
683
+ tracer_tracefunc(PyObject* obj, PyFrameObject* frame, int what, PyObject* arg)
684
+ {
685
+ TracerObject* self = (TracerObject*) obj;
686
+ int ret = 0;
687
+
688
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_IGNORE_C_FUNCTION) &&
689
+ (what == PyTrace_C_CALL || what == PyTrace_C_RETURN || what == PyTrace_C_EXCEPTION)) {
690
+ return 0;
691
+ }
692
+
693
+ PyCodeObject* code = PyFrame_GetCode(frame);
694
+
695
+ switch (what) {
696
+ case PyTrace_CALL:
697
+ ret = tracer_pycall_callback(self, code);
698
+ break;
699
+ case PyTrace_C_CALL:
700
+ ret = tracer_ccall_callback(self, code, arg);
701
+ break;
702
+ case PyTrace_RETURN:
703
+ ret = tracer_pyreturn_callback(self, code, arg);
704
+ break;
705
+ case PyTrace_C_RETURN:
706
+ case PyTrace_C_EXCEPTION:
707
+ ret = tracer_creturn_callback(self, code, arg);
708
+ break;
709
+ default:
710
+ return 0;
711
+ }
712
+
713
+ Py_DECREF(code);
714
+
715
+ return ret;
716
+ }
717
+
718
+ static PyObject*
719
+ tracer_threadtracefunc(PyObject* obj, PyObject* args)
720
+ {
721
+ PyFrameObject* frame = NULL;
722
+ char* event = NULL;
723
+ PyObject* trace_args = NULL;
724
+ int what = 0;
725
+ if (!PyArg_ParseTuple(args, "OsO", &frame, &event, &trace_args)) {
726
+ printf("Error when parsing arguments!\n");
727
+ exit(1);
728
+ }
729
+ PyEval_SetProfile(tracer_tracefunc, obj);
730
+ if (!strcmp(event, "call")) {
731
+ what = PyTrace_CALL;
732
+ } else if (!strcmp(event, "c_call")) {
733
+ what = PyTrace_C_CALL;
734
+ } else if (!strcmp(event, "return")) {
735
+ what = PyTrace_RETURN;
736
+ } else if (!strcmp(event, "c_return")) {
737
+ what = PyTrace_C_RETURN;
738
+ } else if (!strcmp(event, "c_exception")) {
739
+ what = PyTrace_C_EXCEPTION;
740
+ } else {
741
+ printf("Unexpected event type: %s\n", event);
742
+ }
743
+ tracer_tracefunc(obj, frame, what, trace_args);
744
+ Py_RETURN_NONE;
745
+ }
746
+
747
+ // sys.monitoring mechanism
748
+
749
+ PyObject* get_cfunc_from_callable(PyObject* callable, PyObject* self_arg)
750
+ {
751
+ // return a new reference
752
+ if (PyCFunction_Check(callable)) {
753
+ Py_INCREF(callable);
754
+ return (PyObject*)((PyCFunctionObject*)callable);
755
+ }
756
+ if (Py_TYPE(callable) == &PyMethodDescr_Type) {
757
+ /* For backwards compatibility need to
758
+ * convert to builtin method */
759
+
760
+ /* If no arg, skip */
761
+ if (self_arg == sys_monitoring_missing) {
762
+ return NULL;
763
+ }
764
+ PyObject* meth = Py_TYPE(callable)->tp_descr_get(
765
+ callable, self_arg, (PyObject*)Py_TYPE(self_arg));
766
+ if (meth == NULL) {
767
+ return NULL;
768
+ }
769
+ if (PyCFunction_Check(meth)) {
770
+ return (PyObject*)((PyCFunctionObject*)meth);
771
+ }
772
+ } else if (Py_TYPE(callable) == &PyMethod_Type) {
773
+ PyObject* func = PyMethod_GET_FUNCTION(callable);
774
+ if (func && PyCFunction_Check(func)) {
775
+ Py_INCREF(func);
776
+ return func;
777
+ }
778
+ }
779
+ return NULL;
780
+ }
781
+
782
+ PyObject*
783
+ _pystart_callback(PyObject* self, PyObject *const *args, Py_ssize_t nargs)
784
+ {
785
+ PyCodeObject* code = (PyCodeObject*)args[0];
786
+ int ret = tracer_pycall_callback((TracerObject*)self, code);
787
+ if (ret != 0) {
788
+ return NULL;
789
+ }
790
+ Py_RETURN_NONE;
791
+ }
792
+
793
+ PyObject*
794
+ _pyreturn_callback(PyObject* self, PyObject *const *args, Py_ssize_t nargs)
795
+ {
796
+ PyCodeObject* code = (PyCodeObject*)args[0];
797
+ PyObject* arg = args[2];
798
+ int ret = tracer_pyreturn_callback((TracerObject*)self, code, arg);
799
+ if (ret != 0) {
800
+ return NULL;
801
+ }
802
+ Py_RETURN_NONE;
803
+ }
804
+
805
+ PyObject*
806
+ _ccall_callback(PyObject* self, PyObject *const *args, Py_ssize_t nargs)
807
+ {
808
+ PyCodeObject* code = (PyCodeObject*)args[0];
809
+ PyObject* cfunc = get_cfunc_from_callable(args[2], args[3]);
810
+ if (!cfunc) {
811
+ Py_RETURN_NONE;
812
+ }
813
+ int ret = tracer_ccall_callback((TracerObject*)self, code, cfunc);
814
+ Py_DECREF(cfunc);
815
+ if (ret != 0) {
816
+ return NULL;
817
+ }
818
+ Py_RETURN_NONE;
819
+ }
820
+
821
+ PyObject*
822
+ _creturn_callback(PyObject* self, PyObject *const *args, Py_ssize_t nargs)
823
+ {
824
+ PyCodeObject* code = (PyCodeObject*)args[0];
825
+ PyObject* cfunc = get_cfunc_from_callable(args[2], args[3]);
826
+ if (!cfunc) {
827
+ Py_RETURN_NONE;
828
+ }
829
+ int ret = tracer_creturn_callback((TracerObject*)self, code, cfunc);
830
+ Py_DECREF(cfunc);
831
+ if (ret != 0) {
832
+ return NULL;
833
+ }
834
+ Py_RETURN_NONE;
835
+ }
836
+
837
+ static struct {
838
+ unsigned int event;
839
+ PyMethodDef callback_method;
840
+ } callback_table[] = {
841
+ {PY_MONITORING_EVENT_PY_START,
842
+ {"_pystart_callback", (PyCFunction)_pystart_callback, METH_FASTCALL, NULL}},
843
+ {PY_MONITORING_EVENT_PY_RESUME,
844
+ {"_pystart_callback", (PyCFunction)_pystart_callback, METH_FASTCALL, NULL}},
845
+ {PY_MONITORING_EVENT_PY_THROW,
846
+ {"_pystart_callback", (PyCFunction)_pystart_callback, METH_FASTCALL, NULL}},
847
+ {PY_MONITORING_EVENT_PY_RETURN,
848
+ {"_pyreturn_callback", (PyCFunction)_pyreturn_callback, METH_FASTCALL, NULL}},
849
+ {PY_MONITORING_EVENT_PY_YIELD,
850
+ {"_pyreturn_callback", (PyCFunction)_pyreturn_callback, METH_FASTCALL, NULL}},
851
+ {PY_MONITORING_EVENT_PY_UNWIND,
852
+ {"_pyreturn_callback", (PyCFunction)_pyreturn_callback, METH_FASTCALL, NULL}},
853
+ {PY_MONITORING_EVENT_CALL,
854
+ {"_ccall_callback", (PyCFunction)_ccall_callback, METH_FASTCALL, NULL}},
855
+ {PY_MONITORING_EVENT_C_RETURN,
856
+ {"_creturn_callback", (PyCFunction)_creturn_callback, METH_FASTCALL, NULL}},
857
+ {PY_MONITORING_EVENT_C_RAISE,
858
+ {"_creturn_callback", (PyCFunction)_creturn_callback, METH_FASTCALL, NULL}},
859
+ {0,
860
+ {NULL, NULL, 0, NULL}}
861
+ };
862
+
863
+ int
864
+ enable_monitoring(TracerObject* self)
865
+ {
866
+ unsigned int all_events = 0;
867
+ PyObject* monitoring = PyObject_GetAttrString(sys_module, "monitoring");
868
+ if (!monitoring) {
869
+ PyErr_SetString(PyExc_RuntimeError, "Failed to access sys.monitoring");
870
+ goto cleanup;
871
+ }
872
+
873
+ PyObject* ret = PyObject_CallMethod(monitoring, "use_tool_id",
874
+ "is", SNAPTRACE_TOOL_ID, "viztracer");
875
+ if (!ret) {
876
+ PyErr_Clear();
877
+ PyObject_CallMethod(monitoring, "free_tool_id", "i", SNAPTRACE_TOOL_ID);
878
+ ret = PyObject_CallMethod(monitoring, "use_tool_id",
879
+ "is", SNAPTRACE_TOOL_ID, "viztracer");
880
+ if (!ret) {
881
+ goto cleanup;
882
+ }
883
+ }
884
+ Py_DECREF(ret);
885
+
886
+ for (int i = 0; callback_table[i].callback_method.ml_meth != 0; i++) {
887
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_IGNORE_C_FUNCTION) &&
888
+ (callback_table[i].event == PY_MONITORING_EVENT_CALL ||
889
+ callback_table[i].event == PY_MONITORING_EVENT_C_RETURN ||
890
+ callback_table[i].event == PY_MONITORING_EVENT_C_RAISE)) {
891
+ continue;
892
+ }
893
+ unsigned int event = (1 << callback_table[i].event);
894
+ PyObject* callback = PyCFunction_New(&callback_table[i].callback_method, (PyObject*)self);
895
+
896
+ PyObject* regsiter_result = PyObject_CallMethod(monitoring, "register_callback",
897
+ "iiO", SNAPTRACE_TOOL_ID, event, callback);
898
+ Py_DECREF(callback);
899
+
900
+ if (!regsiter_result) {
901
+ goto cleanup;
902
+ }
903
+ Py_DECREF(regsiter_result);
904
+ all_events |= event;
905
+ }
906
+
907
+ PyObject* event_result = PyObject_CallMethod(monitoring, "set_events",
908
+ "ii", SNAPTRACE_TOOL_ID, all_events);
909
+ if (!event_result) {
910
+ goto cleanup;
911
+ }
912
+ Py_DECREF(event_result);
913
+
914
+ cleanup:
915
+
916
+ Py_XDECREF(monitoring);
917
+
918
+ if (PyErr_Occurred()) {
919
+ return -1;
920
+ }
921
+
922
+ return 0;
923
+ }
924
+
925
+ int disable_monitoring(TracerObject* self)
926
+ {
927
+ PyObject* monitoring = PyObject_GetAttrString(sys_module, "monitoring");
928
+ if (!monitoring) {
929
+ PyErr_SetString(PyExc_RuntimeError, "Failed to access sys.monitoring");
930
+ goto cleanup;
931
+ }
932
+
933
+ PyObject* curr_tool = PyObject_CallMethod(monitoring, "get_tool", "i", SNAPTRACE_TOOL_ID);
934
+
935
+ if (!curr_tool) {
936
+ goto cleanup;
937
+ }
938
+
939
+ if (curr_tool == Py_None) {
940
+ // No current tool, nothing to do
941
+ Py_DECREF(curr_tool);
942
+ goto cleanup;
943
+ }
944
+
945
+ PyObject* event_result = PyObject_CallMethod(monitoring, "set_events",
946
+ "ii", SNAPTRACE_TOOL_ID, 0);
947
+ if (!event_result) {
948
+ goto cleanup;
949
+ }
950
+ Py_DECREF(event_result);
951
+
952
+ PyObject* ret = PyObject_CallMethod(monitoring, "free_tool_id",
953
+ "i", SNAPTRACE_TOOL_ID);
954
+ if (!ret) {
955
+ goto cleanup;
956
+ }
957
+ Py_DECREF(ret);
958
+
959
+ cleanup:
960
+
961
+ Py_XDECREF(monitoring);
962
+
963
+ if (PyErr_Occurred()) {
964
+ return -1;
965
+ }
966
+
967
+ return 0;
968
+ }
969
+
970
+ // =============================================================================
971
+ // snaptrace.Tracer methods
972
+ // =============================================================================
973
+
974
+ static void
975
+ tracer__flush_unfinished(TracerObject* self, int flush_as_finish)
976
+ {
977
+ SNAPTRACE_THREAD_PROTECT_START(self);
978
+
979
+ struct MetadataNode* meta_node = self->metadata_head;
980
+ while(meta_node) {
981
+ struct ThreadInfo* info = meta_node->thread_info;
982
+
983
+ if (info == NULL) {
984
+ meta_node = meta_node->next;
985
+ continue;
986
+ }
987
+
988
+ struct FunctionNode* func_node = info->stack_top;
989
+
990
+ while (func_node->prev && info->curr_stack_depth > 0) {
991
+ // Fake a FEE node to get the name
992
+ struct EventNode* fee_node = get_next_node(self);
993
+
994
+ fee_node->ntype = FEE_NODE;
995
+ fee_node->ts = func_node->ts;
996
+ fee_node->tid = meta_node->tid;
997
+
998
+ if (flush_as_finish) {
999
+ fee_node->data.fee.dur = get_ts() - func_node->ts;
1000
+ } else {
1001
+ fee_node->data.fee.dur = 0;
1002
+ }
1003
+
1004
+ if (PyCode_Check(func_node->func)) {
1005
+ PyCodeObject* code = (PyCodeObject*) func_node->func;
1006
+ if (flush_as_finish) {
1007
+ fee_node->data.fee.type = PyTrace_RETURN;
1008
+ } else {
1009
+ fee_node->data.fee.type = PyTrace_CALL;
1010
+ }
1011
+ fee_node->data.fee.code = (PyCodeObject*)Py_NewRef(code);
1012
+ } else if (PyCFunction_Check(func_node->func)) {
1013
+ PyCFunctionObject* cfunc = (PyCFunctionObject*) func_node->func;
1014
+ if (flush_as_finish) {
1015
+ fee_node->data.fee.type = PyTrace_C_RETURN;
1016
+ } else {
1017
+ fee_node->data.fee.type = PyTrace_C_CALL;
1018
+ }
1019
+ fee_node->data.fee.ml_name = cfunc->m_ml->ml_name;
1020
+ if (cfunc->m_module) {
1021
+ // The function belongs to a module
1022
+ fee_node->data.fee.m_module = Py_NewRef(cfunc->m_module);
1023
+ } else {
1024
+ // The function is a class method
1025
+ fee_node->data.fee.m_module = NULL;
1026
+ if (cfunc->m_self) {
1027
+ // It's not a static method, has __self__
1028
+ fee_node->data.fee.tp_name = cfunc->m_self->ob_type->tp_name;
1029
+ } else {
1030
+ // It's a static method, does not have __self__
1031
+ fee_node->data.fee.tp_name = NULL;
1032
+ }
1033
+ }
1034
+ }
1035
+
1036
+ // Clean up the node
1037
+ Py_CLEAR(func_node->args);
1038
+ Py_CLEAR(func_node->func);
1039
+
1040
+ func_node = func_node->prev;
1041
+ info->curr_stack_depth -= 1;
1042
+ }
1043
+ info->stack_top = func_node;
1044
+ meta_node = meta_node->next;
1045
+ }
1046
+
1047
+ SNAPTRACE_THREAD_PROTECT_END(self);
1048
+ }
1049
+
1050
+ static PyObject*
1051
+ tracer_start(TracerObject* self, PyObject* Py_UNUSED(unused))
1052
+ {
1053
+ if (curr_tracer) {
1054
+ printf("Warning! Overwrite tracer! You should not have two VizTracer recording at the same time!\n");
1055
+ } else {
1056
+ curr_tracer = self;
1057
+ }
1058
+
1059
+ self->collecting = 1;
1060
+ #if PY_VERSION_HEX >= 0x030C0000
1061
+ if (enable_monitoring(self) != 0) {
1062
+ return NULL;
1063
+ };
1064
+ #else
1065
+ PyEval_SetProfile(tracer_tracefunc, (PyObject*) self);
1066
+ #endif
1067
+
1068
+ Py_RETURN_NONE;
1069
+ }
1070
+
1071
+ static PyObject*
1072
+ tracer_stop(TracerObject* self, PyObject* stop_option)
1073
+ {
1074
+ if (self) {
1075
+ struct ThreadInfo* info = get_thread_info(self);
1076
+ if (!info) {
1077
+ self->collecting = 0;
1078
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to get thread info. This should not happen.");
1079
+ return NULL;
1080
+ }
1081
+ self->collecting = 0;
1082
+
1083
+ if (PyUnicode_CheckExact(stop_option) &&
1084
+ strcmp(PyUnicode_AsUTF8(stop_option), "flush_as_finish") == 0) {
1085
+ tracer__flush_unfinished(self, 1);
1086
+ } else {
1087
+ tracer__flush_unfinished(self, 0);
1088
+ }
1089
+ info->curr_stack_depth = 0;
1090
+ info->ignore_stack_depth = 0;
1091
+ info->paused = 0;
1092
+ }
1093
+
1094
+ curr_tracer = NULL;
1095
+ #if PY_VERSION_HEX >= 0x030C0000
1096
+ if (disable_monitoring(self) != 0) {
1097
+ return NULL;
1098
+ }
1099
+ #else
1100
+ PyEval_SetProfile(NULL, NULL);
1101
+ #endif
1102
+
1103
+ Py_RETURN_NONE;
1104
+ }
1105
+
1106
+ static PyObject*
1107
+ tracer_pause(TracerObject* self, PyObject* Py_UNUSED(unused))
1108
+ {
1109
+ if (self->collecting) {
1110
+ struct ThreadInfo* info = get_thread_info((TracerObject*)self);
1111
+ if (!info) {
1112
+ self->collecting = 0;
1113
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to get thread info. This should not happen.");
1114
+ return NULL;
1115
+ }
1116
+
1117
+ if (!info->paused) {
1118
+ // When we enter this function, tracer.pause has been called.
1119
+ // We need to reduce the ignore_stack_depth to simulate the
1120
+ // returns from these two functions
1121
+ info->ignore_stack_depth -= 1;
1122
+ info->paused = 1;
1123
+ #if PY_VERSION_HEX >= 0x030C0000
1124
+ if (disable_monitoring(self) != 0) {
1125
+ return NULL;
1126
+ }
1127
+ #else
1128
+ PyEval_SetProfile(NULL, NULL);
1129
+ #endif
1130
+ }
1131
+ }
1132
+
1133
+ Py_RETURN_NONE;
1134
+ }
1135
+
1136
+ static PyObject*
1137
+ tracer_resume(TracerObject* self, PyObject* Py_UNUSED(unused))
1138
+ {
1139
+ if (self->collecting) {
1140
+ struct ThreadInfo* info = get_thread_info(self);
1141
+ if (!info) {
1142
+ self->collecting = 0;
1143
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to get thread info. This should not happen.");
1144
+ return NULL;
1145
+ }
1146
+
1147
+ if (info->paused) {
1148
+ info->paused = 0;
1149
+ #if PY_VERSION_HEX >= 0x030C0000
1150
+ if (enable_monitoring(self) != 0) {
1151
+ return NULL;
1152
+ }
1153
+ #else
1154
+ PyEval_SetProfile(tracer_tracefunc, (PyObject*)self);
1155
+ #endif
1156
+ }
1157
+ }
1158
+
1159
+ Py_RETURN_NONE;
1160
+ }
1161
+
1162
+ static PyObject*
1163
+ tracer_load(TracerObject* self, PyObject* Py_UNUSED(unused))
1164
+ {
1165
+ PyObject* lst = PyList_New(0);
1166
+
1167
+ SNAPTRACE_THREAD_PROTECT_START(self);
1168
+ struct EventNode* curr = self->buffer + self->buffer_head_idx;
1169
+ PyObject* pid = NULL;
1170
+ PyObject* cat_fee = PyUnicode_FromString("FEE");
1171
+ PyObject* cat_instant = PyUnicode_FromString("INSTANT");
1172
+ PyObject* ph_B = PyUnicode_FromString("B");
1173
+ PyObject* ph_i = PyUnicode_FromString("i");
1174
+ PyObject* ph_X = PyUnicode_FromString("X");
1175
+ PyObject* ph_C = PyUnicode_FromString("C");
1176
+ PyObject* ph_M = PyUnicode_FromString("M");
1177
+ unsigned long counter = 0;
1178
+ unsigned long prev_counter = 0;
1179
+ struct MetadataNode* metadata_node = NULL;
1180
+ PyObject* task_dict = NULL;
1181
+ PyObject* func_name_dict = PyDict_New();
1182
+
1183
+ if (self->fix_pid > 0) {
1184
+ pid = PyLong_FromLong(self->fix_pid);
1185
+ } else {
1186
+ #if _WIN32
1187
+ pid = PyLong_FromLong(GetCurrentProcessId());
1188
+ #else
1189
+ pid = PyLong_FromLong(getpid());
1190
+ #endif
1191
+ }
1192
+
1193
+ // == Load the metadata first ==
1194
+ // Process Name
1195
+ {
1196
+ PyObject* dict = PyDict_New();
1197
+ PyObject* args = PyDict_New();
1198
+ PyObject* process_name_string = PyUnicode_FromString("process_name");
1199
+ PyObject* process_name = NULL;
1200
+
1201
+ if (self->process_name) {
1202
+ process_name = Py_NewRef(self->process_name);
1203
+ } else {
1204
+ PyObject* current_process_method = PyObject_GetAttrString(multiprocessing_module, "current_process");
1205
+ if (!current_process_method) {
1206
+ perror("Failed to access multiprocessing.current_process()");
1207
+ exit(-1);
1208
+ }
1209
+ PyObject* current_process = PyObject_CallNoArgs(current_process_method);
1210
+ if (!current_process_method) {
1211
+ perror("Failed to access multiprocessing.current_process()");
1212
+ exit(-1);
1213
+ }
1214
+ process_name = PyObject_GetAttrString(current_process, "name");
1215
+ Py_DECREF(current_process_method);
1216
+ Py_DECREF(current_process);
1217
+ }
1218
+
1219
+ PyDict_SetItemString(dict, "ph", ph_M);
1220
+ PyDict_SetItemString(dict, "pid", pid);
1221
+ PyDict_SetItemString(dict, "tid", pid);
1222
+ PyDict_SetItemString(dict, "name", process_name_string);
1223
+ Py_DECREF(process_name_string);
1224
+ PyDict_SetItemString(args, "name", process_name);
1225
+ PyDict_SetItemString(dict, "args", args);
1226
+ Py_DECREF(args);
1227
+ Py_DECREF(process_name);
1228
+ PyList_Append(lst, dict);
1229
+ }
1230
+
1231
+
1232
+ // Thread Name
1233
+ metadata_node = self->metadata_head;
1234
+ while (metadata_node) {
1235
+ PyObject* dict = PyDict_New();
1236
+ PyObject* args = PyDict_New();
1237
+ PyObject* tid = PyLong_FromLong(metadata_node->tid);
1238
+ PyObject* thread_name_string = PyUnicode_FromString("thread_name");
1239
+
1240
+ PyDict_SetItemString(dict, "ph", ph_M);
1241
+ PyDict_SetItemString(dict, "pid", pid);
1242
+ PyDict_SetItemString(dict, "tid", tid);
1243
+ Py_DECREF(tid);
1244
+ PyDict_SetItemString(dict, "name", thread_name_string);
1245
+ Py_DECREF(thread_name_string);
1246
+ PyDict_SetItemString(args, "name", metadata_node->name);
1247
+ PyDict_SetItemString(dict, "args", args);
1248
+ Py_DECREF(args);
1249
+ metadata_node = metadata_node->next;
1250
+ PyList_Append(lst, dict);
1251
+ }
1252
+
1253
+ // Task Name if using LOG_ASYNC
1254
+ // We need to make up some thread id for the task
1255
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC)) {
1256
+ task_dict = PyDict_New();
1257
+ }
1258
+
1259
+ while (curr != self->buffer + self->buffer_tail_idx) {
1260
+ struct EventNode* node = curr;
1261
+ PyObject* dict = PyDict_New();
1262
+ PyObject* name = NULL;
1263
+ PyObject* tid = PyLong_FromLong(node->tid);
1264
+ PyObject* ts = PyFloat_FromDouble(system_ts_to_us(node->ts));
1265
+
1266
+ PyDict_SetItemString(dict, "pid", pid);
1267
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC)) {
1268
+ if (curr->data.fee.asyncio_task == NULL) {
1269
+ PyDict_SetItemString(dict, "tid", tid);
1270
+ } else {
1271
+ PyObject* task_id = PyLong_FromUnsignedLongLong(((uintptr_t)curr->data.fee.asyncio_task) & 0xffffff);
1272
+ PyDict_SetItemString(dict, "tid", task_id);
1273
+ if (!PyDict_Contains(task_dict, task_id)) {
1274
+ PyObject* task_name = NULL;
1275
+ if (PyObject_HasAttrString(curr->data.fee.asyncio_task, "get_name")) {
1276
+ PyObject* task_name_method = PyObject_GetAttrString(curr->data.fee.asyncio_task, "get_name");
1277
+ task_name = PyObject_CallNoArgs(task_name_method);
1278
+ Py_DECREF(task_name_method);
1279
+ } else if (PyObject_HasAttrString(curr->data.fee.asyncio_task, "name")) {
1280
+ task_name = PyObject_GetAttrString(curr->data.fee.asyncio_task, "name");
1281
+ } else {
1282
+ task_name = PyUnicode_FromString("Task");
1283
+ }
1284
+ PyDict_SetItem(task_dict, task_id, task_name);
1285
+ Py_DECREF(task_name);
1286
+ }
1287
+ Py_DECREF(task_id);
1288
+ }
1289
+ } else {
1290
+ PyDict_SetItemString(dict, "tid", tid);
1291
+ }
1292
+ Py_DECREF(tid);
1293
+ PyDict_SetItemString(dict, "ts", ts);
1294
+ Py_DECREF(ts);
1295
+
1296
+ switch (node->ntype) {
1297
+ case FEE_NODE:
1298
+ name = get_name_from_fee_node(node, func_name_dict);
1299
+
1300
+ if (node->data.fee.type == PyTrace_CALL || node->data.fee.type == PyTrace_C_CALL) {
1301
+ PyDict_SetItemString(dict, "ph", ph_B);
1302
+ } else {
1303
+ PyDict_SetItemString(dict, "ph", ph_X);
1304
+ PyObject* dur = PyFloat_FromDouble(dur_ts_to_us(node->data.fee.dur));
1305
+ PyDict_SetItemString(dict, "dur", dur);
1306
+ Py_DECREF(dur);
1307
+ }
1308
+ PyDict_SetItemString(dict, "name", name);
1309
+ Py_DECREF(name);
1310
+
1311
+ PyObject* arg_dict = Py_XNewRef(node->data.fee.args);
1312
+ if (node->data.fee.retval) {
1313
+ if (!arg_dict) {
1314
+ arg_dict = PyDict_New();
1315
+ }
1316
+ PyDict_SetItemString(arg_dict, "return_value", node->data.fee.retval);
1317
+ }
1318
+ if (arg_dict) {
1319
+ PyDict_SetItemString(dict, "args", arg_dict);
1320
+ Py_DECREF(arg_dict);
1321
+ }
1322
+
1323
+ PyDict_SetItemString(dict, "cat", cat_fee);
1324
+ break;
1325
+ case INSTANT_NODE:
1326
+ PyDict_SetItemString(dict, "ph", ph_i);
1327
+ PyDict_SetItemString(dict, "cat", cat_instant);
1328
+ PyDict_SetItemString(dict, "name", node->data.instant.name);
1329
+ PyDict_SetItemString(dict, "args", node->data.instant.args);
1330
+ PyDict_SetItemString(dict, "s", node->data.instant.scope);
1331
+ break;
1332
+ case COUNTER_NODE:
1333
+ PyDict_SetItemString(dict, "ph", ph_C);
1334
+ PyDict_SetItemString(dict, "name", node->data.counter.name);
1335
+ PyDict_SetItemString(dict, "args", node->data.counter.args);
1336
+ break;
1337
+ case OBJECT_NODE:
1338
+ PyDict_SetItemString(dict, "ph", node->data.object.ph);
1339
+ PyDict_SetItemString(dict, "id", node->data.object.id);
1340
+ PyDict_SetItemString(dict, "name", node->data.object.name);
1341
+ if (!(node->data.object.args == Py_None)) {
1342
+ PyDict_SetItemString(dict, "args", node->data.object.args);
1343
+ }
1344
+ break;
1345
+ case RAW_NODE:
1346
+ // We still need to tid from node and we need the pid
1347
+ tid = PyLong_FromLong(node->tid);
1348
+
1349
+ Py_DECREF(dict);
1350
+ dict = node->data.raw;
1351
+
1352
+ PyDict_SetItemString(dict, "pid", pid);
1353
+ PyDict_SetItemString(dict, "tid", tid);
1354
+ Py_DECREF(tid);
1355
+
1356
+ Py_INCREF(dict);
1357
+ break;
1358
+ default:
1359
+ printf("Unknown Node Type!\n");
1360
+ exit(1);
1361
+ }
1362
+ clear_node(node);
1363
+ PyList_Append(lst, dict);
1364
+ Py_DECREF(dict);
1365
+ curr = curr + 1;
1366
+ if (curr == self->buffer + self->buffer_size) {
1367
+ curr = self->buffer;
1368
+ }
1369
+
1370
+ counter += 1;
1371
+ if (counter - prev_counter > 10000 && (counter - prev_counter) / ((1 + self->total_entries)/100) > 0) {
1372
+ verbose_printf(self, 1, "Loading data, %lu / %lu\r", counter, self->total_entries);
1373
+ prev_counter = counter;
1374
+ }
1375
+ }
1376
+
1377
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC)) {
1378
+ Py_ssize_t pos = 0;
1379
+ PyObject* key = NULL;
1380
+ PyObject* value = NULL;
1381
+ while (PyDict_Next(task_dict, &pos, &key, &value)) {
1382
+ PyObject* dict = PyDict_New();
1383
+ PyObject* args = PyDict_New();
1384
+ PyObject* tid = key;
1385
+ PyObject* thread_name_string = PyUnicode_FromString("thread_name");
1386
+
1387
+ PyDict_SetItemString(dict, "ph", ph_M);
1388
+ PyDict_SetItemString(dict, "pid", pid);
1389
+ PyDict_SetItemString(dict, "tid", tid);
1390
+ PyDict_SetItemString(dict, "name", thread_name_string);
1391
+ Py_DECREF(thread_name_string);
1392
+ PyDict_SetItemString(args, "name", value);
1393
+ PyDict_SetItemString(dict, "args", args);
1394
+ Py_DECREF(args);
1395
+ PyList_Append(lst, dict);
1396
+ }
1397
+ }
1398
+
1399
+ verbose_printf(self, 1, "Loading finish \n");
1400
+ Py_DECREF(pid);
1401
+ Py_DECREF(cat_fee);
1402
+ Py_DECREF(cat_instant);
1403
+ Py_DECREF(ph_B);
1404
+ Py_DECREF(ph_i);
1405
+ Py_DECREF(ph_X);
1406
+ Py_DECREF(ph_C);
1407
+ Py_DECREF(ph_M);
1408
+ Py_DECREF(func_name_dict);
1409
+ self->buffer_tail_idx = self->buffer_head_idx;
1410
+ SNAPTRACE_THREAD_PROTECT_END(self);
1411
+ return lst;
1412
+ }
1413
+
1414
+ static PyObject*
1415
+ tracer_dump(TracerObject* self, PyObject* args, PyObject* kw)
1416
+ {
1417
+ const char* filename = NULL;
1418
+ int sanitize_function_name = 0;
1419
+ static char* kwlist[] = {"filename", "sanitize_function_name", NULL};
1420
+ FILE* fptr = NULL;
1421
+
1422
+ if (!PyArg_ParseTupleAndKeywords(args, kw, "s|p", kwlist,
1423
+ &filename, &sanitize_function_name)) {
1424
+ return NULL;
1425
+ }
1426
+ fptr = fopen(filename, "w");
1427
+ if (!fptr) {
1428
+ PyErr_Format(PyExc_ValueError, "Can't open file %s to write", filename);
1429
+ return NULL;
1430
+ }
1431
+
1432
+ fprintf(fptr, "{\"traceEvents\":[");
1433
+
1434
+ SNAPTRACE_THREAD_PROTECT_START(self);
1435
+ struct EventNode* curr = self->buffer + self->buffer_head_idx;
1436
+ unsigned long pid = 0;
1437
+ uint8_t overflowed = ((self->buffer_tail_idx + 1) % self->buffer_size) == self->buffer_head_idx;
1438
+ struct MetadataNode* metadata_node = NULL;
1439
+ PyObject* task_dict = NULL;
1440
+
1441
+ if (self->fix_pid > 0) {
1442
+ pid = self->fix_pid;
1443
+ } else {
1444
+ #if _WIN32
1445
+ pid = GetCurrentProcessId();
1446
+ #else
1447
+ pid = getpid();
1448
+ #endif
1449
+ }
1450
+
1451
+ // == Load the metadata first ==
1452
+ // Process Name
1453
+ {
1454
+ PyObject* process_name = NULL;
1455
+ if (self->process_name) {
1456
+ process_name = Py_NewRef(self->process_name);
1457
+ } else {
1458
+ PyObject* current_process_method = PyObject_GetAttrString(multiprocessing_module, "current_process");
1459
+ if (!current_process_method) {
1460
+ perror("Failed to access multiprocessing.current_process()");
1461
+ exit(-1);
1462
+ }
1463
+ PyObject* current_process = PyObject_CallNoArgs(current_process_method);
1464
+ if (!current_process_method) {
1465
+ perror("Failed to access multiprocessing.current_process()");
1466
+ exit(-1);
1467
+ }
1468
+ process_name = PyObject_GetAttrString(current_process, "name");
1469
+ Py_DECREF(current_process_method);
1470
+ Py_DECREF(current_process);
1471
+ }
1472
+
1473
+ fprintf(fptr, "{\"ph\":\"M\",\"pid\":%lu,\"tid\":%lu,\"name\":\"process_name\",\"args\":{\"name\":\"",
1474
+ pid, pid);
1475
+ fprint_escape(fptr, PyUnicode_AsUTF8(process_name));
1476
+ fprintf(fptr, "\"}},");
1477
+ Py_DECREF(process_name);
1478
+ }
1479
+
1480
+ // Thread Name
1481
+ metadata_node = self->metadata_head;
1482
+ while (metadata_node) {
1483
+ fprintf(fptr, "{\"ph\":\"M\",\"pid\":%lu,\"tid\":%lu,\"name\":\"thread_name\",\"args\":{\"name\":\"",
1484
+ pid, metadata_node->tid);
1485
+ fprint_escape(fptr, PyUnicode_AsUTF8(metadata_node->name));
1486
+ fprintf(fptr, "\"}},");
1487
+ metadata_node = metadata_node->next;
1488
+ }
1489
+
1490
+ // Task Name if using LOG_ASYNC
1491
+ // We need to make up some thread id for the task
1492
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC)) {
1493
+ task_dict = PyDict_New();
1494
+ }
1495
+
1496
+ while (curr != self->buffer + self->buffer_tail_idx) {
1497
+ struct EventNode* node = curr;
1498
+ long long ts_long = system_ts_to_ns(node->ts);
1499
+ unsigned long tid = node->tid;
1500
+
1501
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC)) {
1502
+ if (curr->data.fee.asyncio_task != NULL) {
1503
+ tid = (unsigned long)(((uintptr_t)curr->data.fee.asyncio_task) & 0xffffff);
1504
+ PyObject* task_id = PyLong_FromLong(tid);
1505
+ if (!PyDict_Contains(task_dict, task_id)) {
1506
+ PyObject* task_name = NULL;
1507
+ if (PyObject_HasAttrString(curr->data.fee.asyncio_task, "get_name")) {
1508
+ PyObject* task_name_method = PyObject_GetAttrString(curr->data.fee.asyncio_task, "get_name");
1509
+ task_name = PyObject_CallNoArgs(task_name_method);
1510
+ Py_DECREF(task_name_method);
1511
+ } else if (PyObject_HasAttrString(curr->data.fee.asyncio_task, "name")) {
1512
+ task_name = PyObject_GetAttrString(curr->data.fee.asyncio_task, "name");
1513
+ } else {
1514
+ task_name = PyUnicode_FromString("Task");
1515
+ }
1516
+ PyDict_SetItem(task_dict, task_id, task_name);
1517
+ Py_DECREF(task_name);
1518
+ }
1519
+ Py_DECREF(task_id);
1520
+ }
1521
+ }
1522
+ if (node->ntype != RAW_NODE) {
1523
+ // printf("%f") is about 10x slower than print("%d")
1524
+ fprintf(fptr, "{\"pid\":%lu,\"tid\":%lu,\"ts\":%lld.%03lld,", pid, tid, ts_long / 1000, ts_long % 1000);
1525
+ }
1526
+
1527
+ switch (node->ntype) {
1528
+ case FEE_NODE:
1529
+ ;
1530
+ long long dur_long = dur_ts_to_ns(node->data.fee.dur);
1531
+ char ph = 'X';
1532
+ if (node->data.fee.type == PyTrace_CALL || node->data.fee.type == PyTrace_C_CALL) {
1533
+ ph = 'B';
1534
+ }
1535
+ fprintf(fptr, "\"ph\":\"%c\",\"cat\":\"fee\",\"dur\":%lld.%03lld,\"name\":\"", ph, dur_long / 1000, dur_long % 1000);
1536
+ fprintfeename(fptr, node, sanitize_function_name);
1537
+ fputc('\"', fptr);
1538
+
1539
+ PyObject* arg_dict = NULL;
1540
+ if (node->data.fee.args) {
1541
+ arg_dict = node->data.fee.args;
1542
+ Py_INCREF(arg_dict);
1543
+ }
1544
+ if (node->data.fee.retval) {
1545
+ if (!arg_dict) {
1546
+ arg_dict = PyDict_New();
1547
+ }
1548
+ PyDict_SetItemString(arg_dict, "return_value", node->data.fee.retval);
1549
+ }
1550
+ if (arg_dict) {
1551
+ fprintf(fptr, ",\"args\":");
1552
+ fprintjson(fptr, arg_dict);
1553
+ Py_DECREF(arg_dict);
1554
+ }
1555
+ break;
1556
+ case INSTANT_NODE:
1557
+ fprintf(fptr, "\"ph\":\"i\",\"cat\":\"instant\",\"name\":\"");
1558
+ fprint_escape(fptr, PyUnicode_AsUTF8(node->data.instant.name));
1559
+ if (node->data.instant.args == Py_None) {
1560
+ fprintf(fptr, "\",\"s\":\"%s\"", PyUnicode_AsUTF8(node->data.instant.scope));
1561
+ } else {
1562
+ fprintf(fptr, "\",\"s\":\"%s\",\"args\":", PyUnicode_AsUTF8(node->data.instant.scope));
1563
+ fprintjson(fptr, node->data.instant.args);
1564
+ }
1565
+ break;
1566
+ case COUNTER_NODE:
1567
+ fprintf(fptr, "\"ph\":\"C\",\"name\":\"");
1568
+ fprint_escape(fptr, PyUnicode_AsUTF8(node->data.counter.name));
1569
+ fprintf(fptr, "\",\"args\":");
1570
+ fprintjson(fptr, node->data.counter.args);
1571
+ break;
1572
+ case OBJECT_NODE:
1573
+ fprintf(fptr, "\"ph\":\"%s\",\"id\":\"%s\",\"name\":\"",
1574
+ PyUnicode_AsUTF8(node->data.object.ph), PyUnicode_AsUTF8(node->data.object.id));
1575
+ fprint_escape(fptr, PyUnicode_AsUTF8(node->data.object.name));
1576
+ fputc('\"', fptr);
1577
+ if (!(node->data.object.args == Py_None)) {
1578
+ fprintf(fptr, ",\"args\":");
1579
+ fprintjson(fptr, node->data.object.args);
1580
+ }
1581
+ break;
1582
+ case RAW_NODE:
1583
+ // We still need to tid from node and we need the pid
1584
+ ;
1585
+ PyObject* py_pid = PyLong_FromLong(pid);
1586
+ PyObject* py_tid = PyLong_FromLong(node->tid);
1587
+ PyObject* dict = node->data.raw;
1588
+
1589
+ PyDict_SetItemString(dict, "pid", py_pid);
1590
+ PyDict_SetItemString(dict, "tid", py_tid);
1591
+ fprintjson(fptr, dict);
1592
+ fputc(',', fptr);
1593
+ Py_DECREF(py_pid);
1594
+ Py_DECREF(py_tid);
1595
+ break;
1596
+ default:
1597
+ printf("Unknown Node Type!\n");
1598
+ exit(1);
1599
+ }
1600
+ if (node->ntype != RAW_NODE) {
1601
+ fputs("},", fptr);
1602
+ }
1603
+ clear_node(node);
1604
+ curr = curr + 1;
1605
+ if (curr == self->buffer + self->buffer_size) {
1606
+ curr = self->buffer;
1607
+ }
1608
+ }
1609
+
1610
+ if (CHECK_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC)) {
1611
+ Py_ssize_t pos = 0;
1612
+ PyObject* key = NULL;
1613
+ PyObject* value = NULL;
1614
+ while (PyDict_Next(task_dict, &pos, &key, &value)) {
1615
+ PyObject* tid_repr = PyObject_Repr(key);
1616
+ fprintf(fptr, "{\"ph\":\"M\",\"pid\":%lu,\"tid\":%s,\"name\":\"thread_name\",\"args\":{\"name\":\"%s\"}},",
1617
+ pid, PyUnicode_AsUTF8(tid_repr), PyUnicode_AsUTF8(value));
1618
+ Py_DECREF(tid_repr);
1619
+ }
1620
+ Py_DECREF(task_dict);
1621
+ }
1622
+
1623
+ self->buffer_tail_idx = self->buffer_head_idx;
1624
+ fseek(fptr, -1, SEEK_CUR);
1625
+ fprintf(fptr, "], \"viztracer_metadata\": {\"overflow\":%s", overflowed? "true": "false");
1626
+
1627
+ if (self->sync_marker > 0)
1628
+ {
1629
+ long long ts_sync_marker = system_ts_to_ns(self->sync_marker);
1630
+ fprintf(fptr, ",\"sync_marker\":%lld.%03lld", ts_sync_marker / 1000, ts_sync_marker % 1000);
1631
+ }
1632
+
1633
+ fprintf(fptr, "}}");
1634
+ fclose(fptr);
1635
+ SNAPTRACE_THREAD_PROTECT_END(self);
1636
+ Py_RETURN_NONE;
1637
+ }
1638
+
1639
+ static PyObject*
1640
+ tracer_clear(TracerObject* self, PyObject* Py_UNUSED(unused))
1641
+ {
1642
+ struct EventNode* curr = self->buffer + self->buffer_head_idx;
1643
+ while (curr != self->buffer + self->buffer_tail_idx) {
1644
+ struct EventNode* node = curr;
1645
+ clear_node(node);
1646
+ curr = curr + 1;
1647
+ if (curr == self->buffer + self->buffer_size) {
1648
+ curr = self->buffer;
1649
+ }
1650
+ }
1651
+ self->buffer_tail_idx = self->buffer_head_idx;
1652
+
1653
+ Py_RETURN_NONE;
1654
+ }
1655
+
1656
+ static PyObject*
1657
+ tracer_setpid(TracerObject* self, PyObject* args)
1658
+ {
1659
+ long input_pid = -1;
1660
+ if (!PyArg_ParseTuple(args, "|l", &input_pid)) {
1661
+ printf("Parsing error on setpid\n");
1662
+ }
1663
+
1664
+ if (input_pid >= 0) {
1665
+ self->fix_pid = input_pid;
1666
+ } else {
1667
+ #if _WIN32
1668
+ self->fix_pid = GetCurrentProcessId();
1669
+ #else
1670
+ self->fix_pid = getpid();
1671
+ #endif
1672
+ }
1673
+
1674
+ Py_RETURN_NONE;
1675
+ }
1676
+
1677
+ static PyObject*
1678
+ tracer_getts(TracerObject* self, PyObject* Py_UNUSED(unused))
1679
+ {
1680
+ int64_t ts = get_ts();
1681
+ double us = system_ts_to_us(ts);
1682
+
1683
+ return PyFloat_FromDouble(us);
1684
+ }
1685
+
1686
+ static PyObject*
1687
+ tracer_getbasetime(TracerObject* self, PyObject* Py_UNUSED(unused))
1688
+ {
1689
+ return PyLong_FromLongLong(get_base_time_ns());
1690
+ }
1691
+
1692
+ static PyObject*
1693
+ tracer_resetstack(TracerObject* self, PyObject* Py_UNUSED(unused))
1694
+ {
1695
+ struct ThreadInfo* info = get_thread_info(self);
1696
+ if (!info) {
1697
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to get thread info. This should not happen.");
1698
+ return NULL;
1699
+ }
1700
+
1701
+ info->curr_stack_depth = 0;
1702
+ info->ignore_stack_depth = 0;
1703
+
1704
+ struct FunctionNode* stack_top = info->stack_top;
1705
+ clear_stack(&stack_top);
1706
+ info->stack_top = stack_top;
1707
+
1708
+ Py_RETURN_NONE;
1709
+ }
1710
+
1711
+ static PyObject*
1712
+ tracer_addinstant(TracerObject* self, PyObject* args, PyObject* kw)
1713
+ {
1714
+ PyObject* name = NULL;
1715
+ PyObject* instant_args = NULL;
1716
+ PyObject* scope = NULL;
1717
+ struct EventNode* node = NULL;
1718
+ static char* kwlist[] = {"name", "args", "scope", NULL};
1719
+ const char* allowed_scope[] = {"g", "p", "t"};
1720
+
1721
+ if (!self->collecting) {
1722
+ Py_RETURN_NONE;
1723
+ }
1724
+
1725
+ if (!PyArg_ParseTupleAndKeywords(args, kw, "O|OO", kwlist,
1726
+ &name, &instant_args, &scope)) {
1727
+ return NULL;
1728
+ }
1729
+
1730
+ struct ThreadInfo* info = get_thread_info(self);
1731
+ if (!info) {
1732
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to get thread info. This should not happen.");
1733
+ return NULL;
1734
+ }
1735
+
1736
+ if (!instant_args) {
1737
+ instant_args = Py_None;
1738
+ }
1739
+
1740
+ if (!scope) {
1741
+ scope = PyUnicode_FromString("g");
1742
+ } else {
1743
+ if (!PyUnicode_CheckExact(scope)) {
1744
+ PyErr_SetString(PyExc_TypeError, "Scope must be a string");
1745
+ return NULL;
1746
+ }
1747
+ for (int i = 0; i < 3; i++) {
1748
+ if (strcmp(PyUnicode_AsUTF8(scope), allowed_scope[i]) == 0) {
1749
+ break;
1750
+ }
1751
+ if (i == 2) {
1752
+ PyErr_SetString(PyExc_ValueError, "Scope must be one of 'g', 'p', 't'");
1753
+ return NULL;
1754
+ }
1755
+ }
1756
+ Py_INCREF(scope);
1757
+ }
1758
+
1759
+ node = get_next_node(self);
1760
+ node->ntype = INSTANT_NODE;
1761
+ node->tid = info->tid;
1762
+ node->ts = get_ts();
1763
+ node->data.instant.name = Py_NewRef(name);
1764
+ node->data.instant.args = Py_NewRef(instant_args);
1765
+ node->data.instant.scope = scope;
1766
+
1767
+ Py_RETURN_NONE;
1768
+ }
1769
+
1770
+ static PyObject*
1771
+ tracer_addfunctionarg(TracerObject* self, PyObject* args, PyObject* kw)
1772
+ {
1773
+ PyObject* key = NULL;
1774
+ PyObject* value = NULL;
1775
+ static char* kwlist[] = {"key", "value", NULL};
1776
+
1777
+ if (!self->collecting) {
1778
+ Py_RETURN_NONE;
1779
+ }
1780
+
1781
+ if (!PyArg_ParseTupleAndKeywords(args, kw, "OO", kwlist, &key, &value)) {
1782
+ return NULL;
1783
+ }
1784
+
1785
+ struct ThreadInfo* info = get_thread_info(self);
1786
+ if (!info) {
1787
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to get thread info. This should not happen.");
1788
+ return NULL;
1789
+ }
1790
+
1791
+ struct FunctionNode* fnode = info->stack_top;
1792
+ if (!fnode->args) {
1793
+ fnode->args = PyDict_New();
1794
+ }
1795
+
1796
+ PyDict_SetItem(fnode->args, key, value);
1797
+
1798
+ Py_RETURN_NONE;
1799
+ }
1800
+
1801
+ static PyObject*
1802
+ tracer_addcounter(TracerObject* self, PyObject* args, PyObject* kw)
1803
+ {
1804
+ PyObject* name = NULL;
1805
+ PyObject* counter_args = NULL;
1806
+ static char* kwlist[] = {"name", "args", NULL};
1807
+ struct EventNode* node = NULL;
1808
+
1809
+ if (!self->collecting) {
1810
+ Py_RETURN_NONE;
1811
+ }
1812
+
1813
+ if (!PyArg_ParseTupleAndKeywords(args, kw, "OO", kwlist,
1814
+ &name, &counter_args)) {
1815
+ return NULL;
1816
+ }
1817
+
1818
+ struct ThreadInfo* info = get_thread_info(self);
1819
+
1820
+ if (!info) {
1821
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to get thread info. This should not happen.");
1822
+ return NULL;
1823
+ }
1824
+
1825
+ node = get_next_node(self);
1826
+ node->ntype = COUNTER_NODE;
1827
+ node->tid = info->tid;
1828
+ node->ts = get_ts();
1829
+ node->data.counter.name = Py_NewRef(name);
1830
+ node->data.counter.args = Py_NewRef(counter_args);
1831
+
1832
+ Py_RETURN_NONE;
1833
+ }
1834
+
1835
+ static PyObject*
1836
+ tracer_addobject(TracerObject* self, PyObject* args, PyObject* kw)
1837
+ {
1838
+ PyObject* ph = NULL;
1839
+ PyObject* id = NULL;
1840
+ PyObject* name = NULL;
1841
+ PyObject* object_args = NULL;
1842
+ static char* kwlist[] = {"ph", "obj_id", "name", "args", NULL};
1843
+ struct EventNode* node = NULL;
1844
+
1845
+ if (!self->collecting) {
1846
+ Py_RETURN_NONE;
1847
+ }
1848
+
1849
+ if (!PyArg_ParseTupleAndKeywords(args, kw, "OOO|O", kwlist,
1850
+ &ph, &id, &name, &object_args)) {
1851
+ return NULL;
1852
+ }
1853
+
1854
+ struct ThreadInfo* info = get_thread_info(self);
1855
+
1856
+ if (!info) {
1857
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to get thread info. This should not happen.");
1858
+ return NULL;
1859
+ }
1860
+
1861
+ if (!object_args) {
1862
+ object_args = Py_None;
1863
+ }
1864
+
1865
+ node = get_next_node(self);
1866
+ node->ntype = OBJECT_NODE;
1867
+ node->tid = info->tid;
1868
+ node->ts = get_ts();
1869
+ node->data.object.ph = Py_NewRef(ph);
1870
+ node->data.object.id = Py_NewRef(id);
1871
+ node->data.object.name = Py_NewRef(name);
1872
+ node->data.object.args = Py_NewRef(object_args);
1873
+
1874
+ Py_RETURN_NONE;
1875
+ }
1876
+
1877
+ static PyObject*
1878
+ tracer_addraw(TracerObject* self, PyObject* args, PyObject* kw)
1879
+ {
1880
+ PyObject* raw = NULL;
1881
+ static char* kwlist[] = {"raw", NULL};
1882
+ struct EventNode* node = NULL;
1883
+
1884
+ if (!PyArg_ParseTupleAndKeywords(args, kw, "O", kwlist, &raw)) {
1885
+ return NULL;
1886
+ }
1887
+
1888
+ struct ThreadInfo* info = get_thread_info(self);
1889
+
1890
+ if (!info) {
1891
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to get thread info. This should not happen.");
1892
+ return NULL;
1893
+ }
1894
+
1895
+ node = get_next_node(self);
1896
+ node->tid = info->tid;
1897
+ node->ntype = RAW_NODE;
1898
+ node->data.raw = Py_NewRef(raw);
1899
+
1900
+ Py_RETURN_NONE;
1901
+ }
1902
+
1903
+ static PyObject*
1904
+ tracer_setignorestackcounter(TracerObject* self, PyObject* value)
1905
+ {
1906
+ int current_value = 0;
1907
+
1908
+ if (!PyLong_Check(value)) {
1909
+ PyErr_SetString(PyExc_TypeError, "value must be an integer");
1910
+ return NULL;
1911
+ }
1912
+
1913
+ struct ThreadInfo* info = get_thread_info(self);
1914
+
1915
+ if (!info) {
1916
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to get thread info. This should not happen.");
1917
+ return NULL;
1918
+ }
1919
+
1920
+ current_value = info->ignore_stack_depth;
1921
+ // +1 to compensate for this call so when it returns, the value is correct
1922
+ info->ignore_stack_depth = PyLong_AsLong(value) + 1;
1923
+
1924
+ // -1 is the actual ignore stack depth before this call
1925
+ return Py_BuildValue("i", current_value - 1);
1926
+ }
1927
+
1928
+ static PyObject*
1929
+ tracer_getfunctionarg(TracerObject* self, PyObject* Py_UNUSED(unused))
1930
+ {
1931
+ struct ThreadInfo* info = get_thread_info(self);
1932
+
1933
+ if (!info) {
1934
+ PyErr_SetString(PyExc_RuntimeError, "VizTracer: Failed to get thread info. This should not happen.");
1935
+ return NULL;
1936
+ }
1937
+
1938
+ struct FunctionNode* fnode = info->stack_top;
1939
+ if (!fnode->args) {
1940
+ Py_RETURN_NONE;
1941
+ }
1942
+
1943
+ return Py_NewRef(fnode->args);
1944
+ }
1945
+
1946
+ static PyObject*
1947
+ tracer_set_sync_marker(TracerObject* self, PyObject* Py_UNUSED(unused))
1948
+ {
1949
+ if (self->sync_marker != 0)
1950
+ {
1951
+ PyErr_WarnFormat(PyExc_RuntimeWarning, 1, "Synchronization marker already set");
1952
+ }
1953
+ self->sync_marker = get_ts();
1954
+ Py_RETURN_NONE;
1955
+ }
1956
+
1957
+ static PyObject*
1958
+ tracer_get_sync_marker(TracerObject* self, PyObject* Py_UNUSED(unused))
1959
+ {
1960
+ if (self->sync_marker == 0)
1961
+ {
1962
+ Py_RETURN_NONE;
1963
+ }
1964
+
1965
+ double ts_sync_marker = system_ts_to_us(self->sync_marker);
1966
+ return PyFloat_FromDouble(ts_sync_marker);
1967
+ }
1968
+
1969
+ static PyMethodDef Tracer_methods[] = {
1970
+ {"threadtracefunc", (PyCFunction)tracer_threadtracefunc, METH_VARARGS, "trace function"},
1971
+ {"start", (PyCFunction)tracer_start, METH_NOARGS, "start profiling"},
1972
+ {"stop", (PyCFunction)tracer_stop, METH_O, "stop profiling"},
1973
+ {"load", (PyCFunction)tracer_load, METH_NOARGS, "load buffer"},
1974
+ {"dump", (PyCFunction)tracer_dump, METH_VARARGS|METH_KEYWORDS, "dump buffer to file"},
1975
+ {"clear", (PyCFunction)tracer_clear, METH_NOARGS, "clear buffer"},
1976
+ {"setpid", (PyCFunction)tracer_setpid, METH_VARARGS, "set fixed pid"},
1977
+ {"add_instant", (PyCFunction)tracer_addinstant, METH_VARARGS|METH_KEYWORDS, "add instant event"},
1978
+ {"add_counter", (PyCFunction)tracer_addcounter, METH_VARARGS|METH_KEYWORDS, "add counter event"},
1979
+ {"add_object", (PyCFunction)tracer_addobject, METH_VARARGS|METH_KEYWORDS, "add object event"},
1980
+ {"add_raw", (PyCFunction)tracer_addraw, METH_VARARGS|METH_KEYWORDS, "add raw event"},
1981
+ {"add_func_args", (PyCFunction)tracer_addfunctionarg, METH_VARARGS|METH_KEYWORDS, "add function arg"},
1982
+ {"get_func_args", (PyCFunction)tracer_getfunctionarg, METH_NOARGS, "get current function arg"},
1983
+ {"getts", (PyCFunction)tracer_getts, METH_NOARGS, "get timestamp"},
1984
+ {"get_base_time", (PyCFunction)tracer_getbasetime, METH_NOARGS, "get base time in nanoseconds"},
1985
+ {"reset_stack", (PyCFunction)tracer_resetstack, METH_NOARGS, "reset stack"},
1986
+ {"pause", (PyCFunction)tracer_pause, METH_NOARGS, "pause profiling"},
1987
+ {"resume", (PyCFunction)tracer_resume, METH_NOARGS, "resume profiling"},
1988
+ {"setignorestackcounter", (PyCFunction)tracer_setignorestackcounter, METH_O, "reset ignore stack depth"},
1989
+ {"set_sync_marker", (PyCFunction)tracer_set_sync_marker, METH_NOARGS, "set current timestamp to synchronization marker"},
1990
+ {"get_sync_marker", (PyCFunction)tracer_get_sync_marker, METH_NOARGS, "get synchronization marker or None if not set"},
1991
+ {NULL, NULL, 0, NULL}
1992
+ };
1993
+
1994
+ // ===========================================================================
1995
+ // snaptrace.Tracer internals
1996
+ // ===========================================================================
1997
+
1998
+ static PyObject*
1999
+ Tracer_New(PyTypeObject* type, PyObject* args, PyObject* kwargs)
2000
+ {
2001
+ TracerObject* self = (TracerObject*) type->tp_alloc(type, 0);
2002
+ if (self) {
2003
+ self->collecting = 0;
2004
+ self->fix_pid = 0;
2005
+ self->total_entries = 0;
2006
+ self->check_flags = 0;
2007
+ self->verbose = 0;
2008
+ self->lib_file_path = NULL;
2009
+ self->max_stack_depth = 0;
2010
+ self->include_files = NULL;
2011
+ self->exclude_files = NULL;
2012
+ self->min_duration = 0;
2013
+ self->buffer = NULL;
2014
+ self->buffer_head_idx = 0;
2015
+ self->buffer_tail_idx = 0;
2016
+ self->sync_marker = 0;
2017
+ self->metadata_head = NULL;
2018
+ }
2019
+
2020
+ return (PyObject*) self;
2021
+ }
2022
+
2023
+ static int
2024
+ Tracer_Init(TracerObject* self, PyObject* args, PyObject* kwargs)
2025
+ {
2026
+ if (!PyArg_ParseTuple(args, "l", &self->buffer_size)) {
2027
+ PyErr_SetString(PyExc_TypeError, "You need to specify buffer size when initializing Tracer");
2028
+ return -1;
2029
+ }
2030
+
2031
+ // We need an extra slot for circular buffer
2032
+ self->buffer_size += 1;
2033
+ self->buffer = (struct EventNode*) PyMem_Calloc(self->buffer_size, sizeof(struct EventNode));
2034
+ if (!self->buffer) {
2035
+ PyErr_NoMemory();
2036
+ return -1;
2037
+ }
2038
+
2039
+ #if _WIN32
2040
+ if ((self->dwTlsIndex = TlsAlloc()) == TLS_OUT_OF_INDEXES) {
2041
+ printf("Error on TLS!\n");
2042
+ exit(-1);
2043
+ }
2044
+ #else
2045
+ if (pthread_key_create(&self->thread_key, snaptrace_threaddestructor)) {
2046
+ perror("Failed to create Tss_Key");
2047
+ exit(-1);
2048
+ }
2049
+ #endif
2050
+
2051
+ #if PY_VERSION_HEX >= 0x030C0000
2052
+ #else
2053
+ // Python: threading.setprofile(tracefunc)
2054
+ {
2055
+ PyObject* handler = PyCFunction_New(&Tracer_methods[0], (PyObject*)self);
2056
+
2057
+ if (PyObject_CallMethod(threading_module, "setprofile", "N", handler) == NULL) {
2058
+ perror("Failed to call threading.setprofile() properly");
2059
+ exit(-1);
2060
+ }
2061
+ }
2062
+ #endif
2063
+
2064
+ #if PY_VERSION_HEX >= 0x030C0000
2065
+ #else
2066
+ PyEval_SetProfile(tracer_tracefunc, (PyObject*)self);
2067
+ #endif
2068
+
2069
+ return 0;
2070
+ }
2071
+
2072
+ static void
2073
+ Tracer_dealloc(TracerObject* self)
2074
+ {
2075
+ tracer_clear(self, NULL);
2076
+ if (self->lib_file_path) {
2077
+ PyMem_FREE(self->lib_file_path);
2078
+ }
2079
+ Py_XDECREF(self->include_files);
2080
+ Py_XDECREF(self->exclude_files);
2081
+ PyMem_FREE(self->buffer);
2082
+
2083
+ struct MetadataNode* node = self->metadata_head;
2084
+ struct MetadataNode* prev = NULL;
2085
+ while (node) {
2086
+ prev = node;
2087
+ Py_CLEAR(node->name);
2088
+ node = node->next;
2089
+ PyMem_FREE(prev);
2090
+ }
2091
+
2092
+ #if PY_VERSION_HEX >= 0x030C0000
2093
+ #else
2094
+ // threading.setprofile(None)
2095
+ // It's possible that during deallocation phase threading module has released setprofile
2096
+ // and we should be okay with that.
2097
+ PyObject* result = PyObject_CallMethod(threading_module, "setprofile", "O", Py_None);
2098
+ if (result != NULL) {
2099
+ Py_DECREF(result);
2100
+ }
2101
+ #endif
2102
+
2103
+ Py_TYPE(self)->tp_free((PyObject*) self);
2104
+ }
2105
+
2106
+ // ================================================================
2107
+ // snaptrace.Tracer Definition
2108
+ // ================================================================
2109
+
2110
+ // We define getsetters in another file
2111
+ extern PyGetSetDef Tracer_getsetters[];
2112
+
2113
+ static PyTypeObject TracerType = {
2114
+ PyVarObject_HEAD_INIT(NULL, 0)
2115
+ .tp_name = "snaptrace.Tracer",
2116
+ .tp_doc = "Tracer",
2117
+ .tp_basicsize = sizeof(TracerObject),
2118
+ .tp_itemsize = 0,
2119
+ .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
2120
+ .tp_new = Tracer_New,
2121
+ .tp_init = (initproc) Tracer_Init,
2122
+ .tp_dealloc = (destructor) Tracer_dealloc,
2123
+ .tp_methods = Tracer_methods,
2124
+ .tp_getset = Tracer_getsetters,
2125
+ };
2126
+
2127
+ // ================================================================
2128
+ // snaptrace Module Functions
2129
+ // ================================================================
2130
+
2131
+ void
2132
+ snaptrace_free(void* Py_UNUSED(unused)) {
2133
+ quicktime_free();
2134
+ Py_CLEAR(threading_module);
2135
+ Py_CLEAR(multiprocessing_module);
2136
+ Py_CLEAR(asyncio_module);
2137
+ Py_CLEAR(asyncio_tasks_module);
2138
+ Py_CLEAR(curr_task_getters[0]);
2139
+ Py_CLEAR(trio_lowlevel_module);
2140
+ Py_CLEAR(curr_task_getters[1]);
2141
+ Py_CLEAR(json_module);
2142
+ Py_CLEAR(sys_module);
2143
+ }
2144
+
2145
+ // ================================================================
2146
+ // snaptrace Module Definition
2147
+ // ================================================================
2148
+
2149
+ static struct PyModuleDef snaptracemodule = {
2150
+ .m_base = PyModuleDef_HEAD_INIT,
2151
+ .m_name = "viztracer.snaptrace",
2152
+ .m_size = -1,
2153
+ .m_free = snaptrace_free,
2154
+ };
2155
+
2156
+ // ================================================================
2157
+ // Python Interface
2158
+ // ================================================================
2159
+
2160
+ PyMODINIT_FUNC
2161
+ PyInit_snaptrace(void)
2162
+ {
2163
+ // Tracer Module
2164
+ PyObject* m = NULL;
2165
+
2166
+ if (PyType_Ready(&TracerType) < 0) {
2167
+ return NULL;
2168
+ }
2169
+
2170
+ m = PyModule_Create(&snaptracemodule);
2171
+
2172
+ if (!m) {
2173
+ return NULL;
2174
+ }
2175
+
2176
+ #ifdef Py_GIL_DISABLED
2177
+ PyUnstable_Module_SetGIL(m, Py_MOD_GIL_NOT_USED);
2178
+ #endif
2179
+
2180
+ Py_INCREF(&TracerType);
2181
+ if (PyModule_AddObject(m, "Tracer", (PyObject*) &TracerType) < 0) {
2182
+ Py_DECREF(&TracerType);
2183
+ Py_DECREF(m);
2184
+ return NULL;
2185
+ }
2186
+
2187
+ threading_module = PyImport_ImportModule("threading");
2188
+ multiprocessing_module = PyImport_ImportModule("multiprocessing");
2189
+ if ((trio_module = PyImport_ImportModule("trio"))) {
2190
+ trio_lowlevel_module = PyImport_AddModule("trio.lowlevel");
2191
+ curr_task_getters[1] = PyObject_GetAttrString(trio_lowlevel_module, "current_task");
2192
+ } else {
2193
+ PyErr_Clear();
2194
+ }
2195
+ json_module = PyImport_ImportModule("json");
2196
+
2197
+ #if PY_VERSION_HEX >= 0x030C0000
2198
+ sys_module = PyImport_ImportModule("sys");
2199
+ PyObject* monitoring = PyObject_GetAttrString(sys_module, "monitoring");
2200
+ sys_monitoring_missing = PyObject_GetAttrString(monitoring, "MISSING");
2201
+ Py_DECREF(monitoring);
2202
+ #endif
2203
+
2204
+ quicktime_init();
2205
+
2206
+ return m;
2207
+ }