Allow tracing kernel threads

src/process_info.ml

@@ -27,3 +27,81 @@ let read_all_proc_info () =
 ;;
 
 let cmdline_of_pid pid = Hashtbl.find state pid
+
+let kthreadd_pid = Pid.of_int 2
+
+let ppid_of_stat_line stat_line =
+  match String.rsplit2 stat_line ~on:')' with
+  | None -> None
+  | Some (_, rest) ->
+    (match
+       rest
+       |> String.lstrip
+       |> String.split ~on:' '
+       |> List.filter ~f:(Fn.non String.is_empty)
+     with
+     | _state :: ppid :: _ -> Some (Pid.of_string ppid)
+     | _ -> None)
+;;
+
+let ppid_of_pid pid =
+  try
+    In_channel.read_all [%string "/proc/%{pid#Pid}/stat"] |> ppid_of_stat_line
+  with
+  | _ -> None
+;;
+
+let is_kernel_thread pid =
+  match ppid_of_pid pid with
+  | Some ppid -> Pid.(ppid = kthreadd_pid)
+  | None -> false
+;;
+
+let vmlinux_candidates () =
+  let release = (Core_unix.uname ()).release in
+  [ [%string "/usr/lib/debug/boot/vmlinux-%{release}"]
+  ; [%string "/boot/vmlinux-%{release}"]
+  ; [%string "/usr/lib/debug/lib/modules/%{release}/vmlinux"]
+  ]
+;;
+
+let find_vmlinux () =
+  List.find_map (vmlinux_candidates ()) ~f:(fun path ->
+    if Sys_unix.file_exists path then Some path else None)
+;;
+
+let executable_of_pid pid =
+  let exe_path = [%string "/proc/%{pid#Pid}/exe"] in
+  match Or_error.try_with (fun () -> Core_unix.readlink exe_path) with
+  | Ok path -> Ok path
+  | Error _ when is_kernel_thread pid ->
+    (match find_vmlinux () with
+     | Some vmlinux -> Ok vmlinux
+     | None ->
+       Or_error.error_string
+         "Cannot trace kernel thread: vmlinux not found. Install your distro's kernel \
+          debug package (for example, linux-image-$(uname -r)-dbgsym on Debian/Ubuntu) \
+          and retry with [-trace-include-kernel] or [-trace-kernel-only].")
+  | Error error -> Error error
+;;
+
+let%expect_test "ppid_of_stat_line parses comm with spaces" =
+  Expect_test_helpers_core.require_equal
+    (module struct
+      type t = Pid.t option
+
+      let equal = Option.equal Pid.equal
+      let sexp_of_t = [%sexp_of: Pid.t option]
+    end)
+    (ppid_of_stat_line "(migration/0) S 2 0 0 0 -1 2129984 0 0 0 0 0 0 0 0 20 0 1 0")
+    (Some (Pid.of_int 2))
+  |> Deferred.return
+;;
+
+let%expect_test "is_kernel_thread is false for the current process" =
+  Expect_test_helpers_core.require_equal
+    (module Bool)
+    (is_kernel_thread (Pid.of_int (Core_unix.getpid ())))
+    false
+  |> Deferred.return
+;;

src/process_info.mli

@@ -9,3 +9,6 @@ end
 val read_proc_info : Pid.t -> unit
 val read_all_proc_info : unit -> unit
 val cmdline_of_pid : Pid.t -> Entry.Cmdline.t option
+val is_kernel_thread : Pid.t -> bool
+val find_vmlinux : unit -> string option
+val executable_of_pid : Pid.t -> string Or_error.t

src/trace.ml

@@ -670,26 +670,11 @@ module Make_commands (Backend : Backend_intf.S) = struct
   let select_pid () =
     if force supports_fzf
     then (
-      let deselect_pid_args pid =
-        let pid = Pid.to_string pid in
-        [ "--ppid"; pid; "-p"; pid; "--deselect" ]
-      in
       (* There are no Linux APIs, or OCaml libraries that I've found, for enumerating
          running processes. The [ps] command uses the /proc/ filesystem and is much easier
          than walking the /proc/ system and filtering ourselves. *)
       let process_lines =
-        [ [ "x"; "-w"; "--no-headers" ]
-        ; [ "-o"; "pid,args" ]
-          (* If running as root, allow tracing all processes, including those owned
-             by non-root users.
-
-             Hide kernel threads (PID 2 and children), since though we can trace them in
-             theory, in practice they don't have their image under /proc/$pid/exe, which
-             we currently rely on. *)
-        ; (if Core_unix.geteuid () = 0 then deselect_pid_args (Pid.of_int 2) else [])
-        ]
-        |> List.concat
-        |> Shell.run_lines "ps"
+        Shell.run_lines "ps" [ "x"; "-w"; "--no-headers"; "-o"; "pid,args" ]
       in
       let%bind.Deferred.Or_error sel_line =
         Fzf.pick_one (Fzf.Pick_from.Inputs process_lines)
@@ -745,14 +730,25 @@ module Make_commands (Backend : Backend_intf.S) = struct
          else (
            (* Always use the head PID for locating triggers since only a single
               trigger can be passed currently. *)
-           let executable =
-             List.hd_exn pids
-             |> fun pid -> Core_unix.readlink [%string "/proc/%{pid#Pid}/exe"]
+           let head_pid = List.hd_exn pids in
+           let%bind executable =
+             Process_info.executable_of_pid head_pid |> Deferred.Or_error.of_or_error
            in
            record_opt_fn ~executable ~f:(fun opts ->
-             let { Record_opts.executable; when_to_snapshot; collection_mode; _ } =
+             let { Record_opts.executable; when_to_snapshot; collection_mode; trace_scope; _ } =
                opts
              in
+             let%bind () =
+               if Process_info.is_kernel_thread head_pid
+               then (
+                 match trace_scope with
+                 | Userspace ->
+                   Deferred.Or_error.error_string
+                     "Cannot trace kernel thread without kernel tracing enabled. Pass \
+                      [-trace-include-kernel] or [-trace-kernel-only]."
+                 | Kernel | Userspace_and_kernel -> return (Ok ()))
+               else return (Ok ())
+             in
              let%bind elf = create_elf ~executable ~when_to_snapshot in
              let%bind range_symbols =
                (* TODO Use LLVM to load the symbol table, because Owee can't handle executables that use DWARF5. *)