vhd_qcow_parsing: Add parse_header_interval for interval-based headers

Since the runtime feature flag vhd_legacy_blocks_format determines which block
format is used to describe allocated VHD clusters, this requires duplicate
parse_header_interval functions for VHD and QCOW.

The right functions are selected in stream_vdi based on the feature flag.

Signed-off-by: Andrii Sultanov <andriy.sultanov@vates.tech>

Author: last-genius

ocaml/xapi/qcow_tool_wrapper.ml

@@ -39,6 +39,10 @@ let parse_header qcow_path =
   let pipe_reader = read_header qcow_path in
   Vhd_qcow_parsing.parse_header pipe_reader
 
+let parse_header_interval qcow_path =
+  let pipe_reader = read_header qcow_path in
+  Vhd_qcow_parsing.parse_header_interval pipe_reader
+
 let send ?relative_to (progress_cb : int -> unit) (unix_fd : Unix.file_descr)
     (path : string) (_size : Int64.t) =
   let qcow_of_device =

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

ocaml/xapi/stream_vdi.ml

@@ -306,37 +306,99 @@ let send_one ofd (__context : Context.t) rpc session_id progress refresh_session
           | Ok (Some (driver, path)) when driver = "vhd" || driver = "qcow2"
             -> (
             try
-              (* Read backing file headers, then only read and write
+              let last_chunk = Int64.((to_int size - 1) / to_int chunk_size) in
+              if !Xapi_globs.vhd_legacy_blocks_format then
+                (* Read backing file headers, then only read and write
                  allocated clusters from the bitmap *)
-              let cluster_size, cluster_list =
-                match driver with
-                | "vhd" ->
-                    Vhd_tool_wrapper.parse_header path
-                | "qcow2" ->
-                    Qcow_tool_wrapper.parse_header path
-                | _ ->
-                    failwith (Printf.sprintf "%s: unreachable" __FUNCTION__)
-              in
-              let set =
-                get_allocated_chunks_from_clusters cluster_size cluster_list
-              in
-              (* First and last chunks are always written - it's a limitation
+                let cluster_size, cluster_list =
+                  match driver with
+                  | "vhd" ->
+                      Vhd_tool_wrapper.parse_header path
+                  | "qcow2" ->
+                      Qcow_tool_wrapper.parse_header path
+                  | _ ->
+                      failwith (Printf.sprintf "%s: unreachable" __FUNCTION__)
+                in
+                let set =
+                  get_allocated_chunks_from_clusters cluster_size cluster_list
+                in
+                (* First and last chunks are always written - it's a limitation
                  of the XVA format *)
-              let last_chunk =
-                Int64.((to_int size - to_int chunk_size + 1) / to_int chunk_size)
-              in
-              let set = set |> ChunkSet.add 0 |> ChunkSet.add last_chunk in
-              ChunkSet.iter
-                (fun this_chunk_no ->
-                  let offset = Int64.(mul (of_int this_chunk_no) chunk_size) in
-                  let _ =
-                    write_chunk this_chunk_no offset
-                      ~write_check:(fun _ _ -> true)
-                      ~seek:true ~timeout_workaround:false
-                  in
-                  ()
-                )
-                set
+                let set = set |> ChunkSet.add 0 |> ChunkSet.add last_chunk in
+                ChunkSet.iter
+                  (fun this_chunk_no ->
+                    let offset =
+                      Int64.(mul (of_int this_chunk_no) chunk_size)
+                    in
+                    let _ =
+                      write_chunk this_chunk_no offset
+                        ~write_check:(fun _ _ -> true)
+                        ~seek:true ~timeout_workaround:false
+                    in
+                    ()
+                  )
+                  set
+              else
+                let cluster_size, cluster_list =
+                  match driver with
+                  | "vhd" ->
+                      Vhd_tool_wrapper.parse_header_interval path
+                  | "qcow2" ->
+                      Qcow_tool_wrapper.parse_header_interval path
+                  | _ ->
+                      failwith (Printf.sprintf "%s: unreachable" __FUNCTION__)
+                in
+                let process_chunk chunk_no ~force =
+                  if force || (chunk_no <> 0 && chunk_no <> last_chunk) then
+                    let offset = Int64.(mul (of_int chunk_no) chunk_size) in
+                    let _ =
+                      write_chunk chunk_no offset
+                        ~write_check:(fun _ _ -> true)
+                        ~seek:true ~timeout_workaround:false
+                    in
+                    ()
+                in
+
+                process_chunk 0 ~force:true ;
+
+                let chunk_size = Int64.to_int chunk_size in
+                let chunks_in_cluster =
+                  (cluster_size + chunk_size - 1) / chunk_size
+                in
+                (* Iterate over allocated intervals, copying every cluster inside *)
+                let _ =
+                  List.fold_left
+                    (fun prev_chunk (cluster_no_left, cluster_no_right) ->
+                      let calc_chunk cluster =
+                        let cluster_offset = cluster * cluster_size in
+                        let chunk_no = cluster_offset / chunk_size in
+                        chunk_no
+                      in
+                      let left_chunk_no = calc_chunk cluster_no_left in
+                      let right_chunk_no =
+                        calc_chunk cluster_no_right + chunks_in_cluster - 1
+                      in
+
+                      (* If a chunk contains multiple clusters, we could have
+                       already copied it. In that case, start with the
+                       following chunk. *)
+                      let left_chunk_no =
+                        if left_chunk_no = prev_chunk then
+                          left_chunk_no + 1
+                        else
+                          left_chunk_no
+                      in
+
+                      for i = left_chunk_no to right_chunk_no do
+                        process_chunk i ~force:false
+                      done ;
+
+                      right_chunk_no
+                    )
+                    (-1) cluster_list
+                in
+
+                process_chunk last_chunk ~force:true
             with e ->
               debug "%s: Falling back to reading the whole raw disk after %s"
                 __FUNCTION__ (Printexc.to_string e) ;

ocaml/xapi/vhd_qcow_parsing.ml

@@ -44,15 +44,36 @@ let run_tool tool ?(replace_fds = []) ?input_fd ?output_fd
       error "%s output: %s" tool out ;
       raise (Api_errors.Server_error (Api_errors.vdi_io_error, [out]))
 
-let parse_header pipe_reader =
+let parse_header_aux pipe_reader =
   let ic = Unix.in_channel_of_descr pipe_reader in
   let buf = Buffer.create 4096 in
   let json = Yojson.Basic.from_channel ~buf ~fname:"header.json" ic in
   In_channel.close ic ;
   let cluster_size =
     1 lsl Yojson.Basic.Util.(member "cluster_bits" json |> to_int)
   in
+  (cluster_size, json)
+
+let parse_header pipe_reader =
+  let cluster_size, json = parse_header_aux pipe_reader in
   let cluster_list =
     Yojson.Basic.Util.(member "data_clusters" json |> to_list |> List.map to_int)
   in
   (cluster_size, cluster_list)
+
+let parse_header_interval pipe_reader =
+  let cluster_size, json = parse_header_aux pipe_reader in
+  let cluster_list =
+    Yojson.Basic.Util.(
+      member "data_clusters" json
+      |> to_list
+      |> List.map (fun x ->
+          match to_list x with
+          | x :: y :: _ ->
+              (to_int x, to_int y)
+          | _ ->
+              raise (Invalid_argument "Invalid JSON")
+      )
+    )
+  in
+  (cluster_size, cluster_list)

ocaml/xapi/vhd_qcow_parsing.mli

@@ -22,3 +22,5 @@ val run_tool :
   -> unit
 
 val parse_header : Unix.file_descr -> int * int list
+
+val parse_header_interval : Unix.file_descr -> int * (int * int) list

ocaml/xapi/vhd_tool_wrapper.ml

@@ -116,9 +116,14 @@ let receive progress_cb format protocol (s : Unix.file_descr)
   in
   run_vhd_tool progress_cb args s s' path
 
-let read_vhd_header path =
+let read_vhd_header path ~legacy =
   let vhd_tool = !Xapi_globs.vhd_tool in
-  let args = ["read_headers"; path] in
+  let args =
+    if legacy then
+      ["read_headers"; path]
+    else
+      ["read_headers_interval"; path]
+  in
   let pipe_reader, pipe_writer = Unix.pipe ~cloexec:true () in
 
   let progress_cb _ = () in
@@ -130,9 +135,13 @@ let read_vhd_header path =
   pipe_reader
 
 let parse_header vhd_path =
-  let pipe_reader = read_vhd_header vhd_path in
+  let pipe_reader = read_vhd_header vhd_path ~legacy:true in
   Vhd_qcow_parsing.parse_header pipe_reader
 
+let parse_header_interval vhd_path =
+  let pipe_reader = read_vhd_header vhd_path ~legacy:false in
+  Vhd_qcow_parsing.parse_header_interval pipe_reader
+
 let send progress_cb ?relative_to (protocol : string) (dest_format : string)
     (s : Unix.file_descr) (path : string) (size : Int64.t) (prefix : string) =
   let __FUN = __FUNCTION__ in