Reduced disk I/O when writing out chunk blocks during download (#309)

1 files changed, 81 insertions, 70 deletions

diff --git a/src/zen/cmds/builds_cmd.cpp b/src/zen/cmds/builds_cmd.cpp
index 61e3c0fab..8e3d50790 100644
--- a/src/zen/cmds/builds_cmd.cpp
+++ b/src/zen/cmds/builds_cmd.cpp
@@ -3955,35 +3955,67 @@ namespace {
 		}
 	}
 
+	IoBuffer MakeBufferMemoryBased(const CompositeBuffer& PartialBlockBuffer)
+	{
+		ZEN_TRACE_CPU("MakeBufferMemoryBased");
+		IoBuffer					  BlockMemoryBuffer;
+		std::span<const SharedBuffer> Segments = PartialBlockBuffer.GetSegments();
+		if (Segments.size() == 1)
+		{
+			IoBufferFileReference FileRef = {};
+			if (PartialBlockBuffer.GetSegments().front().AsIoBuffer().GetFileReference(FileRef))
+			{
+				BlockMemoryBuffer = UniqueBuffer::Alloc(FileRef.FileChunkSize).MoveToShared().AsIoBuffer();
+				BasicFile Reader;
+				Reader.Attach(FileRef.FileHandle);
+				auto			  _		  = MakeGuard([&Reader]() { Reader.Detach(); });
+				MutableMemoryView ReadMem = BlockMemoryBuffer.GetMutableView();
+				Reader.Read(ReadMem.GetData(), FileRef.FileChunkSize, FileRef.FileChunkOffset);
+				return BlockMemoryBuffer;
+			}
+			else
+			{
+				return PartialBlockBuffer.GetSegments().front().AsIoBuffer();
+			}
+		}
+		else
+		{
+			// Not a homogenous memory buffer, read all to memory
+
+			BlockMemoryBuffer		  = UniqueBuffer::Alloc(PartialBlockBuffer.GetSize()).MoveToShared().AsIoBuffer();
+			MutableMemoryView ReadMem = BlockMemoryBuffer.GetMutableView();
+			for (const SharedBuffer& Segment : Segments)
+			{
+				IoBufferFileReference FileRef = {};
+				if (Segment.AsIoBuffer().GetFileReference(FileRef))
+				{
+					BasicFile Reader;
+					Reader.Attach(FileRef.FileHandle);
+					Reader.Read(ReadMem.GetData(), FileRef.FileChunkSize, FileRef.FileChunkOffset);
+					Reader.Detach();
+					ReadMem = ReadMem.Mid(FileRef.FileChunkSize);
+				}
+				else
+				{
+					ReadMem = ReadMem.CopyFrom(Segment.AsIoBuffer().GetView());
+				}
+			}
+			return BlockMemoryBuffer;
+		}
+	}
+
 	bool GetBlockWriteOps(const ChunkedFolderContent&	   RemoteContent,
 						  const ChunkedContentLookup&	   Lookup,
 						  std::span<const IoHash>		   ChunkRawHashes,
 						  std::span<const uint32_t>		   ChunkCompressedLengths,
-						  std::span<const uint32_t>		   ChunkRawLengths,
 						  std::span<std::atomic<uint32_t>> SequenceIndexChunksLeftToWriteCounters,
 						  std::span<std::atomic<bool>>	   RemoteChunkIndexNeedsCopyFromSourceFlags,
-						  CompositeBuffer&&				   PartialBlockBuffer,
+						  const MemoryView				   BlockView,
 						  uint32_t						   FirstIncludedBlockChunkIndex,
 						  uint32_t						   LastIncludedBlockChunkIndex,
 						  BlockWriteOps&				   OutOps)
 	{
 		ZEN_TRACE_CPU("GetBlockWriteOps");
-		MemoryView			  BlockMemoryView;
-		UniqueBuffer		  BlockMemoryBuffer;
-		IoBufferFileReference FileRef = {};
-		if (PartialBlockBuffer.GetSegments().size() == 1 && PartialBlockBuffer.GetSegments()[0].AsIoBuffer().GetFileReference(FileRef))
-		{
-			BlockMemoryBuffer = UniqueBuffer::Alloc(FileRef.FileChunkSize);
-			BasicFile Reader;
-			Reader.Attach(FileRef.FileHandle);
-			Reader.Read(BlockMemoryBuffer.GetData(), FileRef.FileChunkSize, FileRef.FileChunkOffset);
-			BlockMemoryView = BlockMemoryBuffer.GetView();
-			Reader.Detach();
-		}
-		else
-		{
-			BlockMemoryView = PartialBlockBuffer.ViewOrCopyRange(0, PartialBlockBuffer.GetSize(), BlockMemoryBuffer);
-		}
 		uint32_t OffsetInBlock = 0;
 		for (uint32_t ChunkBlockIndex = FirstIncludedBlockChunkIndex; ChunkBlockIndex <= LastIncludedBlockChunkIndex; ChunkBlockIndex++)
 		{
@@ -4000,32 +4032,9 @@ namespace {
 					bool NeedsWrite = true;
 					if (RemoteChunkIndexNeedsCopyFromSourceFlags[ChunkIndex].compare_exchange_strong(NeedsWrite, false))
 					{
-						// CompositeBuffer	 Chunk = PartialBlockBuffer.Mid(OffsetInBlock, ChunkCompressedSize);
-						MemoryView		 ChunkMemory = BlockMemoryView.Mid(OffsetInBlock, ChunkCompressedSize);
-						CompositeBuffer	 Chunk = CompositeBuffer(IoBuffer(IoBuffer::Wrap, ChunkMemory.GetData(), ChunkMemory.GetSize()));
-						IoHash			 VerifyChunkHash;
-						uint64_t		 VerifyRawSize;
-						CompressedBuffer Compressed = CompressedBuffer::FromCompressed(Chunk, VerifyChunkHash, VerifyRawSize);
-						if (!Compressed)
-						{
-							ZEN_ASSERT(false);
-						}
-						if (VerifyChunkHash != ChunkHash)
-						{
-							ZEN_ASSERT(false);
-						}
-						if (!ChunkRawLengths.empty())
-						{
-							if (VerifyRawSize != ChunkRawLengths[ChunkBlockIndex])
-							{
-								ZEN_ASSERT(false);
-							}
-						}
-						CompositeBuffer Decompressed = Compressed.DecompressToComposite();
-						if (!Decompressed)
-						{
-							throw std::runtime_error(fmt::format("Decompression of build blob {} failed", ChunkHash));
-						}
+						MemoryView ChunkMemoryView = BlockView.Mid(OffsetInBlock + CompressedBuffer::GetHeaderSizeForNoneEncoder(),
+																   ChunkCompressedSize - CompressedBuffer::GetHeaderSizeForNoneEncoder());
+						IoBuffer   Decompressed(IoBuffer::Wrap, ChunkMemoryView.GetData(), ChunkMemoryView.GetSize());
 						ZEN_ASSERT_SLOW(ChunkHash == IoHash::HashBuffer(Decompressed));
 						ZEN_ASSERT(Decompressed.GetSize() == RemoteContent.ChunkedContent.ChunkRawSizes[ChunkIndex]);
 						for (const ChunkedContentLookup::ChunkSequenceLocation* Target : ChunkTargetPtrs)
@@ -4040,19 +4049,22 @@ namespace {
 
 			OffsetInBlock += ChunkCompressedSize;
 		}
-		std::sort(OutOps.WriteOps.begin(),
-				  OutOps.WriteOps.end(),
-				  [](const BlockWriteOps::WriteOpData& Lhs, const BlockWriteOps::WriteOpData& Rhs) {
-					  if (Lhs.Target->SequenceIndex < Rhs.Target->SequenceIndex)
-					  {
-						  return true;
-					  }
-					  if (Lhs.Target->SequenceIndex > Rhs.Target->SequenceIndex)
-					  {
-						  return false;
-					  }
-					  return Lhs.Target->Offset < Rhs.Target->Offset;
-				  });
+		{
+			ZEN_TRACE_CPU("GetBlockWriteOps_sort");
+			std::sort(OutOps.WriteOps.begin(),
+					  OutOps.WriteOps.end(),
+					  [](const BlockWriteOps::WriteOpData& Lhs, const BlockWriteOps::WriteOpData& Rhs) {
+						  if (Lhs.Target->SequenceIndex < Rhs.Target->SequenceIndex)
+						  {
+							  return true;
+						  }
+						  if (Lhs.Target->SequenceIndex > Rhs.Target->SequenceIndex)
+						  {
+							  return false;
+						  }
+						  return Lhs.Target->Offset < Rhs.Target->Offset;
+					  });
+		}
 		return true;
 	}
 
@@ -4068,26 +4080,25 @@ namespace {
 	{
 		ZEN_TRACE_CPU("WriteBlockToDisk");
 
+		IoBuffer		 BlockMemoryBuffer = MakeBufferMemoryBased(BlockBuffer);
+		const MemoryView BlockView		   = BlockMemoryBuffer.GetView();
+
 		BlockWriteOps Ops;
 		if ((BlockDescription.HeaderSize == 0) || BlockDescription.ChunkCompressedLengths.empty())
 		{
 			ZEN_TRACE_CPU("WriteBlockToDisk_Legacy");
 
-			UniqueBuffer				CopyBuffer;
-			const MemoryView			BlockView = BlockBuffer.ViewOrCopyRange(0, BlockBuffer.GetSize(), CopyBuffer);
 			uint64_t					HeaderSize;
-			const std::vector<uint32_t> ChunkCompressedLengths = ReadChunkBlockHeader(BlockView, HeaderSize);
-
-			CompositeBuffer PartialBlockBuffer = std::move(BlockBuffer).Mid(CompressedBuffer::GetHeaderSizeForNoneEncoder() + HeaderSize);
+			const std::vector<uint32_t> ChunkCompressedLengths =
+				ReadChunkBlockHeader(BlockView.Mid(CompressedBuffer::GetHeaderSizeForNoneEncoder()), HeaderSize);
 
 			if (GetBlockWriteOps(RemoteContent,
 								 Lookup,
 								 BlockDescription.ChunkRawHashes,
 								 ChunkCompressedLengths,
-								 BlockDescription.ChunkRawLengths,
 								 SequenceIndexChunksLeftToWriteCounters,
 								 RemoteChunkIndexNeedsCopyFromSourceFlags,
-								 std::move(PartialBlockBuffer),
+								 BlockView.Mid(CompressedBuffer::GetHeaderSizeForNoneEncoder() + HeaderSize),
 								 0,
 								 gsl::narrow<uint32_t>(BlockDescription.ChunkRawHashes.size() - 1),
 								 Ops))
@@ -4104,16 +4115,13 @@ namespace {
 			return false;
 		}
 
-		CompositeBuffer PartialBlockBuffer =
-			std::move(BlockBuffer).Mid(CompressedBuffer::GetHeaderSizeForNoneEncoder() + BlockDescription.HeaderSize);
 		if (GetBlockWriteOps(RemoteContent,
 							 Lookup,
 							 BlockDescription.ChunkRawHashes,
 							 BlockDescription.ChunkCompressedLengths,
-							 BlockDescription.ChunkRawLengths,
 							 SequenceIndexChunksLeftToWriteCounters,
 							 RemoteChunkIndexNeedsCopyFromSourceFlags,
-							 std::move(PartialBlockBuffer),
+							 BlockView.Mid(CompressedBuffer::GetHeaderSizeForNoneEncoder() + BlockDescription.HeaderSize),
 							 0,
 							 gsl::narrow<uint32_t>(BlockDescription.ChunkRawHashes.size() - 1),
 							 Ops))
@@ -4143,15 +4151,18 @@ namespace {
 								 WriteChunkStatistics&			  WriteChunkStats)
 	{
 		ZEN_TRACE_CPU("WritePartialBlockToDisk");
+
+		IoBuffer		 BlockMemoryBuffer = MakeBufferMemoryBased(PartialBlockBuffer);
+		const MemoryView BlockView		   = BlockMemoryBuffer.GetView();
+
 		BlockWriteOps Ops;
 		if (GetBlockWriteOps(RemoteContent,
 							 Lookup,
 							 BlockDescription.ChunkRawHashes,
 							 BlockDescription.ChunkCompressedLengths,
-							 BlockDescription.ChunkRawLengths,
 							 SequenceIndexChunksLeftToWriteCounters,
 							 RemoteChunkIndexNeedsCopyFromSourceFlags,
-							 std::move(PartialBlockBuffer),
+							 BlockView,
 							 FirstIncludedBlockChunkIndex,
 							 LastIncludedBlockChunkIndex,
 							 Ops))