fix-race

Signed-off-by: Benoit Jacob <jacob.benoit.1@gmail.com>
nod-ai · Jan 20, 2025 · 0675858 · 0675858
1 parent 1f50538
commit 0675858
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/shortfin/CMakeLists.txt b/shortfin/CMakeLists.txt
@@ -118,11 +118,16 @@ option(SHORTFIN_ENABLE_ASAN "Enable ASAN" OFF)
 if(SHORTFIN_ENABLE_ASAN)
   add_compile_options(-fsanitize=address)
   add_link_options(-fsanitize=address)
-
-  # Enable more ASAN checks.
   add_compile_definitions(IREE_SANITIZER_ADDRESS)
 endif()
 
+option(SHORTFIN_ENABLE_TSAN "Enable TSAN" OFF)
+if(SHORTFIN_ENABLE_TSAN)
+  add_compile_options(-fsanitize=thread)
+  add_link_options(-fsanitize=thread)
+  add_compile_definitions(IREE_SANITIZER_THREAD)
+endif()
+
 # Thread safety annotations: Enabled if the compiler supports it.
 check_cxx_compiler_flag("-Wthread-safety" SHORTFIN_HAS_THREAD_SAFETY_ANNOTATIONS)
 if(SHORTFIN_HAS_THREAD_SAFETY)

diff --git a/shortfin/src/shortfin/support/blocking_executor.cc b/shortfin/src/shortfin/support/blocking_executor.cc
@@ -157,7 +157,11 @@ int BlockingExecutor::ThreadInstance::RunOnThread() noexcept {
     auto status = iree_wait_source_wait_one(signal_transact.await(),
                                             iree_infinite_timeout());
     IREE_CHECK_OK(status);
-    Task exec_task = std::move(current_task);
+    Task exec_task;
+    {
+      iree::slim_mutex_lock_guard g(executor->control_mu_);
+      exec_task = std::move(current_task);
+    }
     if (exec_task) {
       exec_task();
     }