From 214093e6c4a3409aed82f16cebbef9ec8a048333 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 23 Jan 2025 16:01:39 -0700
Subject: [PATCH] fix: prevent `Task.get` deadlocks from threadpool starvation
 (#6758)

This PR prevents deadlocks from non-cyclical task waits that may
otherwise occur during parallel elaboration with small threadpool sizes.
---
 CMakePresets.json      |  2 +-
 src/Init/Core.lean     | 13 +++++++++++--
 src/runtime/object.cpp | 24 ++++++++++++++++++++----
 3 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/CMakePresets.json b/CMakePresets.json
index a7523ac7f2..19060b1c73 100644
--- a/CMakePresets.json
+++ b/CMakePresets.json
@@ -26,7 +26,7 @@
       "displayName": "Sanitize build config",
       "cacheVariables": {
         "LEAN_EXTRA_CXX_FLAGS": "-fsanitize=address,undefined",
-        "LEANC_EXTRA_FLAGS": "-fsanitize=address,undefined -fsanitize-link-c++-runtime",
+        "LEANC_EXTRA_CC_FLAGS": "-fsanitize=address,undefined -fsanitize-link-c++-runtime",
         "SMALL_ALLOCATOR": "OFF",
         "BSYMBOLIC": "OFF"
       },
diff --git a/src/Init/Core.lean b/src/Init/Core.lean
index 814c4dbb3b..be79514fef 100644
--- a/src/Init/Core.lean
+++ b/src/Init/Core.lean
@@ -516,8 +516,17 @@ The tasks have an overridden representation in the runtime.
 structure Task (α : Type u) : Type u where
   /-- `Task.pure (a : α)` constructs a task that is already resolved with value `a`. -/
   pure ::
-  /-- If `task : Task α` then `task.get : α` blocks the current thread until the
-  value is available, and then returns the result of the task. -/
+  /--
+  Blocks the current thread until the given task has finished execution, and then returns the result
+  of the task. If the current thread is itself executing a (non-dedicated) task, the maximum
+  threadpool size is temporarily increased by one while waiting so as to ensure the process cannot
+  be deadlocked by threadpool starvation. Note that when the current thread is unblocked, more tasks
+  than the configured threadpool size may temporarily be running at the same time until sufficiently
+  many tasks have finished.
+
+  `Task.map` and `Task.bind` should be preferred over `Task.get` for setting up task dependencies
+  where possible as they do not require temporarily growing the threadpool in this way.
+  -/
   get : α
   deriving Inhabited, Nonempty
 
diff --git a/src/runtime/object.cpp b/src/runtime/object.cpp
index 654e1b59f2..2bb6d81b63 100644
--- a/src/runtime/object.cpp
+++ b/src/runtime/object.cpp
@@ -693,10 +693,14 @@ class task_manager {
             unique_lock<mutex> lock(m_mutex);
             m_idle_std_workers++;
             while (true) {
-                if (m_queues_size == 0) {
-                    if (m_shutting_down) {
-                        break;
-                    }
+                if (m_queues_size == 0 && m_shutting_down) {
+                    break;
+                }
+                if (m_queues_size == 0 ||
+                        // If we have reached the maximum number of standard workers (because the
+                        // maximum was decreased by `task_get`), wait for someone else to become
+                        // idle before picking up new work.
+                        m_std_workers.size() - m_idle_std_workers >= m_max_std_workers) {
                     m_queue_cv.wait(lock);
                     continue;
                 }
@@ -859,7 +863,19 @@ public:
         unique_lock<mutex> lock(m_mutex);
         if (t->m_value)
             return;
+        // see `Task.get`
+        bool in_pool = g_current_task_object && g_current_task_object->m_imp->m_prio <= LEAN_MAX_PRIO;
+        if (in_pool) {
+            m_max_std_workers++;
+            if (m_idle_std_workers == 0)
+                spawn_worker();
+            else
+                m_queue_cv.notify_one();
+        }
         m_task_finished_cv.wait(lock, [&]() { return t->m_value != nullptr; });
+        if (in_pool) {
+            m_max_std_workers--;
+        }
     }
 
     object * wait_any(object * task_list) {