Merge branch 'de/zen-service-command' of https://github.ol.epicgames.net/ue-foundation/zen into de/zen-service-command

1 files changed, 71 insertions, 19 deletions

diff --git a/src/zencore/basicfile.cpp b/src/zencore/basicfile.cpp
index b3bffd34d..95876cff4 100644
--- a/src/zencore/basicfile.cpp
+++ b/src/zencore/basicfile.cpp
@@ -281,6 +281,20 @@ BasicFile::StreamByteRange(uint64_t FileOffset, uint64_t Size, std::function<voi
 }
 
 uint64_t
+BasicFile::Write(const CompositeBuffer& Data, uint64_t FileOffset)
+{
+	std::error_code Ec;
+	uint64_t		WrittenBytes = Write(Data, FileOffset, Ec);
+
+	if (Ec)
+	{
+		std::error_code Dummy;
+		throw std::system_error(Ec, fmt::format("Failed to write to file '{}'", zen::PathFromHandle(m_FileHandle, Dummy)));
+	}
+	return WrittenBytes;
+}
+
+uint64_t
 BasicFile::Write(const CompositeBuffer& Data, uint64_t FileOffset, std::error_code& Ec)
 {
 	uint64_t WrittenBytes = 0;
@@ -309,6 +323,8 @@ BasicFile::Write(MemoryView Data, uint64_t FileOffset, std::error_code& Ec)
 void
 BasicFile::Write(const void* Data, uint64_t Size, uint64_t FileOffset, std::error_code& Ec)
 {
+	ZEN_ASSERT(m_FileHandle != nullptr);
+
 	Ec.clear();
 
 	const uint64_t MaxChunkSize = 2u * 1024 * 1024 * 1024;
@@ -817,6 +833,17 @@ BasicFileWriter::Write(const void* Data, uint64_t Size, uint64_t FileOffset)
 }
 
 void
+BasicFileWriter::Write(const CompositeBuffer& Data, uint64_t FileOffset)
+{
+	for (const SharedBuffer& Segment : Data.GetSegments())
+	{
+		const uint64_t SegmentSize = Segment.GetSize();
+		Write(Segment.GetData(), SegmentSize, FileOffset);
+		FileOffset += SegmentSize;
+	}
+}
+
+void
 BasicFileWriter::Flush()
 {
 	const uint64_t BufferedBytes = m_BufferEnd - m_BufferStart;
@@ -831,23 +858,21 @@ BasicFileWriter::Flush()
 }
 
 IoBuffer
-WriteToTempFile(CompositeBuffer&& Buffer, const std::filesystem::path& Path, std::function<bool(std::error_code& Ec)>&& RetryCallback)
+WriteToTempFile(CompositeBuffer&& Buffer, const std::filesystem::path& Path)
 {
-	if (std::filesystem::is_regular_file(Path))
+	TemporaryFile	Temp;
+	std::error_code Ec;
+	Temp.CreateTemporary(Path.parent_path(), Ec);
+	if (Ec)
 	{
-		IoBuffer ExistingTempFile = IoBuffer(IoBufferBuilder::MakeFromFile(Path));
-		if (ExistingTempFile && ExistingTempFile.GetSize() == Buffer.GetSize())
-		{
-			ExistingTempFile.SetDeleteOnClose(true);
-			return ExistingTempFile;
-		}
+		throw std::system_error(Ec, fmt::format("Failed to create temp file for blob at '{}'", Path));
 	}
-	BasicFile BlockFile;
-	BlockFile.Open(Path, BasicFile::Mode::kTruncateDelete, std::move(RetryCallback));
-	uint64_t Offset = 0;
+
+	uint64_t BufferSize = Buffer.GetSize();
 	{
+		uint64_t			  Offset		= 0;
 		static const uint64_t BufferingSize = 256u * 1024u;
-		BasicFileWriter		  BufferedOutput(BlockFile, BufferingSize / 2);
+		//		BasicFileWriter		  BufferedOutput(BlockFile, BufferingSize / 2);
 		for (const SharedBuffer& Segment : Buffer.GetSegments())
 		{
 			size_t SegmentSize = Segment.GetSize();
@@ -859,22 +884,49 @@ WriteToTempFile(CompositeBuffer&& Buffer, const std::filesystem::path& Path, std
 						 FileRef.FileChunkOffset,
 						 FileRef.FileChunkSize,
 						 BufferingSize,
-						 [&BufferedOutput, &Offset](const void* Data, size_t Size) {
-							 BufferedOutput.Write(Data, Size, Offset);
+						 [&Temp, &Offset](const void* Data, size_t Size) {
+							 Temp.Write(Data, Size, Offset);
 							 Offset += Size;
 						 });
 			}
 			else
 			{
-				BufferedOutput.Write(Segment.GetData(), SegmentSize, Offset);
+				Temp.Write(Segment.GetData(), SegmentSize, Offset);
 				Offset += SegmentSize;
 			}
 		}
 	}
-	void*	 FileHandle	 = BlockFile.Detach();
-	IoBuffer BlockBuffer = IoBuffer(IoBuffer::File, FileHandle, 0, Offset, /*IsWholeFile*/ true);
-	BlockBuffer.SetDeleteOnClose(true);
-	return BlockBuffer;
+
+	Temp.MoveTemporaryIntoPlace(Path, Ec);
+	if (Ec)
+	{
+		Ec.clear();
+		BasicFile OpenTemp(Path, BasicFile::Mode::kDelete, Ec);
+		if (Ec)
+		{
+			throw std::system_error(Ec, fmt::format("Failed to move temp file to '{}'", Path));
+		}
+		if (OpenTemp.FileSize() != BufferSize)
+		{
+			throw std::runtime_error(fmt::format("Failed to move temp file to '{}' - mismatching file size already exists", Path));
+		}
+		IoBuffer TmpBuffer(IoBuffer::File, OpenTemp.Detach(), 0, BufferSize, true);
+
+		IoHash		 ExistingHash = IoHash::HashBuffer(TmpBuffer);
+		const IoHash ExpectedHash = IoHash::HashBuffer(Buffer);
+		if (ExistingHash != ExpectedHash)
+		{
+			throw std::runtime_error(fmt::format("Failed to move temp file to '{}' - mismatching file hash already exists", Path));
+		}
+		Buffer = CompositeBuffer{};
+		TmpBuffer.SetDeleteOnClose(true);
+		return TmpBuffer;
+	}
+	Buffer = CompositeBuffer{};
+	BasicFile OpenTemp(Path, BasicFile::Mode::kDelete);
+	IoBuffer  TmpBuffer(IoBuffer::File, OpenTemp.Detach(), 0, BufferSize, true);
+	TmpBuffer.SetDeleteOnClose(true);
+	return TmpBuffer;
 }
 
 //////////////////////////////////////////////////////////////////////////