Merge pull request #17934 from MinaProtocol/cjjdespres/gradual-ledger-migration

Add more backing-agnostic methods to Root

Author: cjjdespres

src/lib/mina_ledger/root.ml

@@ -172,14 +172,14 @@ struct
     | Converting_db db ->
         Converting_ledger.merkle_root db
 
-  let create ~logger ~config ~depth () =
+  let create ~logger ~config ~depth ?(assert_synced = false) () =
     match config with
     | Config.Stable_db_config directory_name ->
         Stable_db (Stable_db.create ~directory_name ~depth ())
     | Converting_db_config config ->
         Converting_db
           (Converting_ledger.create ~config:(In_directories config) ~logger
-             ~depth () )
+             ~depth ~assert_synced () )
 
   let create_temporary ~logger ~backing_type ~depth () =
     match backing_type with
@@ -215,6 +215,64 @@ struct
              ; backing_2 = Config.backing_of_config config
              } )
 
+  let create_checkpoint_with_directory t ~directory_name =
+    let backing_type =
+      match t with
+      | Stable_db _db ->
+          Config.Stable_db
+      | Converting_db _db ->
+          Config.Converting_db
+    in
+    let config = Config.with_directory ~backing_type ~directory_name in
+    create_checkpoint t ~config ()
+
+  (** Migrate the accounts in the ledger database [stable_db] and store them in
+      [empty_hardfork_db]. The accounts are set in the target database in chunks
+      so the daemon is still responsive during this operation; the daemon would
+      otherwise stop everything as it hashed every account in the list. *)
+  let chunked_migration ?(chunk_size = 1 lsl 6) stable_locations_and_accounts
+      empty_migrated_db =
+    let open Async.Deferred.Let_syntax in
+    let ledger_depth = Migrated_db.depth empty_migrated_db in
+    let addrs_and_accounts =
+      List.mapi stable_locations_and_accounts ~f:(fun i acct ->
+          ( Migrated_db.Addr.of_int_exn ~ledger_depth i
+          , Account.Hardfork.of_stable acct ) )
+    in
+    let rec set_chunks accounts =
+      let%bind () = Async_unix.Scheduler.yield () in
+      let chunk, accounts' = List.split_n accounts chunk_size in
+      if List.is_empty chunk then return empty_migrated_db
+      else (
+        Migrated_db.set_batch_accounts empty_migrated_db chunk ;
+        set_chunks accounts' )
+    in
+    set_chunks addrs_and_accounts
+
+  let make_converting t =
+    let open Async.Deferred.Let_syntax in
+    match t with
+    | Converting_db _db ->
+        return t
+    | Stable_db db ->
+        let directory_name =
+          Stable_db.get_directory db
+          |> Option.value_exn
+               ~message:"Invariant: database must be in a directory"
+        in
+        let converting_config =
+          Converting_ledger.Config.with_primary ~directory_name
+        in
+        let migrated_db =
+          Migrated_db.create
+            ~directory_name:converting_config.converting_directory
+            ~depth:(Stable_db.depth db) ()
+        in
+        let%map migrated_db =
+          chunked_migration (Stable_db.to_list_sequential db) migrated_db
+        in
+        Converting_db (Converting_ledger.of_ledgers db migrated_db)
+
   let as_unmasked t =
     match t with
     | Stable_db db ->
@@ -270,4 +328,12 @@ struct
         Stable_db.get_all_accounts_rooted_at_exn db
     | Converting_db db ->
         Converting_ledger.get_all_accounts_rooted_at_exn db
+
+  let unsafely_decompose_root t =
+    match t with
+    | Stable_db db ->
+        (db, None)
+    | Converting_db db ->
+        ( Converting_ledger.primary_ledger db
+        , Some (Converting_ledger.converting_ledger db) )
 end

src/lib/mina_ledger/root.mli

@@ -114,7 +114,13 @@ module Make
 
   (** Create a root ledger backed by a single database in the given
       directory. *)
-  val create : logger:Logger.t -> config:Config.t -> depth:int -> unit -> t
+  val create :
+       logger:Logger.t
+    -> config:Config.t
+    -> depth:int
+    -> ?assert_synced:bool
+    -> unit
+    -> t
 
   val create_temporary :
        logger:Logger.t
@@ -124,12 +130,22 @@ module Make
     -> t
 
   (** Make a checkpoint of the root ledger and return a new root ledger backed
-      by that checkpoint *)
+      by that checkpoint. Throws an exception if the config does not match the
+      backing type of the root.  *)
   val create_checkpoint : t -> config:Config.t -> unit -> t
 
-  (** Make a checkpoint of the root ledger *)
+  (** Make a checkpoint of the root ledger. Throws an exception if the config
+      does not match the backing type of the root. *)
   val make_checkpoint : t -> config:Config.t -> unit
 
+  (** Make a checkpoint of the root ledger [t] of the same backing type using
+      [directory_name] as a template for its location. Return a new root ledger
+      backed by that checkpoint. *)
+  val create_checkpoint_with_directory : t -> directory_name:string -> t
+
+  (** Convert a root backed by a [Config.Stable_db] to *)
+  val make_converting : t -> t Async.Deferred.t
+
   (** View the root ledger as an unmasked [Any_ledger] so it can be used by code
       that does not need to know how the root is implemented *)
   val as_unmasked : t -> Any_ledger.witness
@@ -160,4 +176,10 @@ module Make
   (** Get all of the accounts that are in a subtree of the underlying Merkle
     tree rooted at `address`. The accounts are ordered by their addresses. *)
   val get_all_accounts_rooted_at_exn : t -> addr -> (addr * account) list
+
+  (** Decompose a root into its components parts. Users of this method must be
+      careful to ensure that either the underlying databases remain in sync, or
+      that they are not later used to back a root ledger. Use this on temporary
+      copies of root ledgers if possible. *)
+  val unsafely_decompose_root : t -> Stable_db.t * Migrated_db.t option
 end

src/lib/mina_ledger/test/test_mina_ledger.ml

@@ -180,6 +180,36 @@ module Root_test = struct
               ~logger ~depth:ledger_depth ~assert_synced:true ()
             |> close) ) ;
         Deferred.unit )
+
+  (** Test that a root created with a stable backing and then made converting
+      has the expected database states *)
+  let test_root_make_converting ~random () =
+    Mina_stdlib_unix.File_system.with_temp_dir "root_gradual_migration"
+      ~f:(fun cwd ->
+        let cfg =
+          L.Root.Config.with_directory ~directory_name:(cwd ^/ "ledger")
+        in
+        let root =
+          L.Root.create ~logger
+            ~config:(cfg ~backing_type:Stable_db)
+            ~depth:ledger_depth ()
+        in
+        let loc_with_accounts =
+          populate_with_random_accounts ~num:num_accounts ~root ~random
+        in
+        let%bind root = L.Root.make_converting root in
+        (* Make sure the stable accounts are all still present *)
+        assert_accounts ~loc_with_accounts ~root ;
+        L.Root.close root ;
+        (* Re-open the root as converting to check that the databases are in
+           sync *)
+        let converting_root =
+          L.Root.create ~logger
+            ~config:(cfg ~backing_type:Converting_db)
+            ~depth:ledger_depth ~assert_synced:true ()
+        in
+        L.Root.close converting_root ;
+        Deferred.unit )
 end
 
 let () =
@@ -195,5 +225,7 @@ let () =
                 (Root_test.test_root_moving ~random)
             ; Alcotest_async.test_case "make checkpointing a root" `Quick
                 (Root_test.test_root_make_checkpointing ~random)
+            ; Alcotest_async.test_case "make converting a root" `Quick
+                (Root_test.test_root_make_converting ~random)
             ] )
         ] )