CP-310612: optimize get_by_uuid (xapi-project#6757)

I was profiling XAPI recently and noticed this: <img width="1202"
height="834" alt="Screenshot 2025-11-13 at 18 41 38"
src="https://github.com/user-attachments/assets/c4828554-d364-459b-9a34-070414d3b146"
/>

I was fairly sure I fixed it already, and indeed I did back in January,
but the commit got lost in my branches and never made it to master.

Needs some testing, might've bitrotted, hence the draft PR.

Author: edwintorok

ocaml/database/database_test.ml

@@ -463,6 +463,7 @@ functor
       (* reference which we create *)
       let valid_ref = "ref1" in
       let valid_uuid = "uuid1" in
+      let new_uuid = "uuid2" in
       let invalid_ref = "foo" in
       let invalid_uuid = "bar" in
       let t =
@@ -626,6 +627,32 @@ functor
           "read_field_where <valid table> <valid return> <valid field> <valid \
            value>" ;
       test_invalid_where_record "read_field_where" (Client.read_field_where t) ;
+
+      (* before changing the UUID, the new UUID should be missing *)
+      expect_missing_uuid "VM" new_uuid (fun () ->
+          let (_ : string) = Client.db_get_by_uuid t "VM" new_uuid in
+          ()
+      ) ;
+      (* change UUID, can happen during VM import *)
+      Client.write_field t "VM" valid_ref Db_names.uuid new_uuid ;
+      let old_uuid = valid_uuid in
+      (* new UUID should be found *)
+      let r = Client.db_get_by_uuid t "VM" new_uuid in
+      if r <> valid_ref then
+        failwith_fmt "db_get_by_uuid <new uuid>: got %s; expected %s" r
+          valid_ref ;
+      let r = Client.db_get_by_uuid_opt t "VM" new_uuid in
+      ( if r <> Some valid_ref then
+          let rs = Option.value ~default:"None" r in
+          failwith_fmt "db_get_by_uuid_opt <new uuid>: got %s; expected %s" rs
+            valid_ref
+      ) ;
+      (* old UUID should not be found anymore *)
+      expect_missing_uuid "VM" old_uuid (fun () ->
+          let (_ : string) = Client.db_get_by_uuid t "VM" old_uuid in
+          ()
+      ) ;
+
       Printf.printf "write_field <invalid table>\n" ;
       expect_missing_tbl "Vm" (fun () ->
           let (_ : unit) = Client.write_field t "Vm" "" "" "" in
@@ -842,5 +869,23 @@ functor
           )
         in
         Printf.printf "bad sequence: %.2f calls/sec\n%!" malign_time
-      )
+      ) ;
+      Client.delete_row t "VM" valid_ref ;
+      (* after deleting the row, both old and new uuid must be missing *)
+      expect_missing_uuid "VM" new_uuid (fun () ->
+          let (_ : string) = Client.db_get_by_uuid t "VM" new_uuid in
+          ()
+      ) ;
+      expect_missing_uuid "VM" old_uuid (fun () ->
+          let (_ : string) = Client.db_get_by_uuid t "VM" old_uuid in
+          ()
+      ) ;
+      let r = Client.db_get_by_uuid_opt t "VM" old_uuid in
+      if not (Option.is_none r) then
+        failwith_fmt "db_get_by_uuid_opt <old uuid>: got %s; expected None"
+          valid_ref ;
+      let r = Client.db_get_by_uuid_opt t "VM" new_uuid in
+      if not (Option.is_none r) then
+        failwith_fmt "db_get_by_uuid_opt <old uuid>: got %s; expected None"
+          valid_ref
   end

ocaml/database/db_cache_impl.ml

@@ -277,21 +277,13 @@ let read_field_where' conv t rcd =
 let read_field_where t rcd = read_field_where' Fun.id t rcd
 
 let db_get_by_uuid t tbl uuid_val =
-  match
-    read_field_where' Schema.CachedValue.string_of t
-      {
-        table= tbl
-      ; return= Db_names.ref
-      ; where_field= Db_names.uuid
-      ; where_value= uuid_val
-      }
-  with
-  | [] ->
-      raise (Read_missing_uuid (tbl, "", uuid_val))
-  | [r] ->
+  let db = get_database t in
+  match Database.lookup_uuid uuid_val db with
+  | Some (tbl', r) when String.equal tbl tbl' ->
       r
   | _ ->
-      raise (Too_many_values (tbl, "", uuid_val))
+      (* we didn't find the UUID, or it belonged to another table *)
+      raise (Read_missing_uuid (tbl, "", uuid_val))
 
 let db_get_by_uuid_opt t tbl uuid_val =
   match

ocaml/database/db_cache_types.ml

@@ -508,6 +508,8 @@ module Database = struct
 
   let lookup_key key db = KeyMap.find_opt (Ref key) db.keymap
 
+  let lookup_uuid key db = KeyMap.find_opt (Uuid key) db.keymap
+
   let make schema =
     {
       tables= TableSet.empty
@@ -615,6 +617,33 @@ let update_many_to_many g tblname objref f db =
       db
       (Schema.many_to_many tblname (Database.schema db))
 
+let uuid_of ~tblname ~objref db =
+  try
+    Some
+      (Schema.Value.Unsafe_cast.string
+         (Row.find Db_names.uuid
+            (Table.find objref (TableSet.find tblname (Database.tableset db)))
+         )
+      )
+  with _ -> None
+
+let maybe_update_uuid_keymap ~tblname ~objref ~fldname ~newval db =
+  if fldname = Db_names.uuid then
+    db
+    |> Database.update_keymap @@ fun keymap ->
+       let keymap =
+         match uuid_of ~tblname ~objref db with
+         | None ->
+             keymap
+         | Some uuid ->
+             KeyMap.remove (Uuid uuid) keymap
+       in
+       KeyMap.add_unique tblname Db_names.uuid
+         (Uuid (Schema.Value.Unsafe_cast.string newval))
+         (tblname, objref) keymap
+  else
+    db
+
 let set_field tblname objref fldname newval db =
   if fldname = Db_names.ref then
     failwith (Printf.sprintf "Cannot safely update field: %s" fldname) ;
@@ -632,6 +661,7 @@ let set_field tblname objref fldname newval db =
   if need_other_table_update then
     let g = Manifest.generation (Database.manifest db) in
     db
+    |> maybe_update_uuid_keymap ~tblname ~objref ~fldname ~newval
     |> update_many_to_many g tblname objref remove_from_set
     |> update_one_to_many g tblname objref remove_from_set
     |> Database.update
@@ -646,6 +676,7 @@ let set_field tblname objref fldname newval db =
   else
     let g = Manifest.generation (Database.manifest db) in
     db
+    |> maybe_update_uuid_keymap ~tblname ~objref ~fldname ~newval
     |> ((fun _ -> newval)
        |> Row.update g fldname empty
        |> Table.update g objref Row.empty
@@ -696,16 +727,7 @@ let add_row tblname objref newval db =
   |> Database.increment
 
 let remove_row tblname objref db =
-  let uuid =
-    try
-      Some
-        (Schema.Value.Unsafe_cast.string
-           (Row.find Db_names.uuid
-              (Table.find objref (TableSet.find tblname (Database.tableset db)))
-           )
-        )
-    with _ -> None
-  in
+  let uuid = uuid_of ~tblname ~objref db in
   let g = db.Database.manifest.Manifest.generation_count in
   db
   |> Database.update_keymap (fun m ->

ocaml/database/db_cache_types.mli

@@ -195,6 +195,8 @@ module Database : sig
 
   val lookup_key : string -> t -> (string * string) option
 
+  val lookup_uuid : string -> t -> (string * string) option
+
   val reindex : t -> t
 
   val register_callback : string -> (update -> t -> unit) -> t -> t