added --find-max-block-count option to builds upload (#337)

8 files changed, 33 insertions, 13 deletions

diff --git a/src/zen/cmds/builds_cmd.cpp b/src/zen/cmds/builds_cmd.cpp
index 1c24a3c5b..8e8fd480a 100644
--- a/src/zen/cmds/builds_cmd.cpp
+++ b/src/zen/cmds/builds_cmd.cpp
@@ -2900,6 +2900,7 @@ namespace {
 					  const std::filesystem::path& Path,
 					  const std::filesystem::path& ZenFolderPath,
 					  const std::filesystem::path& ManifestPath,
+					  const uint64_t			   FindBlockMaxCount,
 					  const uint8_t				   BlockReuseMinPercentLimit,
 					  bool						   AllowMultiparts,
 					  const CbObject&			   MetaData,
@@ -2939,7 +2940,7 @@ namespace {
 		FindBlocksStatistics FindBlocksStats;
 
 		std::future<PrepareBuildResult> PrepBuildResultFuture = GetNetworkPool().EnqueueTask(std::packaged_task<PrepareBuildResult()>{
-			[&Storage, BuildId, &MetaData, CreateBuild, AllowMultiparts, IgnoreExistingBlocks, &FindBlocksStats] {
+			[&Storage, BuildId, FindBlockMaxCount, &MetaData, CreateBuild, AllowMultiparts, IgnoreExistingBlocks, &FindBlocksStats] {
 				ZEN_TRACE_CPU("PrepareBuild");
 
 				PrepareBuildResult Result;
@@ -2976,7 +2977,7 @@ namespace {
 				{
 					ZEN_TRACE_CPU("FindBlocks");
 					Stopwatch KnownBlocksTimer;
-					Result.KnownBlocks				= ParseChunkBlockDescriptionList(Storage.BuildStorage->FindBlocks(BuildId));
+					Result.KnownBlocks = ParseChunkBlockDescriptionList(Storage.BuildStorage->FindBlocks(BuildId, FindBlockMaxCount));
 					FindBlocksStats.FindBlockTimeMS = KnownBlocksTimer.GetElapsedTimeMs();
 					FindBlocksStats.FoundBlockCount = Result.KnownBlocks.size();
 					Result.FindBlocksTimeMs			= KnownBlocksTimer.GetElapsedTimeMs();
@@ -7417,7 +7418,7 @@ namespace {
 					std::vector<ChunkBlockDescription> AugmentedBlockDescriptions;
 					AugmentedBlockDescriptions.reserve(BlockRawHashes.size());
 					std::vector<ChunkBlockDescription> FoundBlocks =
-						ParseChunkBlockDescriptionList(Storage.BuildStorage->FindBlocks(BuildId));
+						ParseChunkBlockDescriptionList(Storage.BuildStorage->FindBlocks(BuildId, (uint64_t)-1));
 
 					for (const IoHash& BlockHash : BlockRawHashes)
 					{
@@ -8496,6 +8497,12 @@ BuildsCommand::BuildsCommand()
 							   "<manifestpath>");
 	m_UploadOptions
 		.add_option("", "", "verify", "Enable post upload verify of all uploaded data", cxxopts::value(m_PostUploadVerify), "<verify>");
+	m_UploadOptions.add_option("",
+							   "",
+							   "find-max-block-count",
+							   "The maximum number of blocks we search for in the build context",
+							   cxxopts::value(m_FindBlockMaxCount),
+							   "<maxblockcount>");
 
 	m_UploadOptions.parse_positional({"local-path", "build-id"});
 	m_UploadOptions.positional_help("local-path build-id");
@@ -9221,6 +9228,7 @@ BuildsCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
 						 Path,
 						 ZenFolderPath,
 						 MakeSafeAbsolutePath(m_ManifestPath),
+						 m_FindBlockMaxCount,
 						 m_BlockReuseMinPercentLimit,
 						 m_AllowMultiparts,
 						 MetaData,
@@ -9489,6 +9497,7 @@ BuildsCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
 						 Path,
 						 ZenFolderPath,
 						 {},
+						 m_FindBlockMaxCount,
 						 m_BlockReuseMinPercentLimit,
 						 m_AllowMultiparts,
 						 MetaData,
@@ -9677,6 +9686,7 @@ BuildsCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
 						 DownloadPath,
 						 ZenFolderPath,
 						 {},
+						 m_FindBlockMaxCount,
 						 m_BlockReuseMinPercentLimit,
 						 m_AllowMultiparts,
 						 MetaData2,
diff --git a/src/zen/cmds/builds_cmd.h b/src/zen/cmds/builds_cmd.h
index bc5cd235f..b40c3e08f 100644
--- a/src/zen/cmds/builds_cmd.h
+++ b/src/zen/cmds/builds_cmd.h
@@ -89,7 +89,8 @@ private:
 	std::string m_Path;
 
 	cxxopts::Options m_UploadOptions{"upload", "Upload a folder"};
-	bool			 m_PostUploadVerify = false;
+	uint64_t		 m_FindBlockMaxCount = 10000;
+	bool			 m_PostUploadVerify	 = false;
 
 	cxxopts::Options		 m_DownloadOptions{"download", "Download a folder"};
 	std::vector<std::string> m_BuildPartNames;
diff --git a/src/zenserver/projectstore/buildsremoteprojectstore.cpp b/src/zenserver/projectstore/buildsremoteprojectstore.cpp
index fbb9bc344..a6583b722 100644
--- a/src/zenserver/projectstore/buildsremoteprojectstore.cpp
+++ b/src/zenserver/projectstore/buildsremoteprojectstore.cpp
@@ -308,7 +308,7 @@ public:
 	{
 		ZEN_ASSERT(m_OplogBuildPartId != Oid::Zero);
 		JupiterSession Session(m_JupiterClient->Logger(), m_JupiterClient->Client());
-		JupiterResult  FindResult = Session.FindBlocks(m_Namespace, m_Bucket, m_BuildId);
+		JupiterResult  FindResult = Session.FindBlocks(m_Namespace, m_Bucket, m_BuildId, (uint64_t)-1);
 		AddStats(FindResult);
 		GetKnownBlocksResult Result{ConvertResult(FindResult)};
 		if (Result.ErrorCode)
diff --git a/src/zenutil/filebuildstorage.cpp b/src/zenutil/filebuildstorage.cpp
index 0fa05194f..7aa252e44 100644
--- a/src/zenutil/filebuildstorage.cpp
+++ b/src/zenutil/filebuildstorage.cpp
@@ -442,7 +442,7 @@ public:
 		SimulateLatency(0, 0);
 	}
 
-	virtual CbObject FindBlocks(const Oid& BuildId) override
+	virtual CbObject FindBlocks(const Oid& BuildId, uint64_t MaxBlockCount) override
 	{
 		ZEN_TRACE_CPU("FileBuildStorage::FindBlocks");
 		ZEN_UNUSED(BuildId);
@@ -451,6 +451,8 @@ public:
 		auto	  _ = MakeGuard([&]() { m_Stats.TotalExecutionTimeUs += ExecutionTimer.GetElapsedTimeUs(); });
 		m_Stats.TotalRequestCount++;
 
+		uint64_t FoundCount = 0;
+
 		DirectoryContent Content;
 		GetDirectoryContent(GetBlobsMetadataFolder(), DirectoryContentFlags::IncludeFiles, Content);
 		CbObjectWriter Writer;
@@ -469,6 +471,11 @@ public:
 
 					CbObject BlockObject = CbObject(SharedBuffer(BlockMetaDataPayload));
 					Writer.AddObject(BlockObject);
+					FoundCount++;
+					if (FoundCount == MaxBlockCount)
+					{
+						break;
+					}
 				}
 			}
 		}
diff --git a/src/zenutil/include/zenutil/buildstorage.h b/src/zenutil/include/zenutil/buildstorage.h
index f8c7c012c..b0665dbf8 100644
--- a/src/zenutil/include/zenutil/buildstorage.h
+++ b/src/zenutil/include/zenutil/buildstorage.h
@@ -55,7 +55,7 @@ public:
 		std::function<void(uint64_t Offset, const IoBuffer& Chunk, uint64_t BytesRemaining)>&& Receiver) = 0;
 
 	virtual void	 PutBlockMetadata(const Oid& BuildId, const IoHash& BlockRawHash, const CbObject& MetaData) = 0;
-	virtual CbObject FindBlocks(const Oid& BuildId)																= 0;
+	virtual CbObject FindBlocks(const Oid& BuildId, uint64_t MaxBlockCount)										= 0;
 	virtual CbObject GetBlockMetadatas(const Oid& BuildId, std::span<const IoHash> BlockHashes)					= 0;
 
 	virtual void PutBuildPartStats(const Oid& BuildId, const Oid& BuildPartId, const tsl::robin_map<std::string, double>& FloatStats) = 0;
diff --git a/src/zenutil/include/zenutil/jupiter/jupitersession.h b/src/zenutil/include/zenutil/jupiter/jupitersession.h
index fda4a7bfe..417ed7384 100644
--- a/src/zenutil/include/zenutil/jupiter/jupitersession.h
+++ b/src/zenutil/include/zenutil/jupiter/jupitersession.h
@@ -152,7 +152,7 @@ public:
 											  const Oid&	   BuildId,
 											  const Oid&	   PartId,
 											  const IoHash&	   RawHash);
-	JupiterResult			FindBlocks(std::string_view Namespace, std::string_view BucketId, const Oid& BuildId);
+	JupiterResult			FindBlocks(std::string_view Namespace, std::string_view BucketId, const Oid& BuildId, uint64_t MaxBlockCount);
 	JupiterResult			GetBlockMetadata(std::string_view Namespace, std::string_view BucketId, const Oid& BuildId, IoBuffer Payload);
 
 	JupiterResult PutBuildPartStats(std::string_view Namespace,
diff --git a/src/zenutil/jupiter/jupiterbuildstorage.cpp b/src/zenutil/jupiter/jupiterbuildstorage.cpp
index b6d9e3990..f2d190408 100644
--- a/src/zenutil/jupiter/jupiterbuildstorage.cpp
+++ b/src/zenutil/jupiter/jupiterbuildstorage.cpp
@@ -289,13 +289,13 @@ public:
 		}
 	}
 
-	virtual CbObject FindBlocks(const Oid& BuildId) override
+	virtual CbObject FindBlocks(const Oid& BuildId, uint64_t MaxBlockCount) override
 	{
 		ZEN_TRACE_CPU("Jupiter::FindBlocks");
 
 		Stopwatch	  ExecutionTimer;
 		auto		  _			 = MakeGuard([&]() { m_Stats.TotalExecutionTimeUs += ExecutionTimer.GetElapsedTimeUs(); });
-		JupiterResult FindResult = m_Session.FindBlocks(m_Namespace, m_Bucket, BuildId);
+		JupiterResult FindResult = m_Session.FindBlocks(m_Namespace, m_Bucket, BuildId, MaxBlockCount);
 		AddStatistic(FindResult);
 		if (!FindResult.Success)
 		{
diff --git a/src/zenutil/jupiter/jupitersession.cpp b/src/zenutil/jupiter/jupitersession.cpp
index 2e4fe5258..fde86a478 100644
--- a/src/zenutil/jupiter/jupitersession.cpp
+++ b/src/zenutil/jupiter/jupitersession.cpp
@@ -758,10 +758,12 @@ JupiterSession::FinalizeBuildPart(std::string_view Namespace,
 }
 
 JupiterResult
-JupiterSession::FindBlocks(std::string_view Namespace, std::string_view BucketId, const Oid& BuildId)
+JupiterSession::FindBlocks(std::string_view Namespace, std::string_view BucketId, const Oid& BuildId, uint64_t MaxBlockCount)
 {
-	HttpClient::Response Response = m_HttpClient.Get(fmt::format("/api/v2/builds/{}/{}/{}/blocks/listBlocks", Namespace, BucketId, BuildId),
-													 HttpClient::Accept(ZenContentType::kCbObject));
+	const std::string	 Parameters = MaxBlockCount == (uint64_t)-1 ? "" : fmt::format("?count={}", MaxBlockCount);
+	HttpClient::Response Response =
+		m_HttpClient.Get(fmt::format("/api/v2/builds/{}/{}/{}/blocks/listBlocks{}", Namespace, BucketId, BuildId, Parameters),
+						 HttpClient::Accept(ZenContentType::kCbObject));
 	return detail::ConvertResponse(Response, "JupiterSession::FindBlocks"sv);
 }