Map based snark pool tables

Author: nholland94

src/lib/network_pool/snark_pool.ml

@@ -2,16 +2,15 @@ open Core_kernel
 open Async
 open Pipe_lib
 open Network_peer
-module Statement_table = Transaction_snark_work.Statement.Table
 
 module Snark_tables = struct
   type t =
     { all :
         Ledger_proof.t One_or_two.t Priced_proof.t
-        Transaction_snark_work.Statement.Table.t
+        Transaction_snark_work.Statement.Map.t
     ; rebroadcastable :
         (Ledger_proof.t One_or_two.t Priced_proof.t * Core.Time.t)
-        Transaction_snark_work.Statement.Table.t
+        Transaction_snark_work.Statement.Map.t
     }
   [@@deriving sexp]
 end
@@ -107,7 +106,7 @@ struct
         | `New_best_tip of Base_ledger.t ]
 
       type t =
-        { snark_tables : Snark_tables.t
+        { snark_tables : Snark_tables.t ref
         ; snark_table_lock : (unit Throttle.Sequencer.t[@sexp.opaque])
         ; frontier : (unit -> Transition_frontier.t option[@sexp.opaque])
         ; config : Config.t
@@ -136,7 +135,7 @@ struct
 
       let snark_pool_json t : Yojson.Safe.t =
         `List
-          (Statement_table.fold ~init:[] t.snark_tables.all
+          (Map.fold ~init:[] !(t.snark_tables).all
              ~f:(fun ~key ~data:{ proof = _; fee = { fee; prover } } acc ->
                let work_ids =
                  Transaction_snark_work.Statement.compact_json key
@@ -151,7 +150,7 @@ struct
                :: acc ) )
 
       let all_completed_work (t : t) : Transaction_snark_work.Info.t list =
-        Statement_table.fold ~init:[] t.snark_tables.all
+        Map.fold ~init:[] !(t.snark_tables).all
           ~f:(fun ~key ~data:{ proof = _; fee = { fee; prover } } acc ->
             let work_ids = Transaction_snark_work.Statement.work_ids key in
             { Transaction_snark_work.Info.statements = key
@@ -177,7 +176,7 @@ struct
                  in
                  let%bind prover_account_ids =
                    let account_ids =
-                     t.snark_tables.all |> Statement_table.data
+                     !(t.snark_tables).all |> Map.data
                      |> List.map
                           ~f:(fun { Priced_proof.fee = { prover; _ }; _ } ->
                             prover )
@@ -202,7 +201,7 @@ struct
                  let yield = Staged.unstage (Scheduler.yield_every ~n:50) in
                  let%map () =
                    let open Deferred.Let_syntax in
-                   Statement_table.fold ~init:Deferred.unit t.snark_tables.all
+                   Map.fold ~init:Deferred.unit !(t.snark_tables).all
                      ~f:(fun ~key ~data:{ fee = { fee; prover }; _ } acc ->
                        let%bind () = acc in
                        let%map () = yield () in
@@ -216,9 +215,12 @@ struct
                          fee_is_sufficient t ~fee
                            ~account_exists:prover_account_exists
                        in
-                       if not keep then (
-                         Hashtbl.remove t.snark_tables.all key ;
-                         Hashtbl.remove t.snark_tables.rebroadcastable key ) )
+                       if not keep then
+                         t.snark_tables :=
+                           { all = Map.remove !(t.snark_tables).all key
+                           ; rebroadcastable =
+                               Map.remove !(t.snark_tables).rebroadcastable key
+                           } )
                  in
                  () )
             in
@@ -229,12 +231,16 @@ struct
         [%log' trace t.logger]
           "Handling snark refcount table: $num_removed works were removed"
           ~metadata:[ ("num_removed", `Int (List.length removed_work)) ] ;
-        List.iter removed_work ~f:(fun work ->
-            Hashtbl.remove t.snark_tables.rebroadcastable work ;
-            Hashtbl.remove t.snark_tables.all work ) ;
+        t.snark_tables :=
+          { all =
+              List.fold ~init:!(t.snark_tables).all ~f:Map.remove removed_work
+          ; rebroadcastable =
+              List.fold ~init:!(t.snark_tables).rebroadcastable ~f:Map.remove
+                removed_work
+          } ;
         Mina_metrics.(
           Gauge.set Snark_work.snark_pool_size
-            (Float.of_int @@ Hashtbl.length t.snark_tables.all))
+            (Float.of_int @@ Map.length !(t.snark_tables).all))
 
       let handle_transition_frontier_diff u t =
         match u with
@@ -276,9 +282,10 @@ struct
           ~frontier_broadcast_pipe ~config ~logger ~tf_diff_writer =
         let t =
           { snark_tables =
-              { all = Statement_table.create ()
-              ; rebroadcastable = Statement_table.create ()
-              }
+              ref
+                { Snark_tables.all = Transaction_snark_work.Statement.Map.empty
+                ; rebroadcastable = Transaction_snark_work.Statement.Map.empty
+                }
           ; snark_table_lock = Throttle.Sequencer.create ()
           ; frontier =
               (fun () -> Broadcast_pipe.Reader.peek frontier_broadcast_pipe)
@@ -296,7 +303,7 @@ struct
 
       let get_logger t = t.logger
 
-      let request_proof t = Statement_table.find t.snark_tables.all
+      let request_proof t = Map.find !(t.snark_tables).all
 
       let add_snark ?(is_local = false) t ~work
           ~(proof : Ledger_proof.t One_or_two.t) ~fee =
@@ -305,23 +312,28 @@ struct
               ( if work_is_referenced t work then (
                 (*Note: fee against existing proofs and the new proofs are checked in
                   Diff.unsafe_apply which calls this function*)
-                Hashtbl.set t.snark_tables.all ~key:work ~data:{ proof; fee } ;
-                if is_local then
-                  Hashtbl.set t.snark_tables.rebroadcastable ~key:work
-                    ~data:({ proof; fee }, Time.now ())
-                else
-                  (* Stop rebroadcasting locally generated snarks if they are
-                     overwritten. No-op if there is no rebroadcastable SNARK with that
-                     statement. *)
-                  Hashtbl.remove t.snark_tables.rebroadcastable work ;
+                t.snark_tables :=
+                  { all =
+                      Map.set !(t.snark_tables).all ~key:work
+                        ~data:{ proof; fee }
+                  ; rebroadcastable =
+                      ( if is_local then
+                        Map.set !(t.snark_tables).rebroadcastable ~key:work
+                          ~data:({ proof; fee }, Time.now ())
+                      else
+                        (* Stop rebroadcasting locally generated snarks if they are
+                           overwritten. No-op if there is no rebroadcastable SNARK with that
+                           statement. *)
+                        Map.remove !(t.snark_tables).rebroadcastable work )
+                  } ;
                 (*when snark work is added to the pool*)
                 Mina_metrics.(
                   Gauge.set Snark_work.useful_snark_work_received_time_sec
                     Time.(
                       let x = now () |> to_span_since_epoch |> Span.to_sec in
                       x -. Mina_metrics.time_offset_sec) ;
                   Gauge.set Snark_work.snark_pool_size
-                    (Float.of_int @@ Hashtbl.length t.snark_tables.all) ;
+                    (Float.of_int @@ Map.length !(t.snark_tables).all) ;
                   Snark_work.Snark_fee_histogram.observe Snark_work.snark_fee
                     ( fee.Mina_base.Fee_with_prover.fee |> Currency.Fee.to_int
                     |> Float.of_int )) ;
@@ -486,15 +498,17 @@ struct
       | None ->
           []
       | Some best_tips ->
-          Hashtbl.to_alist t.snark_tables.rebroadcastable
+          Map.to_alist !(t.snark_tables).rebroadcastable
           |> List.filter_map ~f:(fun (stmt, (snark, _time)) ->
                  if Hash_set.mem best_tips stmt then
                    Some (Diff.Add_solved_work (stmt, snark))
                  else None )
 
     let remove_solved_work t work =
-      Statement_table.remove t.snark_tables.all work ;
-      Statement_table.remove t.snark_tables.rebroadcastable work
+      t.snark_tables :=
+        { all = Map.remove !(t.snark_tables).all work
+        ; rebroadcastable = Map.remove !(t.snark_tables).rebroadcastable work
+        }
   end
 
   include Network_pool_base.Make (Transition_frontier) (Resource_pool)

src/lib/transaction_snark_work/transaction_snark_work.ml

@@ -27,12 +27,14 @@ module Statement = struct
 
       type _unused = unit constraint t = Arg.Stable.V1.t
 
+      include Comparable.Make_binable (Arg.Stable.V1)
       include Hashable.Make_binable (Arg.Stable.V1)
     end
   end]
 
   type t = Stable.Latest.t [@@deriving sexp, hash, compare, yojson, equal]
 
+  include Comparable.Make_binable (Stable.Latest)
   include Hashable.Make (Stable.Latest)
 
   let gen = One_or_two.gen Transaction_snark.Statement.gen

src/lib/transaction_snark_work/transaction_snark_work.mli

@@ -6,12 +6,16 @@ module Statement : sig
   type t = Transaction_snark.Statement.t One_or_two.t
   [@@deriving compare, sexp, yojson, equal]
 
+  include Comparable.S with type t := t
+
   include Hashable.S with type t := t
 
   module Stable : sig
     module V1 : sig
       type t [@@deriving bin_io, compare, sexp, version, yojson, equal]
 
+      include Comparable.S with type t := t
+
       include Hashable.S_binable with type t := t
     end
   end