Source file tally_rusage.ml

type ru = {
  utime : int;
  stime : int;
  maxrss : int;
  ixrss : int;
  idrss : int;
  isrss : int;
  minflt : int;
  majflt : int;
  nswap : int;
  inblock : int;
  outblock : int;
  msgsnd : int;
  msgrcv : int;
  nsignals : int;
  nvcsw : int;
  nivcsw : int;
  vsize : int;
  rss : int;
  tsize : int;
  dsize : int;
  ssize : int;
  runtime : int;
  cow : int;
  start : int;
}

let ( let* ) = Result.bind

external sysconf_clock_tick : unit -> int = "tally_sysconf_clock_tick"
external sysctl_kinfo_proc : int -> ru = "tally_sysctl_kinfo_proc"
external uname : unit -> string = "tally_uname"

let wrap f arg = try Ok (f arg) with e -> Error (`Msg (Printexc.to_string e))

let string_of_file filename =
  try
    let fh = open_in filename in
    let content = input_line fh in
    close_in_noerr fh ;
    Ok content
  with _ -> Error (`Msg (Fmt.str "Error reading file %S" filename))

let parse_proc_stat s =
  let stats_opt =
    match String.index_opt s '(', String.rindex_opt s ')' with
    | Some idxa, Some idxb ->
      let pid = String.sub s 0 (idxa - 1)
      and tcomm = String.sub s (idxa + 1) (idxb - idxa - 1)
      and rest = String.sub s (idxb + 2) (String.length s - (idxb + 2))
      in
      let rest = String.split_on_char ' ' rest in
      Some (pid :: tcomm :: rest)
    | _ -> None
  in
  Option.to_result ~none:(`Msg "unable to parse /proc/<pid>/stat") stats_opt

let read_proc_status pid =
  try
    let fh = open_in ("/proc/" ^ string_of_int pid ^ "/status") in
    let lines =
      let rec read_lines acc = try
          read_lines (input_line fh :: acc)
        with End_of_file -> acc in
      read_lines []
    in
    close_in_noerr fh ;
    List.map (String.split_on_char ':') lines |>
    List.fold_left (fun acc x -> match acc, x with
        | Some acc, k :: v ->
          (* strip leading tab character and further possible whitespace *)
          let v = String.concat ":" v |> String.trim in
          Some ((k, v) :: acc)
        | _ -> None) (Some []) |>
    Option.to_result ~none:(`Msg "failed to parse /proc/<pid>/status")
  with _ -> Error (`Msg (Fmt.str "error reading file /proc/%d/status" pid))

let linux_rusage pid =
  let* start =
    match Unix.stat ("/proc/" ^ string_of_int pid) with
    | { Unix.st_ctime = start; _ } ->
      let frac = Float.rem start 1. in
      Ok (int_of_float start * 1_000_000 + int_of_float (frac *. 1_000_000.))
    | exception Unix.Unix_error (Unix.ENOENT,_,_) -> Error (`Msg "failed to stat process")
  in
  (* reading /proc/<pid>/stat - since it may disappear mid-time,
     best to have it in memory *)
  let* data = string_of_file ("/proc/" ^ string_of_int pid ^ "/stat") in
  let* stat_vals = parse_proc_stat data in
  let* data = string_of_file ("/proc/" ^ string_of_int pid ^ "/statm") in
  let statm_vals = String.split_on_char ' ' data in
  let* status = read_proc_status pid in
  let assoc_i key : (int, _) result =
    let e x = Option.to_result ~none:(`Msg "error parsing /proc/<pid>/status") x in
    let* v = e (List.assoc_opt key status) in
    e (int_of_string_opt v)
  in
  let to_usec t =
    let clock_tick = sysconf_clock_tick () in
    let s, usec =
      t / clock_tick, ((Int.rem t clock_tick) * 1_000_000) / clock_tick
    in
    s * 1_000_000 + usec
  in
  let i s = try Ok (int_of_string s) with
      Failure _ -> Error (`Msg "couldn't parse integer")
  in
  if List.length stat_vals >= 52 && List.length statm_vals >= 7 then
    let* minflt = i (List.nth stat_vals 9) in
    let* majflt = i (List.nth stat_vals 11) in
    let* utime = Result.map to_usec (i (List.nth stat_vals 13)) in (* divide by sysconf(_SC_CLK_TCK) *)
    let* stime = Result.map to_usec (i (List.nth stat_vals 14)) in (* divide by sysconf(_SC_CLK_TCK) *)
    let runtime = utime + stime in
    let* vsize = i (List.nth stat_vals 22) in (* in bytes *)
    let* rss = i (List.nth stat_vals 23) in (* in pages *)
    let* nswap = i (List.nth stat_vals 35) in (* not maintained, 0 *)
    let* tsize = i (List.nth statm_vals 3) in
    let* dsize = i (List.nth statm_vals 5) in (* data + stack *)
    let* ssize = i (List.nth statm_vals 5) in (* data + stack *)
    let* nvcsw = assoc_i "voluntary_ctxt_switches" in
    let* nivcsw = assoc_i "nonvoluntary_ctxt_switches" in
    Ok { utime ; stime ; maxrss = rss ; ixrss = 0 ;
         idrss = 0 ; isrss = 0 ; minflt ; majflt ; nswap ; inblock = 0 ; outblock = 0 ;
         msgsnd = 0 ; msgrcv = 0 ; nsignals = 0 ; nvcsw ; nivcsw ;
         vsize; rss; tsize; dsize; ssize; runtime; cow = 0; start }
  else
    Error (`Msg "couldn't read /proc/<pid>/stat")

let rusage pid =
  match uname () with
  | "FreeBSD" -> wrap sysctl_kinfo_proc pid
  | "Linux" -> linux_rusage pid
  | os -> Error (`Msg ("unkown operating system " ^ os))

let to_fields ru =
  [
    "utime", ru.utime;
    "stime", ru.stime;
    "maxrss", ru.maxrss;
    "ixrss", ru.ixrss;
    "idrss", ru.idrss;
    "isrss", ru.isrss;
    "minflt", ru.minflt;
    "majflt", ru.majflt;
    "nswap", ru.nswap;
    "inblock", ru.inblock;
    "outblock", ru.outblock;
    "msgsnd", ru.msgsnd;
    "msgrcv", ru.msgrcv;
    "nsignals", ru.nsignals;
    "nvcsw", ru.nvcsw;
    "nivcsw", ru.nivcsw;
    "vsize", ru.vsize;
    "rss", ru.rss;
    "tsize", ru.tsize;
    "dsize", ru.dsize;
    "ssize", ru.ssize;
    "runtime", ru.runtime;
    "cow", ru.cow;
    "start", ru.start;
  ]

let v pid =
  let measure () =
    match rusage pid with
    | Error (`Msg msg) ->
      Logs.err (fun m -> m "error measuring rusage: %s" msg);
      [ [], [] ]
    | Ok ru ->
      [ [], to_fields ru ]
  in
  Tally.v "rusage" measure