WeakValueHashMap: add sharded storage to reduce mutex contention

Author: TheMostDiligent

Common/interface/WeakValueHashMap.hpp

@@ -32,6 +32,7 @@
 #include <unordered_map>
 #include <memory>
 #include <mutex>
+#include <vector>
 
 #include "../../Platforms/Basic/interface/DebugUtilities.hpp"
 
@@ -70,7 +71,14 @@ class WeakValueHashMap
     class Impl;
 
 public:
-    WeakValueHashMap() = default;
+    explicit WeakValueHashMap(size_t NumShards = 1) :
+        m_pImpl(NumShards != 0 ? NumShards : 1)
+    {
+        for (std::shared_ptr<Impl>& pImpl : m_pImpl)
+        {
+            pImpl = std::make_shared<Impl>();
+        }
+    }
 
     // clang-format off
     WeakValueHashMap           (const WeakValueHashMap&) = delete;
@@ -136,16 +144,23 @@ class WeakValueHashMap
 
     ValueHandle Get(const KeyType& Key) const
     {
-        return m_pImpl->Get(Key);
+        const size_t ShardIdx = GetShardIndex(Key);
+        return m_pImpl[ShardIdx]->Get(Key);
     }
 
     template <typename... ArgsType>
     ValueHandle GetOrInsert(const KeyType& Key, ArgsType&&... Args) const
     {
-        return m_pImpl->GetOrInsert(Key, std::forward<ArgsType>(Args)...);
+        const size_t ShardIdx = GetShardIndex(Key);
+        return m_pImpl[ShardIdx]->GetOrInsert(Key, std::forward<ArgsType>(Args)...);
     }
 
 private:
+    size_t GetShardIndex(const KeyType& Key) const
+    {
+        return m_pImpl.size() > 1 ? m_Hasher(Key) % m_pImpl.size() : 0;
+    }
+
     class Impl : public std::enable_shared_from_this<Impl>
     {
     public:
@@ -184,7 +199,8 @@ class WeakValueHashMap
             // Create shared_ptr with deleter that erases the Key from the map when
             // the last reference to the value is destroyed
             std::shared_ptr<ValueType> pNewValue{
-                new ValueType{std::forward<ArgsType>(Args)...},
+                // Use parentheses to avoid initializer_list surprises.
+                new ValueType(std::forward<ArgsType>(Args)...),
                 [WeakSelf = std::move(WeakSelf), Key](ValueType* pValue) {
                     // Get a pointer before deleting the value as the value keeps the map alive.
                     std::shared_ptr<Impl> Self = WeakSelf.lock();
@@ -249,7 +265,8 @@ class WeakValueHashMap
         std::mutex                                                           m_Mtx;
         std::unordered_map<KeyType, std::weak_ptr<ValueType>, Hasher, Keyeq> m_Map;
     };
-    std::shared_ptr<Impl> m_pImpl = std::make_shared<Impl>();
+    std::vector<std::shared_ptr<Impl>> m_pImpl;
+    Hasher                             m_Hasher;
 };
 
 } // namespace Diligent

Tests/DiligentCoreTest/src/Common/WeakValueHashMapTest.cpp

@@ -58,7 +58,7 @@ TEST(Common_WeakValueHashMap, GetOrInsert)
         auto Handle1 = Map.GetOrInsert(1, "Value");
 
         // Release map while the handle is still alive
-        Map = {};
+        Map = WeakValueHashMap<int, std::string>{};
 
         EXPECT_TRUE(Handle1);
         EXPECT_STREQ(Handle1->c_str(), "Value");
@@ -71,7 +71,7 @@ TEST(Common_WeakValueHashMap, GetOrInsert)
         auto Handle1 = Map.GetOrInsert(1, "Value");
 
         // Release map while the handle is still alive
-        Map = {};
+        Map = WeakValueHashMap<int, std::string>{};
 
         WeakValueHashMap<int, std::string>::ValueHandle Handle2{std::move(Handle1)};
         EXPECT_FALSE(Handle1);
@@ -86,7 +86,7 @@ TEST(Common_WeakValueHashMap, GetOrInsert)
         auto Handle1 = Map.GetOrInsert(1, "Value");
 
         // Release map while the handle is still alive
-        Map = {};
+        Map = WeakValueHashMap<int, std::string>{};
 
         WeakValueHashMap<int, std::string>::ValueHandle Handle2;
         Handle2 = std::move(Handle1);
@@ -207,72 +207,78 @@ static constexpr int kNumParallelKeys = 16384;
 // Test that multiple threads can concurrently get or insert values into the map
 TEST(Common_WeakValueHashMap, ParallelGetOrInsert1)
 {
-    std::vector<std::thread> Threads(kNumThreads);
-
-    Threading::Signal StartSignal;
-
-    WeakValueHashMap<int, std::string> Map;
-    for (size_t t = 0; t < kNumThreads; ++t)
+    for (size_t NumShards : {1, 2, 4})
     {
-        Threads[t] = std::thread{
-            [&Map, &StartSignal]() //
-            {
-                StartSignal.Wait(true, kNumThreads);
+        std::vector<std::thread> Threads(kNumThreads);
 
-                for (int k = 0; k < kNumParallelKeys; ++k)
+        Threading::Signal StartSignal;
+
+        WeakValueHashMap<int, std::string> Map{NumShards};
+        for (size_t t = 0; t < kNumThreads; ++t)
+        {
+            Threads[t] = std::thread{
+                [&Map, &StartSignal]() //
                 {
-                    std::string Value = "Value" + std::to_string(k);
+                    StartSignal.Wait(true, kNumThreads);
 
-                    auto Handle = Map.GetOrInsert(k, Value);
-                    EXPECT_TRUE(Handle);
-                    EXPECT_EQ(*Handle, Value);
-                }
-            }};
-    }
+                    for (int k = 0; k < kNumParallelKeys; ++k)
+                    {
+                        std::string Value = "Value" + std::to_string(k);
 
-    StartSignal.Trigger(true);
-    for (auto& Thread : Threads)
-    {
-        Thread.join();
+                        auto Handle = Map.GetOrInsert(k, Value);
+                        EXPECT_TRUE(Handle);
+                        EXPECT_EQ(*Handle, Value);
+                    }
+                }};
+        }
+
+        StartSignal.Trigger(true);
+        for (auto& Thread : Threads)
+        {
+            Thread.join();
+        }
     }
 }
 
 // Similar to the previous test, but all values are kept alive
 TEST(Common_WeakValueHashMap, ParallelGetOrInsert2)
 {
-    std::vector<std::thread> Threads(kNumThreads);
-
-    Threading::Signal StartSignal;
+    for (size_t NumShards : {1, 2, 4})
+    {
+        std::vector<std::thread> Threads(kNumThreads);
 
-    std::vector<WeakValueHashMap<int, std::string>::ValueHandle> Handles(kNumThreads * kNumParallelKeys);
+        Threading::Signal StartSignal;
 
-    WeakValueHashMap<int, std::string> Map;
-    for (size_t t = 0; t < kNumThreads; ++t)
-    {
-        Threads[t] = std::thread{
-            [&Map, &StartSignal, &Handles](size_t ThreadId) //
-            {
-                StartSignal.Wait(true, kNumThreads);
+        std::vector<WeakValueHashMap<int, std::string>::ValueHandle> Handles(kNumThreads * kNumParallelKeys);
 
-                for (int k = 0; k < kNumParallelKeys; ++k)
+        WeakValueHashMap<int, std::string> Map{NumShards};
+        for (size_t t = 0; t < kNumThreads; ++t)
+        {
+            Threads[t] = std::thread{
+                [&Map, &StartSignal, &Handles](size_t ThreadId) //
                 {
-                    std::string Value = "Value" + std::to_string(k);
+                    StartSignal.Wait(true, kNumThreads);
 
-                    auto Handle = Map.GetOrInsert(k, Value);
-                    EXPECT_TRUE(Handle);
-                    EXPECT_EQ(*Handle, Value);
+                    for (int k = 0; k < kNumParallelKeys; ++k)
+                    {
+                        std::string Value = "Value" + std::to_string(k);
 
-                    Handles[ThreadId * kNumParallelKeys + k] = std::move(Handle);
-                }
-            },
-            t,
-        };
-    }
+                        auto Handle = Map.GetOrInsert(k, Value);
+                        EXPECT_TRUE(Handle);
+                        EXPECT_EQ(*Handle, Value);
 
-    StartSignal.Trigger(true);
-    for (auto& Thread : Threads)
-    {
-        Thread.join();
+                        Handles[ThreadId * kNumParallelKeys + k] = std::move(Handle);
+                    }
+                },
+                t,
+            };
+        }
+
+        StartSignal.Trigger(true);
+        for (auto& Thread : Threads)
+        {
+            Thread.join();
+        }
     }
 }