From dec3f27c488a1dda8a2f1133361e2fda9315e0d2 Mon Sep 17 00:00:00 2001
From: Dan Engelbrecht <dan.engelbrecht@epicgames.com>
Date: Fri, 11 Apr 2025 19:22:21 +0200
Subject: fix race condition in multipart download (#358)

---
 src/zenutil/filebuildstorage.cpp | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

(limited to 'src/zenutil/filebuildstorage.cpp')

diff --git a/src/zenutil/filebuildstorage.cpp b/src/zenutil/filebuildstorage.cpp
index 7aa252e44..f335a03a3 100644
--- a/src/zenutil/filebuildstorage.cpp
+++ b/src/zenutil/filebuildstorage.cpp
@@ -369,11 +369,11 @@ public:
 		return IoBuffer{};
 	}
 
-	virtual std::vector<std::function<void()>> GetLargeBuildBlob(
-		const Oid&																			   BuildId,
-		const IoHash&																		   RawHash,
-		uint64_t																			   ChunkSize,
-		std::function<void(uint64_t Offset, const IoBuffer& Chunk, uint64_t BytesRemaining)>&& Receiver) override
+	virtual std::vector<std::function<void()>> GetLargeBuildBlob(const Oid&													   BuildId,
+																 const IoHash&												   RawHash,
+																 uint64_t													   ChunkSize,
+																 std::function<void(uint64_t Offset, const IoBuffer& Chunk)>&& OnReceive,
+																 std::function<void()>&& OnComplete) override
 	{
 		ZEN_TRACE_CPU("FileBuildStorage::GetLargeBuildBlob");
 		ZEN_UNUSED(BuildId);
@@ -387,16 +387,18 @@ public:
 		{
 			struct WorkloadData
 			{
-				std::atomic<uint64_t>																 BytesRemaining;
-				BasicFile																			 BlobFile;
-				std::function<void(uint64_t Offset, const IoBuffer& Chunk, uint64_t BytesRemaining)> Receiver;
+				std::atomic<uint64_t>										BytesRemaining;
+				BasicFile													BlobFile;
+				std::function<void(uint64_t Offset, const IoBuffer& Chunk)> OnReceive;
+				std::function<void()>										OnComplete;
 			};
 
 			std::shared_ptr<WorkloadData> Workload(std::make_shared<WorkloadData>());
 			Workload->BlobFile.Open(BlockPath, BasicFile::Mode::kRead);
 			const uint64_t BlobSize = Workload->BlobFile.FileSize();
 
-			Workload->Receiver		 = std::move(Receiver);
+			Workload->OnReceive		 = std::move(OnReceive);
+			Workload->OnComplete	 = std::move(OnComplete);
 			Workload->BytesRemaining = BlobSize;
 
 			std::vector<std::function<void()>> WorkItems;
@@ -410,8 +412,12 @@ public:
 					IoBuffer PartPayload(Size);
 					Workload->BlobFile.Read(PartPayload.GetMutableView().GetData(), Size, Offset);
 					m_Stats.TotalBytesRead += PartPayload.GetSize();
+					Workload->OnReceive(Offset, PartPayload);
 					uint64_t ByteRemaning = Workload->BytesRemaining.fetch_sub(Size);
-					Workload->Receiver(Offset, PartPayload, ByteRemaning);
+					if (ByteRemaning == Size)
+					{
+						Workload->OnComplete();
+					}
 					SimulateLatency(Size, PartPayload.GetSize());
 				});
 
-- 
cgit v1.2.3