Async Bundle Loading and API to block base game's loading process

GTFO-API/API/AssetAPI.AssetLoadHandle.cs

@@ -0,0 +1,36 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace GTFO.API.API
+{
+    /// <summary>
+    /// Handle that allocated when you called <see cref="AssetAPI.WantToWorkForStartupAssets"/>
+    /// This blocks game loading until every allocated AssetLoadHandle as marked as completed
+    /// </summary>
+    public sealed class AssetLoadHandle
+    {
+        /// <summary>
+        /// true if current load job has ended
+        /// </summary>
+        public bool IsCompleted { get; private set; }
+
+        /// <summary>
+        /// Mark this specific Asset loading job as completed
+        /// </summary>
+        public void SetCompleted()
+        {
+            IsCompleted = true;
+        }
+
+        /// <summary>
+        /// Add Line to loading text 
+        /// </summary>
+        public void AddLoadingText(string text)
+        {
+            MainMenuGuiLayer.Current.PageIntro.m_textCenter.AddLine($"<size=35%>{text}</size>");
+        }
+    }
+}

GTFO-API/API/AssetAPI.cs

@@ -4,6 +4,7 @@
 using System.Linq;
 using AssetShards;
 using BepInEx;
+using GTFO.API.API;
 using GTFO.API.Attributes;
 using GTFO.API.Impl;
 using GTFO.API.Resources;
@@ -20,7 +21,7 @@ public static class AssetAPI
         public static ApiStatusInfo Status => APIStatus.Asset;
 
         /// <summary>
-        /// Invoked when the game's startup assets have been fully loaded
+        /// Invoked when the base game's startup assets have been fully loaded
         /// </summary>
         public static event Action OnStartupAssetsLoaded;
 
@@ -29,11 +30,43 @@ public static class AssetAPI
         /// </summary>
         public static event Action OnAssetBundlesLoaded;
 
+        /// <summary>
+        /// Invoken when loading custom assets are about to start
+        /// </summary>
+        public static event Action OnCustomAssetsLoading;
+
+        /// <summary>
+        /// Invoked when startup asset has fully loaded (Including custom bundles and base game assets)
+        /// </summary>
+        public static event Action OnStartupAssetsFullyLoaded;
+
         /// <summary>
         /// Invoked when the internal handler is ready
         /// </summary>
         public static event Action OnImplReady;
 
+        /// <summary>
+        /// Return true If every assetbundle and startup assets are loaded
+        /// </summary>
+        public static bool IsReadyForStartup
+        {
+            get
+            {
+                if (!Status.Ready)
+                    return false;
+
+                if (s_LoadBlockerHandles.Any(x => !x.IsCompleted))
+                    return false;
+
+                if (!s_StartupAssetsFullyLoaded)
+                {
+                    s_StartupAssetsFullyLoaded = true;
+                    OnStartupAssetsFullyLoaded?.Invoke();
+                }
+                return true;
+            }
+        }
+
         /// <summary>
         /// Checks if an asset is already registered in the <see cref="AssetAPI"/>
         /// </summary>
@@ -188,6 +221,26 @@ public static bool TryInstantiateAsset<TAsset>(string assetName, string copyName
             return clonedObj != null;
         }
 
+        /// <summary>
+        /// Create new Load Job Handle
+        /// </summary>
+        /// <param name="newLoadHandle">Created Load Job Handle</param>
+        /// <returns>true if handle has created, false if custom asset loading process is already done (after <see cref="AssetAPI.OnStartupAssetsFullyLoaded"/> has invoked)</returns>
+        public static bool WantToWorkForStartupAssets(out AssetLoadHandle newLoadHandle)
+        {
+            if (s_StartupAssetsFullyLoaded)
+            {
+                APILogger.Error($"Asset", "Startup Assets are already loaded! Try load them before ");
+                newLoadHandle = null;
+                return false;
+            }
+
+
+            newLoadHandle = new();
+            s_LoadBlockerHandles.Add(newLoadHandle);
+            return true;
+        }
+
         private static void OnAssetsLoaded()
         {
             if (!APIStatus.Asset.Created)
@@ -199,20 +252,21 @@ private static void OnAssetsLoaded()
 
         internal static void InvokeImplReady() => OnImplReady?.Invoke();
 
+        internal static void InvokeBundleLoaded() => OnAssetBundlesLoaded?.Invoke();
+
         internal static void Setup()
         {
             EventAPI.OnAssetsLoaded += OnAssetsLoaded;
-            OnImplReady += LoadAssetBundles;
+            OnImplReady += LoadCustomStartupAssets;
         }
 
-        private static void LoadAssetBundles()
+        private static void LoadCustomStartupAssets()
         {
+            OnCustomAssetsLoading?.Invoke();
             string assetBundleDir = Path.Combine(Paths.BepInExRootPath, "Assets", "AssetBundles");
             string assetBundlesDirOld = Path.Combine(Paths.ConfigPath, "Assets", "AssetBundles");
-            bool anyLoaded = LoadAssetBundles(assetBundleDir);
-            anyLoaded |= LoadAssetBundles(assetBundlesDirOld, outdated: true);
-            if (anyLoaded)
-                OnAssetBundlesLoaded?.Invoke();
+            LoadAssetBundles(assetBundleDir);
+            LoadAssetBundles(assetBundlesDirOld, outdated: true);
         }
 
         private static bool LoadAssetBundles(string assetBundlesDir, bool outdated = false)
@@ -238,19 +292,16 @@ private static bool LoadAssetBundles(string assetBundlesDir, bool outdated = fal
 
             for (int i = 0; i < bundlePaths.Length; i++)
             {
-                try
-                {
-                    LoadAndRegisterAssetBundle(bundlePaths[i]);
-                }
-                catch (Exception ex)
-                {
-                    APILogger.Warn(nameof(AssetAPI), $"Failed to load asset bundle '{bundlePaths[i]}' ({ex.Message})");
-                }
+                WantToWorkForStartupAssets(out var assetLoadHandle);
+                AssetAPI_Impl.Instance.LoadAssetBundle(bundlePaths[i], assetLoadHandle);
             }
+            AssetAPI_Impl.Instance.DEBUG_BundleLoadingStarted(bundlePaths.Length);
 
             return true;
         }
 
-        internal static ConcurrentDictionary<string, UnityEngine.Object> s_RegistryCache = new();
+        internal static readonly ConcurrentDictionary<string, UnityEngine.Object> s_RegistryCache = new();
+        internal static readonly ConcurrentBag<AssetLoadHandle> s_LoadBlockerHandles = new();
+        internal static bool s_StartupAssetsFullyLoaded = false;
     }
 }

GTFO-API/API/Impl/AssetAPI_Impl.cs

@@ -1,6 +1,12 @@
 using System;
+using System.Collections;
+using System.Diagnostics;
+using System.IO;
 using AssetShards;
+using BepInEx.Unity.IL2CPP.Utils;
+using GTFO.API.API;
 using GTFO.API.Resources;
+using Il2CppInterop.Runtime.Attributes;
 using UnityEngine;
 
 namespace GTFO.API.Impl
@@ -21,6 +27,103 @@ public static AssetAPI_Impl Instance
                 return s_Instance;
             }
         }
+
+        private int _LoadingCount = 0;
+
+        private readonly Stopwatch _SW = new();
+        private int _DebugLoadingCount;
+
+        private static readonly string COMPRESSED_STR = "-compressed";
+        private static readonly int COMPRESSED_STR_LENGTH = COMPRESSED_STR.Length;
+
+        [Conditional("DEBUG")]
+        [HideFromIl2Cpp]
+        public void DEBUG_BundleLoadingStarted(int loadingCount)
+        {
+            _SW.Restart();
+            _DebugLoadingCount = loadingCount;
+        }
+
+        [Conditional("DEBUG")]
+        [HideFromIl2Cpp]
+        public void DEBUG_BundleLoadingFinished()
+        {
+            _SW.Stop();
+            APILogger.Verbose($"Asset", $"Elapsed Time to loading {_DebugLoadingCount} bundles:");
+            APILogger.Verbose($"Asset", $" - {_SW.Elapsed}! (or {_SW.ElapsedMilliseconds}ms)");
+        }
+
+        [HideFromIl2Cpp]
+        public void LoadAssetBundle(string filePath, AssetLoadHandle loadHandle)
+        {
+            _LoadingCount++;
+            this.StartCoroutine(DoLoadAssetBundle(filePath, loadHandle));
+        }
+
+        [HideFromIl2Cpp]
+        private IEnumerator DoLoadAssetBundle(string filePath, AssetLoadHandle loadHandle)
+        {
+            AssetBundleCreateRequest loadReq = AssetBundle.LoadFromFileAsync(filePath);
+
+            yield return loadReq;
+
+            AssetBundle loadedBundle = loadReq.assetBundle;
+            if (loadedBundle == null)
+            {
+                _LoadingCount--;
+                APILogger.Warn($"Asset", $"Failed to load asset bundle: [{filePath}]");
+                yield break;
+            }
+
+
+            APILogger.Warn($"Asset", $"Start Loading Bundle!");
+            string[] assetNames = loadedBundle.AllAssetNames();
+            int remainingAssets = assetNames.Length;
+            int loadedCount = 0;
+
+            foreach (string assetName in assetNames)
+            {
+                AssetBundleRequest loadAssetReq = loadedBundle.LoadAssetAsync(assetName);
+                loadAssetReq.add_completed((Action<AsyncOperation>)((x) =>
+                {
+                    remainingAssets--;
+                    loadedCount++;
+
+                    UnityEngine.Object loadedAsset = loadAssetReq.asset;
+                    if (loadedAsset == null)
+                    {
+                        APILogger.Warn("Asset", $"Skipping asset {assetName}");
+                    }
+
+                    RegisterAsset(assetName, loadedAsset);
+                }));
+                yield return null;
+            }
+
+            yield return new WaitUntil((Il2CppSystem.Func<bool>)(() =>
+            {
+                return remainingAssets <= 0;
+            }));
+
+            _LoadingCount--;
+            loadHandle.SetCompleted();
+
+            string fileName = Path.GetFileNameWithoutExtension(filePath);
+            if (fileName.EndsWith(COMPRESSED_STR, StringComparison.InvariantCultureIgnoreCase))
+                fileName = fileName[^COMPRESSED_STR_LENGTH..]; //Trim "-Compressed" from bundle name
+
+            loadHandle.AddLoadingText($"[Assets] <color=orange>{fileName}</color> has loaded!");
+
+            if (_LoadingCount <= 0)
+            {
+                _LoadingCount = 0;
+                loadHandle.AddLoadingText($"[Assets] All bundle has loaded!");
+                AssetAPI.InvokeBundleLoaded();
+                DEBUG_BundleLoadingFinished();
+            }
+
+            yield return null;
+        }
 
         private void Awake()
         {
@@ -34,6 +137,7 @@ private void Awake()
             AssetAPI.InvokeImplReady();
         }
 
+        [HideFromIl2Cpp]
         public void RegisterAsset(string name, UnityEngine.Object gameObject)
         {
             string upperName = name.ToUpper();

GTFO-API/Patches/GS_Offline_Patches.cs

@@ -0,0 +1,25 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using HarmonyLib;
+
+namespace GTFO.API.Patches
+{
+    [HarmonyPatch(typeof(GS_Offline))]
+    internal class GS_Offline_Patches
+    {
+        [HarmonyPrefix]
+        [HarmonyPatch(nameof(GS_Offline.Update))]
+        static bool Prefix()
+        {
+            if (AssetAPI.IsReadyForStartup)
+            {
+                return true; //Run Original
+            }
+
+            return false; //Skip Original
+        }
+    }
+}