use partial blocks for oplog import (#780)

Feature: Add --allow-partial-block-requests to zen oplog-import
Improvement: zen oplog-import now uses partial block requests to reduce download size
Improvement: Use latency to Cloud Storage host and Zen Cache host when calculating partial block requests

9 files changed, 93 insertions, 46 deletions

diff --git a/src/zenremotestore/include/zenremotestore/builds/buildstoragecache.h b/src/zenremotestore/include/zenremotestore/builds/buildstoragecache.h
index bb5b1c5f4..f25ce5b5e 100644
--- a/src/zenremotestore/include/zenremotestore/builds/buildstoragecache.h
+++ b/src/zenremotestore/include/zenremotestore/builds/buildstoragecache.h
@@ -65,6 +65,7 @@ struct ZenCacheEndpointTestResult
 {
 	bool		Success = false;
 	std::string FailureReason;
+	double		LatencySeconds = -1.0;
 };
 
 ZenCacheEndpointTestResult TestZenCacheEndpoint(std::string_view BaseUrl, const bool AssumeHttp2, const bool HttpVerbose);
diff --git a/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h b/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h
index 6800444e0..31733569e 100644
--- a/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h
+++ b/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h
@@ -9,6 +9,7 @@
 #include <zenremotestore/builds/buildstoragecache.h>
 #include <zenremotestore/chunking/chunkblock.h>
 #include <zenremotestore/chunking/chunkedcontent.h>
+#include <zenremotestore/partialblockrequestmode.h>
 #include <zenutil/bufferedwritefilecache.h>
 
 #include <atomic>
@@ -109,17 +110,6 @@ struct RebuildFolderStateStatistics
 	uint64_t			  FinalizeTreeElapsedWallTimeUs = 0;
 };
 
-enum EPartialBlockRequestMode
-{
-	Off,
-	ZenCacheOnly,
-	Mixed,
-	All,
-	Invalid
-};
-
-EPartialBlockRequestMode PartialBlockRequestModeFromString(const std::string_view ModeString);
-
 std::filesystem::path ZenStateFilePath(const std::filesystem::path& ZenFolderPath);
 std::filesystem::path ZenTempFolderPath(const std::filesystem::path& ZenFolderPath);
 
diff --git a/src/zenremotestore/include/zenremotestore/builds/buildstorageutil.h b/src/zenremotestore/include/zenremotestore/builds/buildstorageutil.h
index ab3037c89..4b85d8f1e 100644
--- a/src/zenremotestore/include/zenremotestore/builds/buildstorageutil.h
+++ b/src/zenremotestore/include/zenremotestore/builds/buildstorageutil.h
@@ -17,10 +17,12 @@ struct BuildStorageResolveResult
 	std::string HostUrl;
 	std::string HostName;
 	bool		HostAssumeHttp2 = false;
+	double		HostLatencySec	= -1.0;
 
 	std::string CacheUrl;
 	std::string CacheName;
 	bool		CacheAssumeHttp2 = false;
+	double		CacheLatencySec	 = -1.0;
 };
 
 enum class ZenCacheResolveMode
@@ -54,9 +56,11 @@ struct StorageInstance
 	std::unique_ptr<HttpClient>		   BuildStorageHttp;
 	std::unique_ptr<BuildStorageBase>  BuildStorage;
 	std::string						   StorageName;
+	double							   BuildStorageLatencySec = -1.0;
 	std::unique_ptr<HttpClient>		   CacheHttp;
 	std::unique_ptr<BuildStorageCache> BuildCacheStorage;
 	std::string						   CacheName;
+	double							   CacheLatencySec = -1.0;
 };
 
 }  // namespace zen
diff --git a/src/zenremotestore/include/zenremotestore/chunking/chunkblock.h b/src/zenremotestore/include/zenremotestore/chunking/chunkblock.h
index 57710fcf5..5a17ef79c 100644
--- a/src/zenremotestore/include/zenremotestore/chunking/chunkblock.h
+++ b/src/zenremotestore/include/zenremotestore/chunking/chunkblock.h
@@ -82,8 +82,12 @@ class ChunkBlockAnalyser
 public:
 	struct Options
 	{
-		bool IsQuiet   = false;
-		bool IsVerbose = false;
+		bool	 IsQuiet				  = false;
+		bool	 IsVerbose				  = false;
+		double	 HostLatencySec			  = -1.0;
+		double	 HostHighSpeedLatencySec  = -1.0;
+		uint64_t HostSpeedBytesPerSec	  = (1u * 1024u * 1024u * 1024u) / 8u;	// 1GBit
+		uint64_t HostHighSpeedBytesPerSec = (2u * 1024u * 1024u * 1024u) / 8u;	// 2GBit
 	};
 
 	ChunkBlockAnalyser(OperationLogOutput& LogOutput, std::span<const ChunkBlockDescription> BlockDescriptions, const Options& Options);
@@ -110,7 +114,8 @@ public:
 	{
 		Off,
 		SingleRange,
-		On,
+		MultiRange,
+		MultiRangeHighSpeed,
 		Exact
 	};
 
@@ -130,14 +135,14 @@ private:
 		uint16_t MaxRangeCount;
 	};
 
-	static constexpr uint16_t FullBlockRangePercentLimit = 95;
+	static constexpr uint16_t FullBlockRangePercentLimit = 98;
 
 	static constexpr BlockRangeLimit ForceMergeLimits[] = {{.SizePercent = FullBlockRangePercentLimit, .MaxRangeCount = 1},
-														   {.SizePercent = 90, .MaxRangeCount = 2},
-														   {.SizePercent = 85, .MaxRangeCount = 8},
-														   {.SizePercent = 80, .MaxRangeCount = 16},
-														   {.SizePercent = 75, .MaxRangeCount = 32},
-														   {.SizePercent = 70, .MaxRangeCount = 48},
+														   {.SizePercent = 90, .MaxRangeCount = 4},
+														   {.SizePercent = 85, .MaxRangeCount = 16},
+														   {.SizePercent = 80, .MaxRangeCount = 32},
+														   {.SizePercent = 75, .MaxRangeCount = 48},
+														   {.SizePercent = 70, .MaxRangeCount = 64},
 														   {.SizePercent = 4, .MaxRangeCount = 82},
 														   {.SizePercent = 0, .MaxRangeCount = 96}};
 
@@ -149,7 +154,7 @@ private:
 																				std::span<const BlockRangeDescriptor> Ranges);
 	std::vector<BlockRangeDescriptor>				 CollapseBlockRanges(const uint64_t						   AlwaysAcceptableGap,
 																		 std::span<const BlockRangeDescriptor> BlockRanges);
-	uint64_t										 CalculateNextGap(std::span<const BlockRangeDescriptor> BlockRanges);
+	uint64_t CalculateNextGap(const uint64_t AlwaysAcceptableGap, std::span<const BlockRangeDescriptor> BlockRanges);
 	std::optional<std::vector<BlockRangeDescriptor>> CalculateBlockRanges(uint32_t					   BlockIndex,
 																		  const ChunkBlockDescription& BlockDescription,
 																		  std::span<const uint32_t>	   BlockChunkIndexNeeded,
diff --git a/src/zenremotestore/include/zenremotestore/jupiter/jupiterhost.h b/src/zenremotestore/include/zenremotestore/jupiter/jupiterhost.h
index 432496bc1..7bbf40dfa 100644
--- a/src/zenremotestore/include/zenremotestore/jupiter/jupiterhost.h
+++ b/src/zenremotestore/include/zenremotestore/jupiter/jupiterhost.h
@@ -28,6 +28,7 @@ struct JupiterEndpointTestResult
 {
 	bool		Success = false;
 	std::string FailureReason;
+	double		LatencySeconds = -1.0;
 };
 
 JupiterEndpointTestResult TestJupiterEndpoint(std::string_view BaseUrl, const bool AssumeHttp2, const bool HttpVerbose);
diff --git a/src/zenremotestore/include/zenremotestore/operationlogoutput.h b/src/zenremotestore/include/zenremotestore/operationlogoutput.h
index 9693e69cf..6f10ab156 100644
--- a/src/zenremotestore/include/zenremotestore/operationlogoutput.h
+++ b/src/zenremotestore/include/zenremotestore/operationlogoutput.h
@@ -3,6 +3,7 @@
 #pragma once
 
 #include <zencore/fmtutils.h>
+#include <zencore/logbase.h>
 
 namespace zen {
 
@@ -57,9 +58,7 @@ public:
 	virtual ProgressBar* CreateProgressBar(std::string_view InSubTask) = 0;
 };
 
-struct LoggerRef;
-
-OperationLogOutput* CreateStandardLogOutput(LoggerRef& Log);
+OperationLogOutput* CreateStandardLogOutput(LoggerRef Log);
 
 #define ZEN_OPERATION_LOG(OutputTarget, InLevel, fmtstr, ...)                                         \
 	do                                                                                                \
diff --git a/src/zenremotestore/include/zenremotestore/partialblockrequestmode.h b/src/zenremotestore/include/zenremotestore/partialblockrequestmode.h
new file mode 100644
index 000000000..54adea2b2
--- /dev/null
+++ b/src/zenremotestore/include/zenremotestore/partialblockrequestmode.h
@@ -0,0 +1,20 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#pragma once
+
+#include <string_view>
+
+namespace zen {
+
+enum EPartialBlockRequestMode
+{
+	Off,
+	ZenCacheOnly,
+	Mixed,
+	All,
+	Invalid
+};
+
+EPartialBlockRequestMode PartialBlockRequestModeFromString(const std::string_view ModeString);
+
+}  // namespace zen
diff --git a/src/zenremotestore/include/zenremotestore/projectstore/buildsremoteprojectstore.h b/src/zenremotestore/include/zenremotestore/projectstore/buildsremoteprojectstore.h
index e8b7c15c0..66dfcc62d 100644
--- a/src/zenremotestore/include/zenremotestore/projectstore/buildsremoteprojectstore.h
+++ b/src/zenremotestore/include/zenremotestore/projectstore/buildsremoteprojectstore.h
@@ -34,6 +34,8 @@ std::shared_ptr<RemoteProjectStore> CreateJupiterBuildsRemoteStore(LoggerRef
 																   bool							   Quiet,
 																   bool							   Unattended,
 																   bool							   Hidden,
-																   WorkerThreadPool&			   CacheBackgroundWorkerPool);
+																   WorkerThreadPool&			   CacheBackgroundWorkerPool,
+																   double&						   OutHostLatencySec,
+																   double&						   OutCacheLatencySec);
 
 }  // namespace zen
diff --git a/src/zenremotestore/include/zenremotestore/projectstore/remoteprojectstore.h b/src/zenremotestore/include/zenremotestore/projectstore/remoteprojectstore.h
index 008f94351..152c02ee2 100644
--- a/src/zenremotestore/include/zenremotestore/projectstore/remoteprojectstore.h
+++ b/src/zenremotestore/include/zenremotestore/projectstore/remoteprojectstore.h
@@ -6,6 +6,7 @@
 #include <zenstore/projectstore.h>
 
 #include <zenremotestore/chunking/chunkblock.h>
+#include <zenremotestore/partialblockrequestmode.h>
 
 #include <unordered_set>
 
@@ -73,6 +74,16 @@ public:
 		std::vector<ChunkBlockDescription> Blocks;
 	};
 
+	struct GetBlockDescriptionsResult : public Result
+	{
+		std::vector<ChunkBlockDescription> Blocks;
+	};
+
+	struct AttachmentExistsInCacheResult : public Result
+	{
+		std::vector<bool> HasBody;
+	};
+
 	struct RemoteStoreInfo
 	{
 		bool		CreateBlocks;
@@ -111,10 +122,20 @@ public:
 	virtual FinalizeResult		  FinalizeContainer(const IoHash& RawHash)															   = 0;
 	virtual SaveAttachmentsResult SaveAttachments(const std::vector<SharedBuffer>& Payloads)										   = 0;
 
-	virtual LoadContainerResult	  LoadContainer()										= 0;
-	virtual GetKnownBlocksResult  GetKnownBlocks()										= 0;
-	virtual LoadAttachmentResult  LoadAttachment(const IoHash& RawHash)					= 0;
-	virtual LoadAttachmentsResult LoadAttachments(const std::vector<IoHash>& RawHashes) = 0;
+	virtual LoadContainerResult			  LoadContainer()											 = 0;
+	virtual GetKnownBlocksResult		  GetKnownBlocks()											 = 0;
+	virtual GetBlockDescriptionsResult	  GetBlockDescriptions(std::span<const IoHash> BlockHashes)	 = 0;
+	virtual AttachmentExistsInCacheResult AttachmentExistsInCache(std::span<const IoHash> RawHashes) = 0;
+
+	struct AttachmentRange
+	{
+		uint64_t Offset = 0;
+		uint64_t Bytes	= (uint64_t)-1;
+
+		inline operator bool() const { return Offset != 0 || Bytes != (uint64_t)-1; }
+	};
+	virtual LoadAttachmentResult  LoadAttachment(const IoHash& RawHash, const AttachmentRange& Range) = 0;
+	virtual LoadAttachmentsResult LoadAttachments(const std::vector<IoHash>& RawHashes)				  = 0;
 
 	virtual void Flush() = 0;
 };
@@ -153,14 +174,15 @@ RemoteProjectStore::LoadContainerResult BuildContainer(
 
 class JobContext;
 
-RemoteProjectStore::Result SaveOplogContainer(ProjectStore::Oplog&									  Oplog,
-											  const CbObject&										  ContainerObject,
-											  const std::function<void(std::span<IoHash> RawHashes)>& OnReferencedAttachments,
-											  const std::function<bool(const IoHash& RawHash)>&		  HasAttachment,
-											  const std::function<void(const IoHash& BlockHash, std::vector<IoHash>&& Chunks)>& OnNeedBlock,
-											  const std::function<void(const IoHash& RawHash)>&		 OnNeedAttachment,
-											  const std::function<void(const ChunkedInfo& Chunked)>& OnChunkedAttachment,
-											  JobContext*											 OptionalContext);
+RemoteProjectStore::Result SaveOplogContainer(
+	ProjectStore::Oplog&									Oplog,
+	const CbObject&											ContainerObject,
+	const std::function<void(std::span<IoHash> RawHashes)>& OnReferencedAttachments,
+	const std::function<bool(const IoHash& RawHash)>&		HasAttachment,
+	const std::function<void(ThinChunkBlockDescription&& ThinBlockDescription, std::vector<uint32_t>&& NeededChunkIndexes)>& OnNeedBlock,
+	const std::function<void(const IoHash& RawHash)>&	   OnNeedAttachment,
+	const std::function<void(const ChunkedInfo& Chunked)>& OnChunkedAttachment,
+	JobContext*											   OptionalContext);
 
 RemoteProjectStore::Result SaveOplog(CidStore&				ChunkStore,
 									 RemoteProjectStore&	RemoteStore,
@@ -177,15 +199,18 @@ RemoteProjectStore::Result SaveOplog(CidStore&				ChunkStore,
 									 bool					IgnoreMissingAttachments,
 									 JobContext*			OptionalContext);
 
-RemoteProjectStore::Result LoadOplog(CidStore&			  ChunkStore,
-									 RemoteProjectStore&  RemoteStore,
-									 ProjectStore::Oplog& Oplog,
-									 WorkerThreadPool&	  NetworkWorkerPool,
-									 WorkerThreadPool&	  WorkerPool,
-									 bool				  ForceDownload,
-									 bool				  IgnoreMissingAttachments,
-									 bool				  CleanOplog,
-									 JobContext*		  OptionalContext);
+RemoteProjectStore::Result LoadOplog(CidStore&				  ChunkStore,
+									 RemoteProjectStore&	  RemoteStore,
+									 ProjectStore::Oplog&	  Oplog,
+									 WorkerThreadPool&		  NetworkWorkerPool,
+									 WorkerThreadPool&		  WorkerPool,
+									 bool					  ForceDownload,
+									 bool					  IgnoreMissingAttachments,
+									 bool					  CleanOplog,
+									 EPartialBlockRequestMode PartialBlockRequestMode,
+									 double					  HostLatencySec,
+									 double					  CacheLatencySec,
+									 JobContext*			  OptionalContext);
 
 std::vector<IoHash>					   GetBlockHashesFromOplog(CbObjectView ContainerObject);
 std::vector<ThinChunkBlockDescription> GetBlocksFromOplog(CbObjectView ContainerObject, std::span<const IoHash> IncludeBlockHashes);