iterate chunks crash fix (#376)

* Bugfix: Add explicit lambda capture in CasContainer::IterateChunks to avoid accessing state data references

1 files changed, 46 insertions, 30 deletions

diff --git a/src/zenstore/buildstore/buildstore.cpp b/src/zenstore/buildstore/buildstore.cpp
index cf518c06f..b4891a742 100644
--- a/src/zenstore/buildstore/buildstore.cpp
+++ b/src/zenstore/buildstore/buildstore.cpp
@@ -442,7 +442,10 @@ BuildStore::GetMetadatas(std::span<const IoHash> BlobHashes, WorkerThreadPool* O
 		}
 	}
 
-	auto DoOneBlock = [&](std::span<const size_t> ChunkIndexes) {
+	auto DoOneBlock = [this](std::span<const BlockStoreLocation> MetaLocations,
+							 std::span<const size_t>			 MetaLocationResultIndexes,
+							 std::span<const size_t>			 ChunkIndexes,
+							 std::vector<IoBuffer>&				 Result) {
 		if (ChunkIndexes.size() < 4)
 		{
 			for (size_t ChunkIndex : ChunkIndexes)
@@ -459,7 +462,7 @@ BuildStore::GetMetadatas(std::span<const IoHash> BlobHashes, WorkerThreadPool* O
 		return m_MetadataBlockStore.IterateBlock(
 			MetaLocations,
 			ChunkIndexes,
-			[&](size_t ChunkIndex, const void* Data, uint64_t Size) {
+			[&MetaLocationResultIndexes, &Result](size_t ChunkIndex, const void* Data, uint64_t Size) {
 				if (Data != nullptr)
 				{
 					size_t ResultIndex	= MetaLocationResultIndexes[ChunkIndex];
@@ -479,38 +482,51 @@ BuildStore::GetMetadatas(std::span<const IoHash> BlobHashes, WorkerThreadPool* O
 	{
 		Latch WorkLatch(1);
 
-		m_MetadataBlockStore.IterateChunks(MetaLocations, [&](uint32_t BlockIndex, std::span<const size_t> ChunkIndexes) -> bool {
-			ZEN_UNUSED(BlockIndex);
-			if (ChunkIndexes.size() == MetaLocations.size() || OptionalWorkerPool == nullptr || ReferencedBlocks.size() == 1)
-			{
-				return DoOneBlock(ChunkIndexes);
-			}
-			else
-			{
-				ZEN_ASSERT(OptionalWorkerPool != nullptr);
-				WorkLatch.AddCount(1);
-				try
+		m_MetadataBlockStore.IterateChunks(
+			MetaLocations,
+			[this, OptionalWorkerPool, &Result, &MetaLocations, &MetaLocationResultIndexes, &ReferencedBlocks, DoOneBlock, &WorkLatch](
+				uint32_t				BlockIndex,
+				std::span<const size_t> ChunkIndexes) -> bool {
+				ZEN_UNUSED(BlockIndex);
+				if (ChunkIndexes.size() == MetaLocations.size() || OptionalWorkerPool == nullptr || ReferencedBlocks.size() == 1)
 				{
-					OptionalWorkerPool->ScheduleWork([&, ChunkIndexes = std::vector<size_t>(ChunkIndexes.begin(), ChunkIndexes.end())]() {
-						auto _ = MakeGuard([&WorkLatch]() { WorkLatch.CountDown(); });
-						try
-						{
-							DoOneBlock(ChunkIndexes);
-						}
-						catch (const std::exception& Ex)
-						{
-							ZEN_WARN("Failed getting metadata for {} chunks. Reason: {}", ChunkIndexes.size(), Ex.what());
-						}
-					});
+					return DoOneBlock(MetaLocations, MetaLocationResultIndexes, ChunkIndexes, Result);
 				}
-				catch (const std::exception& Ex)
+				else
 				{
-					WorkLatch.CountDown();
-					ZEN_ERROR("Failed dispatching async work to fetch metadata for {} chunks. Reason: {}", ChunkIndexes.size(), Ex.what());
+					ZEN_ASSERT(OptionalWorkerPool != nullptr);
+					std::vector<size_t> TmpChunkIndexes(ChunkIndexes.begin(), ChunkIndexes.end());
+					WorkLatch.AddCount(1);
+					try
+					{
+						OptionalWorkerPool->ScheduleWork([this,
+														  &Result,
+														  &MetaLocations,
+														  &MetaLocationResultIndexes,
+														  DoOneBlock,
+														  &WorkLatch,
+														  ChunkIndexes = std::move(TmpChunkIndexes)]() {
+							auto _ = MakeGuard([&WorkLatch]() { WorkLatch.CountDown(); });
+							try
+							{
+								DoOneBlock(MetaLocations, MetaLocationResultIndexes, ChunkIndexes, Result);
+							}
+							catch (const std::exception& Ex)
+							{
+								ZEN_WARN("Failed getting metadata for {} chunks. Reason: {}", ChunkIndexes.size(), Ex.what());
+							}
+						});
+					}
+					catch (const std::exception& Ex)
+					{
+						WorkLatch.CountDown();
+						ZEN_ERROR("Failed dispatching async work to fetch metadata for {} chunks. Reason: {}",
+								  ChunkIndexes.size(),
+								  Ex.what());
+					}
+					return true;
 				}
-				return true;
-			}
-		});
+			});
 
 		WorkLatch.CountDown();
 		WorkLatch.Wait();