add check that we have enough free space to complete the builds download (#640)

* add check that we have enough free space to complete the builds download

2 files changed, 56 insertions, 6 deletions

diff --git a/src/zenremotestore/builds/buildstorageoperations.cpp b/src/zenremotestore/builds/buildstorageoperations.cpp
index d8ae9c5a7..fb1b3018a 100644
--- a/src/zenremotestore/builds/buildstorageoperations.cpp
+++ b/src/zenremotestore/builds/buildstorageoperations.cpp
@@ -1013,6 +1013,15 @@ BuildsOperationUpdateFolder::Execute(FolderContent& OutLocalFolderState)
 		uint64_t			  TotalPartWriteCount = 0;
 		std::atomic<uint64_t> WritePartsComplete  = 0;
 
+		tsl::robin_map<std::string, uint32_t> RemotePathToRemoteIndex;
+		RemotePathToRemoteIndex.reserve(m_RemoteContent.Paths.size());
+		for (uint32_t RemotePathIndex = 0; RemotePathIndex < m_RemoteContent.Paths.size(); RemotePathIndex++)
+		{
+			RemotePathToRemoteIndex.insert({m_RemoteContent.Paths[RemotePathIndex].generic_string(), RemotePathIndex});
+		}
+
+		CheckRequiredDiskSpace(RemotePathToRemoteIndex);
+
 		{
 			ZEN_TRACE_CPU("WriteChunks");
 
@@ -2042,12 +2051,6 @@ BuildsOperationUpdateFolder::Execute(FolderContent& OutLocalFolderState)
 			ZEN_TRACE_CPU("PrepareTarget");
 
 			tsl::robin_set<IoHash, IoHash::Hasher> CachedRemoteSequences;
-			tsl::robin_map<std::string, uint32_t>  RemotePathToRemoteIndex;
-			RemotePathToRemoteIndex.reserve(m_RemoteContent.Paths.size());
-			for (uint32_t RemotePathIndex = 0; RemotePathIndex < m_RemoteContent.Paths.size(); RemotePathIndex++)
-			{
-				RemotePathToRemoteIndex.insert({m_RemoteContent.Paths[RemotePathIndex].generic_string(), RemotePathIndex});
-			}
 
 			std::vector<uint32_t> FilesToCache;
 
@@ -3066,6 +3069,51 @@ BuildsOperationUpdateFolder::GetChunkWriteCount(std::span<const std::atomic<uint
 };
 
 void
+BuildsOperationUpdateFolder::CheckRequiredDiskSpace(const tsl::robin_map<std::string, uint32_t>& RemotePathToRemoteIndex)
+{
+	tsl::robin_set<uint32_t> ExistingRemotePaths;
+
+	if (m_Options.EnableTargetFolderScavenging)
+	{
+		for (uint32_t LocalPathIndex = 0; LocalPathIndex < m_LocalContent.Paths.size(); LocalPathIndex++)
+		{
+			const IoHash&				 RawHash   = m_LocalContent.RawHashes[LocalPathIndex];
+			const std::filesystem::path& LocalPath = m_LocalContent.Paths[LocalPathIndex];
+
+			if (auto RemotePathIt = RemotePathToRemoteIndex.find(LocalPath.generic_string()); RemotePathIt != RemotePathToRemoteIndex.end())
+			{
+				const uint32_t RemotePathIndex = RemotePathIt->second;
+				if (m_RemoteContent.RawHashes[RemotePathIndex] == RawHash)
+				{
+					ExistingRemotePaths.insert(RemotePathIndex);
+				}
+			}
+		}
+	}
+
+	uint64_t RequiredSpace = 0;
+	for (uint32_t RemotePathIndex = 0; RemotePathIndex < m_RemoteContent.Paths.size(); RemotePathIndex++)
+	{
+		if (!ExistingRemotePaths.contains(RemotePathIndex))
+		{
+			RequiredSpace += m_RemoteContent.RawSizes[RemotePathIndex];
+		}
+	}
+
+	std::error_code Ec;
+	DiskSpace		Space = DiskSpaceInfo(m_Path, Ec);
+	if (Ec)
+	{
+		throw std::runtime_error(fmt::format("Get free disk space for target path '{}' FAILED, reason: {}", m_Path, Ec.message()));
+	}
+	if (Space.Free < (RequiredSpace + 16u * 1024u * 1024u))
+	{
+		throw std::runtime_error(
+			fmt::format("Not enough free space for target path '{}', {} of free space is needed", m_Path, RequiredSpace));
+	}
+}
+
+void
 BuildsOperationUpdateFolder::WriteScavengedSequenceToCache(const std::filesystem::path&			 ScavengeRootPath,
 														   const ChunkedFolderContent&			 ScavengedContent,
 														   const ScavengedSequenceCopyOperation& ScavengeOp)
diff --git a/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h b/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h
index 6eef794fd..c3d8b4c27 100644
--- a/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h
+++ b/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h
@@ -268,6 +268,8 @@ private:
 
 	uint64_t GetChunkWriteCount(std::span<const std::atomic<uint32_t>> SequenceIndexChunksLeftToWriteCounters, uint32_t ChunkIndex);
 
+	void CheckRequiredDiskSpace(const tsl::robin_map<std::string, uint32_t>& RemotePathToRemoteIndex);
+
 	void WriteScavengedSequenceToCache(const std::filesystem::path&			 ScavengeRootPath,
 									   const ChunkedFolderContent&			 ScavengedContent,
 									   const ScavengedSequenceCopyOperation& ScavengeOp);