Refactor Allocator to make it thread safety (strategy) aware

src/umpire/Allocator.cpp

@@ -9,23 +9,37 @@
 #include "umpire/ResourceManager.hpp"
 #include "umpire/strategy/ThreadSafeAllocator.hpp"
 #include "umpire/util/Macros.hpp"
+#include <mutex>
 
 namespace umpire {
 
 Allocator::Allocator(strategy::AllocationStrategy* allocator) noexcept
     : strategy::mixins::Inspector{},
       strategy::mixins::AllocateNull{},
-      m_allocator{allocator},
       m_tracking{allocator->isTracked()}
 {
-  // Hack: If the strategy for this allocator requires thread safety,
-  // we create a mutex to be used during allocation operations
-  //
-  if (dynamic_cast<umpire::strategy::ThreadSafeAllocator*>(allocator) != nullptr) {
-    m_mutex = std::shared_ptr<std::mutex>(new std::mutex);
-  } else {
-    m_mutex = nullptr;
-  }
+  m_threadsafe = ( dynamic_cast<umpire::strategy::ThreadSafeAllocator*>(allocator) != nullptr );
+}
+
+void* Allocator::thread_safe_allocate(std::size_t bytes)
+{
+  umpire::strategy::ThreadSafeAllocator* alloc{dynamic_cast<umpire::strategy::ThreadSafeAllocator*>(m_allocator)};
+  std::lock_guard<std::mutex> lock(alloc->get_mutex());
+  return do_allocate(bytes);
+}
+
+void* Allocator::thread_safe_named_allocate(const std::string& name, std::size_t bytes)
+{
+  umpire::strategy::ThreadSafeAllocator* alloc{dynamic_cast<umpire::strategy::ThreadSafeAllocator*>(m_allocator)};
+  std::lock_guard<std::mutex> lock(alloc->get_mutex());
+  return do_named_allocate(name, bytes);
+}
+
+void Allocator::thread_safe_deallocate(void* ptr)
+{
+  umpire::strategy::ThreadSafeAllocator* alloc{dynamic_cast<umpire::strategy::ThreadSafeAllocator*>(m_allocator)};
+  std::lock_guard<std::mutex> lock(alloc->get_mutex());
+  return do_deallocate(ptr);
 }
 
 void Allocator::release()

src/umpire/Allocator.hpp

@@ -9,7 +9,6 @@
 
 #include <cstddef>
 #include <memory>
-#include <mutex>
 #include <ostream>
 #include <string>
 
@@ -194,18 +193,21 @@ class Allocator : private strategy::mixins::Inspector, strategy::mixins::Allocat
    */
   Allocator(strategy::AllocationStrategy* allocator) noexcept;
 
+  void* thread_safe_allocate(std::size_t bytes);
+  void* thread_safe_named_allocate(const std::string& name, std::size_t bytes);
+  void thread_safe_deallocate(void* ptr);
+
+  inline void* do_allocate(std::size_t bytes);
+  inline void* do_named_allocate(const std::string& name, std::size_t bytes);
+  inline void do_deallocate(void* ptr);
+
   /*!
    * \brief Pointer to the AllocationStrategy used by this Allocator.
    */
   umpire::strategy::AllocationStrategy* m_allocator;
 
   bool m_tracking{true};
-
-  /*!
-   * \brief Mutex to be used for AllocationStrategys that
-   * require thread safety
-   */
-  std::shared_ptr<std::mutex> m_mutex;
+  bool m_threadsafe{false};
 };
 
 inline std::string to_string(const Allocator& a)

src/umpire/Allocator.inl

@@ -8,6 +8,7 @@
 #define UMPIRE_Allocator_INL
 
 #include "umpire/Allocator.hpp"
+#include "umpire/strategy/ThreadSafeAllocator.hpp"
 #include "umpire/config.hpp"
 #include "umpire/event/event.hpp"
 #include "umpire/event/recorder_factory.hpp"
@@ -16,17 +17,14 @@
 
 namespace umpire {
 
-inline void* Allocator::allocate(std::size_t bytes)
+inline void* Allocator::do_allocate(std::size_t bytes)
 {
   void* ret = nullptr;
 
   UMPIRE_ASSERT(UMPIRE_VERSION_OK());
 
   UMPIRE_LOG(Debug, "(" << bytes << ")");
 
-  if (m_mutex != nullptr)
-    m_mutex->lock();
-
   if (0 == bytes) {
     ret = allocateNull();
   } else {
@@ -35,8 +33,6 @@ inline void* Allocator::allocate(std::size_t bytes)
     } catch (umpire::out_of_memory_error& e) {
       e.set_allocator_id(this->getId());
       e.set_requested_size(bytes);
-      if (m_mutex != nullptr)
-        m_mutex->unlock();
       throw;
     }
   }
@@ -45,26 +41,20 @@ inline void* Allocator::allocate(std::size_t bytes)
     registerAllocation(ret, bytes, m_allocator);
   }
 
-  if (m_mutex != nullptr)
-    m_mutex->unlock();
-
   umpire::event::record<umpire::event::allocate>(
       [&](auto& event) { event.size(bytes).ref((void*)m_allocator).ptr(ret); });
 
   return ret;
 }
 
-inline void* Allocator::allocate(const std::string& name, std::size_t bytes)
+inline void* Allocator::do_named_allocate(const std::string& name, std::size_t bytes)
 {
   void* ret = nullptr;
 
   UMPIRE_ASSERT(UMPIRE_VERSION_OK());
 
   UMPIRE_LOG(Debug, "(" << bytes << ")");
 
-  if (m_mutex != nullptr)
-    m_mutex->lock();
-
   if (0 == bytes) {
     ret = allocateNull();
   } else {
@@ -75,15 +65,12 @@ inline void* Allocator::allocate(const std::string& name, std::size_t bytes)
     registerAllocation(ret, bytes, m_allocator, name);
   }
 
-  if (m_mutex != nullptr)
-    m_mutex->unlock();
-
   umpire::event::record<umpire::event::named_allocate>(
       [&](auto& event) { event.name(name).size(bytes).ref((void*)m_allocator).ptr(ret); });
   return ret;
 }
 
-inline void Allocator::deallocate(void* ptr)
+inline void Allocator::do_deallocate(void* ptr)
 {
   umpire::event::record<umpire::event::deallocate>([&](auto& event) { event.ref((void*)m_allocator).ptr(ptr); });
 
@@ -93,8 +80,6 @@ inline void Allocator::deallocate(void* ptr)
     UMPIRE_LOG(Info, "Deallocating a null pointer (This behavior is intentionally allowed and ignored)");
     return;
   } else {
-    if (m_mutex != nullptr)
-      m_mutex->lock();
     if (m_tracking) {
       auto record = deregisterAllocation(ptr, m_allocator);
       if (!deallocateNull(ptr)) {
@@ -105,8 +90,33 @@ inline void Allocator::deallocate(void* ptr)
         m_allocator->deallocate(ptr);
       }
     }
-    if (m_mutex != nullptr)
-      m_mutex->unlock();
+  }
+}
+
+inline void* Allocator::allocate(std::size_t bytes)
+{
+  if (m_threadsafe) {
+    return thread_safe_allocate(bytes);
+  } else {
+    return do_allocate(bytes);
+  }
+}
+
+inline void* Allocator::allocate(const std::string& name, std::size_t bytes)
+{
+  if (m_threadsafe) {
+    return thread_safe_named_allocate(name, bytes);
+  } else {
+    return do_named_allocate(name, bytes);
+  }
+}
+
+inline void Allocator::deallocate(void* ptr)
+{
+  if (m_threadsafe) {
+    thread_safe_deallocate(ptr);
+  } else {
+    do_deallocate(ptr);
   }
 }
 

src/umpire/strategy/ThreadSafeAllocator.cpp

@@ -22,14 +22,12 @@ ThreadSafeAllocator::ThreadSafeAllocator(const std::string& name, int id, Alloca
 
 void* ThreadSafeAllocator::allocate(std::size_t bytes)
 {
-  std::lock_guard<std::mutex> lock(m_mutex);
   void* ret = m_allocator->allocate_internal(bytes);
   return ret;
 }
 
 void ThreadSafeAllocator::deallocate(void* ptr, std::size_t size)
 {
-  std::lock_guard<std::mutex> lock(m_mutex);
   m_allocator->deallocate_internal(ptr, size);
 }
 
@@ -43,5 +41,10 @@ MemoryResourceTraits ThreadSafeAllocator::getTraits() const noexcept
   return m_allocator->getTraits();
 }
 
+std::mutex& ThreadSafeAllocator::get_mutex()
+{
+  return m_mutex;
+}
+
 } // end of namespace strategy
 } // end of namespace umpire

src/umpire/strategy/ThreadSafeAllocator.hpp

@@ -33,6 +33,8 @@ class ThreadSafeAllocator : public AllocationStrategy {
 
   MemoryResourceTraits getTraits() const noexcept override;
 
+  std::mutex& get_mutex();
+
  protected:
   strategy::AllocationStrategy* m_allocator;