PERF: Guard the ThreadPool singleton with an atomic instead of std::once_flag

Modules/Core/Common/src/itkThreadPool.cxx

@@ -32,6 +32,38 @@
 
 namespace itk
 {
+namespace
+{
+// C++20 std::atomic_flag is guaranteed lock-free and, since C++20, has a non-modifying test().
+// C++17 has no atomic_flag::test(), so std::atomic<bool> is used and asserted lock-free instead.
+// The feature-test macro is used rather than __cplusplus because MSVC misreports the latter.
+#if defined(__cpp_lib_atomic_flag_test) && (__cpp_lib_atomic_flag_test >= 201907L)
+using SingletonGuard = std::atomic_flag;
+inline bool
+IsGuardSet(const SingletonGuard & guard, std::memory_order order)
+{
+  return guard.test(order);
+}
+inline void
+SetGuard(SingletonGuard & guard)
+{
+  guard.test_and_set(std::memory_order_release);
+}
+#else
+using SingletonGuard = std::atomic<bool>;
+static_assert(SingletonGuard::is_always_lock_free, "The ThreadPool singleton guard must be lock-free.");
+inline bool
+IsGuardSet(const SingletonGuard & guard, std::memory_order order)
+{
+  return guard.load(order);
+}
+inline void
+SetGuard(SingletonGuard & guard)
+{
+  guard.store(true, std::memory_order_release);
+}
+#endif
+} // namespace
 
 struct ThreadPoolGlobals
 {
@@ -40,8 +72,8 @@ struct ThreadPoolGlobals
   // To lock on the various internal variables.
   std::mutex m_Mutex;
 
-  // To allow singleton creation of ThreadPool.
-  std::once_flag m_ThreadPoolOnceFlag;
+  // Guards singleton creation. Lives here so every ITK library instance shares one guard.
+  SingletonGuard m_IsSingletonCreated{};
 
   // The singleton instance of ThreadPool.
   ThreadPool::Pointer m_ThreadPoolInstance;
@@ -75,17 +107,24 @@ ThreadPool::GetInstance()
   // initialized.
   itkInitGlobalsMacro(PimplGlobals);
 
-  // Create a singleton ThreadPool.
-  std::call_once(m_PimplGlobals->m_ThreadPoolOnceFlag, []() {
-    m_PimplGlobals->m_ThreadPoolInstance = ObjectFactory<Self>::Create();
-    if (m_PimplGlobals->m_ThreadPoolInstance.IsNull())
+  // Acquire pairs with the release in SetGuard, so seeing the guard set implies seeing the pool.
+  if (!IsGuardSet(m_PimplGlobals->m_IsSingletonCreated, std::memory_order_acquire))
+  {
+    const std::lock_guard<std::mutex> lockGuard(m_PimplGlobals->m_Mutex);
+
+    if (!IsGuardSet(m_PimplGlobals->m_IsSingletonCreated, std::memory_order_relaxed))
     {
-      new ThreadPool(); // constructor sets m_PimplGlobals->m_ThreadPoolInstance
-    }
+      m_PimplGlobals->m_ThreadPoolInstance = ObjectFactory<Self>::Create();
+      if (m_PimplGlobals->m_ThreadPoolInstance.IsNull())
+      {
+        new ThreadPool(); // constructor sets m_PimplGlobals->m_ThreadPoolInstance
+      }
 #if defined(ITK_USE_PTHREADS)
-    pthread_atfork(ThreadPool::PrepareForFork, ThreadPool::ResumeFromFork, ThreadPool::ResumeFromFork);
+      pthread_atfork(ThreadPool::PrepareForFork, ThreadPool::ResumeFromFork, ThreadPool::ResumeFromFork);
 #endif
-  });
+      SetGuard(m_PimplGlobals->m_IsSingletonCreated);
+    }
+  }
 
   return m_PimplGlobals->m_ThreadPoolInstance;
 }
@@ -106,8 +145,7 @@ ThreadPool::SetDoNotWaitForThreads(bool doNotWaitForThreads)
 
 ThreadPool::ThreadPool()
 {
-  // m_PimplGlobals->m_Mutex not needed to be acquired here because construction only occurs via GetInstance which is
-  // protected by call_once.
+  // m_PimplGlobals->m_Mutex is already held by GetInstance while this constructor runs.
 
   m_PimplGlobals->m_ThreadPoolInstance = this;        // threads need this
   m_PimplGlobals->m_ThreadPoolInstance->UnRegister(); // Remove extra reference

Modules/Core/Common/test/CMakeLists.txt

@@ -117,6 +117,10 @@ if(ITK_BUILD_SHARED_LIBS AND ITK_DYNAMIC_LOADING)
   list(APPEND ITKCommon2Tests itkDownCastTest.cxx)
 endif()
 
+if(NOT ITK_BUILD_SHARED_LIBS AND ITK_DYNAMIC_LOADING)
+  list(APPEND ITKCommon2Tests itkThreadPoolSingletonSharingTest.cxx)
+endif()
+
 # Note: Tests are arbitrarily split into two drivers to avoid memory issues.
 createtestdriver(
   ITKCommon1
@@ -1410,6 +1414,48 @@ if(ITK_BUILD_SHARED_LIBS AND ITK_DYNAMIC_LOADING)
   )
 endif()
 
+# Each library statically links its own copy of ITKCommon, as every wrapped Python module does.
+if(NOT ITK_BUILD_SHARED_LIBS AND ITK_DYNAMIC_LOADING)
+  foreach(_name IN ITEMS A B)
+    add_library(
+      ThreadPoolSingletonLibrary${_name}
+      MODULE
+      ThreadPoolSingletonLibrary${_name}.cxx
+    )
+    itk_module_target_label(ThreadPoolSingletonLibrary${_name})
+    target_link_libraries(
+      ThreadPoolSingletonLibrary${_name}
+      LINK_PUBLIC
+        ${ITKCommon-Test_LIBRARIES}
+    )
+    set_property(
+      TARGET
+        ThreadPoolSingletonLibrary${_name}
+      PROPERTY
+        LIBRARY_OUTPUT_DIRECTORY
+          ${ITK_TEST_OUTPUT_DIR}
+    )
+    set_target_properties(
+      ThreadPoolSingletonLibrary${_name}
+      PROPERTIES
+        CXX_VISIBILITY_PRESET
+          hidden
+        VISIBILITY_INLINES_HIDDEN
+          1
+    )
+    add_dependencies(ITKCommon2TestDriver ThreadPoolSingletonLibrary${_name})
+  endforeach()
+
+  itk_add_test(
+    NAME itkThreadPoolSingletonSharingTest
+    COMMAND
+      ITKCommon2TestDriver
+      itkThreadPoolSingletonSharingTest
+      $<TARGET_FILE:ThreadPoolSingletonLibraryA>
+      $<TARGET_FILE:ThreadPoolSingletonLibraryB>
+  )
+endif()
+
 set(
   ITKCommonGTests
   itkAbortProcessObjectGTest.cxx

Modules/Core/Common/test/ThreadPoolSingletonLibraryA.cxx

@@ -0,0 +1,35 @@
+/*=========================================================================
+ *
+ *  Copyright NumFOCUS
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *         https://www.apache.org/licenses/LICENSE-2.0.txt
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ *
+ *=========================================================================*/
+#include "itkSingleton.h"
+#include "itkThreadPool.h"
+
+// Each of these libraries links its own copy of a statically built ITKCommon, the way every
+// wrapped Python module does. They share only the SingletonIndex handed to them by the driver.
+extern "C" ITK_ABI_EXPORT void
+ThreadPoolSingletonLibraryA_AdoptSingletonIndex(itk::SingletonIndex * singletonIndex)
+{
+  itk::SingletonIndex::SetInstance(singletonIndex);
+}
+
+extern "C" ITK_ABI_EXPORT void *
+ThreadPoolSingletonLibraryA_GetThreadPool()
+{
+  const itk::ThreadPool::Pointer threadPool = itk::ThreadPool::GetInstance();
+  threadPool->Register(); // Pin it, so a duplicate singleton cannot reuse the freed address.
+  return threadPool.GetPointer();
+}

Modules/Core/Common/test/ThreadPoolSingletonLibraryB.cxx

@@ -0,0 +1,35 @@
+/*=========================================================================
+ *
+ *  Copyright NumFOCUS
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *         https://www.apache.org/licenses/LICENSE-2.0.txt
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ *
+ *=========================================================================*/
+#include "itkSingleton.h"
+#include "itkThreadPool.h"
+
+// Each of these libraries links its own copy of a statically built ITKCommon, the way every
+// wrapped Python module does. They share only the SingletonIndex handed to them by the driver.
+extern "C" ITK_ABI_EXPORT void
+ThreadPoolSingletonLibraryB_AdoptSingletonIndex(itk::SingletonIndex * singletonIndex)
+{
+  itk::SingletonIndex::SetInstance(singletonIndex);
+}
+
+extern "C" ITK_ABI_EXPORT void *
+ThreadPoolSingletonLibraryB_GetThreadPool()
+{
+  const itk::ThreadPool::Pointer threadPool = itk::ThreadPool::GetInstance();
+  threadPool->Register(); // Pin it, so a duplicate singleton cannot reuse the freed address.
+  return threadPool.GetPointer();
+}

Modules/Core/Common/test/itkThreadPoolSingletonSharingTest.cxx

@@ -0,0 +1,77 @@
+/*=========================================================================
+ *
+ *  Copyright NumFOCUS
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *         https://www.apache.org/licenses/LICENSE-2.0.txt
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ *
+ *=========================================================================*/
+#include "itkSingleton.h"
+#include "itkThreadPool.h"
+#include "itksys/DynamicLoader.hxx"
+#include <iostream>
+
+// Two dynamically loaded libraries, each with its own statically linked copy of ITKCommon,
+// sharing one SingletonIndex, must observe one and the same ThreadPool.
+int
+itkThreadPoolSingletonSharingTest(int argc, char * argv[])
+{
+  if (argc < 3)
+  {
+    std::cerr << "Usage: " << argv[0] << " <libraryA> <libraryB>" << std::endl;
+    return EXIT_FAILURE;
+  }
+
+  using AdoptFunction = void (*)(itk::SingletonIndex *);
+  using GetPoolFunction = void * (*)();
+
+  void * pools[2] = { nullptr, nullptr };
+
+  for (int i = 0; i < 2; ++i)
+  {
+    const char * const libraryName = argv[i + 1];
+
+    itksys::DynamicLoader::LibraryHandle handle = itksys::DynamicLoader::OpenLibrary(libraryName);
+    if (handle == nullptr)
+    {
+      std::cerr << "Failed to load " << libraryName << ": " << itksys::DynamicLoader::LastError() << std::endl;
+      return EXIT_FAILURE;
+    }
+
+    const std::string suffix = (i == 0) ? "A" : "B";
+    const auto        adopt = reinterpret_cast<AdoptFunction>(
+      itksys::DynamicLoader::GetSymbolAddress(handle, "ThreadPoolSingletonLibrary" + suffix + "_AdoptSingletonIndex"));
+    const auto getPool = reinterpret_cast<GetPoolFunction>(
+      itksys::DynamicLoader::GetSymbolAddress(handle, "ThreadPoolSingletonLibrary" + suffix + "_GetThreadPool"));
+
+    if (adopt == nullptr || getPool == nullptr)
+    {
+      std::cerr << "Failed to resolve symbols in " << libraryName << std::endl;
+      return EXIT_FAILURE;
+    }
+
+    adopt(itk::SingletonIndex::GetInstance());
+    pools[i] = getPool();
+    std::cout << "Library " << suffix << " ThreadPool: " << pools[i] << std::endl;
+  }
+
+  if (pools[0] != pools[1])
+  {
+    std::cerr << "FAILED: each library created its own ThreadPool. The singleton guard must live in "
+                 "ThreadPoolGlobals, which is shared through the SingletonIndex, not in per-binary storage."
+              << std::endl;
+    return EXIT_FAILURE;
+  }
+
+  std::cout << "Test finished." << std::endl;
+  return EXIT_SUCCESS;
+}