Move the update_snapshot_info_dest to storage mux

ocaml/xapi/storage_migrate.ml

@@ -278,7 +278,7 @@ module MigrateLocal = struct
 
     let (module Remote) = get_remote_backend url verify_dest in
     (* Find the local VDI *)
-    let local_vdi = find_local_vdi ~dbg ~sr ~vdi in
+    let local_vdi, _ = find_vdi ~dbg ~sr ~vdi (module Local) in
     let mirror_id = State.mirror_id_of (sr, local_vdi.vdi) in
     debug "%s: Adding to active local mirrors before sending: id=%s"
       __FUNCTION__ mirror_id ;

ocaml/xapi/storage_migrate_helper.ml

@@ -346,14 +346,13 @@ let get_remote_backend url verify_dest =
   end)) in
   (module Remote : SMAPIv2)
 
-let find_local_vdi ~dbg ~sr ~vdi =
-  (* Find the local VDI *)
-  let vdis, _ = Local.SR.scan2 dbg sr in
+let find_vdi ~dbg ~sr ~vdi (module SMAPIv2 : SMAPIv2) =
+  let vdis, _ = SMAPIv2.SR.scan2 dbg sr in
   match List.find_opt (fun x -> x.vdi = vdi) vdis with
   | None ->
-      failwith "Local VDI not found"
+      failwith_fmt "VDI %s not found" (Storage_interface.Vdi.string_of vdi)
   | Some v ->
-      v
+      (v, vdis)
 
 (** [similar_vdis dbg sr vdi] returns a list of content_ids of vdis
   which are similar to the input [vdi] in [sr] *)

ocaml/xapi/storage_migrate_helper.mli

@@ -261,6 +261,7 @@ module Local : SMAPIv2
 
 val get_remote_backend : string -> bool -> (module SMAPIv2)
 
-val find_local_vdi : dbg:string -> sr:sr -> vdi:vdi -> vdi_info
+val find_vdi :
+  dbg:string -> sr:sr -> vdi:vdi -> (module SMAPIv2) -> vdi_info * vdi_info list
 
 val similar_vdis : dbg:string -> sr:sr -> vdi:vdi -> uuid list

ocaml/xapi/storage_mux.ml

@@ -19,6 +19,7 @@ module D = Debug.Make (struct let name = "mux" end)
 open D
 open Storage_interface
 open Storage_mux_reg
+open Storage_utils
 
 let s_of_sr = Storage_interface.Sr.string_of
 
@@ -330,11 +331,32 @@ module Mux = struct
       Storage_migrate.update_snapshot_info_src ~dbg:(Debug_info.to_string di)
         ~sr ~vdi ~url ~dest ~dest_vdi ~snapshot_pairs
 
+    let set_is_a_snapshot _context ~dbg ~sr ~vdi ~is_a_snapshot =
+      Server_helpers.exec_with_new_task "VDI.set_is_a_snapshot"
+        ~subtask_of:(Ref.of_string dbg) (fun __context ->
+          let vdi, _ = find_vdi ~__context sr vdi in
+          Db.VDI.set_is_a_snapshot ~__context ~self:vdi ~value:is_a_snapshot
+      )
+
+    let set_snapshot_time _context ~dbg ~sr ~vdi ~snapshot_time =
+      let module Date = Clock.Date in
+      Server_helpers.exec_with_new_task "VDI.set_snapshot_time"
+        ~subtask_of:(Ref.of_string dbg) (fun __context ->
+          let vdi, _ = find_vdi ~__context sr vdi in
+          let snapshot_time = Date.of_iso8601 snapshot_time in
+          Db.VDI.set_snapshot_time ~__context ~self:vdi ~value:snapshot_time
+      )
+
+    let set_snapshot_of _context ~dbg ~sr ~vdi ~snapshot_of =
+      Server_helpers.exec_with_new_task "VDI.set_snapshot_of"
+        ~subtask_of:(Ref.of_string dbg) (fun __context ->
+          let vdi, _ = find_vdi ~__context sr vdi in
+          let snapshot_of, _ = find_vdi ~__context sr snapshot_of in
+          Db.VDI.set_snapshot_of ~__context ~self:vdi ~value:snapshot_of
+      )
+
     let update_snapshot_info_dest () ~dbg ~sr ~vdi ~src_vdi ~snapshot_pairs =
-      with_dbg ~name:"SR.update_snapshot_info_dest" ~dbg @@ fun di ->
-      let module C = StorageAPI (Idl.Exn.GenClient (struct
-        let rpc = of_sr sr
-      end)) in
+      with_dbg ~name:"SR.update_snapshot_info_dest" ~dbg @@ fun _di ->
       info
         "SR.update_snapshot_info_dest dbg:%s sr:%s vdi:%s ~src_vdi:%s \
          snapshot_pairs:%s"
@@ -348,8 +370,44 @@ module Mux = struct
         |> String.concat "; "
         |> Printf.sprintf "[%s]"
         ) ;
-      C.SR.update_snapshot_info_dest (Debug_info.to_string di) sr vdi src_vdi
-        snapshot_pairs
+      Server_helpers.exec_with_new_task "SR.update_snapshot_info_dest"
+        ~subtask_of:(Ref.of_string dbg) (fun __context ->
+          let local_vdis, _ = scan2 () ~dbg ~sr in
+          let find_sm_vdi ~vdi ~vdi_info_list =
+            try List.find (fun x -> x.vdi = vdi) vdi_info_list
+            with Not_found ->
+              raise (Storage_error (Vdi_does_not_exist (s_of_vdi vdi)))
+          in
+          let assert_content_ids_match ~vdi_info1 ~vdi_info2 =
+            if vdi_info1.content_id <> vdi_info2.content_id then
+              raise
+                (Storage_error
+                   (Content_ids_do_not_match
+                      (s_of_vdi vdi_info1.vdi, s_of_vdi vdi_info2.vdi)
+                   )
+                )
+          in
+          (* For each (local snapshot vdi, source snapshot vdi) pair:
+             					 * - Check that the content_ids are the same
+             					 * - Copy snapshot_time from the source VDI to the local VDI
+             					 * - Set the local VDI's snapshot_of to vdi
+             					 * - Set is_a_snapshot = true for the local snapshot *)
+          List.iter
+            (fun (local_snapshot, src_snapshot_info) ->
+              let local_snapshot_info =
+                find_sm_vdi ~vdi:local_snapshot ~vdi_info_list:local_vdis
+              in
+              assert_content_ids_match ~vdi_info1:local_snapshot_info
+                ~vdi_info2:src_snapshot_info ;
+              set_snapshot_time __context ~dbg ~sr ~vdi:local_snapshot
+                ~snapshot_time:src_snapshot_info.snapshot_time ;
+              set_snapshot_of __context ~dbg ~sr ~vdi:local_snapshot
+                ~snapshot_of:vdi ;
+              set_is_a_snapshot __context ~dbg ~sr ~vdi:local_snapshot
+                ~is_a_snapshot:true
+            )
+            snapshot_pairs
+      )
   end
 
   module VDI = struct

ocaml/xapi/storage_smapiv1.ml

@@ -18,8 +18,7 @@ open D
 module Date = Clock.Date
 module XenAPI = Client.Client
 open Storage_interface
-
-exception No_VDI
+open Storage_utils
 
 let s_of_vdi = Vdi.string_of
 
@@ -30,26 +29,6 @@ let with_lock = Xapi_stdext_threads.Threadext.Mutex.execute
 let with_dbg ~name ~dbg f =
   Debug_info.with_dbg ~module_name:"SMAPIv1" ~name ~dbg f
 
-(* Find a VDI given a storage-layer SR and VDI *)
-let find_vdi ~__context sr vdi =
-  let sr = s_of_sr sr in
-  let vdi = s_of_vdi vdi in
-  let open Xapi_database.Db_filter_types in
-  let sr = Db.SR.get_by_uuid ~__context ~uuid:sr in
-  match
-    Db.VDI.get_records_where ~__context
-      ~expr:
-        (And
-           ( Eq (Field "location", Literal vdi)
-           , Eq (Field "SR", Literal (Ref.string_of sr))
-           )
-        )
-  with
-  | x :: _ ->
-      x
-  | _ ->
-      raise No_VDI
-
 (* Find a VDI reference given a name *)
 let find_content ~__context ?sr name =
   (* PR-1255: the backend should do this for us *)
@@ -132,32 +111,6 @@ module SMAPIv1 : Server_impl = struct
     let vdi_rec = Db.VDI.get_record ~__context ~self in
     vdi_info_of_vdi_rec __context vdi_rec
 
-  (* For SMAPIv1, is_a_snapshot, snapshot_time and snapshot_of are stored in
-     	 * xapi's database. For SMAPIv2 they should be implemented by the storage
-     	 * backend. *)
-  let set_is_a_snapshot _context ~dbg ~sr ~vdi ~is_a_snapshot =
-    Server_helpers.exec_with_new_task "VDI.set_is_a_snapshot"
-      ~subtask_of:(Ref.of_string dbg) (fun __context ->
-        let vdi, _ = find_vdi ~__context sr vdi in
-        Db.VDI.set_is_a_snapshot ~__context ~self:vdi ~value:is_a_snapshot
-    )
-
-  let set_snapshot_time _context ~dbg ~sr ~vdi ~snapshot_time =
-    Server_helpers.exec_with_new_task "VDI.set_snapshot_time"
-      ~subtask_of:(Ref.of_string dbg) (fun __context ->
-        let vdi, _ = find_vdi ~__context sr vdi in
-        let snapshot_time = Date.of_iso8601 snapshot_time in
-        Db.VDI.set_snapshot_time ~__context ~self:vdi ~value:snapshot_time
-    )
-
-  let set_snapshot_of _context ~dbg ~sr ~vdi ~snapshot_of =
-    Server_helpers.exec_with_new_task "VDI.set_snapshot_of"
-      ~subtask_of:(Ref.of_string dbg) (fun __context ->
-        let vdi, _ = find_vdi ~__context sr vdi in
-        let snapshot_of, _ = find_vdi ~__context sr snapshot_of in
-        Db.VDI.set_snapshot_of ~__context ~self:vdi ~value:snapshot_of
-    )
-
   module Query = struct
     let query _context ~dbg:_ =
       {
@@ -433,46 +386,9 @@ module SMAPIv1 : Server_impl = struct
         ~dest_vdi:_ ~snapshot_pairs:_ =
       assert false
 
-    let update_snapshot_info_dest _context ~dbg ~sr ~vdi ~src_vdi:_
-        ~snapshot_pairs =
-      Server_helpers.exec_with_new_task "SR.update_snapshot_info_dest"
-        ~subtask_of:(Ref.of_string dbg) (fun __context ->
-          let local_vdis = scan __context ~dbg ~sr in
-          let find_sm_vdi ~vdi ~vdi_info_list =
-            try List.find (fun x -> x.vdi = vdi) vdi_info_list
-            with Not_found ->
-              raise (Storage_error (Vdi_does_not_exist (s_of_vdi vdi)))
-          in
-          let assert_content_ids_match ~vdi_info1 ~vdi_info2 =
-            if vdi_info1.content_id <> vdi_info2.content_id then
-              raise
-                (Storage_error
-                   (Content_ids_do_not_match
-                      (s_of_vdi vdi_info1.vdi, s_of_vdi vdi_info2.vdi)
-                   )
-                )
-          in
-          (* For each (local snapshot vdi, source snapshot vdi) pair:
-             					 * - Check that the content_ids are the same
-             					 * - Copy snapshot_time from the source VDI to the local VDI
-             					 * - Set the local VDI's snapshot_of to vdi
-             					 * - Set is_a_snapshot = true for the local snapshot *)
-          List.iter
-            (fun (local_snapshot, src_snapshot_info) ->
-              let local_snapshot_info =
-                find_sm_vdi ~vdi:local_snapshot ~vdi_info_list:local_vdis
-              in
-              assert_content_ids_match ~vdi_info1:local_snapshot_info
-                ~vdi_info2:src_snapshot_info ;
-              set_snapshot_time __context ~dbg ~sr ~vdi:local_snapshot
-                ~snapshot_time:src_snapshot_info.snapshot_time ;
-              set_snapshot_of __context ~dbg ~sr ~vdi:local_snapshot
-                ~snapshot_of:vdi ;
-              set_is_a_snapshot __context ~dbg ~sr ~vdi:local_snapshot
-                ~is_a_snapshot:true
-            )
-            snapshot_pairs
-      )
+    let update_snapshot_info_dest _context ~dbg:_ ~sr:_ ~vdi:_ ~src_vdi:_
+        ~snapshot_pairs:_ =
+      assert false
   end
 
   module VDI = struct

ocaml/xapi/storage_smapiv1.mli

@@ -20,7 +20,4 @@ val vdi_read_write : (Sr.t * Vdi.t, bool) Hashtbl.t
 
 val vdi_info_of_vdi_rec : Context.t -> API.vDI_t -> Storage_interface.vdi_info
 
-val find_vdi : __context:Context.t -> Sr.t -> Vdi.t -> [`VDI] Ref.t * API.vDI_t
-(** Find a VDI given a storage-layer SR and VDI *)
-
 module SMAPIv1 : Server_impl

ocaml/xapi/storage_smapiv1_migrate.ml

@@ -162,26 +162,11 @@ module Copy = struct
         (Printf.sprintf "Remote SR %s not found"
            (Storage_interface.Sr.string_of dest)
         ) ;
-    let vdis = Remote.SR.scan dbg dest in
-    let remote_vdi =
-      try List.find (fun x -> x.vdi = dest_vdi) vdis
-      with Not_found ->
-        failwith
-          (Printf.sprintf "Remote VDI %s not found"
-             (Storage_interface.Vdi.string_of dest_vdi)
-          )
-    in
+
+    let remote_vdi, _ = find_vdi ~dbg ~sr:dest ~vdi:dest_vdi (module Remote) in
     let dest_content_id = remote_vdi.content_id in
     (* Find the local VDI *)
-    let vdis = Local.SR.scan dbg sr in
-    let local_vdi =
-      try List.find (fun x -> x.vdi = vdi) vdis
-      with Not_found ->
-        failwith
-          (Printf.sprintf "Local VDI %s not found"
-             (Storage_interface.Vdi.string_of vdi)
-          )
-    in
+    let local_vdi, vdis = find_vdi ~dbg ~sr ~vdi (module Local) in
     D.debug "copy local content_id=%s" local_vdi.content_id ;
     D.debug "copy remote content_id=%s" dest_content_id ;
     if local_vdi.virtual_size > remote_vdi.virtual_size then (
@@ -293,6 +278,10 @@ module Copy = struct
       (* PR-1255: XXX: this is useful because we don't have content_ids by default *)
       D.debug "setting local content_id <- %s" local_vdi.content_id ;
       Local.VDI.set_content_id dbg sr local_vdi.vdi local_vdi.content_id ;
+      (* Re-find the VDI to get the updated content_id info *)
+      let remote_vdi, _ =
+        find_vdi ~dbg ~sr:dest ~vdi:dest_vdi (module Remote)
+      in
       Some (Vdi_info remote_vdi)
     with e ->
       D.error "Caught %s: performing cleanup actions" (Printexc.to_string e) ;
@@ -312,11 +301,7 @@ module Copy = struct
     let (module Remote) = get_remote_backend url verify_dest in
     (* Find the local VDI *)
     try
-      let vdis = Local.SR.scan dbg sr in
-      let local_vdi =
-        try List.find (fun x -> x.vdi = vdi) vdis
-        with Not_found -> failwith (Printf.sprintf "Local VDI not found")
-      in
+      let local_vdi, _ = find_vdi ~dbg ~sr ~vdi (module Local) in
       try
         let similar_vdis = Local.VDI.similar_content dbg sr vdi in
         let similars = List.map (fun vdi -> vdi.content_id) similar_vdis in

ocaml/xapi/storage_utils.ml

@@ -14,6 +14,10 @@
 
 open Storage_interface
 
+let s_of_sr = Storage_interface.Sr.string_of
+
+let s_of_vdi = Storage_interface.Vdi.string_of
+
 let string_of_vdi_type vdi_type =
   Rpc.string_of_rpc (API.rpc_of_vdi_type vdi_type)
 
@@ -173,3 +177,24 @@ let transform_storage_exn f =
         (Api_errors.Server_error
            (Api_errors.internal_error, [Printexc.to_string e])
         )
+
+exception No_VDI
+
+let find_vdi ~__context sr vdi =
+  let sr = s_of_sr sr in
+  let vdi = s_of_vdi vdi in
+  let open Xapi_database.Db_filter_types in
+  let sr = Db.SR.get_by_uuid ~__context ~uuid:sr in
+  match
+    Db.VDI.get_records_where ~__context
+      ~expr:
+        (And
+           ( Eq (Field "location", Literal vdi)
+           , Eq (Field "SR", Literal (Ref.string_of sr))
+           )
+        )
+  with
+  | x :: _ ->
+      x
+  | _ ->
+      raise No_VDI

ocaml/xapi/storage_utils.mli

@@ -64,3 +64,12 @@ val rpc :
 
 val transform_storage_exn : (unit -> 'a) -> 'a
 (** [transform_storage_exn f] runs [f], rethrowing any storage error as a nice XenAPI error *)
+
+exception No_VDI
+
+val find_vdi :
+     __context:Context.t
+  -> Storage_interface.sr
+  -> Storage_interface.vdi
+  -> [`VDI] Ref.t * API.vDI_t
+(** Find a VDI given a storage-layer SR and VDI *)

ocaml/xapi/xapi_services.ml

@@ -196,7 +196,7 @@ let put_handler (req : Http.Request.t) s _ =
           http_proxy_to_plugin req s name
       | [""; services; "SM"; "data"; sr; vdi] when services = _services ->
           let vdi, _ =
-            Storage_smapiv1.find_vdi ~__context
+            Storage_utils.find_vdi ~__context
               (Storage_interface.Sr.of_string sr)
               (Storage_interface.Vdi.of_string vdi)
           in