Fixed prefab resource cache not storing all required data (SubnauticaNitrox#2665)

Author: Measurity

Nitrox.Server.Subnautica/Models/Resources/Parsers/EntityDistributionsResource.cs

@@ -12,13 +12,12 @@ internal class EntityDistributionsResource(SubnauticaAssetsManager assetsManager
     private readonly SubnauticaAssetsManager assetsManager = assetsManager;
     private readonly IOptions<ServerStartOptions> options = options;
 
-    private ValueTask<LootDistributionData> lootDistribution;
-    public LootDistributionData LootDistribution => GetLootDistributionDataAsync().GetAwaiter().GetResult();
+    private readonly TaskCompletionSource<LootDistributionData> lootDistributionTcs = new();
+    public LootDistributionData LootDistribution => lootDistributionTcs.Task.GetAwaiter().GetResult();
 
-    public Task LoadAsync(CancellationToken cancellationToken)
+    public async Task LoadAsync(CancellationToken cancellationToken)
     {
-        lootDistribution = GetLootDistributionDataAsync(cancellationToken);
-        return Task.CompletedTask;
+        lootDistributionTcs.TrySetResult(await GetLootDistributionDataAsync(cancellationToken));
     }
 
     public Task CleanupAsync()
@@ -27,13 +26,8 @@ public Task CleanupAsync()
         return Task.CompletedTask;
     }
 
-    private async ValueTask<LootDistributionData> GetLootDistributionDataAsync(CancellationToken cancellationToken = default)
+    private async Task<LootDistributionData> GetLootDistributionDataAsync(CancellationToken cancellationToken = default)
     {
-        if (lootDistribution is { IsCompletedSuccessfully : true, Result: not null })
-        {
-            return await lootDistribution;
-        }
-
         // TODO: Do not depend on game code; use custom types to map to game JSON files.
         LootDictionary result = JsonSerializer.Deserialize<LootDictionary>(await GetJsonAsync(cancellationToken),
                                                                            new JsonSerializerOptions

Nitrox.Server.Subnautica/Models/Resources/Parsers/PrefabPlaceholderGroupsResource.cs

@@ -10,6 +10,8 @@
 using Nitrox.Server.Subnautica.Models.Helper;
 using Nitrox.Server.Subnautica.Models.Resources.AddressablesTools.Catalog;
 using Nitrox.Server.Subnautica.Models.Resources.Core;
+using ClassIdByRuntimeKeyDictionary = System.Collections.Generic.Dictionary<string, string>;
+using AddressableCatalogDictionary = System.Collections.Generic.Dictionary<string, string[]>;
 
 namespace Nitrox.Server.Subnautica.Models.Resources.Parsers;
 
@@ -23,18 +25,17 @@ internal sealed class PrefabPlaceholderGroupsResource(SubnauticaAssetsManager as
     ///         the cache is rebuilt
     ///     </para>
     /// </summary>
-    private const int CACHE_VERSION = 3;
+    private const int CACHE_VERSION = 4;
+
     private const string CACHE_FILENAME = "PrefabPlaceholdersGroupAssetsCache.json";
 
-    private readonly ConcurrentDictionary<string, string[]> addressableCatalog = new();
     private readonly SubnauticaAssetsManager assetsManager = assetsManager;
-    private readonly ConcurrentDictionary<string, string> classIdByRuntimeKey = new();
-    private readonly ConcurrentDictionary<string, PrefabPlaceholdersGroupAsset> groupsByClassId = new();
     private readonly ILogger<PrefabPlaceholderGroupsResource> logger = logger;
     private readonly IOptions<ServerStartOptions> options = options;
-    private readonly ConcurrentDictionary<string, PrefabPlaceholderAsset> placeholdersByClassId = [];
     private readonly TaskCompletionSource resourceLoadFinished = new();
     private readonly JsonSerializer serializer = new() { TypeNameHandling = TypeNameHandling.Auto };
+    private ConcurrentDictionary<string, PrefabPlaceholdersGroupAsset> groupsByClassId = [];
+    private ConcurrentDictionary<string, PrefabPlaceholderAsset> placeholdersByClassId = [];
     private ConcurrentDictionary<string, string[]> randomPossibilitiesByClassId = [];
 
     public ConcurrentDictionary<string, PrefabPlaceholdersGroupAsset> GroupsByClassId
@@ -62,10 +63,6 @@ public ConcurrentDictionary<string, string[]> RandomPossibilitiesByClassId
             resourceLoadFinished.Task.GetAwaiter().GetResult();
             return randomPossibilitiesByClassId;
         }
-        private set
-        {
-            randomPossibilitiesByClassId = value;
-        }
     }
 
     public async Task LoadAsync(CancellationToken cancellationToken)
@@ -89,7 +86,7 @@ public void PickRandomClassIdIfRequired(ref string classId)
         }
     }
 
-    private static Dictionary<string, string> LoadPrefabDatabase(string fullFilename)
+    private static ClassIdByRuntimeKeyDictionary LoadPrefabDatabase(string fullFilename)
     {
         Dictionary<string, string> prefabFiles = new();
         if (!File.Exists(fullFilename))
@@ -122,40 +119,29 @@ private static void GetPrefabGameObjectInfoFromBundle(SubnauticaAssetsManager am
         prefabGameObjectInfo = assetFileInst.file.Metadata.GetAssetInfo(rootAssetPathId);
     }
 
-    private Task<Dictionary<string, PrefabPlaceholdersGroupAsset>> LoadPrefabsAndSpawnPossibilitiesAsync(CancellationToken cancellationToken = default)
+    private async Task LoadPrefabsAndSpawnPossibilitiesAsync(CancellationToken cancellationToken = default)
     {
-        string prefabDatabasePath = Path.Combine(options.Value.GetSubnauticaResourcesPath(), "StreamingAssets", "SNUnmanagedData", "prefabs.db");
-
-        // Get all prefab-classIds linked to the (partial) bundle path
-        Dictionary<string, string> prefabDatabase = LoadPrefabDatabase(prefabDatabasePath);
-        cancellationToken.ThrowIfCancellationRequested();
-
         // Loading all prefabs by their classId and file paths (first the path to the prefab then the dependencies)
-        LoadAddressableCatalog(options.Value.GetSubnauticaAaResourcePath(), prefabDatabase);
         cancellationToken.ThrowIfCancellationRequested();
-        Dictionary<string, PrefabPlaceholdersGroupAsset> result = CreateOrLoadPrefabCache(options.Value.GetServerCachePath());
+        await CreateOrLoadPrefabCacheAsync(options.Value.GetServerCachePath());
         cancellationToken.ThrowIfCancellationRequested();
 
         // Select only prefabs with a PrefabPlaceholdersGroups component in the root and link them with their dependencyPaths
         // Do not remove: the internal cache list is slowing down the process more than loading a few assets again. There maybe is a better way in the new AssetToolsNetVersion but, we need a byte to texture library bc ATNs sub-package is only for netstandard.
         assetsManager.UnloadAll(true);
-        // Clear private collections that were used temporarily to parse the files.
-        addressableCatalog.Clear();
-        classIdByRuntimeKey.Clear();
 
         // Get all needed data for the filtered PrefabPlaceholdersGroups to construct PrefabPlaceholdersGroupAssets and add them to the dictionary by classId
         Validate.IsFalse(randomPossibilitiesByClassId.IsEmpty);
-        return Task.FromResult(result);
     }
 
-    private Dictionary<string, PrefabPlaceholdersGroupAsset> CreateOrLoadPrefabCache(string nitroxCachePath)
+    private async Task CreateOrLoadPrefabCacheAsync(string nitroxCachePath)
     {
-        Dictionary<string, PrefabPlaceholdersGroupAsset> prefabPlaceholdersGroupPaths = null;
+        Dictionary<string, PrefabPlaceholdersGroupAsset> prefabPlaceholdersGroupPaths;
         string cacheFilePath = Path.Combine(nitroxCachePath, CACHE_FILENAME);
         Cache? cache = null;
         try
         {
-            cache = Cache.Deserialize(serializer, cacheFilePath);
+            cache = await Cache.DeserializeAsync(serializer, cacheFilePath);
         }
         catch (Exception ex)
         {
@@ -169,24 +155,33 @@ private Dictionary<string, PrefabPlaceholdersGroupAsset> CreateOrLoadPrefabCache
             }
             prefabPlaceholdersGroupPaths = cache.Value.PrefabPlaceholdersGroupPaths;
             randomPossibilitiesByClassId = cache.Value.RandomPossibilitiesByClassId;
+            groupsByClassId = cache.Value.GroupsByClassId;
+            placeholdersByClassId = cache.Value.PlaceholdersByClassId;
             logger.ZLogDebug($"Successfully loaded cache with {prefabPlaceholdersGroupPaths.Count:@PrefabPlaceholdersCount} prefab placeholder groups and {randomPossibilitiesByClassId.Count:@RandomPossibilitiesCount} random spawn behaviours.");
         }
         // Fallback solution
-        if (prefabPlaceholdersGroupPaths is null)
+        else
         {
             logger.ZLogInformation($"Building cache, this may take a while...");
-            prefabPlaceholdersGroupPaths = new(GetPrefabPlaceholderGroupAssetsByGroupClassId(assetsManager, GetAllPrefabPlaceholdersGroupsFast(assetsManager)));
-            Cache.Serialize(serializer, new Cache(CACHE_VERSION, prefabPlaceholdersGroupPaths, randomPossibilitiesByClassId), cacheFilePath);
+            // Get all prefab-classIds linked to the (partial) bundle path
+            string prefabDatabasePath = Path.Combine(options.Value.GetSubnauticaResourcesPath(), "StreamingAssets", "SNUnmanagedData", "prefabs.db");
+            Dictionary<string, string> prefabDatabase = LoadPrefabDatabase(prefabDatabasePath);
+            (AddressableCatalogDictionary addressableCatalog, ClassIdByRuntimeKeyDictionary classIdByRuntimeKey) = LoadAddressableCatalog(options.Value.GetSubnauticaAaResourcePath(), prefabDatabase);
+            prefabPlaceholdersGroupPaths = new(GetPrefabPlaceholderGroupAssetsByGroupClassId(assetsManager, GetAllPrefabPlaceholdersGroupsFast(assetsManager, addressableCatalog, classIdByRuntimeKey), addressableCatalog, classIdByRuntimeKey));
+            await Cache.SerializeAsync(serializer, new Cache(CACHE_VERSION, prefabPlaceholdersGroupPaths, randomPossibilitiesByClassId, groupsByClassId, placeholdersByClassId), cacheFilePath);
             logger.ZLogDebug(
                 $"Successfully built cache with {prefabPlaceholdersGroupPaths.Count:@PrefabPlaceholdersCount} prefab placeholder groups and {randomPossibilitiesByClassId.Count:@RandomPossibilitiesCount} random spawn behaviours. Future server starts will take less time.");
         }
         Validate.IsTrue(prefabPlaceholdersGroupPaths.Count > 0);
         Validate.IsTrue(randomPossibilitiesByClassId.Count > 0);
-        return prefabPlaceholdersGroupPaths;
+        Validate.IsTrue(groupsByClassId.Count > 0);
+        Validate.IsTrue(placeholdersByClassId.Count > 0);
     }
 
-    private void LoadAddressableCatalog(string aaRootPath, Dictionary<string, string> prefabDatabase)
+    private (AddressableCatalogDictionary, ClassIdByRuntimeKeyDictionary) LoadAddressableCatalog(string aaRootPath, Dictionary<string, string> prefabDatabase)
     {
+        ClassIdByRuntimeKeyDictionary classIdByRuntimeKey = [];
+        AddressableCatalogDictionary addressableCatalog = [];
         ContentCatalogData ccd = ContentCatalogData.FromJson(File.ReadAllText(Path.Combine(aaRootPath, "catalog.json")));
         Dictionary<string, string> classIdByPath = prefabDatabase.ToDictionary(m => m.Value, m => m.Key);
 
@@ -216,13 +211,15 @@ private void LoadAddressableCatalog(string aaRootPath, Dictionary<string, string
                 break;
             }
         }
+
+        return (addressableCatalog, classIdByRuntimeKey);
     }
 
     /// <summary>
     ///     Gathers bundle paths by class id for prefab placeholder groups.
     ///     Also fills <see cref="RandomPossibilitiesByClassId" />
     /// </summary>
-    private ConcurrentDictionary<string, string[]> GetAllPrefabPlaceholdersGroupsFast(SubnauticaAssetsManager am)
+    private ConcurrentDictionary<string, string[]> GetAllPrefabPlaceholdersGroupsFast(SubnauticaAssetsManager am, AddressableCatalogDictionary addressableCatalog, ClassIdByRuntimeKeyDictionary classIdByRuntimeKey)
     {
         // First step is to find out about the hash of the types PrefabPlaceholdersGroup and SpawnRandom
         // to be able to recognize them easily later on
@@ -313,7 +310,8 @@ private ConcurrentDictionary<string, string[]> GetAllPrefabPlaceholdersGroupsFas
         return prefabPlaceholdersGroupPaths;
     }
 
-    private ConcurrentDictionary<string, PrefabPlaceholdersGroupAsset> GetPrefabPlaceholderGroupAssetsByGroupClassId(SubnauticaAssetsManager am, ConcurrentDictionary<string, string[]> prefabPlaceholdersGroupPaths)
+    private ConcurrentDictionary<string, PrefabPlaceholdersGroupAsset> GetPrefabPlaceholderGroupAssetsByGroupClassId(SubnauticaAssetsManager am, ConcurrentDictionary<string, string[]> prefabPlaceholdersGroupPaths,
+                                                                                                                     AddressableCatalogDictionary addressableCatalog, ClassIdByRuntimeKeyDictionary classIdByRuntimeKey)
     {
         ConcurrentDictionary<string, PrefabPlaceholdersGroupAsset> prefabPlaceholderGroupsByGroupClassId = new();
 
@@ -323,7 +321,7 @@ private ConcurrentDictionary<string, PrefabPlaceholdersGroupAsset> GetPrefabPlac
             SubnauticaAssetsManager amInnerClone = amClone.Clone();
             AssetsFileInstance assetFileInst = amInnerClone.LoadBundleWithDependencies(keyValuePair.Value);
 
-            PrefabPlaceholdersGroupAsset prefabPlaceholderGroup = GetAndCachePrefabPlaceholdersGroupOfBundle(amInnerClone, assetFileInst, keyValuePair.Key);
+            PrefabPlaceholdersGroupAsset prefabPlaceholderGroup = GetAndCachePrefabPlaceholdersGroupOfBundle(amInnerClone, assetFileInst, keyValuePair.Key, addressableCatalog, classIdByRuntimeKey);
             amInnerClone.UnloadAll();
 
             if (!prefabPlaceholderGroupsByGroupClassId.TryAdd(keyValuePair.Key, prefabPlaceholderGroup))
@@ -334,13 +332,15 @@ private ConcurrentDictionary<string, PrefabPlaceholdersGroupAsset> GetPrefabPlac
         return prefabPlaceholderGroupsByGroupClassId;
     }
 
-    private PrefabPlaceholdersGroupAsset GetAndCachePrefabPlaceholdersGroupOfBundle(SubnauticaAssetsManager amInst, AssetsFileInstance assetFileInst, string classId)
+    private PrefabPlaceholdersGroupAsset GetAndCachePrefabPlaceholdersGroupOfBundle(SubnauticaAssetsManager amInst, AssetsFileInstance assetFileInst, string classId, AddressableCatalogDictionary addressableCatalog,
+                                                                                    ClassIdByRuntimeKeyDictionary classIdByRuntimeKey)
     {
         GetPrefabGameObjectInfoFromBundle(amInst, assetFileInst, out AssetFileInfo prefabGameObjectInfo);
-        return GetAndCachePrefabPlaceholdersGroupGroup(amInst, assetFileInst, prefabGameObjectInfo, classId);
+        return GetAndCachePrefabPlaceholdersGroupGroup(amInst, assetFileInst, prefabGameObjectInfo, classId, addressableCatalog, classIdByRuntimeKey);
     }
 
-    private PrefabPlaceholdersGroupAsset GetAndCachePrefabPlaceholdersGroupGroup(SubnauticaAssetsManager amInst, AssetsFileInstance assetFileInst, AssetFileInfo rootGameObjectInfo, string classId)
+    private PrefabPlaceholdersGroupAsset GetAndCachePrefabPlaceholdersGroupGroup(SubnauticaAssetsManager amInst, AssetsFileInstance assetFileInst, AssetFileInfo rootGameObjectInfo, string classId, AddressableCatalogDictionary addressableCatalog,
+                                                                                 ClassIdByRuntimeKeyDictionary classIdByRuntimeKey)
     {
         if (!string.IsNullOrEmpty(classId) && groupsByClassId.TryGetValue(classId, out PrefabPlaceholdersGroupAsset cachedGroup))
         {
@@ -368,7 +368,7 @@ private PrefabPlaceholdersGroupAsset GetAndCachePrefabPlaceholdersGroupGroup(Sub
 
             AssetTypeValueField gameObjectPtr = prefabPlaceholder["m_GameObject"];
             AssetTypeValueField gameObjectField = amInst.GetExtAsset(assetFileInst, gameObjectPtr).baseField;
-            IPrefabAsset asset = GetAndCacheAsset(amInst, prefabPlaceholder["prefabClassId"].AsString);
+            IPrefabAsset asset = GetAndCacheAsset(amInst, prefabPlaceholder["prefabClassId"].AsString, addressableCatalog, classIdByRuntimeKey);
             bool isEntitySlotAsset = asset is PrefabPlaceholderAsset prefabPlaceholderAsset && prefabPlaceholderAsset.EntitySlot.HasValue;
             NitroxTransform transform = amInst.GetTransformFromGameObject(assetFileInst, gameObjectField, rootGameObjectName, isEntitySlotAsset);
             string prefabAssetClassId = prefabPlaceholder["prefabClassId"].AsString;
@@ -388,7 +388,7 @@ private PrefabPlaceholdersGroupAsset GetAndCachePrefabPlaceholdersGroupGroup(Sub
         return prefabPlaceholdersGroup;
     }
 
-    private IPrefabAsset? GetAndCacheAsset(SubnauticaAssetsManager am, string classId)
+    private IPrefabAsset? GetAndCacheAsset(SubnauticaAssetsManager am, string classId, AddressableCatalogDictionary addressableCatalog, ClassIdByRuntimeKeyDictionary classIdByRuntimeKey)
     {
         if (string.IsNullOrEmpty(classId))
         {
@@ -415,7 +415,7 @@ private PrefabPlaceholdersGroupAsset GetAndCachePrefabPlaceholdersGroupGroup(Sub
         AssetFileInfo placeholdersGroupInfo = am.GetMonoBehaviourFromGameObject(assetFileInst, prefabGameObjectInfo, "PrefabPlaceholdersGroup");
         if (placeholdersGroupInfo != null)
         {
-            PrefabPlaceholdersGroupAsset groupAsset = GetAndCachePrefabPlaceholdersGroupOfBundle(am, assetFileInst, classId);
+            PrefabPlaceholdersGroupAsset groupAsset = GetAndCachePrefabPlaceholdersGroupOfBundle(am, assetFileInst, classId, addressableCatalog, classIdByRuntimeKey);
             groupsByClassId[classId] = groupAsset;
             return groupAsset;
         }
@@ -469,24 +469,30 @@ private PrefabPlaceholdersGroupAsset GetAndCachePrefabPlaceholdersGroupGroup(Sub
         return prefabPlaceholderAsset;
     }
 
-    private record struct Cache(int Version, Dictionary<string, PrefabPlaceholdersGroupAsset> PrefabPlaceholdersGroupPaths, ConcurrentDictionary<string, string[]> RandomPossibilitiesByClassId)
+    private record struct Cache(
+        int Version,
+        Dictionary<string, PrefabPlaceholdersGroupAsset> PrefabPlaceholdersGroupPaths,
+        ConcurrentDictionary<string, string[]> RandomPossibilitiesByClassId,
+        ConcurrentDictionary<string, PrefabPlaceholdersGroupAsset> GroupsByClassId,
+        ConcurrentDictionary<string, PrefabPlaceholderAsset> PlaceholdersByClassId
+    )
     {
-        public static void Serialize(JsonSerializer serializer, Cache cache, string filePath)
+        public static async Task SerializeAsync(JsonSerializer serializer, Cache cache, string filePath)
         {
             Directory.CreateDirectory(Path.GetDirectoryName(filePath) ?? throw new Exception("Failed to get directory path from cache file path"));
-            using StreamWriter stream = File.CreateText(filePath);
+            await using StreamWriter stream = File.CreateText(filePath);
             serializer.Serialize(stream, cache);
         }
 
-        public static Cache? Deserialize(JsonSerializer serializer, string filePath)
+        public static Task<Cache?> DeserializeAsync(JsonSerializer serializer, string filePath)
         {
             if (!File.Exists(filePath))
             {
-                return null;
+                return Task.FromResult<Cache?>(null);
             }
             Directory.CreateDirectory(Path.GetDirectoryName(filePath) ?? throw new Exception("Failed to get directory path from cache file path"));
             using StreamReader reader = File.OpenText(filePath);
-            return (Cache?)serializer.Deserialize(reader, typeof(Cache));
+            return Task.FromResult((Cache?)serializer.Deserialize(reader, typeof(Cache)));
         }
     }
 }