Merge pull request #17801 from MinaProtocol/lyh/refactor-work-selector

Redesign work selector

Author: glyh

src/lib/mina_lib/mina_lib.ml

@@ -920,25 +920,7 @@ let add_complete_work ~logger ~fee ~prover
   Or_error.try_with update_metrics
   |> Result.iter_error ~f:(fun err ->
          [%log warn] "Failed to update metrics on adding work"
-           ~metadata:[ ("error", `String (Error.to_string_hum err)) ] ) ;
-  let cb result =
-    (* remove it from seen jobs after attempting to adding it to the pool to
-           avoid this work being reassigned.
-     * If the diff is accepted then remove it from the seen jobs.
-     * If not then the work should have already been in the pool with a 
-           lower fee or the statement isn't referenced anymore or any other 
-           error. In any case remove it from the seen jobs so that it can be 
-           picked up if needed *)
-    Work_selector.remove t.work_selector stmts ;
-    Result.iter_error result ~f:(fun err ->
-        (* Possible reasons of failure: receiving pipe's capacity exceeded,
-            fee that isn't the lowest, failure in verification or application to the pool *)
-        [%log warn] "Failed to add completed work to the pool"
-          ~metadata:
-            [ ("work_ids", Transaction_snark_work.Statement.compact_json stmts)
-            ; ("error", `String (Error.to_string_hum err))
-            ] )
-  in
+           ~metadata:[ ("error", Error_json.error_to_yojson err) ] ) ;
   Network_pool.Snark_pool.(
     Local_sink.push t.pipes.snark_local_sink
       ( Add_solved_work
@@ -949,7 +931,15 @@ let add_complete_work ~logger ~fee ~prover
                   |> One_or_two.map ~f:Ledger_proof.Cached.read_proof_from_disk
               ; fee = fee_with_prover
               } )
-      , cb ))
+      , Result.iter_error ~f:(fun err ->
+            (* Possible reasons of failure: receiving pipe's capacity exceeded,
+                fee that isn't the lowest, failure in verification or application to the pool *)
+            [%log warn] "Failed to add completed work to the pool"
+              ~metadata:
+                [ ( "work_ids"
+                  , Transaction_snark_work.Statement.compact_json stmts )
+                ; ("error", Error_json.error_to_yojson err)
+                ] ) ))
   |> Deferred.don't_wait_for
 
 let add_work t (work : Snark_work_lib.Result.Partitioned.Stable.Latest.t) =
@@ -1778,9 +1768,8 @@ let create ~commit_id ?wallets (config : Config.t) =
           | Ok () ->
               ()
           | Error err ->
-              [%log' warn config.logger]
-                "Error when setting ITN logger data: %s"
-                (Error.to_string_hum err)
+              [%log' warn config.logger] "Error when setting ITN logger data"
+                ~metadata:[ ("error", Error_json.error_to_yojson err) ]
         else Deferred.unit
       in
       O1trace.thread "mina_lib" (fun () ->
@@ -1866,13 +1855,13 @@ let create ~commit_id ?wallets (config : Config.t) =
               Or_error.iter_error result ~f:(fun error ->
                   [%log' warn config.logger]
                     "Failed to toggle verifier internal tracing: $error"
-                    ~metadata:[ ("error", `String (Error.to_string_hum error)) ] ) ) ;
+                    ~metadata:[ ("error", Error_json.error_to_yojson error) ] ) ) ;
           Internal_tracing.register_toggle_callback (fun enabled ->
               let%map result = Prover.toggle_internal_tracing prover enabled in
               Or_error.iter_error result ~f:(fun error ->
                   [%log' warn config.logger]
                     "Failed to toggle prover internal tracing: $error"
-                    ~metadata:[ ("error", `String (Error.to_string_hum error)) ] ) ) ;
+                    ~metadata:[ ("error", Error_json.error_to_yojson error) ] ) ) ;
           let%bind vrf_evaluator =
             Monitor.try_with ~here:[%here]
               ~rest:
@@ -2137,7 +2126,6 @@ let create ~commit_id ?wallets (config : Config.t) =
           in
           let work_selector =
             Work_selector.State.init
-              ~reassignment_wait:config.work_reassignment_wait
               ~frontier_broadcast_pipe:frontier_broadcast_pipe_r
               ~logger:config.logger
           in

src/lib/work_selector/intf.ml

@@ -100,8 +100,7 @@ module type State_intf = sig
   type transition_frontier
 
   val init :
-       reassignment_wait:int
-    -> frontier_broadcast_pipe:
+       frontier_broadcast_pipe:
          transition_frontier option Pipe_lib.Broadcast_pipe.Reader.t
     -> logger:Logger.t
     -> t
@@ -116,45 +115,30 @@ module type Lib_intf = sig
     include
       State_intf with type transition_frontier := Inputs.Transition_frontier.t
 
-    val remove_old_assignments : t -> logger:Logger.t -> unit
-
-    (**Jobs that have not been assigned yet*)
-    val all_unseen_works :
+    (** [mark_scheduled t work] Mark [work] as scheduled in [t] *)
+    val mark_scheduled :
          t
       -> ( Transaction_witness.t
          , Ledger_proof.Cached.t )
          Snark_work_lib.Work.Single.Spec.t
          One_or_two.t
-         list
-
-    val remove : t -> Transaction_snark.Statement.t One_or_two.t -> unit
+      -> unit
 
-    val set :
-         t
+    (** [all_unscheduled_expensive_works ~snark_pool ~fee t] filters out all
+        works in the list that satisfy the predicate
+        [does_not_have_better_fee ~snark_pool ~fee], and are not scheduled yet
+        *)
+    val all_unscheduled_expensive_works :
+         snark_pool:Snark_pool.t
+      -> fee:Fee.t
+      -> t
       -> ( Transaction_witness.t
          , Ledger_proof.Cached.t )
          Snark_work_lib.Work.Single.Spec.t
          One_or_two.t
-      -> unit
+         list
   end
 
-  (** [get_expensive_work ~snark_pool ~fee works] filters out all works in the
-      list that satisfy the predicate
-      [does_not_have_better_fee ~snark_pool ~fee] *)
-  val get_expensive_work :
-       snark_pool:Snark_pool.t
-    -> fee:Fee.t
-    -> ( Transaction_witness.t
-       , Ledger_proof.Cached.t )
-       Snark_work_lib.Work.Single.Spec.t
-       One_or_two.t
-       list
-    -> ( Transaction_witness.t
-       , Ledger_proof.Cached.t )
-       Snark_work_lib.Work.Single.Spec.t
-       One_or_two.t
-       list
-
   (**jobs that are not in the snark pool yet*)
   val pending_work_statements :
        snark_pool:Snark_pool.t

src/lib/work_selector/random.ml

@@ -1,16 +1,15 @@
 open Core_kernel
 
 module Make (Lib : Intf.Lib_intf) = struct
-  let work ~snark_pool ~fee ~logger (state : Lib.State.t) =
-    Lib.State.remove_old_assignments state ~logger ;
-    let unseen_jobs = Lib.State.all_unseen_works state in
-    match Lib.get_expensive_work ~snark_pool ~fee unseen_jobs with
+  let work ~snark_pool ~fee ~logger:_ (state : Lib.State.t) =
+    match Lib.State.all_unscheduled_expensive_works ~snark_pool ~fee state with
     | [] ->
         None
     | expensive_work ->
         let i = Random.int (List.length expensive_work) in
         let x = List.nth_exn expensive_work i in
-        Lib.State.set state x ; Some x
+        Lib.State.mark_scheduled state x ;
+        Some x
 end
 
 let%test_module "test" =

src/lib/work_selector/random_offset.ml

@@ -42,16 +42,15 @@ module Make (Lib : Intf.Lib_intf) = struct
         List.nth_exn expensive_work !offset
   end
 
-  let work ~snark_pool ~fee ~logger (state : Lib.State.t) =
-    Lib.State.remove_old_assignments state ~logger ;
-    let unseen_jobs = Lib.State.all_unseen_works state in
-    match Lib.get_expensive_work ~snark_pool ~fee unseen_jobs with
+  let work ~snark_pool ~fee ~logger:_ (state : Lib.State.t) =
+    match Lib.State.all_unscheduled_expensive_works ~snark_pool ~fee state with
     | [] ->
         None
     | expensive_work ->
         Offset.update ~new_length:(List.length expensive_work) ;
         let x = Offset.get_nth expensive_work in
-        Lib.State.set state x ; Some x
+        Lib.State.mark_scheduled state x ;
+        Some x
 end
 
 let%test_module "test" =

src/lib/work_selector/sequence.ml

@@ -1,12 +1,11 @@
 module Make (Lib : Intf.Lib_intf) = struct
-  let work ~snark_pool ~fee ~logger (state : Lib.State.t) =
-    Lib.State.remove_old_assignments state ~logger ;
-    let unseen_jobs = Lib.State.all_unseen_works state in
-    match Lib.get_expensive_work ~snark_pool ~fee unseen_jobs with
+  let work ~snark_pool ~fee ~logger:_ (state : Lib.State.t) =
+    match Lib.State.all_unscheduled_expensive_works ~snark_pool ~fee state with
     | [] ->
         None
     | x :: _ ->
-        Lib.State.set state x ; Some x
+        Lib.State.mark_scheduled state x ;
+        Some x
 end
 
 let%test_module "test" =

src/lib/work_selector/test.ml

@@ -6,9 +6,6 @@ open Pipe_lib
 module Make_test (Make_selection_method : Intf.Make_selection_method_intf) =
 struct
   module T = Inputs.Test_inputs
-
-  let reassignment_wait = 2000
-
   module Lib = Work_lib.Make (T)
   module Selection_method = Make_selection_method (Lib)
 
@@ -19,11 +16,10 @@ struct
 
   let precomputed_values = Precomputed_values.for_unit_tests
 
-  let init_state sl reassignment_wait logger =
+  let init_state sl logger =
     let tf_reader, tf_writer = Broadcast_pipe.create None in
     let work_state =
-      Lib.State.init ~reassignment_wait ~frontier_broadcast_pipe:tf_reader
-        ~logger
+      Lib.State.init ~frontier_broadcast_pipe:tf_reader ~logger
     in
     let%map () = Broadcast_pipe.Writer.write tf_writer (Some sl) in
     work_state
@@ -37,7 +33,7 @@ struct
     Quickcheck.test gen_staged_ledger ~trials:100 ~f:(fun sl ->
         Async.Thread_safe.block_on_async_exn (fun () ->
             let open Deferred.Let_syntax in
-            let%bind work_state = init_state sl reassignment_wait logger in
+            let%bind work_state = init_state sl logger in
             let rec go i =
               [%test_result: Bool.t]
                 ~message:"Exceeded time expected to exhaust work" ~expect:true
@@ -49,34 +45,6 @@ struct
             in
             go 0 ) )
 
-  let%test_unit "Reassign work after the wait time" =
-    Backtrace.elide := false ;
-    let snark_pool = T.Snark_pool.create () in
-    let fee = Currency.Fee.zero in
-    let logger = Logger.null () in
-    let send_work work_state =
-      let rec go all_work =
-        let stuff = Selection_method.work ~snark_pool ~fee ~logger work_state in
-        match stuff with
-        | None ->
-            all_work
-        | Some work ->
-            go (One_or_two.to_list work @ all_work)
-      in
-      go []
-    in
-    Quickcheck.test gen_staged_ledger ~trials:10 ~f:(fun sl ->
-        Async.Thread_safe.block_on_async_exn (fun () ->
-            let open Deferred.Let_syntax in
-            let%bind work_state = init_state sl reassignment_wait logger in
-            let work_sent = send_work work_state in
-            (*wait for wait_time after which all the work will be reassigned*)
-            let%map () =
-              Async.after (Time.Span.of_ms (Float.of_int reassignment_wait))
-            in
-            let work_sent_again = send_work work_state in
-            assert (List.length work_sent = List.length work_sent_again) ) )
-
   let gen_snark_pool (works : ('a, 'b) Lib.Work_spec.t One_or_two.t list) fee =
     let open Quickcheck.Generator.Let_syntax in
     let cheap_work_fee = Option.value_exn Fee.(sub fee one) in
@@ -123,7 +91,7 @@ struct
       ~trials:100 ~f:(fun (sl, snark_pool) ->
         Async.Thread_safe.block_on_async_exn (fun () ->
             let open Deferred.Let_syntax in
-            let%bind work_state = init_state sl reassignment_wait logger in
+            let%bind work_state = init_state sl logger in
             let rec go i =
               [%test_result: Bool.t]
                 ~message:"Exceeded time expected to exhaust work" ~expect:true