Merge pull request #16677 from MinaProtocol/dkijania/use_zkapp_vk_in_transaction_pool_dev

[DEV] Introduce zkapp_vk_cache_tag project for caching zkapp verification keys

Author: mrmr1993

src/dune-project

@@ -231,4 +231,5 @@
 (package (name work_selector))
 (package (name zkapp_command_builder))
 (package (name zkapps_examples))
-(package (name zkapp_limits))
+(package (name zkapp_limits))
+(package (name zkapp_vk_cache_tag))

src/lib/mina_lib/mina_lib.ml

@@ -1671,6 +1671,11 @@ let initialize_proof_cache_db (config : Config.t) =
     (config.conf_dir ^/ "proof_cache")
   >>| function Error e -> raise_on_initialization_error e | Ok db -> db
 
+let initialize_zkapp_vk_cache_db (config : Config.t) =
+  Zkapp_vk_cache_tag.create_db ~logger:config.logger
+    (config.conf_dir ^/ "zkapp_vk_cache")
+  >>| function Error e -> raise_on_initialization_error e | Ok db -> db
+
 let create ~commit_id ?wallets (config : Config.t) =
   let commit_id_short = String.sub ~pos:0 ~len:8 commit_id in
   let catchup_mode = if config.super_catchup then `Super else `Normal in
@@ -1705,6 +1710,7 @@ let create ~commit_id ?wallets (config : Config.t) =
                  in_memory_reverse_structured_log_messages_for_integration_test
                  config.start_filtered_logs ;
           let%bind proof_cache_db = initialize_proof_cache_db config in
+          let%bind zkapp_vk_cache_db = initialize_zkapp_vk_cache_db config in
           let module Context =
           (val context ~proof_cache_db ~commit_id config)
           in
@@ -2001,7 +2007,7 @@ let create ~commit_id ?wallets (config : Config.t) =
               ~pool_max_size:
                 config.precomputed_values.genesis_constants.txpool_max_size
               ~genesis_constants:config.precomputed_values.genesis_constants
-              ~slot_tx_end
+              ~slot_tx_end ~vk_cache_db:zkapp_vk_cache_db
           in
           let first_received_message_signal = Ivar.create () in
           let online_status, notify_online_impl =

src/lib/mina_lib/tests/tests.ml

@@ -211,6 +211,7 @@ let%test_module "Epoch ledger sync tests" =
             ~pool_max_size:precomputed_values.genesis_constants.txpool_max_size
             ~genesis_constants:precomputed_values.genesis_constants
             ~slot_tx_end:None
+            ~vk_cache_db:(Zkapp_vk_cache_tag.For_tests.create_db ())
         in
         Network_pool.Transaction_pool.create ~config ~constraint_constants
           ~consensus_constants ~time_controller ~logger

src/lib/network_pool/dune

@@ -77,7 +77,8 @@
   kimchi_pasta.basic
   coda_genesis_ledger
   staged_ledger_diff
-  mina_wire_types)
+  mina_wire_types
+  zkapp_vk_cache_tag)
  (instrumentation
   (backend bisect_ppx))
  (preprocess

src/lib/network_pool/intf.ml

@@ -387,6 +387,7 @@ module type Transaction_resource_pool_intf = sig
     -> verifier:Verifier.t
     -> genesis_constants:Genesis_constants.t
     -> slot_tx_end:Mina_numbers.Global_slot_since_hard_fork.t option
+    -> vk_cache_db:Zkapp_vk_cache_tag.cache_db
     -> Config.t
 
   val member : t -> Transaction_hash.User_command_with_valid_signature.t -> bool

src/lib/network_pool/transaction_pool.ml

@@ -287,18 +287,19 @@ struct
         ; verifier : (Verifier.t[@sexp.opaque])
         ; genesis_constants : Genesis_constants.t
         ; slot_tx_end : Mina_numbers.Global_slot_since_hard_fork.t option
+        ; vk_cache_db : Zkapp_vk_cache_tag.cache_db
         }
-      [@@deriving sexp_of]
 
       (* remove next line if there's a way to force [@@deriving make] write a
          named parameter instead of an optional parameter *)
       let make ~trust_system ~pool_max_size ~verifier ~genesis_constants
-          ~slot_tx_end =
+          ~slot_tx_end ~vk_cache_db =
         { trust_system
         ; pool_max_size
         ; verifier
         ; genesis_constants
         ; slot_tx_end
+        ; vk_cache_db
         }
     end
 
@@ -331,15 +332,17 @@ struct
     module Vk_refcount_table = struct
       type t =
         { verification_keys :
-            (int * Verification_key_wire.t) Zkapp_basic.F_map.Table.t
+            (int * Zkapp_vk_cache_tag.t) Zkapp_basic.F_map.Table.t
         ; account_id_to_vks : int Zkapp_basic.F_map.Map.t Account_id.Table.t
         ; vk_to_account_ids : int Account_id.Map.t Zkapp_basic.F_map.Table.t
+        ; vk_cache_db : Zkapp_vk_cache_tag.cache_db
         }
 
-      let create () =
+      let create vk_cache_db () =
         { verification_keys = Zkapp_basic.F_map.Table.create ()
         ; account_id_to_vks = Account_id.Table.create ()
         ; vk_to_account_ids = Zkapp_basic.F_map.Table.create ()
+        ; vk_cache_db
         }
 
       let find_vk (t : t) = Hashtbl.find t.verification_keys
@@ -365,7 +368,7 @@ struct
         in
         Hashtbl.update t.verification_keys vk.hash ~f:(function
           | None ->
-              (1, vk)
+              (1, Zkapp_vk_cache_tag.write_key_to_disk t.vk_cache_db vk)
           | Some (count, vk) ->
               (count + 1, vk) ) ;
         Hashtbl.update t.account_id_to_vks account_id
@@ -425,13 +428,12 @@ struct
       ; mutable current_batch : int
       ; mutable remaining_in_batch : int
       ; config : Config.t
-      ; logger : (Logger.t[@sexp.opaque])
+      ; logger : Logger.t
       ; batcher : Batcher.t
-      ; mutable best_tip_diff_relay : (unit Deferred.t[@sexp.opaque]) Option.t
-      ; mutable best_tip_ledger : (Base_ledger.t[@sexp.opaque]) Option.t
-      ; verification_key_table : (Vk_refcount_table.t[@sexp.opaque])
+      ; mutable best_tip_diff_relay : unit Deferred.t Option.t
+      ; mutable best_tip_ledger : Base_ledger.t Option.t
+      ; verification_key_table : Vk_refcount_table.t
       }
-    [@@deriving sexp_of]
 
     let member t x =
       Indexed_pool.member t.pool (Transaction_hash.User_command.of_checked x)
@@ -837,7 +839,8 @@ struct
         ; batcher = Batcher.create ~logger config.verifier
         ; best_tip_diff_relay = None
         ; best_tip_ledger = None
-        ; verification_key_table = Vk_refcount_table.create ()
+        ; verification_key_table =
+            Vk_refcount_table.create config.vk_cache_db ()
         }
       in
       don't_wait_for
@@ -1169,7 +1172,12 @@ struct
                             in
                             let vks =
                               vks
-                              |> List.map ~f:(fun vk -> (vk.hash, vk))
+                              |> List.map ~f:(fun vk_cached ->
+                                     let vk =
+                                       Zkapp_vk_cache_tag.read_key_from_disk
+                                         vk_cached
+                                     in
+                                     (vk.hash, vk) )
                               |> Zkapp_basic.F_map.Map.of_alist_exn
                             in
                             (account_id, vks) )
@@ -1930,6 +1938,7 @@ let%test_module _ =
       let config =
         Test.Resource_pool.make_config ~trust_system ~pool_max_size ~verifier
           ~genesis_constants ~slot_tx_end
+          ~vk_cache_db:(Zkapp_vk_cache_tag.For_tests.create_db ())
       in
       let pool_, _, _ =
         Test.create ~config ~logger ~constraint_constants ~consensus_constants

src/lib/transaction_inclusion_status/transaction_inclusion_status.ml

@@ -131,6 +131,7 @@ let%test_module "transaction_status" =
         Transaction_pool.Resource_pool.make_config ~trust_system ~pool_max_size
           ~verifier ~genesis_constants:precomputed_values.genesis_constants
           ~slot_tx_end:None
+          ~vk_cache_db:(Zkapp_vk_cache_tag.For_tests.create_db ())
       in
       let transaction_pool, _, local_sink =
         Transaction_pool.create ~config

src/lib/zkapp_vk_cache_tag/dune

@@ -0,0 +1,14 @@
+(library
+ (public_name zkapp_vk_cache_tag)
+ (libraries
+   ;; opam libraries
+   core_kernel
+   async
+   ;; local libraries
+   logger
+   disk_cache
+   mina_base
+   )
+ (preprocess
+  (pps ppx_mina ppx_version))
+ (instrumentation (backend bisect_ppx)))

src/lib/zkapp_vk_cache_tag/zkapp_vk_cache_tag.ml

@@ -0,0 +1,28 @@
+module Cache = Disk_cache.Make (Mina_base.Verification_key_wire.Stable.Latest)
+
+type cache_db = Lmdb_cache of Cache.t | Identity_cache
+
+type t =
+  | Lmdb of { cache_id : Cache.id; cache_db : Cache.t }
+  | Identity of Mina_base.Verification_key_wire.t
+
+let read_key_from_disk = function
+  | Lmdb t ->
+      Cache.get t.cache_db t.cache_id
+  | Identity key ->
+      key
+
+let write_key_to_disk db key =
+  match db with
+  | Lmdb_cache cache_db ->
+      Lmdb { cache_id = Cache.put cache_db key; cache_db }
+  | Identity_cache ->
+      Identity key
+
+let create_db path ~logger =
+  Cache.initialize ~logger path
+  |> Async.Deferred.Result.map ~f:(fun cache -> Lmdb_cache cache)
+
+module For_tests = struct
+  let create_db () = Identity_cache
+end

src/lib/zkapp_vk_cache_tag/zkapp_vk_cache_tag.mli

@@ -0,0 +1,19 @@
+open Core
+open Async
+
+type t
+
+type cache_db
+
+val create_db :
+     string
+  -> logger:Logger.t
+  -> (cache_db, [> `Initialization_error of Error.t ]) Deferred.Result.t
+
+val read_key_from_disk : t -> Mina_base.Verification_key_wire.t
+
+val write_key_to_disk : cache_db -> Mina_base.Verification_key_wire.t -> t
+
+module For_tests : sig
+  val create_db : unit -> cache_db
+end