From 89d3e36b6d8ae6e6833c79176a89544225dedfa5 Mon Sep 17 00:00:00 2001
From: Rudi Grinberg <me@rgrinberg.com>
Date: Sun, 11 Jun 2023 15:53:42 +0100
Subject: [PATCH] feature: extend traces to include source scans

Source scans can be quite slow. Record them in the trace.

Signed-off-by: Rudi Grinberg <me@rgrinberg.com>

<!-- ps-id: 373fb4fc-304f-43c1-9fe2-c67175f3ee40 -->
---
 CHANGES.md                    |  3 ++
 bin/common.ml                 |  1 +
 src/dune_rules/source_tree.ml | 53 +++++++++++++++++++++++++++--------
 src/dune_stats/dune_stats.ml  |  9 ++++++
 src/dune_stats/dune_stats.mli |  4 +++
 5 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 3796a8f9d744..905102de57ba 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,6 +1,9 @@
 Unreleased
 ----------
 
+- Include source tree scans in the traces produced by `--trace-file` (#7937,
+  @rgrinberg)
+
 - Cinaps: The promotion rules for cinaps would only offer one file at a time no
   matter how many promotions were available. Now we offer all the promotions at
   once (#7901, @rgrinberg)
diff --git a/bin/common.ml b/bin/common.ml
index 33c6f6c5ffbf..8a0507125642 100644
--- a/bin/common.ml
+++ b/bin/common.ml
@@ -1175,6 +1175,7 @@ let build (builder : Builder.t) ~default_root_is_cwd =
             ~extended_build_job_info:builder.stats_trace_extended
             (Out (open_out f))
         in
+        Dune_stats.set_global stats;
         at_exit (fun () -> Dune_stats.close stats);
         stats)
   in
diff --git a/src/dune_rules/source_tree.ml b/src/dune_rules/source_tree.ml
index 89fba2a8f9e3..c2a0ae7ac7f8 100644
--- a/src/dune_rules/source_tree.ml
+++ b/src/dune_rules/source_tree.ml
@@ -671,18 +671,47 @@ module Dir = struct
   module Make_map_reduce (M : Memo.S) (Outcome : Monoid) = struct
     open M.O
 
-    let rec map_reduce t ~traverse ~f =
-      let must_traverse = Sub_dirs.Status.Map.find traverse t.status in
-      match must_traverse with
-      | false -> M.return Outcome.empty
-      | true ->
-        let+ here = f t
-        and+ in_sub_dirs =
-          M.List.map (Filename.Map.values t.contents.sub_dirs) ~f:(fun s ->
-              let* t = M.of_memo (sub_dir_as_t s) in
-              map_reduce t ~traverse ~f)
-        in
-        List.fold_left in_sub_dirs ~init:here ~f:Outcome.combine
+    let map_reduce =
+      let rec map_reduce t ~traverse ~f =
+        let must_traverse = Sub_dirs.Status.Map.find traverse t.status in
+        match must_traverse with
+        | false -> M.return Outcome.empty
+        | true ->
+          let+ here = f t
+          and+ in_sub_dirs =
+            M.List.map (Filename.Map.values t.contents.sub_dirs) ~f:(fun s ->
+                let* t = M.of_memo (sub_dir_as_t s) in
+                map_reduce t ~traverse ~f)
+          in
+          List.fold_left in_sub_dirs ~init:here ~f:Outcome.combine
+      in
+      let impl =
+        lazy
+          (match Dune_stats.global () with
+          | None -> map_reduce
+          | Some stats ->
+            fun t ~traverse ~f ->
+              let start = Unix.gettimeofday () in
+              let+ res = map_reduce t ~traverse ~f in
+              let event =
+                let stop = Unix.gettimeofday () in
+                let module Event = Chrome_trace.Event in
+                let module Timestamp = Event.Timestamp in
+                let dur = Timestamp.of_float_seconds (stop -. start) in
+                let common =
+                  Event.common_fields ~name:"Source tree scan"
+                    ~ts:(Timestamp.of_float_seconds start)
+                    ()
+                in
+                let args =
+                  [ ("dir", `String (Path.Source.to_string t.path)) ]
+                in
+                Event.complete common ~args ~dur
+              in
+              Dune_stats.emit stats event;
+              res)
+      in
+      fun t ~traverse ~f -> (Lazy.force impl) t ~traverse ~f
   end
 
   let cram_tests (t : t) =
diff --git a/src/dune_stats/dune_stats.ml b/src/dune_stats/dune_stats.ml
index 5d95d33037a8..dad8dcc567ce 100644
--- a/src/dune_stats/dune_stats.ml
+++ b/src/dune_stats/dune_stats.ml
@@ -118,6 +118,15 @@ let close { print; close; _ } =
   print "]\n";
   close ()
 
+let global = ref None
+
+let set_global t =
+  if Option.is_some !global then
+    Code_error.raise "global stats have been set" [];
+  global := Some t
+
+let global () = !global
+
 let create ~extended_build_job_info dst =
   let print =
     match dst with
diff --git a/src/dune_stats/dune_stats.mli b/src/dune_stats/dune_stats.mli
index 4d16d7237c48..4bc8347aabd2 100644
--- a/src/dune_stats/dune_stats.mli
+++ b/src/dune_stats/dune_stats.mli
@@ -12,6 +12,10 @@ type dst =
       ; flush : unit -> unit
       }
 
+val global : unit -> t option
+
+val set_global : t -> unit
+
 val create : extended_build_job_info:bool -> dst -> t
 
 val emit : t -> Chrome_trace.Event.t -> unit