Implement a sync marker

docs/source/basic_usage.rst

@@ -271,6 +271,14 @@ this usage.
 
     viztracer --align_combine run1.json run2.json -o compare_report.json
 
+You can also set a sync-marker from your source code, and VizTracer will align both reports to this particular timestamp.
+
+.. code-block::
+
+    from viztracer import get_tracer
+    get_tracer().set_sync_marker()
+
+
 Compress Your Report
 --------------------
 

src/viztracer/modules/snaptrace.c

@@ -1612,7 +1612,15 @@ tracer_dump(TracerObject* self, PyObject* args, PyObject* kw)
 
     self->buffer_tail_idx = self->buffer_head_idx;
     fseek(fptr, -1, SEEK_CUR);
-    fprintf(fptr, "], \"viztracer_metadata\": {\"overflow\":%s}}", overflowed? "true": "false");
+    fprintf(fptr, "], \"viztracer_metadata\": {\"overflow\":%s", overflowed? "true": "false");
+
+    if (self->sync_marker > 0)
+    {
+        long long ts_sync_marker = system_ts_to_ns(self->sync_marker);
+        fprintf(fptr, ",\"sync_marker\":%lld.%03lld", ts_sync_marker / 1000, ts_sync_marker % 1000);
+    }
+
+    fprintf(fptr, "}}");
     fclose(fptr);
     SNAPTRACE_THREAD_PROTECT_END(self);
     Py_RETURN_NONE;
@@ -1925,6 +1933,29 @@ tracer_getfunctionarg(TracerObject* self, PyObject* Py_UNUSED(unused))
     return Py_NewRef(fnode->args);
 }
 
+static PyObject*
+tracer_set_sync_marker(TracerObject* self, PyObject* Py_UNUSED(unused))
+{
+    if (self->sync_marker != 0)
+    {
+        PyErr_WarnFormat(PyExc_RuntimeWarning, 1, "Synchronization marker already set");
+    }
+    self->sync_marker = get_ts();
+    Py_RETURN_NONE;
+}
+
+static PyObject*
+tracer_get_sync_marker(TracerObject* self, PyObject* Py_UNUSED(unused))
+{
+    if (self->sync_marker == 0)
+    {
+        Py_RETURN_NONE;
+    }
+
+    double ts_sync_marker = system_ts_to_us(self->sync_marker);
+    return PyFloat_FromDouble(ts_sync_marker);
+}
+
 static PyMethodDef Tracer_methods[] = {
     {"threadtracefunc", (PyCFunction)tracer_threadtracefunc, METH_VARARGS, "trace function"},
     {"start", (PyCFunction)tracer_start, METH_NOARGS, "start profiling"},
@@ -1945,6 +1976,8 @@ static PyMethodDef Tracer_methods[] = {
     {"pause", (PyCFunction)tracer_pause, METH_NOARGS, "pause profiling"},
     {"resume", (PyCFunction)tracer_resume, METH_NOARGS, "resume profiling"},
     {"setignorestackcounter", (PyCFunction)tracer_setignorestackcounter, METH_O, "reset ignore stack depth"},
+    {"set_sync_marker", (PyCFunction)tracer_set_sync_marker, METH_NOARGS, "set current timestamp to synchronization marker"},
+    {"get_sync_marker", (PyCFunction)tracer_get_sync_marker, METH_NOARGS, "get synchronization marker or None if not set"},
     {NULL, NULL, 0, NULL}
 };
 
@@ -1970,6 +2003,7 @@ Tracer_New(PyTypeObject* type, PyObject* args, PyObject* kwargs)
         self->buffer = NULL;
         self->buffer_head_idx = 0;
         self->buffer_tail_idx = 0;
+        self->sync_marker = 0;
         self->metadata_head = NULL;
     }
 

src/viztracer/modules/snaptrace.h

@@ -121,6 +121,7 @@ typedef struct {
     long buffer_size;
     long buffer_head_idx;
     long buffer_tail_idx;
+    int64_t sync_marker;
     struct MetadataNode* metadata_head;
 } TracerObject;
 

src/viztracer/report_builder.py

@@ -125,7 +125,7 @@ def combine_json(self) -> None:
                 raise ValueError("Can't get report of nothing")
         if self.align:
             for one in self.jsons:
-                self.align_events(one["traceEvents"])
+                self.align_events(one["traceEvents"], one['viztracer_metadata'].get('sync_marker', None))
         self.combined_json = self.jsons[0]
         if "viztracer_metadata" not in self.combined_json:
             self.combined_json["viztracer_metadata"] = {}
@@ -143,14 +143,20 @@ def combine_json(self) -> None:
                 self.combined_json["file_info"]["files"].update(one["file_info"]["files"])
                 self.combined_json["file_info"]["functions"].update(one["file_info"]["functions"])
 
-    def align_events(self, original_events: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    def align_events(self, original_events: list[dict[str, Any]], sync_marker: Optional[float] = None) -> list[dict[str, Any]]:
         """
         Apply an offset to all the trace events, making the start timestamp 0
         This is useful when comparing multiple runs of the same script
 
+        If sync_marker is not None then sync_marker  be used as an offset
+
         This function will change the timestamp in place, and return the original list
         """
-        offset_ts = min((event["ts"] for event in original_events if "ts" in event))
+        if sync_marker is None:
+            offset_ts = min((event["ts"] for event in original_events if "ts" in event))
+        else:
+            offset_ts = sync_marker
+
         for event in original_events:
             if "ts" in event:
                 event["ts"] -= offset_ts

src/viztracer/snaptrace.pyi

@@ -67,3 +67,11 @@ class Tracer:
 
     def add_instant(self, name: str, args: Any = None, scope: Literal["g", "p", "t"] = "g") -> None:
         ...
+
+    def set_sync_marker(self) -> None:
+        """set current timestamp to synchronization marker"""
+        ...
+
+    def get_sync_marker(self) -> Optional[float]:
+        """get synchronization marker or None if not set"""
+        ...

src/viztracer/viztracer.py

@@ -254,6 +254,10 @@ def parse(self) -> int:
                     "overflow": False,
                 },
             }
+            sync_marker = self.get_sync_marker()
+            if sync_marker is not None:
+                self.data['viztracer_metadata']['sync_marker'] = sync_marker
+
             metadata_count = 0
             for d in self.data["traceEvents"]:
                 if d["ph"] == "M":

tests/test_cmdline.py

@@ -347,6 +347,76 @@ def test_combine(self):
                        os.path.join(example_json_dir, "different_sorts.json")],
                       expected_output_file="result.json")
 
+    def test_set_sync_marker(self):
+        test_script = textwrap.dedent("""
+            from viztracer import get_tracer
+
+            get_tracer().set_sync_marker()
+            get_tracer().set_sync_marker()
+
+        """)
+
+        def expect_sync_marker(data):
+            self.assertGreater(data['viztracer_metadata'].get('sync_marker'), 0)
+
+        self.template(
+            [sys.executable, "-m", "viztracer", "--ignore_frozen", "-o", "result.json", "cmdline_test.py"],
+            expected_output_file="result.json",
+            expected_stderr="Synchronization marker already set",
+            script=test_script,
+            check_func=expect_sync_marker,
+        )
+
+    def test_align_combine_sync_marker(self):
+        test_script = textwrap.dedent("""
+            import random
+            import time
+            from viztracer import get_tracer
+
+            def test_func():
+                return 2 * 4
+
+            # should sleep from 50ms to 100ms
+            time.sleep(random.uniform(0.05, 0.1))
+
+            get_tracer().set_sync_marker()
+            test_func()
+        """)
+
+        def expect_aligned_to_sync_marker(data):
+
+            funcs = [event for event in data['traceEvents'] if 'ts' in event and event['name'].startswith('test_func')]
+            self.assertEqual(len(funcs), 2)
+
+            # we expect that aligned events shifted not more than 1ms
+            aligned_diff = abs(funcs[1]['ts'] - funcs[0]['ts'])
+            self.assertLessEqual(aligned_diff, 1000.0, str(data))
+
+        for extra_args in [[], ["--dump_raw"]]:
+            with tempfile.TemporaryDirectory() as tmpdir:
+                res1_filename = os.path.join(tmpdir, 'res1.json')
+                res2_filename = os.path.join(tmpdir, 'res2.json')
+
+                common_cmd_line = [sys.executable, "-m", "viztracer", "--ignore_frozen"] + extra_args
+                self.template(
+                    common_cmd_line + ["-o", res1_filename, "cmdline_test.py"],
+                    expected_output_file=res1_filename,
+                    script=test_script,
+                    cleanup=False,
+                )
+                self.template(
+                    common_cmd_line + ["-o", res2_filename, "cmdline_test.py"],
+                    expected_output_file=res2_filename,
+                    script=test_script,
+                    cleanup=False,
+                )
+
+                self.template(
+                    [sys.executable, "-m", "viztracer", "--align_combine", res1_filename, res2_filename],
+                    expected_output_file="result.json",
+                    check_func=expect_aligned_to_sync_marker,
+                )
+
     def test_tracer_entries(self):
         self.template([sys.executable, "-m", "viztracer", "--tracer_entries", "1000", "cmdline_test.py"])
         self.template([sys.executable, "-m", "viztracer", "--tracer_entries", "50", "cmdline_test.py"])