chore: Task.get block profiling (#7016)

* `--profile` now reports `blocking` time spent in `Task.get` inside
other profiling categories
* environment variable `LEAN_TRACE_TASK_GET_BLOCKED` when set makes
`lean` dump stack traces of `Task.get` blocks

src/library/time_task.cpp

@@ -15,10 +15,17 @@ static std::map<std::string, second_duration> * g_cum_times;
 static mutex * g_cum_times_mutex;
 LEAN_THREAD_PTR(time_task, g_current_time_task);
 
+bool has_profiling_task() {
+    return g_current_time_task != nullptr;
+}
 void report_profiling_time(std::string const & category, second_duration time) {
     lock_guard<mutex> _(*g_cum_times_mutex);
     (*g_cum_times)[category] += time;
 }
+void exclude_profiling_time_from_current_task(second_duration time) {
+    if (g_current_time_task)
+        g_current_time_task->exclude_duration(time);
+}
 
 void display_cumulative_profiling_times(std::ostream & out) {
     if (g_cum_times->empty())

src/library/time_task.h

@@ -11,8 +11,10 @@ Author: Sebastian Ullrich
 #include "util/message_definitions.h"
 
 namespace lean {
+LEAN_EXPORT bool has_profiling_task();
 LEAN_EXPORT void report_profiling_time(std::string const & category, second_duration time);
 LEAN_EXPORT void display_cumulative_profiling_times(std::ostream & out);
+LEAN_EXPORT void exclude_profiling_time_from_current_task(second_duration time);
 
 /** Measure time of some task and report it for the final cumulative profile. */
 class LEAN_EXPORT time_task {
@@ -22,6 +24,7 @@ class LEAN_EXPORT time_task {
 public:
     time_task(std::string const & category, options const & opts, name decl = name());
     ~time_task();
+    void exclude_duration(second_duration duration) { m_timeit->exclude_duration(duration); }
 };
 
 void initialize_time_task();

src/runtime/object.cpp

@@ -1063,10 +1063,21 @@ extern "C" LEAN_EXPORT obj_res lean_task_map_core(obj_arg f, obj_arg t, unsigned
     }
 }
 
+// We don't use `time_task` here as it's outside runtime/, and we wouldn't have access to `options`
+// anyway
+LEAN_EXPORT void (*g_lean_report_task_get_blocked_time)(std::chrono::nanoseconds) = nullptr;
+
 extern "C" LEAN_EXPORT b_obj_res lean_task_get(b_obj_arg t) {
     if (object * v = lean_to_task(t)->m_value)
         return v;
-    g_task_manager->wait_for(lean_to_task(t));
+    if (g_lean_report_task_get_blocked_time) {
+        std::chrono::steady_clock::time_point start = std::chrono::steady_clock::now();
+        g_task_manager->wait_for(lean_to_task(t));
+        std::chrono::steady_clock::time_point end = std::chrono::steady_clock::now();
+        g_lean_report_task_get_blocked_time(std::chrono::nanoseconds(end - start));
+    } else {
+        g_task_manager->wait_for(lean_to_task(t));
+    }
     lean_assert(lean_to_task(t)->m_value != nullptr);
     object * r = lean_to_task(t)->m_value;
     return r;

src/util/shell.cpp

@@ -434,6 +434,24 @@ void check_optarg(char const * option_name) {
 
 extern "C" object * lean_enable_initializer_execution(object * w);
 
+namespace lean {
+extern void (*g_lean_report_task_get_blocked_time)(std::chrono::nanoseconds);
+}
+static bool trace_task_get_blocked = getenv("LEAN_TRACE_TASK_GET_BLOCKED") != nullptr;
+static void report_task_get_blocked_time(std::chrono::nanoseconds d) {
+    if (has_profiling_task()) {
+        report_profiling_time("blocked", d);
+        exclude_profiling_time_from_current_task(d);
+        if (trace_task_get_blocked) {
+            sstream ss;
+            ss << "Task.get blocked for " << std::chrono::duration_cast<std::chrono::seconds>(d).count() << "s";
+            // using a panic for reporting is a bit of a hack, but good enough for this
+            // `lean`-specific use case
+            lean_panic(ss.str().c_str(), /* force stderr */ true);
+        }
+    }
+}
+
 extern "C" LEAN_EXPORT int lean_main(int argc, char ** argv) {
 #ifdef LEAN_EMSCRIPTEN
     // When running in command-line mode under Node.js, we make system directories available in the virtual filesystem.
@@ -651,6 +669,7 @@ extern "C" LEAN_EXPORT int lean_main(int argc, char ** argv) {
     }
 
     if (get_profiler(opts)) {
+        g_lean_report_task_get_blocked_time = report_task_get_blocked_time;
         report_profiling_time("initialization", init_time);
     }