Source file smaps.ml

open Stdint
open Lwt.Syntax

type value = Stdint.uint64

module Fields = struct
  type key = string
  type 'a t = (key * 'a) list

  let empty = []
  let add k v l = (k, v) :: l
  let find_exn = List.assoc
  let find = List.assoc_opt

  let fold f = List.fold_left (fun s (k, v) -> f k v s)

  let map_opt f mp =
    fold (fun k v m -> match f k v with Some v -> add k v m | None -> m) empty mp

  let to_list l = l

  let key_of_string x = x

  (* RSS = Shared_Clean + Shared_Dirty + Private_Clean + Private_Dirty *)
  let size = "Size"
  let kernel_page_size ="KernelPageSize"
  let mmu_page_size = "MMUPageSize"
  let rss = "Rss"
  let pss = "Pss"
  let shared_clean = "Shared_Clean"
  let shared_dirty = "Shared_Dirty"
  let private_clean = "Private_Clean"
  let private_dirty = "Private_Dirty"
  let referenced = "Referenced"
  let anonymous = "Anonymous"
  let lazy_free = "LazyFree"
  let anon_huge_pages = "AnonHugePages"
  let shmem_pmd_mapped = "ShmemPmdMapped"
  let file_pmd_mapped = "FilePmdMapped"
  let shared_hugetlb = "Shared_Hugetlb"
  let private_hugetlb = "Private_Hugetlb"
  let swap = "Swap"
  let swap_pss = "SwapPss"
  let locked = "Locked"
end

type t =
  { address           : value * value
  ; perms             : int
  ; offset            : value
  ; inode             : int
  ; pathname          : string
  ; size_fields       : value Fields.t
  }

let get_address { address; _ } = address
let get_perms { perms; _ } = perms
let get_offset { offset; _ } = offset
let get_inode { inode; _ } = inode
let get_pathname { pathname; _ } = pathname
let get_size_exn mapping field = Fields.find_exn field mapping.size_fields
let get_size mapping field = Fields.find field mapping.size_fields

let pp_perms ppf perms =
  let f perm b c =
    if (b land perm) > 0 then c
    else '-' in
  let r = f perms 16 'r' in
  let w = f perms 8 'w' in
  let x = f perms 4 'x' in
  let s = f perms 2 's' in
  let p = f perms 1 'p' in
  Format.fprintf ppf "%c%c%c%c%c" r w x s p

let pp_mapping ppf
    { address = (address_l, address_r)
    ; perms
    ; offset
    ; inode
    ; pathname
    ; size_fields } =
  let pp_sz_field ppf (fld, sz) =
    let sz = Uint64.to_string sz in
    Format.fprintf ppf "%s: %s kB" fld sz in

  Format.fprintf ppf (
    "Pathname: %s\n" ^^
    "Address: %s - %s\n" ^^
    "Offset: %s\n" ^^
    "Permission: %a\n" ^^
    "Inode: %d\n" ^^
    "%a\n"
  )
    pathname
    (Uint64.to_string_hex address_l)
    (Uint64.to_string_hex address_r)
    (Uint64.to_string_hex offset)
    pp_perms perms
    inode
    (Format.pp_print_list
       ~pp_sep:(fun ppf () -> Format.fprintf ppf "\n")
       pp_sz_field)
    (Fields.to_list size_fields)

let pp_mappings ppf mappings =
  Format.fprintf ppf "%a"
    (Format.pp_print_list
       ~pp_sep:(fun ppf () -> Format.fprintf ppf "\n")
       pp_mapping)
    mappings

let bind x f = Lwt.bind x (function Ok x -> f x | Error e -> Lwt.return_error e)
let bind_opt x f = Lwt.bind x (function Ok (Some x) -> f x | Ok None -> Lwt.return_ok None | Error e -> Lwt.return_error e)
let return = Lwt.return_ok
let (>>=?) = bind
let (>>>=?) = bind_opt

let parse_perms str =
  if String.length str < 4 then
    Error "Failed to parse permissions"
  else
    let r = if String.get str 0 = 'r' then 16 else 0 in
    let w = if String.get str 1 = 'w' then 8 else 0 in
    let x = if String.get str 2 = 'x' then 4 else 0 in

    (* shared flag and private flag are represented in the same position *)
    let s = if String.get str 3 = 's' then 2 else 0 in
    let p = if String.get str 3 = 'p' then 1 else 0 in
    Ok (r + w + x + s + p)

let is_shared { perms; _ } =
  (perms land 2) > 0

let parse_mapping stream =
  let rec loop stream flds =
    let* str = Lwt_stream.peek stream in
    match str with
    | Some s ->
        let ss = Str.(split (regexp "[ \t]+") s) in
        (match ss with
         | ssh :: sst ->
             if String.length ssh < 1 then Lwt.return (Error "")
             else if String.get ssh (String.length ssh -1) = ':' then
               let* _ = Lwt_stream.get stream in
               let key = String.sub ssh 0 (String.length ssh - 1) in
               let flds = Fields.add key sst flds in
               loop stream flds
             else
               return (Some flds) (* May reach next mapping header *)
         | [] -> loop stream flds (* Empty line *))
    | None -> return None in
  Lwt.bind (Lwt_stream.get stream) @@ return >>>=? fun str ->
  loop stream Fields.empty >>>=? fun flds ->
  begin
    let (let*) = Result.bind in
    let ss = Str.(split (regexp "[ \t]+") str) in
    let* (address, perms, offset, _dev, inode, pathname) =
      match ss with
      | [ s0; s1; s2; s3; s4 ] -> Ok (s0, s1, s2, s3, s4, "[ anon ]")
      | s0 :: s1 :: s2 :: s3 :: s4 :: tl ->
          let pathname = String.concat " " tl in
          Ok (s0, s1, s2, s3, s4, pathname)
      | _ -> Error "The first line of a mapping is not found" in
    let* address = match Str.(split (regexp "-") address) with
      | [l ; r] -> Ok (Uint64.of_string ("0x" ^ l), Uint64.of_string ("0x" ^ r))
      | _ -> Error "Failed to parse address" in
    let* perms = parse_perms perms in
    let* offset = try Ok (Uint64.of_string ("0x" ^ offset)) with Failure err -> Error err in
    let* inode = match int_of_string_opt inode with Some v -> Ok v | None -> Error "Failed to parse inode" in
    let size_fields =
      Fields.map_opt (fun _k v ->
          match v with
          | sz :: un :: _ when un = "kB" ->
              Some (Uint64.of_string sz)
          | _ -> None) flds in
    Ok (Some { address
             ; perms
             ; offset
             ; inode
             ; pathname
             ; size_fields
             })
  end |> Lwt.return

let parse_stream stream =
  let rec loop stream l =
    parse_mapping stream >>=? function
    | Some m -> loop stream (m::l)
    | None -> Lwt.return_ok l in
  loop stream []

let get_smaps fn =
  Lwt_io.with_file ~mode:Input fn
    (fun ch ->
       let stream = Lwt_io.read_lines ch in
       parse_stream stream)

let get_self_smaps () =
  let fn = Format.sprintf "/proc/self/smaps" in
  get_smaps fn

let get_smaps pid =
  let fn = Format.sprintf "/proc/%d/smaps" pid in
  get_smaps fn

let sum_rss =
  List.fold_left (fun s m -> Uint64.add s (get_size_exn m Fields.rss)) (Uint64.of_int 0)

(** Summalize RSS of shared mappings *)
let sum_shared_rss mappings =
  List.filter is_shared mappings
  |> sum_rss