Add thread_setup callback to tasking::Executor (#5581)

Signed-off-by: Michał Zientkiewicz <[email protected]>

Author: mzient

dali/core/exec/tasking_test.cc

@@ -35,6 +35,30 @@ TEST(TaskingTest, ExecutorShutdown) {
   });
 }
 
+TEST(TaskingTest, ExecutorSetup) {
+  /** Check that setup has effect on all threads */
+  static thread_local int tls = 0;
+  int num_threads = 32;
+  Executor ex(num_threads);
+  ex.Start([]() { tls = 42; });  // set thread-local value
+  std::atomic_int correct{0}, incorrect{0};
+  auto complete = Task::Create([](){});
+  int num_tasks = num_threads * 8;  // launch a lot of tasks - all taks must see the expected value
+  for (int i = 0; i < num_tasks; i++) {
+    auto task = Task::Create([&](){
+      if (tls == 42)
+        ++correct;
+      else
+        ++incorrect;
+    });
+    complete->Succeed(task);
+    ex.AddSilentTask(task);
+  }
+  ex.AddSilentTask(complete);
+  ex.Wait(complete);
+  EXPECT_EQ(correct, num_tasks);
+  EXPECT_EQ(incorrect, 0);
+}
 
 TEST(TaskingTest, IndependentTasksAreParallel) {
   int num_threads = 4;

include/dali/core/exec/tasking/executor.h

@@ -15,6 +15,7 @@
 #ifndef DALI_CORE_EXEC_TASKING_EXECUTOR_H_
 #define DALI_CORE_EXEC_TASKING_EXECUTOR_H_
 
+#include <functional>
 #include <memory>
 #include <thread>
 #include <vector>
@@ -38,13 +39,19 @@ class Executor : public Scheduler {
   /** Launches the worker threads.
    *
    * Multiple calls to Start have no effect. The function is not thread safe.
+   *
+   * @param thread_setup A function executed before the main loop.
    */
-  void Start() {
+  void Start(std::function<void()> thread_setup = {}) {
     if (started_)
       return;
     assert(workers_.empty());
     for (int i = 0; i < num_threads_; i++)
-      workers_.emplace_back([this]() { RunWorker(); });
+      workers_.emplace_back([thread_setup, this]() {
+        if (thread_setup)
+          thread_setup();
+        RunWorker();
+      });
     started_ = true;
   }