Fix TTbbLocalExecutor::GetWorkerThreadId. Enable test_fit_on_scipy_sparse_spmatrix on Windows.

commit_hash:c52cdc5529aff5cda1cbd11be2852647736ceb49

2 files changed, 23 insertions, 4 deletions

diff --git a/library/cpp/threading/local_executor/tbb_local_executor.cpp b/library/cpp/threading/local_executor/tbb_local_executor.cpp
index 65d66594438..91a8460b0eb 100644
--- a/library/cpp/threading/local_executor/tbb_local_executor.cpp
+++ b/library/cpp/threading/local_executor/tbb_local_executor.cpp
@@ -14,9 +14,21 @@ int NPar::TTbbLocalExecutor<RespectTls>::GetThreadCount() const noexcept {
 
 template <bool RespectTls>
 int NPar::TTbbLocalExecutor<RespectTls>::GetWorkerThreadId() const noexcept {
-    return TbbArena.execute([] {
-        return tbb::this_task_arena::current_thread_index();
-    });
+    static thread_local int WorkerThreadId = -1;
+    if (WorkerThreadId == -1) {
+        // Can't rely on return value except checking that it is 'not_initialized' because of
+        //  "Since a thread may exit the arena at any time if it does not execute a task, the index of
+        //   a thread may change between any two tasks"
+        //  (https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/task_scheduler/task_arena/this_task_arena_ns#_CPPv4N3tbb15this_task_arena20current_thread_indexEv)
+        const auto tbbThreadIndex = tbb::this_task_arena::current_thread_index();
+        if (tbbThreadIndex == tbb::task_arena::not_initialized) {
+            // This thread does not belong to TBB worker threads
+            WorkerThreadId = 0;
+        } else {
+            WorkerThreadId = ++RegisteredThreadCounter;
+        }
+    }
+    return WorkerThreadId;
 }
 
 template <bool RespectTls>
diff --git a/library/cpp/threading/local_executor/tbb_local_executor.h b/library/cpp/threading/local_executor/tbb_local_executor.h
index 751a37e2d47..f67c07349d4 100644
--- a/library/cpp/threading/local_executor/tbb_local_executor.h
+++ b/library/cpp/threading/local_executor/tbb_local_executor.h
@@ -9,6 +9,9 @@
 #include <contrib/libs/tbb/include/tbb/task_arena.h>
 #include <contrib/libs/tbb/include/tbb/task_group.h>
 
+#include <atomic>
+
+
 namespace NPar {
     template <bool RespectTls = false>
     class TTbbLocalExecutor final: public ILocalExecutor  {
@@ -16,7 +19,9 @@ namespace NPar {
         TTbbLocalExecutor(int nThreads)
             : ILocalExecutor()
             , TbbArena(nThreads)
-            , NumberOfTbbThreads(nThreads) {}
+            , NumberOfTbbThreads(nThreads)
+            , RegisteredThreadCounter(0)
+        {}
         ~TTbbLocalExecutor() noexcept override {}
 
         virtual int GetWorkerThreadId() const noexcept override;
@@ -44,5 +49,7 @@ namespace NPar {
         mutable tbb::task_arena TbbArena;
         tbb::task_group Group;
         int NumberOfTbbThreads;
+
+        mutable std::atomic_int RegisteredThreadCounter;
     };
 }