Merge pull request #17730 from MinaProtocol/georgeee/use-stricter-error-type-in-work-paritioner

Use stricter error type in Work partitioner

Author: georgeee

src/lib/snark_worker/prod.ml

@@ -136,18 +136,21 @@ module Impl = struct
 
   let perform_single_untimed ~(m : (module Worker_state.S)) ~logger
       ~proof_cache_db ~single_spec ~signature_kind ~sok_digest () =
-    let open Deferred.Or_error.Let_syntax in
+    let open Deferred.Result.Let_syntax in
     let (module M) = m in
     match single_spec with
     | Work.Work.Single.Spec.Transition
         (input, (w : Transaction_witness.Stable.Latest.t)) -> (
+        (* TODO: remove this case after delivering the full snark worker
+           optimization PR series *)
         match w.transaction with
         | Command (Zkapp_command zkapp_command) -> (
             let%bind witnesses_specs_stmts =
               extract_zkapp_segment_works ~m ~input ~witness:w
                 ~zkapp_command:
                   (Zkapp_command.write_all_proofs_to_disk ~signature_kind
                      ~proof_cache_db zkapp_command )
+              |> Result.map_error ~f:Failed_to_generate_inputs.error_of_t
               |> Deferred.return
             in
             match Mina_stdlib.Nonempty_list.uncons witnesses_specs_stmts with

src/lib/work_partitioner/snark_worker_shared.ml

@@ -40,10 +40,24 @@ module Zkapp_command_inputs = struct
         , stmt ) )
 end
 
+module Failed_to_generate_inputs = struct
+  type t = [ `Failed_to_generate_inputs of Zkapp_command.t * Error.t ]
+
+  (* TODO: remove this after fully delivering the snark worker
+     optimization PR series *)
+  let error_of_t (`Failed_to_generate_inputs (tx, e) : t) =
+    Error.tag_s e
+      ~tag:
+        ( Zkapp_command.Stable.Latest.sexp_of_t
+        @@ Zkapp_command.read_all_proofs_from_disk tx )
+    |> Error.tag ~tag:"failed to generate inputs for zkapp command"
+end
+
 let extract_zkapp_segment_works ~m:(module M : S)
     ~(input : Mina_state.Snarked_ledger_state.t)
     ~(witness : Transaction_witness.Stable.Latest.t)
-    ~(zkapp_command : Zkapp_command.t) : Zkapp_command_inputs.t Or_error.t =
+    ~(zkapp_command : Zkapp_command.t) :
+    (Zkapp_command_inputs.t, Failed_to_generate_inputs.t) Result.t =
   let inputs =
     Or_error.try_with (fun () ->
         Transaction_snark.zkapp_command_witnesses_exn
@@ -67,12 +81,7 @@ let extract_zkapp_segment_works ~m:(module M : S)
   in
   match inputs with
   | Error e ->
-      Error
-        (Error.createf
-           !"Failed to generate inputs for zkapp_command : %s: %s"
-           ( Zkapp_command.read_all_proofs_from_disk zkapp_command
-           |> Zkapp_command.Stable.Latest.to_yojson |> Yojson.Safe.to_string )
-           (Error.to_string_hum e) )
+      Result.fail (`Failed_to_generate_inputs (zkapp_command, e))
   | Ok (first_segment :: rest_segments) ->
       Ok (Nonempty_list.init first_segment rest_segments)
   | Ok [] ->

src/lib/work_partitioner/work_partitioner.ml

@@ -70,7 +70,7 @@ let reschedule_if_old ~reassignment_timeout
 let reschedule_old_zkapp_job
     ~partitioner:
       ({ reassignment_timeout; zkapp_jobs_sent_by_partitioner; _ } : t) :
-    Work.Spec.Partitioned.Stable.Latest.t Or_error.t option =
+    (Work.Spec.Partitioned.Stable.Latest.t, _) Result.t option =
   let%map.Option job =
     Sent_zkapp_job_pool.remove_until_reschedule
       ~f:(reschedule_if_old ~reassignment_timeout)
@@ -81,7 +81,7 @@ let reschedule_old_zkapp_job
 let reschedule_old_single_job
     ~partitioner:
       ({ reassignment_timeout; single_jobs_sent_by_partitioner; _ } : t) :
-    Work.Spec.Partitioned.Stable.Latest.t Or_error.t option =
+    (Work.Spec.Partitioned.Stable.Latest.t, _) Result.t option =
   let%map.Option job =
     Sent_single_job_pool.remove_until_reschedule
       ~f:(reschedule_if_old ~reassignment_timeout)
@@ -118,7 +118,7 @@ let schedule_from_pending_zkapp_command ~(id : Work.Id.Single.t)
     ~pending
 
 let schedule_from_any_pending_zkapp_command ~(partitioner : t) :
-    Work.Spec.Partitioned.Stable.Latest.t Or_error.t option =
+    (Work.Spec.Partitioned.Stable.Latest.t, _) Result.t option =
   let spec_generated = ref None in
   (* TODO: Consider remove all works no longer relevant for current frontier,
      this may need changes from underlying work selector. *)
@@ -155,7 +155,7 @@ let convert_zkapp_command_from_selector ~partitioner ~job ~pairing
 
 let convert_single_work_from_selector ~(partitioner : t)
     ~(single_spec : Work.Spec.Single.t) ~sok_message ~pairing :
-    Work.Spec.Partitioned.Stable.Latest.t Or_error.t =
+    (Work.Spec.Partitioned.Stable.Latest.t, _) Result.t =
   let job =
     Work.With_job_meta.
       { spec = single_spec
@@ -208,7 +208,7 @@ let schedule_from_tmp_slot ~(partitioner : t) =
     ~sok_message
 
 let schedule_job_from_partitioner ~(partitioner : t) :
-    Work.Spec.Partitioned.Stable.Latest.t Or_error.t option =
+    (Work.Spec.Partitioned.Stable.Latest.t, _) Result.t option =
   List.find_map ~f:Lazy.force
     [ lazy (reschedule_old_zkapp_job ~partitioner)
     ; lazy (reschedule_old_single_job ~partitioner)
@@ -220,7 +220,7 @@ let schedule_job_from_partitioner ~(partitioner : t) :
 let consume_job_from_selector ~(partitioner : t)
     ~(sok_message : Mina_base.Sok_message.t)
     ~(instances : Work.Spec.Single.t One_or_two.t) :
-    Work.Spec.Partitioned.Stable.Latest.t Or_error.t =
+    (Work.Spec.Partitioned.Stable.Latest.t, _) Result.t =
   let pairing_id = Id_generator.next_id partitioner.single_id_gen () in
   Hashtbl.add_exn partitioner.pairing_pool ~key:pairing_id
     ~data:(Spec_only { spec = instances; sok_message }) ;
@@ -248,7 +248,7 @@ let request_from_selector_and_consume_by_partitioner ~(partitioner : t)
 
 let request_partitioned_work ~(sok_message : Mina_base.Sok_message.t)
     ~(work_from_selector : work_from_selector) ~(partitioner : t) :
-    Work.Spec.Partitioned.Stable.Latest.t Or_error.t option =
+    (Work.Spec.Partitioned.Stable.Latest.t, _) Result.t option =
   List.find_map ~f:Lazy.force
     [ lazy (schedule_job_from_partitioner ~partitioner)
     ; lazy

src/lib/work_partitioner/work_partitioner.mli

@@ -56,7 +56,10 @@ val request_partitioned_work :
      sok_message:Mina_base.Sok_message.t
   -> work_from_selector:work_from_selector
   -> partitioner:t
-  -> Snark_work_lib.Spec.Partitioned.Stable.V1.t Or_error.t option
+  -> ( Snark_work_lib.Spec.Partitioned.Stable.V1.t
+     , Snark_worker_shared.Failed_to_generate_inputs.t )
+     Result.t
+     option
 
 type submit_result =
   | SpecUnmatched  (** Submitted work doesn't match the scheme in pool *)