Merge pull request #5778 from contificate/IH-642

IH-642: Restructure xs-trace to use Cmdliner

Author: Colin

ocaml/xs-trace/dune

@@ -4,7 +4,20 @@
   (public_name xs-trace)
   (package xapi)
   (libraries
+    cmdliner
     tracing_export
     xapi-stdext-unix
   )
 )
+
+(rule
+  (targets xs-trace.1)
+  (deps (:exe xs_trace.exe))
+  (action (with-stdout-to %{targets} (run %{exe} --help=groff)))
+)
+
+(install
+  (section man)
+  (package xapi)
+  (files (xs-trace.1 as man1/xs-trace.1))
+)

ocaml/xs-trace/xs_trace.ml

@@ -12,48 +12,92 @@
  * GNU Lesser General Public License for more details.
  *)
 
-let _ =
-  match Sys.argv with
-  | [|_; action; origin; url_string|] ->
-      let url = Uri.of_string url_string in
-      let submit_json json =
-        if json <> "" then
-          let result = Tracing_export.Destination.Http.export ~url json in
-          match result with
-          | Ok _ ->
-              ()
-          | Error err ->
-              Printf.eprintf "Error: %s" (Printexc.to_string err) ;
-              exit 1
-      in
-      let rec export_file orig =
-        if Sys.is_directory orig then
-          let files = Sys.readdir orig in
-          let file_paths = Array.map (Filename.concat orig) files in
-          Array.iter export_file file_paths
-        else if Filename.check_suffix orig ".zst" then
-          Xapi_stdext_unix.Unixext.with_file orig [O_RDONLY] 0o000
-          @@ fun compressed_file ->
-          Zstd.Fast.decompress_passive compressed_file @@ fun decompressed ->
-          if Filename.check_suffix orig ".ndjson.zst" then
+module Exporter = struct
+  module Unixext = Xapi_stdext_unix.Unixext
+
+  (* Submit JSON to a specified endpoint. *)
+  let submit_json url json =
+    if json <> "" then
+      match Tracing_export.Destination.Http.export ~url json with
+      | Error err ->
+          Printf.eprintf "Error: %s" (Printexc.to_string err) ;
+          exit 1
+      | _ ->
+          ()
+
+  (** Export traces from file system to a remote endpoint. *)
+  let export erase src dst =
+    let dst = Uri.of_string dst in
+    let submit_json = submit_json dst in
+    let rec export_file = function
+      | path when Sys.is_directory path ->
+          (* Recursively export trace files. *)
+          Sys.readdir path
+          |> Array.iter (fun f -> Filename.concat path f |> export_file)
+      | path when Filename.check_suffix path ".zst" ->
+          (* Decompress compressed trace file and decide whether to
+             treat it as line-delimited or not. *)
+          let ( let@ ) = ( @@ ) in
+          let@ compressed = Unixext.with_file path [O_RDONLY] 0o000 in
+          let@ decompressed = Zstd.Fast.decompress_passive compressed in
+          if Filename.check_suffix path ".ndjson.zst" then
             let ic = Unix.in_channel_of_descr decompressed in
-            Xapi_stdext_unix.Unixext.lines_iter
-              (fun line -> submit_json line)
-              ic
+            Unixext.lines_iter submit_json ic
           else
-            let json = Xapi_stdext_unix.Unixext.string_of_fd decompressed in
+            let json = Unixext.string_of_fd decompressed in
             submit_json json
-        else if Filename.check_suffix orig ".ndjson" then
-          Xapi_stdext_unix.Unixext.readfile_line
-            (fun line -> submit_json line)
-            orig
-        else
-          let json = Xapi_stdext_unix.Unixext.string_of_file orig in
+      | path when Filename.check_suffix path ".ndjson" ->
+          (* Submit traces line by line. *)
+          Unixext.readfile_line submit_json path
+      | path ->
+          (* Assume any other extension is a valid JSON file. *)
+          let json = Unixext.string_of_file path in
           submit_json json
-      in
-      export_file origin ;
-      if action = "mv" then
-        Xapi_stdext_unix.Unixext.rm_rec ~rm_top:true origin
-  | _ ->
-      Printf.eprintf "Usage: %s cp/mv <origin> <destination>\n" Sys.argv.(0) ;
-      exit 1
+    in
+    export_file src ;
+    if erase then
+      Unixext.rm_rec ~rm_top:true src
+end
+
+module Cli = struct
+  open Cmdliner
+
+  let src =
+    let doc = "The trace file, e.g. /path/to/trace.ndjson" in
+    Arg.(required & pos 0 (some string) None (info [] ~docv:"SRC" ~doc))
+
+  let dst =
+    let doc =
+      "The destination endpoint URL, e.g. http://localhost:9411/api/v2/spans"
+    in
+    Arg.(required & pos 1 (some string) None (info [] ~docv:"DST" ~doc))
+
+  let export_term ~erase = Term.(const Exporter.export $ const erase $ src $ dst)
+
+  let cp_cmd =
+    let term = export_term ~erase:false in
+    let doc = "copy a trace to an endpoint" in
+    Cmd.(v (info "cp" ~doc) term)
+
+  let mv_cmd =
+    let term = export_term ~erase:true in
+    let doc = "copy a trace to an endpoint and erase it afterwards" in
+    Cmd.(v (info "mv" ~doc) term)
+
+  let xs_trace_cmd =
+    let man =
+      [
+        `S "DESCRIPTION"
+      ; `P "$(mname) is a utility for working with local trace files"
+      ]
+    in
+    let desc =
+      let doc = "utility for working with local trace files" in
+      Cmd.info "xs-trace" ~doc ~version:"0.1" ~man
+    in
+    Cmd.group desc [cp_cmd; mv_cmd]
+
+  let main () = Cmd.eval xs_trace_cmd
+end
+
+let () = exit (Cli.main ())