Fix crash due to async content data streaming

Properly checks for asset data unloading before taking lock on asset chunks.
#3358 #3085 #3515

Source/Engine/Content/Storage/FlaxChunk.h

@@ -87,6 +87,10 @@ public:
     /// </summary>
     double LastAccessTime = 0.0;
 
+    /// <summary>
+    /// Flag set to indicate that chunk is during loading (atomic access to sync multiple reading threads).
+    /// </summary>
+    int64 IsLoading = 0;
     /// <summary>
     /// The chunk data.
     /// </summary>
@@ -146,7 +150,7 @@ public:
     /// </summary>
     FORCE_INLINE bool IsLoaded() const
     {
-        return Data.IsValid();
+        return Data.IsValid() && Platform::AtomicRead(&IsLoading) == 0;
     }
 
     /// <summary>
@@ -154,7 +158,7 @@ public:
     /// </summary>
     FORCE_INLINE bool IsMissing() const
     {
-        return Data.IsInvalid();
+        return !IsLoaded();
     }
 
     /// <summary>

Source/Engine/Content/Storage/FlaxStorage.cpp

@@ -5,6 +5,7 @@
 #include "FlaxPackage.h"
 #include "ContentStorageManager.h"
 #include "Engine/Core/Log.h"
+#include "Engine/Core/ScopeExit.h"
 #include "Engine/Core/Types/TimeSpan.h"
 #include "Engine/Platform/File.h"
 #include "Engine/Profiler/ProfilerCPU.h"
@@ -246,6 +247,7 @@ FlaxStorage::~FlaxStorage()
     ASSERT(IsDisposed());
     CHECK(_chunksLock == 0);
     CHECK(_refCount == 0);
+    CHECK(_isUnloadingData == 0);
     ASSERT(_chunks.IsEmpty());
 
 #if USE_EDITOR
@@ -261,6 +263,22 @@ FlaxStorage::~FlaxStorage()
 #endif
 }
 
+void FlaxStorage::LockChunks()
+{
+RETRY:
+    Platform::InterlockedIncrement(&_chunksLock);
+    if (Platform::AtomicRead(&_isUnloadingData) != 0)
+    {
+        // Someone else is closing file handles or freeing chunks so wait for it to finish and retry
+        Platform::InterlockedDecrement(&_chunksLock);
+        do
+        {
+            Platform::Sleep(1);
+        } while (Platform::AtomicRead(&_isUnloadingData) != 0);
+        goto RETRY;
+    }
+}
+
 FlaxStorage::LockData FlaxStorage::LockSafe()
 {
     auto lock = LockData(this);
@@ -689,7 +707,6 @@ bool FlaxStorage::LoadAssetHeader(const Guid& id, AssetInitData& data)
         return true;
     }
 
-    // Load header
     return LoadAssetHeader(e, data);
 }
 
@@ -699,7 +716,10 @@ bool FlaxStorage::LoadAssetChunk(FlaxChunk* chunk)
     ASSERT(IsLoaded());
     ASSERT(chunk != nullptr && _chunks.Contains(chunk));
 
-    // Check if already loaded
+    // Protect against loading the same chunk from multiple threads at once
+    while (Platform::InterlockedCompareExchange(&chunk->IsLoading, 1, 0) != 0)
+        Platform::Sleep(1);
+    SCOPE_EXIT{ Platform::AtomicStore(&chunk->IsLoading, 0); };
     if (chunk->IsLoaded())
         return false;
 
@@ -776,12 +796,10 @@ bool FlaxStorage::LoadAssetChunk(FlaxChunk* chunk)
             // Raw data
             chunk->Data.Read(stream, size);
         }
-        ASSERT(chunk->IsLoaded());
         chunk->RegisterUsage();
     }
 
     UnlockChunks();
-
     return failed;
 }
 
@@ -1420,10 +1438,12 @@ FileReadStream* FlaxStorage::OpenFile()
 
 bool FlaxStorage::CloseFileHandles()
 {
+    // Guard the whole process so if new thread wants to lock the chunks will need to wait for this to end
+    Platform::InterlockedIncrement(&_isUnloadingData);
+    SCOPE_EXIT{ Platform::InterlockedDecrement(&_isUnloadingData); };
+
     if (Platform::AtomicRead(&_chunksLock) == 0 && Platform::AtomicRead(&_files) == 0)
-    {
-        return false;
-    }
+        return false; // Early out when no files are opened
     PROFILE_CPU();
     PROFILE_MEM(ContentFiles);
 
@@ -1496,9 +1516,21 @@ void FlaxStorage::Tick(double time)
     {
         auto chunk = _chunks.Get()[i];
         const bool wasUsed = (time - chunk->LastAccessTime) < unusedDataChunksLifetime;
-        if (!wasUsed && chunk->IsLoaded() && EnumHasNoneFlags(chunk->Flags, FlaxChunkFlags::KeepInMemory))
+        if (!wasUsed && 
+            chunk->IsLoaded() && 
+            EnumHasNoneFlags(chunk->Flags, FlaxChunkFlags::KeepInMemory) && 
+            Platform::AtomicRead(&chunk->IsLoading) == 0)
         {
+            // Guard the unloading so if other thread wants to lock the chunks will need to wait for this to end
+            Platform::InterlockedIncrement(&_isUnloadingData);
+            if (Platform::AtomicRead(&_chunksLock) != 0 || Platform::AtomicRead(&chunk->IsLoading) != 0)
+            {
+                // Someone started loading so skip ticking
+                Platform::InterlockedDecrement(&_isUnloadingData);
+                return;
+            }
             chunk->Unload();
+            Platform::InterlockedDecrement(&_isUnloadingData);
         }
         wasAnyUsed |= wasUsed;
     }

Source/Engine/Content/Storage/FlaxStorage.h

@@ -90,6 +90,7 @@ protected:
     int64 _refCount = 0;
     int64 _chunksLock = 0;
     int64 _files = 0;
+    int64 _isUnloadingData = 0;
     double _lastRefLostTime;
     CriticalSection _loadLocker;
 
@@ -129,10 +130,7 @@ public:
     /// <summary>
     /// Locks the storage chunks data to prevent disposing them. Also ensures that file handles won't be closed while chunks are locked.
     /// </summary>
-    FORCE_INLINE void LockChunks()
-    {
-        Platform::InterlockedIncrement(&_chunksLock);
-    }
+    void LockChunks();
 
     /// <summary>
     /// Unlocks the storage chunks data.