fix jupiterbuildstorage concurrency (#906)

- Bugfix: Fixed concurrency issue in JupiterBuildStorage when updating stats

2 files changed, 35 insertions, 15 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 51ab7bd62..f46adc60f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -42,6 +42,7 @@
 - Bugfix: Fix an issue when a file goes from being internally stored to externally referenced and it could get into an inconsistent state
 - Bugfix: Hub watchdog loop did not check `m_ShutdownFlag`, causing it to spin indefinitely on shutdown
 - Bugfix: Fixed `MakeSafeAbsolutePathInPlace` when a UNC prefix is present but path uses mixed delimiters
+- Bugfix: Fixed concurrency issue in JupiterBuildStorage when updating stats
 
 ## 5.8.0
 - Feature: Hub watchdog automatically deprovisions inactive provisioned and hibernated instances
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>