[SPEC-7644] ParallelDeepAssetReferences is failing intermittently (#5721)

* Fixed race condition caused by trying to handle asset ready event before asset container has finished filling out all the data structures.

Added check to only handle asset ready once init is complete
Added unit test to verify fix

Signed-off-by: amzn-mike <80125227+amzn-mike@users.noreply.github.com>

* Re-enable test

Signed-off-by: amzn-mike <80125227+amzn-mike@users.noreply.github.com>

* Add missing space to error message

Signed-off-by: amzn-mike <80125227+amzn-mike@users.noreply.github.com>

* Add comment on sleep

Signed-off-by: amzn-mike <80125227+amzn-mike@users.noreply.github.com>

* Collapse nested namespace

Signed-off-by: amzn-mike <80125227+amzn-mike@users.noreply.github.com>

* Collapse nested namespace

Signed-off-by: amzn-mike <80125227+amzn-mike@users.noreply.github.com>

Code/Framework/AzCore/AzCore/Asset/AssetContainer.cpp

@@ -35,6 +35,43 @@ namespace AZ::Data
         AssetLoadBus::MultiHandler::BusDisconnect();
     }
 
+    AZStd::vector<AZStd::pair<AssetInfo, Asset<AssetData>>> AssetContainer::CreateAndQueueDependentAssets(
+        const AZStd::vector<AssetInfo>& dependencyInfoList, const AssetLoadParameters& loadParamsCopyWithNoLoadingFilter)
+    {
+        AZStd::vector<AZStd::pair<AssetInfo, Asset<AssetData>>> dependencyAssets;
+
+        for (auto& thisInfo : dependencyInfoList)
+        {
+            auto dependentAsset = AssetManager::Instance().FindOrCreateAsset(
+                thisInfo.m_assetId, thisInfo.m_assetType, AZ::Data::AssetLoadBehavior::Default);
+
+            if (!dependentAsset || !dependentAsset.GetId().IsValid())
+            {
+                AZ_Warning("AssetContainer", false, "Dependency Asset %s (%s) was not found\n",
+                           thisInfo.m_assetId.ToString<AZStd::string>().c_str(), thisInfo.m_relativePath.c_str());
+                RemoveWaitingAsset(thisInfo.m_assetId);
+                continue;
+            }
+            dependencyAssets.emplace_back(thisInfo, AZStd::move(dependentAsset));
+        }
+
+        // Queue the loading of all of the dependent assets before loading the root asset.
+        for (auto& [dependentAssetInfo, dependentAsset] : dependencyAssets)
+        {
+            // Queue each asset to load.
+            auto queuedDependentAsset = AssetManager::Instance().GetAssetInternal(
+                dependentAsset.GetId(), dependentAsset.GetType(),
+                AZ::Data::AssetLoadBehavior::Default, loadParamsCopyWithNoLoadingFilter,
+                dependentAssetInfo, HasPreloads(dependentAsset.GetId()));
+
+            // Verify that the returned asset reference matches the one that we found or created and queued to load.
+            AZ_Assert(dependentAsset == queuedDependentAsset, "GetAssetInternal returned an unexpected asset reference for Asset %s",
+                      dependentAsset.GetId().ToString<AZStd::string>().c_str());
+        }
+
+        return dependencyAssets;
+    }
+
     void AssetContainer::AddDependentAssets(Asset<AssetData> rootAsset, const AssetLoadParameters& loadParams)
     {
         AssetId rootAssetId = rootAsset.GetId();
@@ -181,34 +218,7 @@ namespace AZ::Data
         // Since we've set the load filter to not load dependencies, we need to ensure all the assets are created beforehand
         // so the dependencies can be hooked up as soon as each asset gets serialized in, even if they start getting serialized
         // while we're still in the middle of triggering all of the asset loads below.
-        for (auto& thisInfo : dependencyInfoList)
-        {
-            auto dependentAsset = AssetManager::Instance().FindOrCreateAsset(
-                thisInfo.m_assetId, thisInfo.m_assetType, AZ::Data::AssetLoadBehavior::Default);
-
-            if (!dependentAsset || !dependentAsset.GetId().IsValid())
-            {
-                AZ_Warning("AssetContainer", false, "Dependency Asset %s (%s) was not found\n",
-                    thisInfo.m_assetId.ToString<AZStd::string>().c_str(), thisInfo.m_relativePath.c_str());
-                RemoveWaitingAsset(thisInfo.m_assetId);
-                continue;
-            }
-            dependencyAssets.emplace_back(thisInfo, AZStd::move(dependentAsset));
-        }
-
-        // Queue the loading of all of the dependent assets before loading the root asset.
-        for (auto& [dependentAssetInfo, dependentAsset] : dependencyAssets)
-        {
-            // Queue each asset to load.
-            auto queuedDependentAsset = AssetManager::Instance().GetAssetInternal(
-                dependentAsset.GetId(), dependentAsset.GetType(),
-                AZ::Data::AssetLoadBehavior::Default, loadParamsCopyWithNoLoadingFilter,
-                dependentAssetInfo, HasPreloads(dependentAsset.GetId()));
-
-            // Verify that the returned asset reference matches the one that we found or created and queued to load.
-            AZ_Assert(dependentAsset == queuedDependentAsset, "GetAssetInternal returned an unexpected asset reference for Asset %s",
-                dependentAsset.GetId().ToString<AZStd::string>().c_str());
-        }
+        dependencyAssets = CreateAndQueueDependentAssets(dependencyInfoList, loadParamsCopyWithNoLoadingFilter);
 
         // Add all of the queued dependent assets as dependencies
         {
@@ -347,9 +357,16 @@ namespace AZ::Data
 
     void AssetContainer::HandleReadyAsset(Asset<AssetData> asset)
     {
+        // Wait until we've finished initialization before allowing this
+        // If a ready event happens before we've gotten all the maps/structures set up, there may be some missing data
+        // which can lead to a crash
+        // We'll go through and check the ready status of every dependency immediately after finishing initialization anyway
+        if (m_initComplete)
+        {
         RemoveFromAllWaitingPreloads(asset->GetId());
         RemoveWaitingAsset(asset->GetId());
     }
+    }
 
     void AssetContainer::OnAssetDataLoaded(Asset<AssetData> asset)
     {

Code/Framework/AzCore/AzCore/Asset/AssetContainer.h

@@ -24,8 +24,8 @@ namespace AZ
 
         // AssetContainer loads an asset and all of its dependencies as a collection which is parallellized as much as possible.
         // With the container, the data will all load in parallel.  Dependent asset loads will still obey the expected rules
-        // where PreLoad assets will emit OnAssetReady before the parent does, and QueueLoad assets will emit OnAssetReady in 
-        // no guaranteed order.  However, the OnAssetContainerReady signals will not emit until all PreLoad and QueueLoad assets 
+        // where PreLoad assets will emit OnAssetReady before the parent does, and QueueLoad assets will emit OnAssetReady in
+        // no guaranteed order.  However, the OnAssetContainerReady signals will not emit until all PreLoad and QueueLoad assets
         // are ready.  NoLoad dependencies are not loaded by default but can be loaded along with their dependencies using the
         // same rules as above by using the LoadAll dependency rule.
         class AssetContainer :
@@ -36,7 +36,7 @@ namespace AZ
             AZ_CLASS_ALLOCATOR(AssetContainer, SystemAllocator, 0);
 
             AssetContainer() = default;
-            
+
             AssetContainer(Asset<AssetData> asset, const AssetLoadParameters& loadParams);
             ~AssetContainer();
 
@@ -81,6 +81,10 @@ namespace AZ
             // AssetLoadBus
             void OnAssetDataLoaded(AZ::Data::Asset<AZ::Data::AssetData> asset) override;
         protected:
+
+            virtual AZStd::vector<AZStd::pair<AssetInfo, Asset<AssetData>>> CreateAndQueueDependentAssets(
+                const AZStd::vector<AssetInfo>& dependencyInfoList, const AssetLoadParameters& loadParamsCopyWithNoLoadingFilter);
+
             // Waiting assets are those which have not yet signalled ready.  In the case of PreLoad dependencies the data may have completed the load cycle but
             // the Assets aren't considered "Ready" yet if there are PreLoad dependencies still loading and will still be in the list until the point that asset and
             // All of its preload dependencies have been loaded, when it signals OnAssetReady
@@ -97,7 +101,7 @@ namespace AZ
             void AddDependency(Asset<AssetData>&& addDependency);
 
             // Add a "graph section" to our list of dependencies.  This checks the catalog for all Pre and Queue load assets which are dependents of the requested asset and kicks off loads
-            // NoLoads which are encounted are placed in another list and can be loaded on demand with the LoadDependency call.  
+            // NoLoads which are encounted are placed in another list and can be loaded on demand with the LoadDependency call.
             void AddDependentAssets(Asset<AssetData> rootAsset, const AssetLoadParameters& loadParams);
 
             // If "PreLoad" assets are found in the graph these are cached and tracked with both OnAssetReady and OnAssetDataLoaded messages.
@@ -117,7 +121,7 @@ namespace AZ
             // duringInit if we're coming from the checkReady method - containers that start ready don't need to signal
             void HandleReadyAsset(AZ::Data::Asset<AZ::Data::AssetData> asset);
 
-            // Optimization to save the lookup in the dependencies map 
+            // Optimization to save the lookup in the dependencies map
             AssetInternal::WeakAsset<AssetData> m_rootAsset;
 
             // The root asset id is stored here semi-redundantly on initialization so that we can still refer to it even if the
@@ -136,7 +140,7 @@ namespace AZ
             AZStd::atomic_bool m_finalNotificationSent{false};
 
             mutable AZStd::recursive_mutex m_preloadMutex;
-            // AssetId -> List of assets it is still waiting on 
+            // AssetId -> List of assets it is still waiting on
             PreloadAssetListType m_preloadList;
 
             // AssetId -> List of assets waiting on it

Code/Framework/AzCore/AzCore/Asset/AssetManager.h

@@ -169,14 +169,14 @@ namespace AZ
             /// Register handler with the system for a particular asset type.
             /// A handler should be registered for each asset type it handles.
             /// Please note that all the handlers are registered just once during app startup from the main thread
-            /// and therefore this is not a thread safe method and should not be invoked from different threads. 
+            /// and therefore this is not a thread safe method and should not be invoked from different threads.
             void RegisterHandler(AssetHandler* handler, const AssetType& assetType);
             /// Unregister handler from the asset system.
             /// Please note that all the handlers are unregistered just once during app shutdown from the main thread
             /// and therefore this is not a thread safe method and should not be invoked from different threads.
             void UnregisterHandler(AssetHandler* handler);
             // @}
-            
+
             // @{ Asset catalog management
             /// Register a catalog with the system for a particular asset type.
             /// A catalog should be registered for each asset type it is responsible for.
@@ -295,7 +295,7 @@ namespace AZ
             /**
             * Old 'legacy' assetIds and asset hints can be automatically replaced  with new ones during deserialize / assignment.
             * This operation can be somewhat costly, and its only useful if the program subsequently re-saves the files its loading so that
-            * the asset hints and assetIds actually persist.  Thus, it can be disabled in situations where you know you are not going to be 
+            * the asset hints and assetIds actually persist.  Thus, it can be disabled in situations where you know you are not going to be
             * saving over or creating new source files (for example builders/background apps)
             * By default, it is enabled.
             */
@@ -316,7 +316,7 @@ namespace AZ
             * This method must be invoked before you start unregistering handlers manually and shutting down the asset manager.
             * This method ensures that all jobs in flight are either canceled or completed.
             * This method is automatically called in the destructor but if you are unregistering handlers manually,
-            * you must invoke it yourself. 
+            * you must invoke it yourself.
             */
             void        PrepareShutDown();
 
@@ -366,7 +366,7 @@ namespace AZ
             /**
             * Creates a new shared AssetContainer with an optional loadFilter
             * **/
-            AZStd::shared_ptr<AssetContainer> CreateAssetContainer(Asset<AssetData> asset, const AssetLoadParameters& loadParams = AssetLoadParameters{}) const;
+            virtual AZStd::shared_ptr<AssetContainer> CreateAssetContainer(Asset<AssetData> asset, const AssetLoadParameters& loadParams = AssetLoadParameters{}) const;
 
 
             /**
@@ -452,7 +452,7 @@ namespace AZ
 
 
             // Variant of RegisterAssetLoading used for jobs which have been queued and need to verify the status of the asset
-            // before loading in order to prevent cases where a load is queued, then a blocking load goes through, then the queued 
+            // before loading in order to prevent cases where a load is queued, then a blocking load goes through, then the queued
             // load is processed.  This validation step leaves the loaded (And potentially modified) data as is in that case.
             bool ValidateAndRegisterAssetLoading(const Asset<AssetData>& asset);
 
@@ -482,7 +482,7 @@ namespace AZ
          * the blocking. That will result in a single thread deadlock.
          *
          * If you need to queue work, the logic needs to be similar to this:
-         * 
+         *
          AssetHandler::LoadResult MyAssetHandler::LoadAssetData(const Asset<AssetData>& asset, AZStd::shared_ptr<AssetDataStream> stream,
                                                                 const AZ::Data::AssetFilterCB& assetLoadFilterCB)
          {
@@ -496,13 +496,13 @@ namespace AZ
             }
             else
             {
-                // queue job to load asset in thread identified by m_loadingThreadId 
+                // queue job to load asset in thread identified by m_loadingThreadId
                 auto* queuedJob = QueueLoadingOnOtherThread(...);
 
                 // block waiting for queued job to complete
                 queuedJob->BlockUntilComplete();
             }
-            
+
             .
             .
             .
@@ -525,7 +525,7 @@ namespace AZ
             //! Result from LoadAssetData - it either finished loading, didn't finish and is waiting for more data, or had an error.
             enum class LoadResult : u8
             {
-                
+
                 Error,              // The provided data failed to load correctly
                 MoreDataRequired,   // The provided data loaded correctly, but more data is required to finish the asset load
                 LoadComplete        // The provided data loaded correctly, and the asset has been created

Code/Framework/AzCore/Tests/Asset/AssetManagerLoadingTests.cpp

@@ -2282,6 +2282,45 @@ namespace UnitTest
             AssetManager::Destroy();
         }
 
+        struct MockAssetContainer : AssetContainer
+        {
+            MockAssetContainer(Asset<AssetData> assetData, const AssetLoadParameters& loadParams)
+            {
+                // Copying the code in the original constructor, we can't call that constructor because it will not invoke our virtual method
+                m_rootAsset = AssetInternal::WeakAsset<AssetData>(assetData);
+                m_containerAssetId = m_rootAsset.GetId();
+
+                AddDependentAssets(assetData, loadParams);
+            }
+
+        protected:
+            AZStd::vector<AZStd::pair<AssetInfo, Asset<AssetData>>> CreateAndQueueDependentAssets(
+                const AZStd::vector<AssetInfo>& dependencyInfoList, const AssetLoadParameters& loadParamsCopyWithNoLoadingFilter) override
+            {
+                auto result = AssetContainer::CreateAndQueueDependentAssets(dependencyInfoList, loadParamsCopyWithNoLoadingFilter);
+
+                // Sleep for a long enough time to allow asset loads to complete and start triggering AssetReady events
+                // This forces the race condition to occur
+                AZStd::this_thread::sleep_for(AZStd::chrono::milliseconds(500));
+
+                return result;
+            }
+        };
+
+        struct MockAssetManager : AssetManager
+        {
+            explicit MockAssetManager(const Descriptor& desc)
+                : AssetManager(desc)
+            {
+            }
+
+        protected:
+            AZStd::shared_ptr<AssetContainer> CreateAssetContainer(Asset<AssetData> asset, const AssetLoadParameters& loadParams) const override
+            {
+                return AZStd::shared_ptr<AssetContainer>(aznew MockAssetContainer(asset, loadParams));
+            }
+        };
+
         void ParallelDeepAssetReferences()
         {
             SerializeContext context;
@@ -2289,7 +2328,7 @@ namespace UnitTest
             AssetWithAssetReference::Reflect(context);
 
             AssetManager::Descriptor desc;
-            AssetManager::Create(desc);
+            AssetManager::SetInstance(aznew MockAssetManager(desc));
 
             auto& db = AssetManager::Instance();
 
@@ -2312,17 +2351,17 @@ namespace UnitTest
 
                 // AssetC is MYASSETC
                 AssetWithAssetReference c;
-                c.m_asset = AssetManager::Instance().CreateAsset<AssetWithSerializedData>(AssetId(MyAssetDId)); // point at D
+                c.m_asset = db.CreateAsset<AssetWithSerializedData>(AssetId(MyAssetDId)); // point at D
                 EXPECT_TRUE(AZ::Utils::SaveObjectToFile(GetTestFolderPath() + "TestAsset3.txt", AZ::DataStream::ST_XML, &c, &context));
 
                 // AssetB is MYASSETB
                 AssetWithAssetReference b;
-                b.m_asset = AssetManager::Instance().CreateAsset<AssetWithAssetReference>(AssetId(MyAssetCId)); // point at C
+                b.m_asset = db.CreateAsset<AssetWithAssetReference>(AssetId(MyAssetCId)); // point at C
                 EXPECT_TRUE(AZ::Utils::SaveObjectToFile(GetTestFolderPath() + "TestAsset2.txt", AZ::DataStream::ST_XML, &b, &context));
 
                 // AssetA will be written to disk as MYASSETA
                 AssetWithAssetReference a;
-                a.m_asset = AssetManager::Instance().CreateAsset<AssetWithAssetReference>(AssetId(MyAssetBId)); // point at B
+                a.m_asset = db.CreateAsset<AssetWithAssetReference>(AssetId(MyAssetBId)); // point at B
                 EXPECT_TRUE(AZ::Utils::SaveObjectToFile(GetTestFolderPath() + "TestAsset1.txt", AZ::DataStream::ST_XML, &a, &context));
             }
 
@@ -2531,7 +2570,7 @@ namespace UnitTest
     TEST_F(AssetJobsMultithreadedTest, DISABLED_ParallelDeepAssetReferences)
 #else
     // temporarily disabled until sporadic failures can be root caused
-    TEST_F(AssetJobsMultithreadedTest, DISABLED_ParallelDeepAssetReferences)
+    TEST_F(AssetJobsMultithreadedTest, ParallelDeepAssetReferences)
 #endif // AZ_TRAIT_DISABLE_FAILED_ASSET_MANAGER_TESTS
     {
         ParallelDeepAssetReferences();