Make IoBufferCore::Materialize threadsafe

2 files changed, 58 insertions, 43 deletions

diff --git a/zencore/include/zencore/iobuffer.h b/zencore/include/zencore/iobuffer.h
index 04b3b33dd..a39f442e3 100644
--- a/zencore/include/zencore/iobuffer.h
+++ b/zencore/include/zencore/iobuffer.h
@@ -2,6 +2,7 @@
 
 #pragma once
 
+#include <atomic>
 #include <memory.h>
 #include <zencore/memory.h>
 #include "refcount.h"
@@ -111,39 +112,39 @@ public:
 
 	inline void EnsureDataValid() const
 	{
-		if ((m_Flags & kIsExtended) && !(m_Flags & kIsMaterialized))
+		uint32_t LocalFlags = m_Flags.load(std::memory_order_acquire);
+		if ((LocalFlags & kIsExtended) && !(LocalFlags & kIsMaterialized))
 		{
 			Materialize();
 		}
 	}
 
-	inline bool IsOwnedByThis() const { return !!(m_Flags & kIsOwnedByThis); }
+	inline bool IsOwnedByThis() const { return !!(m_Flags.load(std::memory_order_relaxed) & kIsOwnedByThis); }
 
 	inline void SetIsOwnedByThis(bool NewState)
 	{
 		if (NewState)
 		{
-			m_Flags |= kIsOwnedByThis;
+			m_Flags.fetch_or(kIsOwnedByThis, std::memory_order_relaxed);
 		}
 		else
 		{
-			m_Flags &= ~kIsOwnedByThis;
+			m_Flags.fetch_and(~kIsOwnedByThis, std::memory_order_relaxed);
 		}
 	}
 
 	inline bool IsOwned() const
 	{
-		bool ThisIsOwned = !!(m_Flags & kIsOwnedByThis);
-		if (ThisIsOwned)
+		if (IsOwnedByThis())
 		{
 			return true;
 		}
 		return m_OuterCore && m_OuterCore->IsOwned();
 	}
 
-	inline bool IsImmutable() const { return (m_Flags & kIsMutable) == 0; }
-	inline bool IsWholeFile() const { return (m_Flags & kIsWholeFile) != 0; }
-	inline bool IsNull() const { return (m_Flags & kIsNull) != 0; }
+	inline bool IsImmutable() const { return (m_Flags.load(std::memory_order_relaxed) & kIsMutable) == 0; }
+	inline bool IsWholeFile() const { return (m_Flags.load(std::memory_order_relaxed) & kIsWholeFile) != 0; }
+	inline bool IsNull() const { return (m_Flags.load(std::memory_order_relaxed) & kIsNull) != 0; }
 
 	inline IoBufferExtendedCore*	   ExtendedCore();
 	inline const IoBufferExtendedCore* ExtendedCore() const;
@@ -168,11 +169,11 @@ public:
 	{
 		if (!NewState)
 		{
-			m_Flags |= kIsMutable;
+			m_Flags.fetch_or(kIsMutable, std::memory_order_relaxed);
 		}
 		else
 		{
-			m_Flags &= ~kIsMutable;
+			m_Flags.fetch_and(~kIsMutable, std::memory_order_relaxed);
 		}
 	}
 
@@ -180,27 +181,36 @@ public:
 	{
 		if (NewState)
 		{
-			m_Flags |= kIsWholeFile;
+			m_Flags.fetch_or(kIsWholeFile, std::memory_order_relaxed);
 		}
 		else
 		{
-			m_Flags |= ~kIsWholeFile;
+			m_Flags.fetch_and(~kIsWholeFile, std::memory_order_relaxed);
 		}
 	}
 
 	inline void SetContentType(ZenContentType ContentType)
 	{
 		ZEN_ASSERT_SLOW((uint32_t(ContentType) & kContentTypeMask) == uint32_t(ContentType));
-		m_Flags = (m_Flags & ~(kContentTypeMask << kContentTypeShift)) | (uint32_t(ContentType) << kContentTypeShift);
+		uint32_t OldValue = m_Flags.load(std::memory_order_relaxed);
+		bool	 bSucceeded;
+		do
+		{
+			uint32_t NewValue = (OldValue & ~(kContentTypeMask << kContentTypeShift)) | (uint32_t(ContentType) << kContentTypeShift);
+			bSucceeded		  = m_Flags.compare_exchange_weak(OldValue, NewValue, std::memory_order_relaxed, std::memory_order_relaxed);
+		} while (!bSucceeded);
 	}
 
-	inline ZenContentType GetContentType() const { return ZenContentType((m_Flags >> kContentTypeShift) & kContentTypeMask); }
+	inline ZenContentType GetContentType() const
+	{
+		return ZenContentType((m_Flags.load(std::memory_order_relaxed) >> kContentTypeShift) & kContentTypeMask);
+	}
 
 	inline uint32_t GetRefCount() const { return m_RefCount; }
 
 protected:
 	uint32_t				   m_RefCount  = 0;
-	mutable uint32_t		   m_Flags	   = 0;
+	mutable std::atomic<uint32_t> m_Flags{0};
 	mutable const void*		   m_DataPtr   = nullptr;
 	size_t					   m_DataBytes = 0;
 	RefPtr<const IoBufferCore> m_OuterCore;
@@ -213,7 +223,7 @@ protected:
 
 	static_assert((uint32_t(ZenContentType::kUnknownContentType) & ~kContentTypeMask) == 0);
 
-	enum Flags
+	enum Flags : uint32_t
 	{
 		kIsNull			= 1 << 0,  // This is a null IoBuffer
 		kIsMutable		= 1 << 1,
@@ -266,7 +276,7 @@ private:
 inline IoBufferExtendedCore*
 IoBufferCore::ExtendedCore()
 {
-	if (m_Flags & kIsExtended)
+	if (m_Flags.load(std::memory_order_relaxed) & kIsExtended)
 	{
 		return static_cast<IoBufferExtendedCore*>(this);
 	}
@@ -277,7 +287,7 @@ IoBufferCore::ExtendedCore()
 inline const IoBufferExtendedCore*
 IoBufferCore::ExtendedCore() const
 {
-	if (m_Flags & kIsExtended)
+	if (m_Flags.load(std::memory_order_relaxed) & kIsExtended)
 	{
 		return static_cast<const IoBufferExtendedCore*>(this);
 	}
diff --git a/zencore/iobuffer.cpp b/zencore/iobuffer.cpp
index 7077942bf..9a7c6b30e 100644
--- a/zencore/iobuffer.cpp
+++ b/zencore/iobuffer.cpp
@@ -39,7 +39,7 @@ IoBufferCore::AllocateBuffer(size_t InSize, size_t Alignment)
 #if ZEN_PLATFORM_WINDOWS
 	if (((InSize & 0xffFF) == 0) && (Alignment == 0x10000))
 	{
-		m_Flags |= kLowLevelAlloc;
+		m_Flags.fetch_or(kLowLevelAlloc, std::memory_order_relaxed);
 		m_DataPtr = VirtualAlloc(nullptr, InSize, MEM_COMMIT, PAGE_READWRITE);
 
 		return;
@@ -48,7 +48,7 @@ IoBufferCore::AllocateBuffer(size_t InSize, size_t Alignment)
 
 #if ZEN_USE_MIMALLOC
 	void* Ptr = mi_aligned_alloc(Alignment, RoundUp(InSize, Alignment));
-	m_Flags |= kIoBufferAlloc;
+	m_Flags.fetch_or(kIoBufferAlloc, std::memory_order_relaxed);
 #else
 	void* Ptr = Memory::Alloc(InSize, Alignment);
 #endif
@@ -66,8 +66,9 @@ IoBufferCore::FreeBuffer()
 		return;
 	}
 
+	uint32_t LocalFlags = m_Flags.load(std::memory_order_relaxed);
 #if ZEN_PLATFORM_WINDOWS
-	if (m_Flags & kLowLevelAlloc)
+	if (LocalFlags & kLowLevelAlloc)
 	{
 		VirtualFree(const_cast<void*>(m_DataPtr), 0, MEM_DECOMMIT);
 
@@ -76,7 +77,7 @@ IoBufferCore::FreeBuffer()
 #endif	// ZEN_PLATFORM_WINDOWS
 
 #if ZEN_USE_MIMALLOC
-	if (m_Flags & kIoBufferAlloc)
+	if (LocalFlags & kIoBufferAlloc)
 	{
 		return mi_free(const_cast<void*>(m_DataPtr));
 	}
@@ -170,11 +171,11 @@ IoBufferExtendedCore::IoBufferExtendedCore(void* FileHandle, uint64_t Offset, ui
 , m_FileHandle(FileHandle)
 , m_FileOffset(Offset)
 {
-	m_Flags |= kIsOwnedByThis | kIsExtended;
+	m_Flags.fetch_or(kIsOwnedByThis | kIsExtended, std::memory_order_relaxed);
 
 	if (TransferHandleOwnership)
 	{
-		m_Flags |= kOwnsFile;
+		m_Flags.fetch_or(kOwnsFile, std::memory_order_relaxed);
 	}
 }
 
@@ -183,7 +184,7 @@ IoBufferExtendedCore::IoBufferExtendedCore(const IoBufferExtendedCore* Outer, ui
 , m_FileHandle(Outer->m_FileHandle)
 , m_FileOffset(Outer->m_FileOffset + Offset)
 {
-	m_Flags |= kIsOwnedByThis | kIsExtended;
+	m_Flags.fetch_or(kIsOwnedByThis | kIsExtended, std::memory_order_relaxed);
 }
 
 IoBufferExtendedCore::~IoBufferExtendedCore()
@@ -198,14 +199,15 @@ IoBufferExtendedCore::~IoBufferExtendedCore()
 #endif
 	}
 
+	uint32_t LocalFlags = m_Flags.load(std::memory_order_relaxed);
 #if ZEN_PLATFORM_WINDOWS
-	if (m_Flags & kOwnsMmap)
+	if (LocalFlags & kOwnsMmap)
 	{
 		CloseHandle(m_MmapHandle);
 	}
 #endif
 
-	if (m_Flags & kOwnsFile)
+	if (LocalFlags & kOwnsFile)
 	{
 #if ZEN_PLATFORM_WINDOWS
 		BOOL Success = CloseHandle(m_FileHandle);
@@ -233,43 +235,45 @@ IoBufferExtendedCore::Materialize() const
 	// The synchronization scheme here is very primitive, if we end up with
 	// a lot of contention we can make it more fine-grained
 
-	if (m_MmapHandle)
+	if (m_Flags.load(std::memory_order_acquire) & kIsMaterialized)
 		return;
 
 	RwLock::ExclusiveLockScope _(g_MappingLock);
 
 	// Someone could have gotten here first
-	if (m_MmapHandle)
+	if (m_Flags.load(std::memory_order_acquire) & kIsMaterialized)
 		return;
 
+	uint32_t NewFlags = kIsMaterialized;
+
 	const uint64_t MapOffset				= m_FileOffset & ~0xffffull;
 	const uint64_t MappedOffsetDisplacement = m_FileOffset - MapOffset;
 	const uint64_t MapSize					= m_DataBytes + MappedOffsetDisplacement;
 
 #if ZEN_PLATFORM_WINDOWS
-	m_MmapHandle = CreateFileMapping(m_FileHandle,
-									 /* lpFileMappingAttributes */ nullptr,
-									 /* flProtect */ PAGE_READONLY,
-									 /* dwMaximumSizeLow */ 0,
-									 /* dwMaximumSizeHigh */ 0,
-									 /* lpName */ nullptr);
-
-	if (m_MmapHandle == nullptr)
+	void* NewMmapHandle = CreateFileMapping(m_FileHandle,
+											/* lpFileMappingAttributes */ nullptr,
+											/* flProtect */ PAGE_READONLY,
+											/* dwMaximumSizeLow */ 0,
+											/* dwMaximumSizeHigh */ 0,
+											/* lpName */ nullptr);
+
+	if (NewMmapHandle == nullptr)
 	{
 		throw std::system_error(std::error_code(::GetLastError(), std::system_category()),
 								"CreateFileMapping failed on file '{}'"_format(zen::PathFromHandle(m_FileHandle)));
 	}
 
-	m_Flags |= kOwnsMmap;
+	NewFlags |= kOwnsMmap;
 
-	void* MappedBase = MapViewOfFile(m_MmapHandle,
+	void* MappedBase = MapViewOfFile(NewMmapHandle,
 									 /* dwDesiredAccess */ FILE_MAP_READ,
 									 /* FileOffsetHigh */ uint32_t(MapOffset >> 32),
 									 /* FileOffsetLow */ uint32_t(MapOffset & 0xffFFffFFu),
 									 /* dwNumberOfBytesToMap */ MapSize);
 #else
-	m_MmapHandle = (void*)uintptr_t(~MapSize);	// ~ so it's never null (assuming MapSize >= 0)
-	m_Flags |= kOwnsMmap;
+	NewMmapHandle = (void*)uintptr_t(~MapSize);	 // ~ so it's never null (assuming MapSize >= 0)
+	NewFlags |= kOwnsMmap;
 
 	void* MappedBase = mmap(
 		/* addr */ nullptr,
@@ -289,8 +293,9 @@ IoBufferExtendedCore::Materialize() const
 
 	m_MappedPointer = MappedBase;
 	m_DataPtr		= reinterpret_cast<uint8_t*>(MappedBase) + MappedOffsetDisplacement;
+	m_MmapHandle	= NewMmapHandle;
 
-	m_Flags |= kIsMaterialized;
+	m_Flags.fetch_or(NewFlags, std::memory_order_release);
 }
 
 bool