[master] Optimize data cluster format for VHD and QCOW

ocaml/libs/vhd/vhd_format/f.ml

@@ -2903,24 +2903,10 @@ functor
 
       let raw ?from (vhd : fd Vhd.t) = raw_common ?from vhd
 
-      let vhd_blocks_to_json (t : fd Vhd.t) =
+      let vhd_blocks_to_json_aux (t : fd Vhd.t) blocks =
         let block_size_sectors_shift =
           t.Vhd.header.Header.block_size_sectors_shift
         in
-        let max_table_entries = Vhd.used_max_table_entries t in
-
-        let include_block = include_block None t in
-
-        let blocks =
-          Seq.init max_table_entries Fun.id
-          |> Seq.filter_map (fun i ->
-              if include_block i then
-                Some (`Int i)
-              else
-                None
-          )
-          |> List.of_seq
-        in
         let json =
           `Assoc
             [
@@ -2934,6 +2920,52 @@ functor
         let json_string = Yojson.to_string json in
         print_string json_string ; return ()
 
+      let vhd_blocks_to_json (t : fd Vhd.t) =
+        let max_table_entries = Vhd.used_max_table_entries t in
+        let blocks =
+          Seq.init max_table_entries Fun.id
+          |> Seq.filter_map (fun i ->
+              if include_block None t i then
+                Some (`Int i)
+              else
+                None
+          )
+          |> List.of_seq
+        in
+        vhd_blocks_to_json_aux t blocks
+
+      let vhd_blocks_to_json_interval (t : fd Vhd.t) =
+        let max_table_entries = Vhd.used_max_table_entries t in
+        let blocks, last_block =
+          Seq.init max_table_entries Fun.id
+          |> Seq.fold_left
+               (fun (acc, left_block) i ->
+                 if include_block None t i then
+                   match left_block with
+                   | Some _ ->
+                       (acc, left_block)
+                   | None ->
+                       (acc, Some i)
+                 else
+                   match left_block with
+                   | Some x ->
+                       (`List [`Int x; `Int (i - 1)] :: acc, None)
+                   | None ->
+                       (acc, None)
+               )
+               ([], None)
+        in
+        (* Close off the interval we were tracking we ran off the end of the seq *)
+        let blocks =
+          match last_block with
+          | Some x ->
+              `List [`Int x; `Int (max_table_entries - 1)] :: blocks
+          | None ->
+              blocks
+        in
+        let blocks = List.rev blocks in
+        vhd_blocks_to_json_aux t blocks
+
       let vhd_common ?from ?raw ?(emit_batmap = false) (t : fd Vhd.t) =
         let block_size_sectors_shift =
           t.Vhd.header.Header.block_size_sectors_shift
@@ -3173,6 +3205,8 @@ functor
         Vhd_input.vhd_common ?from ~raw vhd
 
       let blocks_json = Vhd_input.vhd_blocks_to_json
+
+      let blocks_json_interval = Vhd_input.vhd_blocks_to_json_interval
     end
 
     (* Create a VHD stream from data on t, using `include_block` guide us which blocks have data *)

ocaml/libs/vhd/vhd_format/f.mli

@@ -470,6 +470,8 @@ module From_file : functor (F : S.FILE) -> sig
         [from] into [t] *)
 
     val blocks_json : fd Vhd.t -> unit t
+
+    val blocks_json_interval : fd Vhd.t -> unit t
   end
 
   module Raw_input : sig

ocaml/qcow-stream-tool/qcow_stream_tool.ml

@@ -4,52 +4,13 @@ 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
-      (* TODO: List.map becomes tail-recursive in OCaml 5.1, and could be used here instead *)
-      let clusters =
-        data_cluster_map
-        |> Qcow_types.Cluster.Map.to_seq
-        |> Seq.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
-        )
-        |> List.of_seq
-      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
   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
     let man =
@@ -62,18 +23,7 @@ module Cli = struct
       (Cmd.info "stream_decode" ~doc ~man)
       Term.(ret (const Impl.stream_decode $ output "test.raw"))
 
-  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]
+  let cmds = [stream_decode_cmd]
 end
 
 let info =

ocaml/vhd-tool/cli/main.ml

@@ -385,19 +385,34 @@ let stream_cmd =
   , Cmd.info "stream" ~sdocs:_common_options ~doc ~man
   )
 
+let vhd_source =
+  let doc = Printf.sprintf "Path to the VHD file" in
+  Arg.(required & pos 0 (some file) None & info [] ~doc)
+
 let read_headers_cmd =
   let doc =
     {|Parse VHD headers and output allocated blocks information in JSON format \
      like: {"virtual_size": X, "cluster_bits": X, "data_clusters": [1,2,3]}|}
   in
-  let source =
-    let doc = Printf.sprintf "Path to the VHD file" in
-    Arg.(required & pos 0 (some file) None & info [] ~doc)
-  in
-  ( Term.(ret (const Impl.read_headers $ common_options_t $ source))
+  ( Term.(
+      ret
+        (const (Impl.read_headers ~legacy:true) $ common_options_t $ vhd_source)
+    )
   , Cmd.info "read_headers" ~sdocs:_common_options ~doc
   )
 
+let read_headers_interval_cmd =
+  let doc =
+    {|Parse VHD headers and output allocated blocks intervals information in JSON format \
+     like: {"virtual_size": X, "cluster_bits": X, "data_clusters": [[1,13],[17,17],[19,272]]|}
+  in
+  ( Term.(
+      ret
+        (const (Impl.read_headers ~legacy:false) $ common_options_t $ vhd_source)
+    )
+  , Cmd.info "read_headers_interval" ~sdocs:_common_options ~doc
+  )
+
 let cmds =
   [
     info_cmd
@@ -408,6 +423,7 @@ let cmds =
   ; serve_cmd
   ; stream_cmd
   ; read_headers_cmd
+  ; read_headers_interval_cmd
   ]
   |> List.map (fun (t, i) -> Cmd.v i t)
 

ocaml/vhd-tool/src/impl.ml

@@ -1168,11 +1168,15 @@ let stream_t common args ?(progress = no_progress_bar) () =
     args.StreamCommon.tar_filename_prefix args.StreamCommon.good_ciphersuites
     args.StreamCommon.verify_cert
 
-let read_headers common source =
+let read_headers common source ~legacy =
   let path = [Filename.dirname source] in
   let thread =
     retry common 3 (fun () -> Vhd_IO.openchain ~path source false) >>= fun t ->
-    Vhd_IO.close t >>= fun () -> Hybrid_input.blocks_json t
+    Vhd_IO.close t >>= fun () ->
+    if legacy then
+      Hybrid_input.blocks_json t
+    else
+      Hybrid_input.blocks_json_interval t
   in
   Lwt_main.run thread ; `Ok ()
 

ocaml/vhd-tool/src/impl.mli

@@ -36,7 +36,7 @@ val stream :
   Common.t -> StreamCommon.t -> [> `Error of bool * string | `Ok of unit]
 
 val read_headers :
-  Common.t -> string -> [> `Error of bool * string | `Ok of unit]
+  Common.t -> string -> legacy:bool -> [> `Error of bool * string | `Ok of unit]
 
 val serve :
      Common.t

ocaml/xapi/qcow_tool_wrapper.ml

@@ -22,28 +22,43 @@ let receive (progress_cb : int -> unit) (unix_fd : Unix.file_descr)
   Vhd_qcow_parsing.run_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
-  let (_ : Thread.t) =
+  let run_in_thread tool args pipe_writer replace_fds =
     Thread.create
       (fun () ->
         Xapi_stdext_pervasives.Pervasiveext.finally
           (fun () ->
-            Vhd_qcow_parsing.run_tool qcow_tool progress_cb args
-              ~output_fd:pipe_writer
+            Vhd_qcow_parsing.run_tool tool progress_cb args
+              ~output_fd:pipe_writer ~replace_fds
           )
           (fun () -> Unix.close pipe_writer)
       )
       ()
   in
-  pipe_reader
+
+  let map_pipe_reader, map_pipe_writer = Unix.pipe ~cloexec:true () in
+  let (_ : Thread.t) =
+    run_in_thread !Xapi_globs.qemu_img
+      ["map"; qcow_path; "--output=json"]
+      map_pipe_writer []
+  in
+
+  let info_pipe_reader, info_pipe_writer = Unix.pipe ~cloexec:true () in
+  let (_ : Thread.t) =
+    run_in_thread !Xapi_globs.qemu_img
+      ["info"; qcow_path; "--output=json"]
+      info_pipe_writer []
+  in
+
+  (map_pipe_reader, info_pipe_reader)
 
 let parse_header qcow_path =
-  let pipe_reader = read_header qcow_path in
-  Vhd_qcow_parsing.parse_header pipe_reader
+  let pipe, _ = read_header qcow_path in
+  Vhd_qcow_parsing.parse_header pipe
+
+let parse_header_interval qcow_path =
+  let pipes = read_header qcow_path in
+  Vhd_qcow_parsing.parse_header_qemu_img pipes
 
 let send ?relative_to (progress_cb : int -> unit) (unix_fd : Unix.file_descr)
     (path : string) (_size : Int64.t) =
@@ -54,7 +69,12 @@ let send ?relative_to (progress_cb : int -> unit) (unix_fd : Unix.file_descr)
 
   (* If VDI is backed by QCOW, parse the header to determine nonzero clusters
      to avoid reading all of the raw disk *)
-  let input_fd = Result.map read_header qcow_path |> Result.to_option in
+  let input_fds = Result.map read_header qcow_path |> Result.to_option in
+
+  (* TODO: If VHD headers are to be consulted as well, qcow2-to-stdout
+     needs to properly account for cluster_bits. Currently QCOW2 export
+     from VHD-backed VDIs will just revert to raw, without any
+     allocation accounting. *)
 
   (* Parse the header of the VDI we are diffing against as well *)
   let relative_to_qcow_path =
@@ -64,28 +84,60 @@ let send ?relative_to (progress_cb : int -> unit) (unix_fd : Unix.file_descr)
     | None ->
         None
   in
-  let diff_fd = Option.map read_header relative_to_qcow_path in
+  let diff_fds = Option.map read_header relative_to_qcow_path in
+
+  let map_fd_string = Uuidx.(to_string (make ())) in
+  let info_fd_string = Uuidx.(to_string (make ())) in
+  let diff_map_fd_string = Uuidx.(to_string (make ())) in
+  let diff_info_fd_string = Uuidx.(to_string (make ())) in
 
-  let unique_string = Uuidx.(to_string (make ())) in
   let args =
     [path]
     @ (match relative_to with None -> [] | Some vdi -> ["--diff"; vdi])
     @ ( match relative_to_qcow_path with
       | None ->
           []
       | Some _ ->
-          ["--json-header-diff"; unique_string]
+          [
+            "--json-header-diff-map"
+          ; diff_map_fd_string
+          ; "--json-header-diff-info"
+          ; diff_info_fd_string
+          ]
       )
-    @ match qcow_path with Error _ -> [] | Ok _ -> ["--json-header"]
+    @
+    match qcow_path with
+    | Error _ ->
+        []
+    | Ok _ ->
+        [
+          "--json-header-map"
+        ; map_fd_string
+        ; "--json-header-info"
+        ; info_fd_string
+        ]
   in
   let qcow_tool = !Xapi_globs.qcow_to_stdout in
-  let replace_fds = Option.map (fun fd -> [(unique_string, fd)]) diff_fd in
+  let replace_fds =
+    Option.map
+      (fun (map_fd, info_fd) ->
+        let rfds = [(map_fd_string, map_fd); (info_fd_string, info_fd)] in
+        match diff_fds with
+        | Some (diff_map_fd, diff_info_fd) ->
+            (diff_map_fd_string, diff_map_fd)
+            :: (diff_info_fd_string, diff_info_fd)
+            :: rfds
+        | None ->
+            rfds
+      )
+      input_fds
+  in
   Xapi_stdext_pervasives.Pervasiveext.finally
     (fun () ->
-      Vhd_qcow_parsing.run_tool qcow_tool progress_cb args ?input_fd
-        ~output_fd:unix_fd ?replace_fds
+      Vhd_qcow_parsing.run_tool qcow_tool progress_cb args ~output_fd:unix_fd
+        ?replace_fds
     )
     (fun () ->
-      Option.iter Unix.close input_fd ;
-      Option.iter Unix.close diff_fd
+      Option.iter (fun (x, y) -> Unix.close x ; Unix.close y) input_fds ;
+      Option.iter (fun (x, y) -> Unix.close x ; Unix.close y) diff_fds
     )

ocaml/xapi/qcow_tool_wrapper.mli

@@ -25,3 +25,5 @@ val send :
   -> unit
 
 val parse_header : string -> int * int list
+
+val parse_header_interval : string -> int * (int * int) list