pkg: add argument to pkg lock to print perf stats (ocaml#11263)

Signed-off-by: Stephen Sherratt <[email protected]>

Author: gridbugs

bin/lock_dev_tool.ml

@@ -81,6 +81,7 @@ let solve ~dev_tool ~local_packages =
        ~solver_env_from_current_system
        ~version_preference:None
        ~lock_dirs:[ lock_dir ]
+       ~print_perf_stats:false
 ;;
 
 let compiler_package_name = Package_name.of_string "ocaml"

bin/pkg/lock.ml

@@ -70,6 +70,7 @@ let solve_lock_dir
   workspace
   ~local_packages
   ~project_pins
+  ~print_perf_stats
   version_preference
   solver_env_from_current_system
   lock_dir_path
@@ -94,6 +95,7 @@ let solve_lock_dir
       ~unset_solver_vars_from_context:
         (unset_solver_vars_of_workspace workspace ~lock_dir_path)
   in
+  let time_start = Unix.gettimeofday () in
   let* repos =
     let repo_map = repositories_of_workspace workspace in
     let repo_names =
@@ -105,6 +107,7 @@ let solve_lock_dir
     get_repos repo_map ~repositories:(repositories_of_lock_dir workspace ~lock_dir_path)
   in
   let* pins = resolve_project_pins project_pins in
+  let time_solve_start = Unix.gettimeofday () in
   progress_state := Some Progress_indicator.Per_lockdir.State.Solving;
   Dune_pkg.Opam_solver.solve_lock_dir
     solver_env
@@ -119,19 +122,30 @@ let solve_lock_dir
     ~constraints:(constraints_of_workspace workspace ~lock_dir_path)
   >>= function
   | Error (`Diagnostic_message message) -> Fiber.return (Error (lock_dir_path, message))
-  | Ok { lock_dir; files; pinned_packages } ->
+  | Ok { lock_dir; files; pinned_packages; num_expanded_packages } ->
+    let time_end = Unix.gettimeofday () in
+    let maybe_perf_stats =
+      if print_perf_stats
+      then
+        [ Pp.nop
+        ; Pp.textf "Expanded packages: %d" num_expanded_packages
+        ; Pp.textf "Updated repos in: %.2fs" (time_solve_start -. time_start)
+        ; Pp.textf "Solved dependencies in: %.2fs" (time_end -. time_solve_start)
+        ]
+      else []
+    in
     let summary_message =
       User_message.make
-        [ Pp.tag
-            User_message.Style.Success
-            (Pp.textf
-               "Solution for %s:"
-               (Path.Source.to_string_maybe_quoted lock_dir_path))
-        ; (match Package_name.Map.values lock_dir.packages with
-           | [] ->
-             Pp.tag User_message.Style.Warning @@ Pp.text "(no dependencies to lock)"
-           | packages -> pp_packages packages)
-        ]
+        (Pp.tag
+           User_message.Style.Success
+           (Pp.textf
+              "Solution for %s:"
+              (Path.Source.to_string_maybe_quoted lock_dir_path))
+         :: (match Package_name.Map.values lock_dir.packages with
+             | [] ->
+               Pp.tag User_message.Style.Warning @@ Pp.text "(no dependencies to lock)"
+             | packages -> pp_packages packages)
+         :: maybe_perf_stats)
     in
     progress_state := None;
     let+ lock_dir = Lock_dir.compute_missing_checksums ~pinned_packages lock_dir in
@@ -145,6 +159,7 @@ let solve
   ~solver_env_from_current_system
   ~version_preference
   ~lock_dirs
+  ~print_perf_stats
   =
   let open Fiber.O in
   (* a list of thunks that will perform all the file IO side
@@ -166,6 +181,7 @@ let solve
                workspace
                ~local_packages
                ~project_pins
+               ~print_perf_stats
                version_preference
                solver_env_from_current_system
                lockdir_path
@@ -198,7 +214,7 @@ let project_pins =
     Pin_stanza.DB.combine_exn acc (Dune_project.pins project))
 ;;
 
-let lock ~version_preference ~lock_dirs_arg =
+let lock ~version_preference ~lock_dirs_arg ~print_perf_stats =
   let open Fiber.O in
   let* solver_env_from_current_system =
     Dune_pkg.Sys_poll.make ~path:(Env_path.path Stdune.Env.initial)
@@ -223,15 +239,18 @@ let lock ~version_preference ~lock_dirs_arg =
     ~solver_env_from_current_system
     ~version_preference
     ~lock_dirs
+    ~print_perf_stats
 ;;
 
 let term =
   let+ builder = Common.Builder.term
   and+ version_preference = Version_preference.term
-  and+ lock_dirs_arg = Pkg_common.Lock_dirs_arg.term in
+  and+ lock_dirs_arg = Pkg_common.Lock_dirs_arg.term
+  and+ print_perf_stats = Arg.(value & flag & info [ "print-perf-stats" ]) in
   let builder = Common.Builder.forbid_builds builder in
   let common, config = Common.init builder in
-  Scheduler.go ~common ~config (fun () -> lock ~version_preference ~lock_dirs_arg)
+  Scheduler.go ~common ~config (fun () ->
+    lock ~version_preference ~lock_dirs_arg ~print_perf_stats)
 ;;
 
 let info =

bin/pkg/lock.mli

@@ -7,6 +7,7 @@ val solve
   -> solver_env_from_current_system:Dune_pkg.Solver_env.t option
   -> version_preference:Dune_pkg.Version_preference.t option
   -> lock_dirs:Path.Source.t list
+  -> print_perf_stats:bool
   -> unit Fiber.t
 
 (** Command to create lock directory *)

src/dune_pkg/opam_solver.ml

@@ -78,6 +78,9 @@ module Context = struct
          packages for which we've printed a warning. *)
       available_cache : (OpamPackage.t, bool) Table.t
     ; constraints : OpamTypes.filtered_formula Package_name.Map.t
+    ; (* Number of versions of each package whose opam files were read from
+         disk while solving. Used to report performance statistics. *)
+      expanded_packages : (Package_name.t, int) Table.t
     }
 
   let create
@@ -105,6 +108,15 @@ module Context = struct
         end)
         1
     in
+    let expanded_packages =
+      Table.create
+        (module struct
+          include Package_name
+
+          let to_dyn = Package_name.to_dyn
+        end)
+        1
+    in
     { repos
     ; version_preference
     ; local_packages
@@ -115,6 +127,7 @@ module Context = struct
     ; candidates_cache
     ; available_cache
     ; constraints
+    ; expanded_packages
     }
   ;;
 
@@ -180,6 +193,10 @@ module Context = struct
 
   let repo_candidate t name =
     let+ resolved = Opam_repo.load_all_versions t.repos name in
+    Table.add_exn
+      t.expanded_packages
+      (Package_name.of_opam_package_name name)
+      (OpamPackage.Version.Map.cardinal resolved);
     let available =
       OpamPackage.Version.Map.values resolved
       |> List.map ~f:(fun p -> p, Priority.make (Resolved_package.opam_file p))
@@ -252,6 +269,8 @@ module Context = struct
       ~with_test:package_is_local
       filtered_formula
   ;;
+
+  let count_expanded_packages t = Table.fold t.expanded_packages ~init:0 ~f:( + )
 end
 
 module Solver = struct
@@ -1738,6 +1757,7 @@ module Solver_result = struct
     { lock_dir : Lock_dir.t
     ; files : File_entry.t Package_name.Map.Multi.t
     ; pinned_packages : Package_name.Set.t
+    ; num_expanded_packages : int
     }
 end
 
@@ -1965,5 +1985,10 @@ let solve_lock_dir
       >>| List.filter ~f:(fun (_, entries) -> List.is_non_empty entries)
       >>| Package_name.Map.of_list_exn
     in
-    Ok { Solver_result.lock_dir; files; pinned_packages = pinned_package_names }
+    Ok
+      { Solver_result.lock_dir
+      ; files
+      ; pinned_packages = pinned_package_names
+      ; num_expanded_packages = Context.count_expanded_packages context
+      }
 ;;

src/dune_pkg/opam_solver.mli

@@ -5,6 +5,7 @@ module Solver_result : sig
     { lock_dir : Lock_dir.t
     ; files : File_entry.t Package_name.Map.Multi.t
     ; pinned_packages : Package_name.Set.t
+    ; num_expanded_packages : int
     }
 end