scripts: Convert sched ftrace helper scripts to python

Merge the sched_switch ftrace helper scripts into a single python script
that prints the result to stdout.

In this way it's possible to generate a perfetto-compatible trace
running:

 $ sudo ./scripts/sched_ftrace.py > sched.ftrace

Signed-off-by: Andrea Righi <andrea.righi@linux.dev>

scripts/ftrace_trim

@@ -1,26 +0,0 @@
-#!/usr/bin/env python3
-import sys
-
-HEADER = "TASK-PID"
-BUFF_STARTED = "buffer started ###"
-
-def main():
-    nproc = int(sys.argv[2]) - 1
-
-    seen_header = False
-    proc_buffer_started = 0
-    with open(sys.argv[1]) as f:
-        for line in f:
-            l = line.replace("\n", "")
-            if HEADER in l:
-                seen_header = True
-                print(l)
-            if BUFF_STARTED in line:
-                proc_buffer_started += 1
-                continue
-            if proc_buffer_started == nproc or not seen_header:
-                print(l)
-
-
-if __name__=="__main__":
-    main()

scripts/sched_ftrace.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+
+import os
+import sys
+import time
+
+HEADER = "TASK-PID"
+BUFF_STARTED = "buffer started ###"
+TRACING_PATH = "/sys/kernel/debug/tracing"
+TRACE_PIPE_PATH = os.path.join(TRACING_PATH, "trace_pipe")
+
+
+def ftrace_trim(stream, duration, nproc):
+    nproc = nproc - 1
+    seen_header = False
+    proc_buffer_started = 0
+    start_time = time.time()
+
+    for line in stream:
+        if time.time() - start_time >= duration:
+            break
+        l = line.replace("\n", "")
+        if HEADER in l:
+            seen_header = True
+            print(l)
+        if BUFF_STARTED in line:
+            proc_buffer_started += 1
+            continue
+        if proc_buffer_started == nproc or not seen_header:
+            print(l)
+
+
+def run_trace(duration):
+    tracing_on_path = os.path.join(TRACING_PATH, "tracing_on")
+    sched_switch_enable_path = os.path.join(
+        TRACING_PATH, "events/sched/sched_switch/enable"
+    )
+
+    # Enable tracing and sched_switch event
+    with open(tracing_on_path, "w") as f:
+        f.write("1")
+    with open(sched_switch_enable_path, "w") as f:
+        f.write("1")
+
+    # Process the sched_switch events from the trace file
+    try:
+        with open(TRACE_PIPE_PATH, "r") as trace_pipe:
+            ftrace_trim(trace_pipe, duration, os.cpu_count())
+
+    except KeyboardInterrupt:
+        pass  # Allow clean termination with Ctrl+C
+
+    # Disable tracing and sched_switch event after the duration
+    with open(sched_switch_enable_path, "w") as f:
+        f.write("0")
+    with open(tracing_on_path, "w") as f:
+        f.write("0")
+
+
+def main():
+    duration = int(sys.argv[1]) if len(sys.argv) > 1 else 5
+    run_trace(duration)
+
+
+if __name__ == "__main__":
+    main()

scripts/sched_ftrace.sh

@@ -1,12 +0,0 @@
-#!/usr/bin/env bash
-duration=${1:-5}
-out=${2:-"$(pwd)/sched.ftrace"}
-cd /sys/kernel/debug/tracing
-echo 1 > tracing_on
-echo 1 > events/sched/sched_switch/enable
-timeout "${duration}" cat trace_pipe > "${out}.tmp"
-echo 0 > events/sched/sched_switch/enable
-echo 0 > tracing_on
-cd -
-./ftrace_trim "${out}.tmp" `nproc` > "${out}"
-rm "${out}.tmp"