Merge pull request #41 from talex5/cpu

Add --cpu option

Author: talex5

dune-project

@@ -11,6 +11,7 @@
  (synopsis "Trace visualisation for Eio programs")
  (depends
   (ocaml (>= 5.1.0))
+  processor
   (eio_main (>= 0.14))
   (cmdliner (>= 1.2.0))
   (lablgtk3 (>= 3.1.4))

eio-trace.opam

@@ -9,6 +9,7 @@ bug-reports: "https://github.com/ocaml-multicore/eio-trace/issues"
 depends: [
   "dune" {>= "3.11"}
   "ocaml" {>= "5.1.0"}
+  "processor"
   "eio_main" {>= "0.14"}
   "cmdliner" {>= "1.2.0"}
   "lablgtk3" {>= "3.1.4"}

src/dune

@@ -1,4 +1,4 @@
 (executable
   (public_name eio-trace)
   (name main)
-  (libraries eio_main eio_trace cmdliner fxt))
+  (libraries eio_main eio_trace cmdliner fxt processor))

src/main.ml

@@ -27,10 +27,6 @@ let fmt =
   in
   Arg.(value @@ opt (some formats) None @@ info ["T"; "format"] ~docv:"TYPE" ~doc)
 
-let freq =
-  let doc = "How many times per second to check for events." in
-  Arg.(value @@ opt float 100.0 @@ info ["F"; "freq"] ~docv:"RATE" ~doc)
-
 let start_time =
   let doc = "Seconds to skip before the section to display." in
   Arg.(value @@ opt (some time) None @@ info ["s"; "start-time"] ~doc)
@@ -39,10 +35,6 @@ let duration =
   let doc = "Width of the output image in seconds." in
   Arg.(value @@ opt (some time) None @@ info ["d"; "duration"] ~doc)
 
-let child_args =
-  let doc = "The command to be executed and monitored." in
-  Arg.(non_empty @@ pos_all string [] @@ info [] ~docv:"command" ~doc)
-    
 let eio_trace_gtk = "eio-trace-gtk"
 
 let find_eio_trace_gtk () =
@@ -66,16 +58,16 @@ let exec_gtk args =
 
 let show tracefiles = exec_gtk ("show" :: tracefiles)
 
-let run ~fs ~proc_mgr freq args =
+let run ~fs ~proc_mgr config =
   let gtk_exe = find_eio_trace_gtk () in
   let ui tracefile =
-    Eio.Process.run proc_mgr (gtk_exe :: "run" :: tracefile :: args);
+    Eio.Process.run proc_mgr (gtk_exe :: "run" :: tracefile :: config.Record.child_args);
     Ok ()
   in
-  Record.run ~fs ~proc_mgr ~freq ~ui args
+  Record.run ~fs ~proc_mgr ~ui config
 
-let record ~fs ~proc_mgr freq tracefile args =
-  Record.run ~fs ~proc_mgr ~freq ~tracefile args
+let record ~fs ~proc_mgr tracefile config =
+  Record.run ~fs ~proc_mgr ~tracefile config
 
 let ( let* ) = Result.bind
 
@@ -104,9 +96,9 @@ let cmd env =
   let path = Eio.Path.( / ) fs in
   Cmd.group (Cmd.info "eio-trace")
   @@ List.map (fun (name, term) -> Cmd.v (Cmd.info name) term) [
-    "record", record ~fs ~proc_mgr           $$ freq $ (path $$ tracefile) $ child_args;
+    "record", record ~fs ~proc_mgr           $$ (path $$ tracefile) $ Record.cmdliner;
     "dump",   Dump.main Format.std_formatter $$ (List.map path $$ tracefiles);
-    "run",    run ~fs ~proc_mgr              $$ freq $ child_args;
+    "run",    run ~fs ~proc_mgr              $$ Record.cmdliner;
     "show",   show                           $$ tracefiles;
     "render", render                         $$ tracefiles $ imagefile $ fmt $ start_time $ duration;
     "gc-stats", Gc_stats.main Format.std_formatter $$ (List.map path $$ tracefiles);

src/record.ml

@@ -5,6 +5,12 @@ module Fibers = Map.Make(Int)
 
 module Write = Fxt.Write
 
+type config = {
+  freq : float;
+  cpu : int option;
+  child_args : string list;
+}
+
 type fiber = {
   id : int;
   mutable op : string option;   (* If suspended *)
@@ -199,7 +205,7 @@ let rec get_cursor tmp_dir child =
 
 let ( / ) = Eio.Path.( / )
 
-let run ?ui ?tracefile ~proc_mgr ~fs ~freq args =
+let run ?ui ?tracefile ~proc_mgr ~fs { freq; cpu; child_args } =
   let delay = 1. /. freq in
   let fs = (fs :> Eio.Fs.dir_ty Eio.Path.t) in
   let tracefile = (tracefile :> Eio.Fs.dir_ty Eio.Path.t option) in
@@ -216,7 +222,8 @@ let run ?ui ?tracefile ~proc_mgr ~fs ~freq args =
   Eio.Buf_write.with_flow out @@ fun w ->
   let fxt = Write.of_writer w in
   traceln "Recording to %a" Eio.Path.pp tracefile;
-  let child = spawn_child ~sw ~proc_mgr ~tmp_dir args in
+  let child = spawn_child ~sw ~proc_mgr ~tmp_dir child_args in
+  cpu |> Option.iter (fun cpu -> Processor.Affinity.set_cpus [List.nth Processor.Topology.t cpu]);
   let t = {
     fxt;
     pid = Int64.of_int (Eio.Process.pid child);
@@ -246,3 +253,33 @@ let run ?ui ?tracefile ~proc_mgr ~fs ~freq args =
       (fun () -> Promise.await finished)
       (fun () -> Eio_unix.sleep 1.);
     ui (Eio.Path.native_exn tracefile)
+
+open Cmdliner
+
+let ( $ ) = Term.app
+let ( $$ ) f x = Term.const f $ x
+
+let freq =
+  let doc = "How many times per second to check for events." in
+  Arg.(value @@ opt float 100.0 @@ info ["F"; "freq"] ~docv:"RATE" ~doc)
+
+let cpu =
+  let doc = "CPU for eio-trace" in
+  let parse_cpu s =
+    match Arg.(conv_parser int) s with
+    | Error _ as e -> e
+    | Ok x ->
+      let l = List.length Processor.Topology.t in
+      if x >= 0 && x < l then Ok x
+      else Fmt.error_msg "CPU %d not in range 0..%d" x (l - 1)
+  in
+  let cpu_conv = Arg.conv (parse_cpu, Fmt.int) in
+  Arg.(value @@ opt (some cpu_conv) None @@ info ["cpu"] ~docv:"CPU" ~doc)
+
+let child_args =
+  let doc = "The command to be executed and monitored." in
+  Arg.(non_empty @@ pos_all string [] @@ info [] ~docv:"command" ~doc)
+
+let cmdliner =
+  let make freq cpu child_args = { freq; cpu; child_args } in
+  make $$ freq $ cpu $ child_args

src/record.mli

@@ -1,8 +1,15 @@
+type config = {
+  freq : float;
+  cpu : int option;
+  child_args : string list;
+}
+
 val run :
   ?ui : (string -> (unit, string) result) ->
   ?tracefile:_ Eio.Path.t ->
   proc_mgr:_ Eio.Process.mgr ->
   fs:_ Eio.Path.t ->
-  freq:float ->
-  string list ->
+  config ->
   (unit, string) result
+
+val cmdliner : config Cmdliner.Term.t