[SPEC-7644] Cherry Pick - ParallelDeepAssetReferences is failing intermittently (#5797)

* [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>
(cherry picked from commit 56900484fc)

# Conflicts:
#	Code/Framework/AzCore/AzCore/Asset/AssetContainer.cpp

* Fix indentation

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

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

@@ -37,6 +37,43 @@ namespace AZ
             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();
@@ -183,34 +220,7 @@ namespace AZ
             // 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
             {
@@ -348,10 +358,17 @@ namespace AZ
         }
 
         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

@@ -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

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

@@ -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;
 
 
             /**

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

@@ -2297,6 +2297,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;
@@ -2304,7 +2343,7 @@ namespace UnitTest
             AssetWithAssetReference::Reflect(context);
 
             AssetManager::Descriptor desc;
-            AssetManager::Create(desc);
+            AssetManager::SetInstance(aznew MockAssetManager(desc));
 
             auto& db = AssetManager::Instance();
 
@@ -2327,17 +2366,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));
             }
 
@@ -2546,7 +2585,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();