Optimize data cluster format for VHD and QCOW

dune-project

@@ -506,6 +506,7 @@
   cohttp
   cohttp-lwt
   conf-libssl
+  (conf-qemu-img :with-test)
   (cstruct
    (>= "3.0.0"))
   (ezxenstore

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/dune

@@ -71,3 +71,7 @@
   (section libexec_root)
   (files (get_vhd_vsize.exe as xapi/get_vhd_vsize) (sparse_dd.exe as xapi/sparse_dd))
 )
+
+(cram
+ (package vhd-tool)
+ (deps main.exe))

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/cli/read_headers.t

@@ -0,0 +1,30 @@
+Create a sparse raw file
+  $ dd if=/dev/random of=test.raw bs=2097152 count=16 > /dev/null 2>&1
+  $ dd if=/dev/zero of=test.raw bs=2097152 count=10 seek=3 conv=notrunc > /dev/null 2>&1
+
+Convert to .vhd
+  $ qemu-img convert -f raw -O vpc test.raw test.vhd
+
+Check if the right clusters are found to be allocated (legacy JSON format)
+  $ ./main.exe read_headers test.vhd
+  {"virtual_size":33562624,"cluster_bits":21,"data_clusters":[0,1,2,13,14,15]}
+
+Check if the right clusters are found to be allocated (interval-based JSON format)
+  $ ./main.exe read_headers_interval test.vhd
+  {"virtual_size":33562624,"cluster_bits":21,"data_clusters":[[0,2],[13,15]]}
+
+Check fully allocated file
+  $ dd if=/dev/random of=test.raw bs=2097152 count=16 > /dev/null 2>&1
+  $ qemu-img convert -f raw -O vpc test.raw test.vhd
+  $ ./main.exe read_headers test.vhd
+  {"virtual_size":33562624,"cluster_bits":21,"data_clusters":[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]}
+  $ ./main.exe read_headers_interval test.vhd
+  {"virtual_size":33562624,"cluster_bits":21,"data_clusters":[[0,15]]}
+
+Check empty file
+  $ dd if=/dev/zero of=test.raw bs=2097152 count=16 > /dev/null 2>&1
+  $ qemu-img convert -f raw -O vpc test.raw test.vhd
+  $ ./main.exe read_headers test.vhd
+  {"virtual_size":33562624,"cluster_bits":21,"data_clusters":[]}
+  $ ./main.exe read_headers_interval test.vhd
+  {"virtual_size":33562624,"cluster_bits":21,"data_clusters":[]}

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