qcow: Only process allocated clusters on export from raw (#6769)

Implements an optimization similar to the "hybrid" mode in `vhd-tool`:
when exporting to qcow from raw, if the VDI is backed by a QCOW file,
read its header, determine the allocated clusters, and only export
these. This allows skipping over zero clusters in a sparse disk.

Unlike `vhd-tool`, however, this is implemented in a modular way -
`qcow-stream-tool` gets a new `read_headers` command that outputs the
list of allocated clusters (and other info) in JSON format, which allows
it to be consumed by the Python `qcow-to-stdout` script (and by
`vhd-tool` in future stages of this work, see below).

This is the first step of improving handling of sparse VDIs in xapi.
I've got the rest working, but I'll be opening PRs step-by-step for the
following once this PR gets merged:
1. `vhd-tool` gets a `read_headers` command outputting list of allocated
blocks as well
2. `stream_vdi` uses `read_headers` for both VHD and QCOW to avoid
reading zero blocks on XVA export (greatly speeds up handling of sparse
disks and avoids issues with timeouts)
3. `vhd-tool` and `qcow-to-stdout` can read headers of the opposite
format, allowing faster export of sparse VDIs backed by a different
format.

Best reviewed by commit.

Author: last-genius

ocaml/qcow-stream-tool/dune

@@ -7,5 +7,9 @@
     qcow-stream
     cmdliner
     unix
+    lwt.unix
+    lwt
+    qcow-types
+    yojson
   )
 )

ocaml/qcow-stream-tool/qcow_stream_tool.ml

@@ -1,11 +1,53 @@
+open Cmdliner
+
 module Impl = struct
   let stream_decode output =
     Qcow_stream.stream_decode Unix.stdin output ;
     `Ok ()
+
+  let read_headers qcow_path =
+    let open Lwt.Syntax in
+    let t =
+      let* fd = Lwt_unix.openfile qcow_path [Unix.O_RDONLY] 0 in
+      let* virtual_size, cluster_bits, _, data_cluster_map =
+        Qcow_stream.start_stream_decode fd
+      in
+      let clusters = Qcow_types.Cluster.Map.bindings data_cluster_map in
+      let clusters =
+        List.map
+          (fun (_, virt_address) ->
+            let ( >> ) = Int64.shift_right_logical in
+            let address =
+              Int64.to_int (virt_address >> Int32.to_int cluster_bits)
+            in
+            `Int address
+          )
+          clusters
+      in
+      let json =
+        `Assoc
+          [
+            ("virtual_size", `Int (Int64.to_int virtual_size))
+          ; ("cluster_bits", `Int (Int32.to_int cluster_bits))
+          ; ("data_clusters", `List clusters)
+          ]
+      in
+      let json_string = Yojson.to_string json in
+      let* () = Lwt_io.print json_string in
+      let* () = Lwt_io.flush Lwt_io.stdout in
+      Lwt.return_unit
+    in
+    Lwt_main.run t ; `Ok ()
 end
 
 module Cli = struct
-  open Cmdliner
+  let output default =
+    let doc = Printf.sprintf "Path to the output file." in
+    Arg.(value & pos 0 string default & info [] ~doc)
+
+  let input =
+    let doc = Printf.sprintf "Path to the input file." in
+    Arg.(required & pos 0 (some string) None & info [] ~doc)
 
   let stream_decode_cmd =
     let doc = "decode qcow2 formatted data from stdin and write a raw image" in
@@ -15,15 +57,28 @@ module Cli = struct
       ; `P "Decode qcow2 formatted data from stdin and write to a raw file."
       ]
     in
-    let output default =
-      let doc = Printf.sprintf "Path to the output file." in
-      Arg.(value & pos 0 string default & info [] ~doc)
-    in
     Cmd.v
       (Cmd.info "stream_decode" ~doc ~man)
       Term.(ret (const Impl.stream_decode $ output "test.raw"))
 
-  let main () = Cmd.eval stream_decode_cmd
+  let read_headers_cmd =
+    let doc =
+      "Determine allocated clusters by parsing qcow2 file at the provided \
+       path. Returns JSON like the following: {'virtual_size': X, \
+       'cluster_bits': Y, 'data_clusters': [1,2,3]}"
+    in
+    let man = [`S "DESCRIPTION"; `P doc] in
+    Cmd.v
+      (Cmd.info "read_headers" ~doc ~man)
+      Term.(ret (const Impl.read_headers $ input))
+
+  let cmds = [stream_decode_cmd; read_headers_cmd]
 end
 
-let () = exit (Cli.main ())
+let info =
+  let doc = "minimal CLI for qcow-stream" in
+  Cmd.info "qcow-stream-tool" ~version:"1.0.0" ~doc
+
+let () =
+  let cmd = Cmd.group info Cli.cmds in
+  exit (Cmd.eval cmd)

ocaml/xapi/qcow_tool_wrapper.ml

@@ -16,14 +16,15 @@ module D = Debug.Make (struct let name = __MODULE__ end)
 
 open D
 
-let run_qcow_tool qcow_tool ?input_fd ?output_fd (_progress_cb : int -> unit)
-    (args : string list) =
+let run_qcow_tool qcow_tool ?(replace_fds = []) ?input_fd ?output_fd
+    (_progress_cb : int -> unit) (args : string list) =
   info "Executing %s %s" qcow_tool (String.concat " " args) ;
   let open Forkhelpers in
   match
     with_logfile_fd "qcow-tool" (fun log_fd ->
         let pid =
-          safe_close_and_exec input_fd output_fd (Some log_fd) [] qcow_tool args
+          safe_close_and_exec input_fd output_fd (Some log_fd) replace_fds
+            qcow_tool args
         in
         let _, status = waitpid pid in
         if status <> Unix.WEXITED 0 then (
@@ -46,14 +47,70 @@ let update_task_progress (__context : Context.t) (x : int) =
 
 let receive (progress_cb : int -> unit) (unix_fd : Unix.file_descr)
     (path : string) =
-  let args = [path] in
+  let args = ["stream_decode"; path] in
   let qcow_tool = !Xapi_globs.qcow_stream_tool in
   run_qcow_tool qcow_tool progress_cb args ~input_fd:unix_fd
 
+let read_header qcow_path =
+  let args = ["read_headers"; qcow_path] in
+  let qcow_tool = !Xapi_globs.qcow_stream_tool in
+  let pipe_reader, pipe_writer = Unix.pipe ~cloexec:true () in
+
+  let progress_cb _ = () in
+  Xapi_stdext_pervasives.Pervasiveext.finally
+    (fun () -> run_qcow_tool qcow_tool progress_cb args ~output_fd:pipe_writer)
+    (fun () -> Unix.close pipe_writer) ;
+  pipe_reader
+
+let parse_header qcow_path =
+  let pipe_reader = read_header qcow_path in
+  let ic = Unix.in_channel_of_descr pipe_reader in
+  let buf = Buffer.create 4096 in
+  let json = Yojson.Basic.from_channel ~buf ~fname:"qcow_header.json" ic in
+  In_channel.close ic ;
+  let cluster_size =
+    1 lsl Yojson.Basic.Util.(member "cluster_bits" json |> to_int)
+  in
+  let cluster_list =
+    Yojson.Basic.Util.(member "data_clusters" json |> to_list |> List.map to_int)
+  in
+  (cluster_size, cluster_list)
+
 let send ?relative_to (progress_cb : int -> unit) (unix_fd : Unix.file_descr)
     (path : string) (_size : Int64.t) =
+  let qcow_of_device =
+    Vhd_tool_wrapper.backing_file_of_device ~driver:"qcow2"
+  in
+  let qcow_path = qcow_of_device path in
+
+  (* If VDI is backed by QCOW, parse the header to determine nonzero clusters
+     to avoid reading all of the raw disk *)
+  let input_fd = Option.map read_header qcow_path in
+
+  (* Parse the header of the VDI we are diffing against as well *)
+  let relative_to_qcow_path = Option.bind relative_to qcow_of_device in
+  let diff_fd = Option.map read_header relative_to_qcow_path in
+
+  let unique_string = Uuidx.(to_string (make ())) in
   let args =
-    [path] @ match relative_to with None -> [] | Some vdi -> ["--diff"; vdi]
+    [path]
+    @ (match relative_to with None -> [] | Some vdi -> ["--diff"; vdi])
+    @ ( match relative_to_qcow_path with
+      | None ->
+          []
+      | Some _ ->
+          ["--json-header-diff"; unique_string]
+      )
+    @ match qcow_path with None -> [] | Some _ -> ["--json-header"]
   in
   let qcow_tool = !Xapi_globs.qcow_to_stdout in
-  run_qcow_tool qcow_tool progress_cb args ~output_fd:unix_fd
+  let replace_fds = Option.map (fun fd -> [(unique_string, fd)]) diff_fd in
+  Xapi_stdext_pervasives.Pervasiveext.finally
+    (fun () ->
+      run_qcow_tool qcow_tool progress_cb args ?input_fd ~output_fd:unix_fd
+        ?replace_fds
+    )
+    (fun () ->
+      Option.iter Unix.close input_fd ;
+      Option.iter Unix.close diff_fd
+    )

ocaml/xapi/qcow_tool_wrapper.mli

@@ -23,3 +23,5 @@ val send :
   -> string
   -> int64
   -> unit
+
+val parse_header : string -> int * int list

ocaml/xapi/vhd_tool_wrapper.ml

@@ -149,25 +149,27 @@ let find_backend_device path =
         raise Not_found
   with _ -> None
 
-(** [vhd_of_device path] returns (Some vhd) where 'vhd' is the vhd leaf backing a particular device [path] or None.
-    [path] may either be a blktap2 device *or* a blkfront device backed by a blktap2 device. If the latter then
-    the script must be run in the same domain as blkback. *)
-let vhd_of_device path =
+(** [backing_file_of_device path] returns (Some backing_file) where 'backing_file'
+    is the leaf backing a particular device [path] (with a driver of type
+    [driver] or None. [path] may either be a blktap2 device *or* a blkfront
+    device backed by a blktap2 device. If the latter then the script must be
+    run in the same domain as blkback. *)
+let backing_file_of_device path ~driver =
   let tapdisk_of_path path =
     try
       match Tapctl.of_device (Tapctl.create ()) path with
-      | _, _, Some ("vhd", vhd) ->
-          Some vhd
+      | _, _, Some (typ, backing_file) when typ = driver ->
+          Some backing_file
       | _, _, _ ->
           raise Not_found
     with
     | Tapctl.Not_blktap -> (
         debug "Device %s is not controlled by blktap" path ;
-        (* Check if it is a VHD behind a NBD deivce *)
+        (* Check if it is a [driver] behind a NBD device *)
         Stream_vdi.(get_nbd_device path |> image_behind_nbd_device) |> function
-        | Some ("vhd", vhd) ->
-            debug "%s is a VHD behind NBD device %s" vhd path ;
-            Some vhd
+        | Some (typ, backing_file) when typ = driver ->
+            debug "%s is a %s behind NBD device %s" backing_file driver path ;
+            Some backing_file
         | _ ->
             None
       )
@@ -182,6 +184,7 @@ let vhd_of_device path =
 
 let send progress_cb ?relative_to (protocol : string) (dest_format : string)
     (s : Unix.file_descr) (path : string) (size : Int64.t) (prefix : string) =
+  let vhd_of_device = backing_file_of_device ~driver:"vhd" in
   let s' = Uuidx.(to_string (make ())) in
   let source_format, source =
     match (Stream_vdi.get_nbd_device path, vhd_of_device path, relative_to) with