Refactored, process table now loaded centrally.
[virt-mem.git] / lib / virt_mem.ml
index 69a0821..22cbbc5 100644 (file)
@@ -42,12 +42,18 @@ let kernel_size =
 let tools = ref []
 
 (* Registration function used by the tools. *)
-let register ?(external_cmd = true) ?(extra_args = [])
-    ?argcheck ?beforeksyms ?beforeutsname ?run
+let register
+    ?(needs_ksyms = false) ?(needs_utsname = false)
+    ?(needs_tasks = false) ?(needs_everything = false)
+    ~run
+    ?(external_cmd = true)
+    ?(extra_args = [])
+    ?argcheck
     name summary description =
   tools :=
-    (name, (name, summary, description, external_cmd, extra_args,
-           argcheck, beforeksyms, beforeutsname, run))
+    (name, (name, summary, description,
+           needs_ksyms, needs_utsname, needs_tasks, needs_everything,
+           run, external_cmd, extra_args, argcheck))
   :: !tools
 
 (* Main program, called from mem/virt_mem_main.ml when all the
@@ -105,7 +111,7 @@ let main () =
     match tool with
     | None ->                          (* Generic usage message. *)
        let tools = List.map (
-         fun (name, (_, summary, _, external_cmd, _, _, _, _, _)) ->
+         fun (name, (_, summary, _, _, _, _, _, _, external_cmd, _, _)) ->
            if external_cmd then "virt-"^name, summary
            else                 "virt-mem "^name, summary
        ) tools in
@@ -127,12 +133,12 @@ General usage is:
   <tool> [-options] [domains...]
 
 To display extra help for a single tool, do:
-  virt-mem help <tool>
+  virt-mem --help <tool>
 
 Options:") tools
 
                                         (* Tool-specific usage message. *)
-    | Some (name, summary, description, external_cmd, _, _, _, _, _) ->
+    | Some (name, summary, description, _, _, _, _, _, external_cmd, _, _) ->
        let cmd =
          if external_cmd then "virt-" ^ name else "virt-mem " ^ name in
 
@@ -242,7 +248,7 @@ Options:") cmd summary description in
   let argspec =
     let extra_args = match tool with
       | None -> []
-      | Some (_, _, _, _, extra_args, _, _, _, _) -> extra_args in
+      | Some (_, _, _, _, _, _, _, _, _, extra_args, _) -> extra_args in
     let argspec = [
       "-A", Arg.String set_architecture,
         "arch " ^ s_"Set kernel architecture, endianness and word size";
@@ -293,7 +299,9 @@ Options:") cmd summary description in
    * or the user didn't give us a valid tool (eg. "virt-mem foobar").
    * Detect that final case now and give an error.
    *)
-  let name, _, _, _, _, argcheck, beforeksyms,  beforeutsname, run =
+  let name, _, _,
+    needs_ksyms, needs_utsname, needs_tasks, needs_everything,
+    run, external_cmd, extra_args, argcheck =
     match tool with
     | Some t -> t
     | None ->
@@ -480,68 +488,95 @@ Possibly the '-T' command line parameter was used inconsistently.");
       ) testimages
     ) in
 
-  (* Optional callback into the tool before we start looking for
-   * kernel symbols.
-   *)
-  (match beforeksyms with
-   | None -> ()
-   | Some beforeksyms -> beforeksyms debug images
-  );
-
-  (* If there are no more callback functions, then there is no point
-   * continuing with the rest of the program (kernel symbol analysis) ...
-   *)
-  if beforeutsname = None && run = None then exit 0;
-
-  (* Do the kernel symbol analysis. *)
+  (* Now build the kdata, depending on what the tool asked for. *)
   let images =
     List.map (
       fun image ->
-       (* Look for ordinary kernel symbols: *)
-       let image = Virt_mem_ksyms.find_kernel_symbols debug image in
-       (* Look for kallsyms: *)
-       let image = Virt_mem_kallsyms.find_kallsyms debug image in
-
-(*
-       (* Finally, just wrap the lookup_ksym call in something
-        * which prints the query when debug is set.
-        *)
-       let image =
-         if debug then
-           let (domid, name, arch, mem, lookup_ksym) = image in
-           let lookup_ksym sym =
-             try
-               let value = lookup_ksym sym in
-               eprintf "lookup_ksym %S = %Lx\n%!" sym value;
-               value
-             with Not_found ->
-               eprintf "lookup_ksym %S failed\n%!" sym;
-               raise Not_found
-           in
-           (domid, name, arch, mem, lookup_ksym)
-         else
-           image in
-*)
-
-       image
+       let kdata = { ksyms = None; utsname = None; tasks = None } in
+       image, kdata
     ) images in
+  (* Certain needs are dependent on others ... *)
+  let needs_ksyms =
+    if needs_utsname then true
+    else needs_ksyms in
+  let needs_ksyms, needs_utsname =
+    if needs_tasks then true, true
+    else needs_ksyms, needs_utsname in
+  let needs_ksyms, needs_utsname, needs_tasks =
+    if needs_everything then true, true, true
+    else needs_ksyms, needs_utsname, needs_tasks in
 
-  (* Before utsname analysis. *)
-  (match beforeutsname with
-   | None -> ()
-   | Some beforeutsname -> List.iter (beforeutsname debug) images
-  );
-
-  (* If there are no more callback functions, then there is no point
-   * continuing with the rest of the program (kernel version analysis) ...
-   *)
-  if run = None then exit 0;
+  (* Do the kernel symbol analysis. *)
+  let images =
+    if not needs_ksyms then images
+    else
+      List.map (
+       fun (image, kdata) ->
+         (* Look for ordinary kernel symbols: *)
+         let image, ksyms =
+           Virt_mem_ksyms.find_kernel_symbols debug image in
+
+         match ksyms with
+         | None -> image, kdata
+         | Some ksyms ->
+             (* Look for kallsyms: *)
+             let image, kallsyms =
+               Virt_mem_kallsyms.find_kallsyms debug image ksyms in
+
+             let ksyms =
+               match kallsyms with
+               | None -> ksyms (* no kallsyms, just use module symbols *)
+               | Some kallsyms -> kallsyms (* ksyms + kallsyms *) in
+
+             image, { kdata with ksyms = Some ksyms }
+      ) images in
 
   (* Get the kernel version (utsname analysis). *)
-  let images = List.map (Virt_mem_utsname.find_utsname debug) images in
+  let images =
+    if not needs_utsname then images
+    else
+      List.map (
+       fun (image, ({ ksyms = ksyms } as kdata)) ->
+         match ksyms with
+         | None -> image, kdata
+         | Some ksyms ->
+             let image, utsname =
+               Virt_mem_utsname.find_utsname debug image ksyms in
+             let kdata = { kdata with utsname = utsname } in
+             image, kdata
+      ) images in
+
+  (* Get the tasks. *)
+  let images =
+    if not needs_tasks then images
+    else
+      List.map (
+       fun (image, ({ ksyms = ksyms; utsname = utsname } as kdata)) ->
+         match ksyms, utsname with
+         | Some ksyms, Some { kernel_release = kversion } ->
+             let image, tasks =
+               Virt_mem_tasks.find_tasks debug image ksyms kversion in
+             let kdata = { kdata with tasks = tasks } in
+             image, kdata
+         | _, _ -> image, kdata
+      ) images in
 
   (* Run the tool's main function. *)
-  (match run with
-   | None -> ()
-   | Some run -> List.iter (run debug) images
-  )
+  let errors = ref 0 in
+  List.iter (
+    fun (image, kdata) ->
+      try
+       if not needs_everything then (
+         if needs_ksyms && kdata.ksyms = None then
+           failwith (s_"could not read kernel symbols")
+         else if needs_utsname && kdata.utsname = None then
+           failwith (s_"could not read kernel version")
+         else if needs_tasks && kdata.tasks = None then
+           failwith (s_"could not read process table")
+       );
+       run debug image kdata
+      with exn ->
+       eprintf "%s: %s\n" image.domname (Printexc.to_string exn);
+       incr errors
+  ) images;
+  exit (if !errors > 0 then 1 else 0)