fix jupiterbuildstorage concurrency (#906)

- Bugfix: Fixed concurrency issue in JupiterBuildStorage when updating stats

1 files changed, 34 insertions, 15 deletions

diff --git a/src/zenremotestore/builds/jupiterbuildstorage.cpp b/src/zenremotestore/builds/jupiterbuildstorage.cpp
index ad4c4bc89..fa1a69f18 100644
--- a/src/zenremotestore/builds/jupiterbuildstorage.cpp
+++ b/src/zenremotestore/builds/jupiterbuildstorage.cpp
@@ -444,11 +444,13 @@ public:
 
 	virtual bool GetExtendedStatistics(ExtendedStatistics& OutStats) override
 	{
-		OutStats.ReceivedBytesPerSource.reserve(m_ReceivedBytesPerSource.size());
-		for (auto& It : m_ReceivedBytesPerSource)
-		{
-			OutStats.ReceivedBytesPerSource.insert_or_assign(It.first, m_SourceBytes[It.second]);
-		}
+		m_SourceLock.WithSharedLock([this, &OutStats]() {
+			OutStats.ReceivedBytesPerSource.reserve(m_ReceivedBytesPerSource.size());
+			for (auto& It : m_ReceivedBytesPerSource)
+			{
+				OutStats.ReceivedBytesPerSource.insert_or_assign(It.first, m_SourceBytes[It.second].load(std::memory_order_relaxed));
+			}
+		});
 		return true;
 	}
 
@@ -521,15 +523,29 @@ private:
 		}
 		if (!Result.Source.empty())
 		{
-			if (tsl::robin_map<std::string, uint32_t>::const_iterator It = m_ReceivedBytesPerSource.find(Result.Source);
-				It != m_ReceivedBytesPerSource.end())
-			{
-				m_SourceBytes[It->second] += Result.ReceivedBytes;
-			}
-			else
+			if (!m_SourceLock.WithSharedLock([&]() {
+					if (tsl::robin_map<std::string, uint32_t>::const_iterator It = m_ReceivedBytesPerSource.find(Result.Source);
+						It != m_ReceivedBytesPerSource.end())
+					{
+						m_SourceBytes[It->second] += Result.ReceivedBytes;
+						return true;
+					}
+					return false;
+				}))
 			{
-				m_ReceivedBytesPerSource.insert_or_assign(Result.Source, m_SourceBytes.size());
-				m_SourceBytes.push_back(Result.ReceivedBytes);
+				m_SourceLock.WithExclusiveLock([&]() {
+					if (tsl::robin_map<std::string, uint32_t>::const_iterator It = m_ReceivedBytesPerSource.find(Result.Source);
+						It != m_ReceivedBytesPerSource.end())
+					{
+						m_SourceBytes[It->second] += Result.ReceivedBytes;
+					}
+					else if (m_SourceCount < MaxSourceCount)
+					{
+						size_t Index = m_SourceCount++;
+						m_ReceivedBytesPerSource.insert_or_assign(Result.Source, Index);
+						m_SourceBytes[Index] += Result.ReceivedBytes;
+					}
+				});
 			}
 		}
 	}
@@ -540,8 +556,11 @@ private:
 	const std::string			m_Bucket;
 	const std::filesystem::path m_TempFolderPath;
 
-	tsl::robin_map<std::string, uint32_t> m_ReceivedBytesPerSource;
-	std::vector<uint64_t>				  m_SourceBytes;
+	RwLock											  m_SourceLock;
+	tsl::robin_map<std::string, uint32_t>			  m_ReceivedBytesPerSource;
+	static constexpr size_t							  MaxSourceCount = 8u;
+	std::array<std::atomic<uint64_t>, MaxSourceCount> m_SourceBytes;
+	size_t											  m_SourceCount = 0;
 };
 
 std::unique_ptr<BuildStorageBase>