fix race condition in zenserver during batched fetch (#149)

* fix race condition in zenserver duing batched fetch

1 files changed, 92 insertions, 1 deletions

diff --git a/src/zenstore/cache/cachedisklayer.cpp b/src/zenstore/cache/cachedisklayer.cpp
index 37404053e..7c9b56b1b 100644
--- a/src/zenstore/cache/cachedisklayer.cpp
+++ b/src/zenstore/cache/cachedisklayer.cpp
@@ -1340,6 +1340,7 @@ ZenCacheDiskLayer::CacheBucket::EndGetBatch(GetBatchHandle* Batch) noexcept
 		{
 			std::vector<DiskLocation>		StandaloneDiskLocations;
 			std::vector<size_t>				StandaloneKeyIndexes;
+			std::vector<size_t>				MemCachedKeyIndexes;
 			std::vector<DiskLocation>		InlineDiskLocations;
 			std::vector<BlockStoreLocation> InlineBlockLocations;
 			std::vector<size_t>				InlineKeyIndexes;
@@ -1372,6 +1373,7 @@ ZenCacheDiskLayer::CacheBucket::EndGetBatch(GetBatchHandle* Batch) noexcept
 						if (Payload.MemCached)
 						{
 							OutValue.Value = m_MemCachedPayloads[Payload.MemCached].Payload;
+							MemCachedKeyIndexes.push_back(KeyIndex);
 							m_MemoryHitCount++;
 						}
 						else
@@ -1396,11 +1398,17 @@ ZenCacheDiskLayer::CacheBucket::EndGetBatch(GetBatchHandle* Batch) noexcept
 				}
 			}
 
+			// MemCached and MetaData are set independently so we need to check if meta data has been set or if we need to set it even
+			// if we found the data as memcached.
+			// Often we will find the metadata due to the thread setting the mem cached part doing it before us so it is worth
+			// checking if it is present once more before spending time fetching and setting the RawHash and RawSize in metadata
+
 			auto FillOne = [&](const DiskLocation& Location, size_t KeyIndex, IoBuffer&& Value) {
 				if (!Value)
 				{
 					return;
 				}
+				const IoHash&  Key		   = Batch->Keys[KeyIndex];
 				size_t		   ResultIndex = Batch->ResultIndexes[KeyIndex];
 				ZenCacheValue& OutValue	   = Batch->OutResults[ResultIndex];
 				OutValue.Value			   = std::move(Value);
@@ -1420,6 +1428,24 @@ ZenCacheDiskLayer::CacheBucket::EndGetBatch(GetBatchHandle* Batch) noexcept
 
 				if (SetMetaInfo)
 				{
+					// See ZenCacheDiskLayer::CacheBucket::Get - it sets the memcache part first and then if it needs to it set the
+					// metadata sparately, check if it had time to set the metdata
+					RwLock::SharedLockScope UpdateIndexLock(m_IndexLock);
+					if (auto UpdateIt = m_Index.find(Key); UpdateIt != m_Index.end())
+					{
+						BucketPayload& Payload = m_Payloads[UpdateIt->second];
+						if (Payload.MetaData)
+						{
+							const BucketMetaData& MetaData = m_MetaDatas[Payload.MetaData];
+							OutValue.RawHash			   = MetaData.RawHash;
+							OutValue.RawSize			   = MetaData.RawSize;
+							SetMetaInfo					   = false;
+						}
+					}
+				}
+
+				if (SetMetaInfo)
+				{
 					ZEN_TRACE_CPU("Z$::Bucket::EndGetBatch::MetaData");
 					if (Location.IsFlagSet(DiskLocation::kCompressed))
 					{
@@ -1440,7 +1466,6 @@ ZenCacheDiskLayer::CacheBucket::EndGetBatch(GetBatchHandle* Batch) noexcept
 				if (SetMetaInfo || AddToMemCache)
 				{
 					ZEN_TRACE_CPU("Z$::Bucket::EndGetBatch::MemCache");
-					const IoHash&			   Key = Batch->Keys[KeyIndex];
 					RwLock::ExclusiveLockScope UpdateIndexLock(m_IndexLock);
 					{
 						if (auto UpdateIt = m_Index.find(Key); UpdateIt != m_Index.end())
@@ -1505,6 +1530,72 @@ ZenCacheDiskLayer::CacheBucket::EndGetBatch(GetBatchHandle* Batch) noexcept
 				}
 			}
 
+			if (!MemCachedKeyIndexes.empty())
+			{
+				ZEN_TRACE_CPU("Z$::Bucket::EndGetBatch::MemCached");
+				for (size_t KeyIndex : MemCachedKeyIndexes)
+				{
+					const IoHash& Key		  = Batch->Keys[KeyIndex];
+					bool		  SetMetaInfo = FillRawHashAndRawSize[KeyIndex];
+					if (SetMetaInfo)
+					{
+						ZEN_TRACE_CPU("Z$::Bucket::EndGetBatch::MetaData");
+						size_t		   ResultIndex = Batch->ResultIndexes[KeyIndex];
+						ZenCacheValue& OutValue	   = Batch->OutResults[ResultIndex];
+						{
+							// See ZenCacheDiskLayer::CacheBucket::Get - it sets the memcache part first and then if it needs to it set the
+							// metadata sparately, check if it had time to set the metdata
+							RwLock::SharedLockScope UpdateIndexLock(m_IndexLock);
+							if (auto UpdateIt = m_Index.find(Key); UpdateIt != m_Index.end())
+							{
+								BucketPayload& Payload = m_Payloads[UpdateIt->second];
+								if (Payload.MetaData)
+								{
+									const BucketMetaData& MetaData = m_MetaDatas[Payload.MetaData];
+									OutValue.RawHash			   = MetaData.RawHash;
+									OutValue.RawSize			   = MetaData.RawSize;
+									SetMetaInfo					   = false;
+								}
+							}
+						}
+
+						if (SetMetaInfo)
+						{
+							if (OutValue.Value.GetContentType() == ZenContentType::kCompressedBinary)
+							{
+								if (!CompressedBuffer::ValidateCompressedHeader(OutValue.Value, OutValue.RawHash, OutValue.RawSize))
+								{
+									OutValue = ZenCacheValue{};
+								}
+							}
+							else
+							{
+								OutValue.RawHash = IoHash::HashBuffer(OutValue.Value);
+								OutValue.RawSize = OutValue.Value.GetSize();
+							}
+
+							{
+								RwLock::ExclusiveLockScope UpdateIndexLock(m_IndexLock);
+								{
+									if (auto UpdateIt = m_Index.find(Key); UpdateIt != m_Index.end())
+									{
+										BucketPayload& Payload = m_Payloads[UpdateIt->second];
+
+										// Only update if it has not already been updated by other thread
+										if (!Payload.MetaData)
+										{
+											SetMetaData(UpdateIndexLock,
+														Payload,
+														{.RawSize = OutValue.RawSize, .RawHash = OutValue.RawHash});
+										}
+									}
+								}
+							}
+						}
+					}
+				}
+			}
+
 			for (size_t ResultIndex : Batch->ResultIndexes)
 			{
 				bool Hit = !!Batch->OutResults[ResultIndex].Value;