update state when wildcard (#657)

* add --append option and improve state handling when using downloads for `zen builds download`

4 files changed, 143 insertions, 66 deletions

diff --git a/src/zenremotestore/include/zenremotestore/builds/buildcontent.h b/src/zenremotestore/include/zenremotestore/builds/buildcontent.h
new file mode 100644
index 000000000..67aebeda3
--- /dev/null
+++ b/src/zenremotestore/include/zenremotestore/builds/buildcontent.h
@@ -0,0 +1,43 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#pragma once
+
+#include <zenremotestore/chunking/chunkedcontent.h>
+
+namespace zen {
+
+class CbObjectView;
+class CbObjectWriter;
+
+void WriteBuildContentToCompactBinary(CbObjectWriter&						 PartManifestWriter,
+									  const SourcePlatform					 Platform,
+									  std::span<const std::filesystem::path> Paths,
+									  std::span<const IoHash>				 RawHashes,
+									  std::span<const uint64_t>				 RawSizes,
+									  std::span<const uint32_t>				 Attributes,
+									  std::span<const IoHash>				 SequenceRawHashes,
+									  std::span<const uint32_t>				 ChunkCounts,
+									  std::span<const IoHash>				 LocalChunkHashes,
+									  std::span<const uint64_t>				 LocalChunkRawSizes,
+									  const std::vector<uint32_t>&			 AbsoluteChunkOrders,
+									  const std::span<const uint32_t>		 LooseLocalChunkIndexes,
+									  const std::span<IoHash>				 BlockHashes);
+
+void ReadBuildContentFromCompactBinary(CbObjectView						   BuildPartManifest,
+									   SourcePlatform&					   OutPlatform,
+									   std::vector<std::filesystem::path>& OutPaths,
+									   std::vector<IoHash>&				   OutRawHashes,
+									   std::vector<uint64_t>&			   OutRawSizes,
+									   std::vector<uint32_t>&			   OutAttributes,
+									   std::vector<IoHash>&				   OutSequenceRawHashes,
+									   std::vector<uint32_t>&			   OutChunkCounts,
+									   std::vector<uint32_t>&			   OutAbsoluteChunkOrders,
+									   std::vector<IoHash>&				   OutLooseChunkHashes,
+									   std::vector<uint64_t>&			   OutLooseChunkRawSizes,
+									   std::vector<IoHash>&				   OutBlockRawHashes);
+
+#if ZEN_WITH_TESTS
+void buildsavedstate_forcelink();
+#endif	// ZEN_WITH_TESTS
+
+}  // namespace zen
diff --git a/src/zenremotestore/include/zenremotestore/builds/buildsavedstate.h b/src/zenremotestore/include/zenremotestore/builds/buildsavedstate.h
index a11641b0d..f808a7a3b 100644
--- a/src/zenremotestore/include/zenremotestore/builds/buildsavedstate.h
+++ b/src/zenremotestore/include/zenremotestore/builds/buildsavedstate.h
@@ -13,48 +13,67 @@ namespace zen {
 
 class CbObjectView;
 
-CbObject CreateStateObject(const Oid&									   BuildId,
-						   const std::vector<std::pair<Oid, std::string>>& AllBuildParts,
-						   std::span<const ChunkedFolderContent>		   PartContents,
-						   const FolderContent&							   LocalFolderState,
-						   const std::filesystem::path&					   LocalPath);
-
-void ReadBuildContentFromCompactBinary(CbObjectView						   BuildPartManifest,
-									   SourcePlatform&					   OutPlatform,
-									   std::vector<std::filesystem::path>& OutPaths,
-									   std::vector<IoHash>&				   OutRawHashes,
-									   std::vector<uint64_t>&			   OutRawSizes,
-									   std::vector<uint32_t>&			   OutAttributes,
-									   std::vector<IoHash>&				   OutSequenceRawHashes,
-									   std::vector<uint32_t>&			   OutChunkCounts,
-									   std::vector<uint32_t>&			   OutAbsoluteChunkOrders,
-									   std::vector<IoHash>&				   OutLooseChunkHashes,
-									   std::vector<uint64_t>&			   OutLooseChunkRawSizes,
-									   std::vector<IoHash>&				   OutBlockRawHashes);
-
-void CalculateLocalChunkOrders(const std::span<const uint32_t>&				 AbsoluteChunkOrders,
-							   const std::span<const IoHash>				 LooseChunkHashes,
-							   const std::span<const uint64_t>				 LooseChunkRawSizes,
-							   const std::span<const ChunkBlockDescription>& BlockDescriptions,
-							   std::vector<IoHash>&							 OutLocalChunkHashes,
-							   std::vector<uint64_t>&						 OutLocalChunkRawSizes,
-							   std::vector<uint32_t>&						 OutLocalChunkOrders);
-
-void ReadStateObject(CbObjectView						StateView,
-					 Oid&								OutBuildId,
-					 std::vector<Oid>&					BuildPartsIds,
-					 std::vector<std::string>&			BuildPartsNames,
-					 std::vector<ChunkedFolderContent>& OutPartContents,
-					 FolderContent&						OutLocalFolderState);
-
-void ReadStateFile(const std::filesystem::path& StateFilePath, FolderContent& OutLocalFolderState, ChunkedFolderContent& OutLocalContent);
-
-void AddDownloadedPath(const std::filesystem::path&					   SystemRootDir,
-					   const Oid&									   BuildId,
-					   const std::vector<std::pair<Oid, std::string>>& BuildParts,
-					   const std::filesystem::path&					   StateFilePath,
-					   const std::filesystem::path&					   LocalPath);
+struct BuildsSelection
+{
+	struct BuildPart
+	{
+		Oid			Id;
+		std::string Name;
+	};
+	struct Build
+	{
+		Oid						 Id;
+		std::vector<BuildPart>	 Parts;
+		std::vector<std::string> IncludeWildcards;
+		std::vector<std::string> ExcludeWildcards;
+	};
+	std::vector<Build> Builds;
+
+	static void			   Write(const BuildsSelection& Selection, CbWriter& Output);
+	static BuildsSelection Read(CbObjectView& Input, std::vector<ChunkedFolderContent>* OptionalOutLegacyPartsContent);
+};
+
+struct BuildState
+{
+	BuildsSelection		 Selection;
+	ChunkedFolderContent ChunkedContent;
+
+	static void		  Write(const BuildState& Selection, CbWriter& Output);
+	static BuildState Read(CbObjectView& Input);
+};
+
+struct BuildSaveState
+{
+	BuildState			  State;
+	FolderContent		  FolderState;
+	std::filesystem::path LocalPath;
+
+	static void			  Write(const BuildSaveState& SaveState, CbWriter& Output);
+	static BuildSaveState Read(CbObjectView& Input);
+};
+
+struct BuildsDownloadInfo
+{
+	BuildsSelection		  Selection;
+	std::filesystem::path LocalPath;
+	std::filesystem::path StateFilePath;
+	std::string			  Iso8601Date;
+
+	static void				  Write(const BuildsDownloadInfo& Info, CbWriter& Output);
+	static BuildsDownloadInfo Read(CbObjectView& Input);
+};
+
+CbObject	   CreateBuildSaveStateObject(const BuildSaveState& State);
+BuildSaveState ReadBuildSaveStateFile(const std::filesystem::path& StateFilePath);
+
+void AddDownloadedPath(const std::filesystem::path& SystemRootDir, const BuildsDownloadInfo& Info);
+
+BuildsDownloadInfo ReadDownloadedInfoFile(const std::filesystem::path& DownloadInfoPath);
 
 std::vector<std::filesystem::path> GetDownloadedStatePaths(const std::filesystem::path& SystemRootDir);
 
+#if ZEN_WITH_TESTS
+void buildsavedstate_forcelink();
+#endif	// ZEN_WITH_TESTS
+
 }  // namespace zen
diff --git a/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h b/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h
index eba0966f7..3c4535d9c 100644
--- a/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h
+++ b/src/zenremotestore/include/zenremotestore/builds/buildstorageoperations.h
@@ -497,8 +497,9 @@ public:
 
 	struct Options
 	{
-		bool IsQuiet   = false;
-		bool IsVerbose = false;
+		bool IsQuiet				  = false;
+		bool IsVerbose				  = false;
+		bool DoExtraContentValidation = false;
 
 		const uint64_t		  FindBlockMaxCount			= 10000;
 		const uint8_t		  BlockReuseMinPercentLimit = 85;
@@ -579,27 +580,6 @@ private:
 													   const tsl::robin_map<IoHash, uint32_t, IoHash::Hasher>& ChunkHashToLocalChunkIndex,
 													   const std::span<const uint32_t>&						   LooseChunkIndexes,
 													   const std::span<const ChunkBlockDescription>&		   BlockDescriptions);
-	void				  CalculateLocalChunkOrders(const std::span<const uint32_t>&			  AbsoluteChunkOrders,
-													const std::span<const IoHash>				  LooseChunkHashes,
-													const std::span<const uint64_t>				  LooseChunkRawSizes,
-													const std::span<const ChunkBlockDescription>& BlockDescriptions,
-													std::vector<IoHash>&						  OutLocalChunkHashes,
-													std::vector<uint64_t>&						  OutLocalChunkRawSizes,
-													std::vector<uint32_t>&						  OutLocalChunkOrders);
-
-	void WriteBuildContentToCompactBinary(CbObjectWriter&						 PartManifestWriter,
-										  const SourcePlatform					 Platform,
-										  std::span<const std::filesystem::path> Paths,
-										  std::span<const IoHash>				 RawHashes,
-										  std::span<const uint64_t>				 RawSizes,
-										  std::span<const uint32_t>				 Attributes,
-										  std::span<const IoHash>				 SequenceRawHashes,
-										  std::span<const uint32_t>				 ChunkCounts,
-										  std::span<const IoHash>				 LocalChunkHashes,
-										  std::span<const uint64_t>				 LocalChunkRawSizes,
-										  const std::vector<uint32_t>&			 AbsoluteChunkOrders,
-										  const std::span<const uint32_t>		 LooseLocalChunkIndexes,
-										  const std::span<IoHash>				 BlockHashes);
 
 	CompositeBuffer FetchChunk(const ChunkedFolderContent& Content,
 							   const ChunkedContentLookup& Lookup,
diff --git a/src/zenremotestore/include/zenremotestore/chunking/chunkedcontent.h b/src/zenremotestore/include/zenremotestore/chunking/chunkedcontent.h
index 227d877d8..e4be7923a 100644
--- a/src/zenremotestore/include/zenremotestore/chunking/chunkedcontent.h
+++ b/src/zenremotestore/include/zenremotestore/chunking/chunkedcontent.h
@@ -113,7 +113,7 @@ struct ChunkedContentLookup
 };
 
 void				 SaveChunkedFolderContentToCompactBinary(const ChunkedFolderContent& Content, CbWriter& Output);
-ChunkedFolderContent LoadChunkedFolderContentToCompactBinary(CbObjectView Input);
+ChunkedFolderContent LoadChunkedFolderContentFromCompactBinary(CbObjectView Input);
 
 ChunkedFolderContent MergeChunkedFolderContents(const ChunkedFolderContent& Base, std::span<const ChunkedFolderContent> Overlays);
 ChunkedFolderContent DeletePathsFromChunkedContent(const ChunkedFolderContent&			  Base,
@@ -123,6 +123,11 @@ ChunkedFolderContent DeletePathsFromChunkedContent(const ChunkedFolderContent& B
 
 bool CompareChunkedContent(const ChunkedFolderContent& Lhs, const ChunkedFolderContent& Rhs);
 
+ChunkedFolderContent ApplyChunkedContentOverlay(const ChunkedFolderContent&	 Base,
+												const ChunkedFolderContent&	 Overlay,
+												std::span<const std::string> OverlayIncludeWildcards,
+												std::span<const std::string> OverlayExcludeWildcards);
+
 struct ChunkingStatistics
 {
 	std::atomic<uint64_t> FilesProcessed	   = 0;
@@ -198,14 +203,44 @@ GetFirstPathIndexForRawHash(const ChunkedContentLookup& Lookup, const IoHash& Ra
 
 struct ChunkBlockDescription;
 
+void CalculateLocalChunkOrders(const std::span<const uint32_t>&				 AbsoluteChunkOrders,
+							   const std::span<const IoHash>				 LooseChunkHashes,
+							   const std::span<const uint64_t>				 LooseChunkRawSizes,
+							   const std::span<const ChunkBlockDescription>& BlockDescriptions,
+							   std::vector<IoHash>&							 OutLocalChunkHashes,
+							   std::vector<uint64_t>&						 OutLocalChunkRawSizes,
+							   std::vector<uint32_t>&						 OutLocalChunkOrders,
+							   bool											 DoExtraVerify);
+
 void ValidateChunkedFolderContent(const ChunkedFolderContent&			 Content,
 								  std::span<const ChunkBlockDescription> BlockDescriptions,
-								  std::span<const IoHash>				 LooseChunks);
+								  std::span<const IoHash>				 LooseChunks,
+								  std::span<const std::string>			 IncludeWildcards,
+								  std::span<const std::string>			 ExcludeWildcards);
 
 void InlineRemoveUnusedHashes(std::vector<IoHash>& InOutHashes, std::span<const IoHash> UsedHashes);
 
 #if ZEN_WITH_TESTS
 void chunkedcontent_forcelink();
+
+struct FastRandom;
+
+namespace chunkedcontent_testutils {
+	struct ChunkedFile
+	{
+		IoHash				  RawHash;
+		std::vector<IoHash>	  ChunkHashes;
+		std::vector<uint64_t> ChunkSizes;
+		std::vector<IoBuffer> Chunks;
+	};
+
+	ChunkedFile			 CreateChunkedFile(FastRandom& Random, size_t Size, size_t ChunkingSize);
+	ChunkedFolderContent CreateChunkedFolderContent(FastRandom&												Random,
+													std::span<const std::pair<const std::string, uint64_t>> PathAndSizes,
+													uint64_t												ChunkingSize,
+													std::vector<IoBuffer>&									ChunkPayloads);
+}  // namespace chunkedcontent_testutils
+
 #endif	// ZEN_WITH_TESTS
 
 }  // namespace zen