Merge branch 'main' of https://github.ol.epicgames.net/ue-foundation/zen

5 files changed, 17 insertions, 8 deletions

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);
 }