lib/utils: add the "Utils.get_disk_allocated" helper function

Add the "Utils.get_disk_allocated" helper function, which calculates the
number of allocated bytes in an output disk image, through a corresponding
NBD server connection.

Original "NBD.block_status" summation example code by Rich Jones
(thanks!), chunking, metadata context simplification, and virt-v2v
integration by myself. Chunking added because "NBD.block_status" is not
64-bit clean just yet (Eric Blake is working on it).

Bugzilla: https://bugzilla.redhat.com/show_bug.cgi?id=2027598
Signed-off-by: Laszlo Ersek <lersek@redhat.com>
Message-Id: <20211210113537.10907-3-lersek@redhat.com>
Acked-by: Richard W.M. Jones <rjones@redhat.com>

lib/utils.ml

@@ -176,3 +176,33 @@ let with_nbd_connect_unix ~socket ~meta_contexts ~f =
             ~finally:(fun () -> NBD.shutdown nbd)
        )
     ~finally:(fun () -> NBD.close nbd)
+
+let get_disk_allocated ~dir ~disknr =
+  let socket = sprintf "%s/out%d" dir disknr
+  and alloc_ctx = "base:allocation" in
+  with_nbd_connect_unix ~socket ~meta_contexts:[alloc_ctx]
+    ~f:(fun nbd ->
+         if NBD.can_meta_context nbd alloc_ctx then (
+           (* Get the list of extents, using a 2GiB chunk size as hint. *)
+           let size = NBD.get_size nbd
+           and allocated = ref 0_L
+           and fetch_offset = ref 0_L in
+           while !fetch_offset < size do
+             let remaining = size -^ !fetch_offset in
+             let fetch_size = min 0x8000_0000_L remaining in
+             NBD.block_status nbd fetch_size !fetch_offset
+               (fun ctx offset entries err ->
+                  assert (ctx = alloc_ctx);
+                  for i = 0 to Array.length entries / 2 - 1 do
+                    let len = Int64.of_int32 entries.(i * 2)
+                    and typ = entries.(i * 2 + 1) in
+                    if Int32.logand typ 1_l = 0_l then
+                      allocated := !allocated +^ len;
+                    fetch_offset := !fetch_offset +^ len
+                  done;
+                  0
+               )
+           done;
+           Some !allocated
+         ) else None
+       )

lib/utils.mli

@@ -87,3 +87,11 @@ val with_nbd_connect_unix : socket:string ->
     contexts in [meta_contexts] requested (each of which is not necessarily
     supported by the server though). The connection is torn down either on
     normal return or if the function [f] throws an exception. *)
+
+val get_disk_allocated : dir:string -> disknr:int -> int64 option
+(** Callable only in the finalization step. [get_disk_allocated dir disknr]
+    examines output disk [disknr] through the corresponding NBD server socket
+    that resides in [dir]. Returns the number of bytes allocated in the disk
+    image, according to the "base:allocation" metadata context. If the context
+    is not supported by the NBD server behind the socket, the function returns
+    None. *)