Merge pull request #4 from DiFFoZ/feature/async-support

Decompress bundles in background

BepInExFasterLoadAssetBundles/Helpers/AsyncHelper.cs

@@ -0,0 +1,22 @@
+using System;
+using System.Threading.Tasks;
+
+namespace BepInExFasterLoadAssetBundles.Helpers;
+internal static class AsyncHelper
+{
+    public static void Schedule(Func<Task> func)
+    {
+        _ = Task.Run(async () =>
+        {
+            try
+            {
+                await func();
+            }
+
+            catch (Exception ex)
+            {
+                Patcher.Logger.LogError(ex);
+            }
+        });
+    }
+}

BepInExFasterLoadAssetBundles/Helpers/AsyncOperationHelper.cs

@@ -1,4 +1,6 @@
-using System.Threading;
+using System;
+using System.Runtime.CompilerServices;
+using System.Threading;
 using UnityEngine;
 
 namespace BepInExFasterLoadAssetBundles.Helpers;
@@ -11,4 +13,49 @@ public static void WaitUntilOperationComplete<T>(T op) where T : AsyncOperation
             Thread.Sleep(100);
         }
     }
+
+    public static AsyncOperationAwaiter WaitCompletionAsync<T>(this T op) where T : AsyncOperation
+    {
+        return new AsyncOperationAwaiter(op);
+    }
+
+    public struct AsyncOperationAwaiter : ICriticalNotifyCompletion
+    {
+        private AsyncOperation? m_AsyncOperation;
+        private Action? m_ContinuationAction;
+
+        public AsyncOperationAwaiter(AsyncOperation asyncOperation)
+        {
+            m_AsyncOperation = asyncOperation;
+            m_ContinuationAction = null;
+        }
+
+        public readonly AsyncOperationAwaiter GetAwaiter() => this;
+        public readonly bool IsCompleted => m_AsyncOperation!.isDone;
+
+        public void GetResult()
+        {
+            if (m_AsyncOperation != null)
+                m_AsyncOperation.completed -= OnCompleted;
+
+            m_AsyncOperation = null;
+            m_ContinuationAction = null;
+        }
+
+        public void OnCompleted(Action continuation)
+        {
+            UnsafeOnCompleted(continuation);
+        }
+
+        public void UnsafeOnCompleted(Action continuation)
+        {
+            m_ContinuationAction = continuation;
+            m_AsyncOperation!.completed += OnCompleted;
+        }
+
+        private readonly void OnCompleted(AsyncOperation _)
+        {
+            m_ContinuationAction?.Invoke();
+        }
+    }
 }

BepInExFasterLoadAssetBundles/Helpers/HashingHelper.cs

@@ -1,16 +1,34 @@
 using System;
+using System.Buffers;
 using System.Globalization;
 using System.IO;
 using System.Security.Cryptography;
 
 namespace BepInExFasterLoadAssetBundles.Helpers;
 internal class HashingHelper
 {
+    private const int c_BufferSize = 81920;
+
     public static byte[] HashFile(string path)
     {
-        using var fileStream = File.OpenRead(path);
+        using var fileStream = new FileStream(path, FileMode.Open, FileAccess.Read, FileShare.None, c_BufferSize, FileOptions.SequentialScan);
+        return HashStream(fileStream);
+    }
+
+    public static byte[] HashStream(Stream stream)
+    {
         using var sha1 = new SHA1Managed();
-        return sha1.ComputeHash(fileStream);
+
+        var array = ArrayPool<byte>.Shared.Rent(c_BufferSize);
+        int readBytes;
+        while ((readBytes = stream.Read(array, 0, c_BufferSize)) > 0)
+        {
+            sha1.TransformBlock(array, 0, readBytes, array, 0);
+        }
+
+        sha1.TransformFinalBlock([], 0, 0);
+
+        return sha1.Hash;
     }
 
     public static string HashToString(byte[] hash)
@@ -20,7 +38,7 @@ public static string HashToString(byte[] hash)
         for (var i = 0; i < hash.Length; i++)
         {
             var b = hash[i];
-            b.TryFormat(chars[(i * 2)..], out var charsWritten, "X2");
+            b.TryFormat(chars[(i * 2)..], out _, "X2");
         }
 
         return chars.ToString();

BepInExFasterLoadAssetBundles/Managers/AssetBundleManager.cs

@@ -1,6 +1,8 @@
 using System;
 using System.IO;
+using System.Threading.Tasks;
 using BepInExFasterLoadAssetBundles.Helpers;
+using BepInExFasterLoadAssetBundles.Models;
 using UnityEngine;
 
 namespace BepInExFasterLoadAssetBundles.Managers;
@@ -21,8 +23,24 @@ public AssetBundleManager(string cachePath)
 
     public bool TryRecompressAssetBundle(ref string path)
     {
-        var originalFileName = Path.GetFileNameWithoutExtension(path);
-        byte[] hash = HashingHelper.HashFile(path);
+        return TryRecompressAssetBundleInternal(ref path, null);
+    }
+
+    public FileStream? TryRecompressAssetBundle(FileStream stream)
+    {
+        var path = string.Copy(stream.Name);
+
+        if (TryRecompressAssetBundleInternal(ref path, HashingHelper.HashStream(stream)))
+        {
+            return File.OpenRead(path);
+        }
+
+        return null;
+    }
+
+    public bool TryRecompressAssetBundleInternal(ref string path, byte[]? hash)
+    {
+        hash ??= HashingHelper.HashFile(path);
 
         var metadata = Patcher.MetadataManager.FindMetadataByHash(hash);
         if (metadata != null)
@@ -35,8 +53,7 @@ public bool TryRecompressAssetBundle(ref string path)
             var newPath = Path.Combine(CachePath, metadata.UncompressedAssetBundleName);
             if (File.Exists(newPath))
             {
-                Patcher.Logger.LogDebug(
-                    $"Found uncompressed bundle {metadata.UncompressedAssetBundleName}, loading it instead of {originalFileName}");
+                Patcher.Logger.LogDebug($"Loading uncompressed bundle \"{metadata.UncompressedAssetBundleName}\"");
                 path = newPath;
 
                 metadata.LastAccessTime = DateTime.Now;
@@ -45,55 +62,73 @@ public bool TryRecompressAssetBundle(ref string path)
                 return true;
             }
 
-            Patcher.Logger.LogWarning($"Failed to find decompressed assetbundle at {newPath}. Probably it was deleted?");
+            Patcher.Logger.LogWarning($"Failed to find decompressed assetbundle at \"{newPath}\". Probably it was deleted?");
+        }
+
+        var nonRefPath = path;
+        AsyncHelper.Schedule(() => DecompressAssetBundleAsync(nonRefPath, hash));
+
+        return false;
+    }
+
+    public void DeleteCachedAssetBundle(string path)
+    {
+        FileHelper.TryDeleteFile(path, out var fileException);
+        if (fileException != null)
+        {
+            Patcher.Logger.LogError($"Failed to delete uncompressed assetbundle\n{fileException}");
         }
+    }
 
-        metadata = new()
+    private async Task DecompressAssetBundleAsync(string path, byte[] hash)
+    {
+        var metadata = new Metadata()
         {
             OriginalAssetBundleHash = HashingHelper.HashToString(hash),
             LastAccessTime = DateTime.Now,
         };
+        var originalFileName = Path.GetFileNameWithoutExtension(path);
         var outputName = originalFileName + '_' + metadata.GetHashCode() + ".assetbundle";
         var outputPath = Path.Combine(CachePath, outputName);
 
+        // when loading assetbundle async via stream, the file can be still in use. Wait a bit for that
+        var tries = 5;
+        while (--tries > 0)
+        {
+            try
+            {
+                using var tempStream = new FileStream(path, FileMode.Open, FileAccess.Read, FileShare.None);
+            }
+            catch (IOException)
+            {
+                await Task.Delay(100);
+            }
+        }
+
         var op = AssetBundle.RecompressAssetBundleAsync(path, outputPath,
-            BuildCompression.UncompressedRuntime, 0, ThreadPriority.High);
-        AsyncOperationHelper.WaitUntilOperationComplete(op);
+            BuildCompression.UncompressedRuntime, 0, ThreadPriority.Normal);
+
+        await op.WaitCompletionAsync();
 
         if (op.result is not AssetBundleLoadResult.Success)
         {
-            Patcher.Logger.LogWarning($"Failed to decompress a assetbundle at {path}\n{op.humanReadableResult}");
-            return false;
+            Patcher.Logger.LogWarning($"Failed to decompress a assetbundle at \"{path}\"\n{op.humanReadableResult}");
+            return;
         }
 
         // check if unity returned the same assetbundle (means that assetbundle is already decompressed)
         if (hash.AsSpan().SequenceEqual(HashingHelper.HashFile(outputPath)))
         {
-            Patcher.Logger.LogDebug($"Assetbundle {originalFileName} is already uncompressed, adding to ignore list");
+            Patcher.Logger.LogDebug($"Assetbundle \"{originalFileName}\" is already uncompressed, adding to ignore list");
 
             metadata.ShouldNotDecompress = true;
             Patcher.MetadataManager.SaveMetadata(metadata);
 
             DeleteCachedAssetBundle(outputPath);
-            return false;
+            return;
         }
 
-        path = outputPath;
-
         metadata.UncompressedAssetBundleName = outputName;
         Patcher.MetadataManager.SaveMetadata(metadata);
-
-        Patcher.Logger.LogDebug($"Loading uncompressed bundle {outputName} instead of {originalFileName}");
-
-        return true;
-    }
-
-    public void DeleteCachedAssetBundle(string path)
-    {
-        FileHelper.TryDeleteFile(path, out var fileException);
-        if (fileException != null)
-        {
-            Patcher.Logger.LogError($"Failed to delete uncompressed assetbundle\n{fileException}");
-        }
     }
 }

BepInExFasterLoadAssetBundles/Managers/MetadataManager.cs

@@ -40,15 +40,18 @@ public MetadataManager(string metadataFile)
 
     public void SaveMetadata(Metadata metadata)
     {
-        var index = m_Metadata.FindIndex(m => m.OriginalAssetBundleHash == metadata.OriginalAssetBundleHash);
-
-        if (index == -1)
-        {
-            m_Metadata.Add(metadata);
-        }
-        else
+        lock (m_Lock)
         {
-            m_Metadata[index] = metadata;
+            var index = m_Metadata.FindIndex(m => m.OriginalAssetBundleHash == metadata.OriginalAssetBundleHash);
+
+            if (index == -1)
+            {
+                m_Metadata.Add(metadata);
+            }
+            else
+            {
+                m_Metadata[index] = metadata;
+            }
         }
 
         SaveFile();

BepInExFasterLoadAssetBundles/Patcher.cs

@@ -1,10 +1,8 @@
 using System;
 using System.IO;
-using System.Threading;
 using BepInEx;
 using BepInEx.Bootstrap;
 using BepInEx.Logging;
-using BepInExFasterLoadAssetBundles.Helpers;
 using BepInExFasterLoadAssetBundles.Managers;
 using HarmonyLib;
 using UnityEngine;
@@ -37,12 +35,20 @@ public static void ChainloaderInitialized()
         var thisType = typeof(Patcher);
         var harmony = BepInExFasterLoadAssetBundlesPatcher.Harmony;
 
+        // file
         harmony.Patch(AccessTools.Method(typeof(AssetBundle), nameof(AssetBundle.LoadFromFile_Internal)),
             prefix: new(thisType.GetMethod(nameof(LoadAssetBundleFromFileFast))));
 
-        // todo
-        /*harmony.Patch(AccessTools.Method(typeof(AssetBundle), nameof(AssetBundle.LoadFromStreamInternal)),
-            prefix: new(thisType.GetMethod(nameof(LoadAssetBundleFromStreamFast))));*/
+        harmony.Patch(AccessTools.Method(typeof(AssetBundle), nameof(AssetBundle.LoadFromFileAsync_Internal)),
+            prefix: new(thisType.GetMethod(nameof(LoadAssetBundleFromFileFast))));
+
+
+        // streams
+        harmony.Patch(AccessTools.Method(typeof(AssetBundle), nameof(AssetBundle.LoadFromStreamInternal)),
+            prefix: new(thisType.GetMethod(nameof(LoadAssetBundleFromStreamFast))));
+
+        harmony.Patch(AccessTools.Method(typeof(AssetBundle), nameof(AssetBundle.LoadFromStreamAsyncInternal)),
+            prefix: new(thisType.GetMethod(nameof(LoadAssetBundleFromStreamFast))));
     }
 
     public static void LoadAssetBundleFromFileFast(ref string path)
@@ -69,4 +75,30 @@ public static void LoadAssetBundleFromFileFast(ref string path)
             path = tempPath;
         }
     }
+
+    public static void LoadAssetBundleFromStreamFast(ref Stream stream)
+    {
+        if (stream is not FileStream fileStream)
+        {
+            return;
+        }
+
+        var previousPosition = fileStream.Position;
+
+        try
+        {
+            var decompressedStream = AssetBundleManager.TryRecompressAssetBundle(fileStream);
+            if (decompressedStream != null)
+            {
+                stream = decompressedStream;
+                return;
+            }
+        }
+        catch (Exception ex)
+        {
+            Logger.LogError($"Failed to decompress assetbundle\n{ex}");
+        }
+
+        fileStream.Position = previousPosition;
+    }
 }

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.2.0] 2024-03-28
+### Changed
+- Decompression is now happens in background.
+- Decompression thread priority is set to `Normal` instead of `High`.
+- AssetBundle loaded via `FileStream` will be now cached.
+
 ## [0.1.0] 2024-03-25
 ### Added
 - Debug log when decompressed assetbundle is loaded instead.