make subprocess child shows different process name

src/viztracer/main.py

@@ -224,8 +224,10 @@ def parse(self, argv: List[str]) -> VizProcedureResult:
             # It's not practical to cover this line as it requires coverage
             # instrumentation on subprocess.
             output_file = self.ofile  # pragma: no cover
+            process_name = sys.argv[0]
         else:
             output_file = os.path.join(self.multiprocess_output_dir, "result.json")
+            process_name = None
 
         if options.log_multiprocess or options.log_subprocess:  # pragma: no cover
             color_print(
@@ -263,6 +265,7 @@ def parse(self, argv: List[str]) -> VizProcedureResult:
             "sanitize_function_name": options.sanitize_function_name,
             "dump_raw": True,
             "minimize_memory": options.minimize_memory,
+            "process_name": process_name,
         }
 
         return True, None

src/viztracer/modules/snaptrace.c

@@ -650,20 +650,27 @@ snaptrace_load(TracerObject* self, PyObject* args)
         PyObject* dict = PyDict_New();
         PyObject* args = PyDict_New();
         PyObject* process_name_string = PyUnicode_FromString("process_name");
-        PyObject* current_process_method = PyObject_GetAttrString(multiprocessing_module, "current_process");
-        if (!current_process_method) {
-            perror("Failed to access multiprocessing.current_process()");
-            exit(-1);
-        }
-        PyObject* current_process = PyObject_CallObject(current_process_method, NULL);
-        if (!current_process_method) {
-            perror("Failed to access multiprocessing.current_process()");
-            exit(-1);
-        }
-        PyObject* process_name = PyObject_GetAttrString(current_process, "name");
+        PyObject* process_name = NULL;
 
-        Py_DECREF(current_process_method);
-        Py_DECREF(current_process);
+        if (self->process_name) {
+            process_name = self->process_name;
+            Py_INCREF(process_name);
+        } else {
+            PyObject* current_process_method = PyObject_GetAttrString(multiprocessing_module, "current_process");
+            if (!current_process_method) {
+                perror("Failed to access multiprocessing.current_process()");
+                exit(-1);
+            }
+            PyObject* current_process = PyObject_CallObject(current_process_method, NULL);
+            if (!current_process_method) {
+                perror("Failed to access multiprocessing.current_process()");
+                exit(-1);
+            }
+            process_name = PyObject_GetAttrString(current_process, "name");
+            Py_DECREF(current_process_method);
+            Py_DECREF(current_process);
+        }
+
         PyDict_SetItemString(dict, "ph", ph_M);
         PyDict_SetItemString(dict, "pid", pid);
         PyDict_SetItemString(dict, "tid", pid);
@@ -898,22 +905,30 @@ snaptrace_dump(TracerObject* self, PyObject* args)
     // == Load the metadata first ==
     //    Process Name
     {
-        PyObject* current_process_method = PyObject_GetAttrString(multiprocessing_module, "current_process");
-        if (!current_process_method) {
-            perror("Failed to access multiprocessing.current_process()");
-            exit(-1);
-        }
-        PyObject* current_process = PyObject_CallObject(current_process_method, NULL);
-        if (!current_process_method) {
-            perror("Failed to access multiprocessing.current_process()");
-            exit(-1);
+        PyObject* process_name = NULL;
+        if (self->process_name) {
+            process_name = self->process_name;
+            Py_INCREF(process_name);
+        } else {
+            PyObject* current_process_method = PyObject_GetAttrString(multiprocessing_module, "current_process");
+            if (!current_process_method) {
+                perror("Failed to access multiprocessing.current_process()");
+                exit(-1);
+            }
+            PyObject* current_process = PyObject_CallObject(current_process_method, NULL);
+            if (!current_process_method) {
+                perror("Failed to access multiprocessing.current_process()");
+                exit(-1);
+            }
+            process_name = PyObject_GetAttrString(current_process, "name");
+            Py_DECREF(current_process_method);
+            Py_DECREF(current_process);
         }
-        PyObject* process_name = PyObject_GetAttrString(current_process, "name");
 
-        Py_DECREF(current_process_method);
-        Py_DECREF(current_process);
-        fprintf(fptr, "{\"ph\":\"M\",\"pid\":%lu,\"tid\":%lu,\"name\":\"process_name\",\"args\":{\"name\":\"%s\"}},",
-                pid, pid, PyUnicode_AsUTF8(process_name));
+        fprintf(fptr, "{\"ph\":\"M\",\"pid\":%lu,\"tid\":%lu,\"name\":\"process_name\",\"args\":{\"name\":\"",
+                pid, pid);
+        fprint_escape(fptr, PyUnicode_AsUTF8(process_name));
+        fprintf(fptr, "\"}},");
         Py_DECREF(process_name);
     }
 
@@ -1140,11 +1155,12 @@ snaptrace_config(TracerObject* self, PyObject* args, PyObject* kw)
     static char* kwlist[] = {"verbose", "lib_file_path", "max_stack_depth", 
             "include_files", "exclude_files", "ignore_c_function", "ignore_frozen",
             "log_func_retval", "log_func_args", "log_async", "trace_self",
-            "min_duration",
+            "min_duration", "process_name",
             NULL};
     int kw_verbose = -1;
     int kw_max_stack_depth = 0;
     char* kw_lib_file_path = NULL;
+    PyObject* kw_process_name = NULL;
     PyObject* kw_include_files = NULL;
     PyObject* kw_exclude_files = NULL;
     int kw_ignore_c_function = -1;
@@ -1154,7 +1170,7 @@ snaptrace_config(TracerObject* self, PyObject* args, PyObject* kw)
     int kw_log_async = -1;
     int kw_trace_self = -1;
     double kw_min_duration = 0;
-    if (!PyArg_ParseTupleAndKeywords(args, kw, "|isiOOppppppd", kwlist,
+    if (!PyArg_ParseTupleAndKeywords(args, kw, "|isiOOppppppdO", kwlist,
             &kw_verbose,
             &kw_lib_file_path,
             &kw_max_stack_depth,
@@ -1166,7 +1182,8 @@ snaptrace_config(TracerObject* self, PyObject* args, PyObject* kw)
             &kw_log_func_args,
             &kw_log_async,
             &kw_trace_self,
-            &kw_min_duration)) {
+            &kw_min_duration,
+            &kw_process_name)) {
         return NULL;
     }
 
@@ -1190,6 +1207,15 @@ snaptrace_config(TracerObject* self, PyObject* args, PyObject* kw)
         strcpy(self->lib_file_path, kw_lib_file_path);
     }
 
+    if (kw_process_name && !Py_IsNone(kw_process_name)) {
+        if (!PyUnicode_CheckExact(kw_process_name)) {
+            PyErr_SetString(PyExc_TypeError, "process_name must be a string");
+            return NULL;
+        }
+        Py_INCREF(kw_process_name);
+        Py_XSETREF(self->process_name, kw_process_name); 
+    }
+
     if (kw_ignore_c_function == 1) {
         SET_FLAG(self->check_flags, SNAPTRACE_IGNORE_C_FUNCTION);
     } else if (kw_ignore_c_function == 0) {

src/viztracer/modules/snaptrace.h

@@ -69,6 +69,7 @@ typedef struct {
     int verbose;
     char* lib_file_path;
     int max_stack_depth;
+    PyObject* process_name;
     PyObject* include_files;
     PyObject* exclude_files;
     double min_duration;

src/viztracer/tracer.py

@@ -28,7 +28,8 @@ def __init__(
             log_gc: bool = False,
             log_async: bool = False,
             trace_self: bool = False,
-            min_duration: float = 0) -> None:
+            min_duration: float = 0,
+            process_name: Optional[str] = None) -> None:
         self.initialized = False
         self.enable = False
         self.parsed = False
@@ -53,6 +54,7 @@ def __init__(
         self.total_entries = 0
         self.gc_start_args: Dict[str, int] = {}
         self.initialized = True
+        self.process_name = process_name
 
     @property
     def max_stack_depth(self) -> int:
@@ -237,6 +239,7 @@ def config(self) -> None:
             "log_async": self.log_async,
             "trace_self": self.trace_self,
             "min_duration": self.min_duration,
+            "process_name": self.process_name,
         }
 
         self._tracer.config(**cfg)

src/viztracer/viztracer.py

@@ -42,6 +42,7 @@ def __init__(self,
                  minimize_memory: bool = False,
                  dump_raw: bool = False,
                  sanitize_function_name: bool = False,
+                 process_name: Optional[str] = None,
                  output_file: str = "result.json",
                  plugins: Sequence[Union[VizPluginBase, str]] = []) -> None:
         super().__init__(
@@ -58,6 +59,7 @@ def __init__(self,
             log_async=log_async,
             trace_self=trace_self,
             min_duration=min_duration,
+            process_name=process_name,
         )
         self._tracer: Any
         self.verbose = verbose

tests/test_multiprocess.py

@@ -7,6 +7,7 @@
 import sys
 import tempfile
 import unittest
+import json
 
 from .cmdline_tmpl import CmdlineTmpl
 
@@ -44,6 +45,11 @@ def fib(n):
 print(subprocess.call(["python", "-m", "timeit", "-n", "100", "'1+1'"]))
 """
 
+file_subprocess_code = """
+import subprocess
+p = subprocess.Popen(["python", "-c", "print('test')"])
+"""
+
 file_fork = """
 import os
 import time
@@ -240,13 +246,26 @@ def test_child_process(self):
             self.template(["viztracer", "-o", os.path.join(tmpdir, "result.json"), "--subprocess_child", "child.py"],
                           expected_output_file=None)
             self.assertEqual(len(os.listdir(tmpdir)), 1)
+            with open(os.path.join(tmpdir, os.listdir(tmpdir)[0])) as f:
+                trace_events = json.load(f)["traceEvents"]
+                for entry in trace_events:
+                    if entry["name"] == "process_name":
+                        self.assertNotEqual(entry["args"]["name"], "MainProcess")
+                        break
+                else:
+                    self.fail("no process_name event found")
 
     def test_module(self):
         self.template(["viztracer", "-o", "result.json", "cmdline_test.py"],
                       expected_output_file="result.json",
                       expected_stdout=".*100 loops.*",
                       script=file_subprocess_module)
 
+    def test_code(self):
+        self.template(["viztracer", "-o", "result.json", "cmdline_test.py"],
+                      expected_output_file="result.json",
+                      script=file_subprocess_code)
+
     def test_nested(self):
         def check_func(data):
             pids = set()