Skip to content

Commit

Permalink
Merge pull request #7 from DiFFoZ/feature/unmanaged-bundle-caching
Browse files Browse the repository at this point in the history 
Merge `Feature/unmanaged bundle caching`
  • Loading branch information
@DiFFoZ
DiFFoZ authored Apr 24, 2024
2 parents 88b13e3 + ab587f7 commit 5840e19
Show file tree
Hide file tree
Showing 4 changed files with 131 additions and 85 deletions.
2 changes: 2 additions & 0 deletions BepInExFasterLoadAssetBundles/Helpers/HashingHelper.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@ public static byte[] HashStream(Stream stream)

sha1.TransformFinalBlock([], 0, 0);

ArrayPool<byte>.Shared.Return(array);

return sha1.Hash;
}

Expand Down
147 changes: 98 additions & 49 deletions BepInExFasterLoadAssetBundles/Managers/AssetBundleManager.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,67 +51,100 @@ private void DeleteTempFiles()
}
}

public bool TryRecompressAssetBundle(ref string path)
public bool TryRecompressAssetBundle(Stream stream, out string path)
{
return TryRecompressAssetBundleInternal(ref path, HashingHelper.HashFile(path));
}
var hash = HashingHelper.HashStream(stream);

public bool TryRecompressAssetBundle(FileStream stream, out string path)
{
path = string.Copy(stream.Name);
return TryRecompressAssetBundleInternal(ref path, HashingHelper.HashStream(stream));
}
if (FindCachedBundleByHash(hash, out path))
{
return true;
}

public bool TryRecompressAssetBundleInternal(ref string path, byte[] hash)
{
if (!File.Exists(path))
if (stream is FileStream fileStream)
{
path = string.Copy(fileStream.Name);
RecompressAssetBundleInternal(new(path, hash, false));
return false;
}

var metadata = Patcher.MetadataManager.FindMetadataByHash(hash);
if (metadata != null)
// copy stream to temp file
var tempDirectory = Path.Combine(CachePath, "temp");
if (!Directory.Exists(tempDirectory))
{
if (metadata.ShouldNotDecompress || metadata.UncompressedAssetBundleName == null)
{
return false;
}
Directory.CreateDirectory(tempDirectory);
}

var newPath = Path.Combine(CachePath, metadata.UncompressedAssetBundleName);
if (File.Exists(newPath))
{
Patcher.Logger.LogDebug($"Loading uncompressed bundle \"{metadata.UncompressedAssetBundleName}\"");
path = newPath;
var name = Guid.NewGuid().ToString("N") + ".assetbundle";
var tempFile = Path.Combine(tempDirectory, name);

metadata.LastAccessTime = DateTime.Now;
Patcher.MetadataManager.SaveMetadata(metadata);
using (var fs = new FileStream(tempFile, FileMode.CreateNew, FileAccess.Write))
{
stream.Seek(0, SeekOrigin.Begin);
stream.CopyTo(fs);
}

return true;
}
RecompressAssetBundleInternal(new(tempFile, hash, true));
return false;
}

Patcher.Logger.LogWarning($"Failed to find decompressed assetbundle at \"{newPath}\". Probably it was deleted?");
public void DeleteCachedAssetBundle(string path)
{
FileHelper.TryDeleteFile(path, out var fileException);
if (fileException != null)
{
Patcher.Logger.LogError($"Failed to delete uncompressed assetbundle\n{fileException}");
}
}

if (DriveHelper.HasDriveSpaceOnPath(CachePath, 10))
private bool FindCachedBundleByHash(byte[] hash, out string path)
{
path = null!;

var metadata = Patcher.MetadataManager.FindMetadataByHash(hash);
if (metadata == null)
{
m_WorkAssets.Enqueue(new(path, hash));
StartRunner();
return false;
}
else

if (metadata.ShouldNotDecompress || metadata.UncompressedAssetBundleName == null)
{
Patcher.Logger.LogWarning($"Ignoring request of decompressing, because the drive space is less than 10GB");
return false;
}

return false;

var newPath = Path.Combine(CachePath, metadata.UncompressedAssetBundleName);
if (!File.Exists(newPath))
{
Patcher.Logger.LogWarning($"Failed to find decompressed assetbundle at \"{newPath}\". Probably it was deleted?");
return false;
}

Patcher.Logger.LogDebug($"Loading uncompressed bundle \"{metadata.UncompressedAssetBundleName}\"");
path = newPath;

metadata.LastAccessTime = DateTime.Now;
Patcher.MetadataManager.SaveMetadata(metadata);

return true;
}

public void DeleteCachedAssetBundle(string path)
private void RecompressAssetBundleInternal(WorkAsset workAsset)
{
FileHelper.TryDeleteFile(path, out var fileException);
if (fileException != null)
if (!File.Exists(workAsset.Path))
{
Patcher.Logger.LogError($"Failed to delete uncompressed assetbundle\n{fileException}");
return;
}

if (DriveHelper.HasDriveSpaceOnPath(CachePath, 10))
{
Patcher.Logger.LogDebug($"Queued recompress of \"{Path.GetFileName(workAsset.Path)}\" assetbundle");

m_WorkAssets.Enqueue(workAsset);
StartRunner();
return;
}

Patcher.Logger.LogWarning($"Ignoring request of decompressing, because the free drive space is less than 10GB");
return;
}

private void StartRunner()
Expand Down Expand Up @@ -140,7 +173,7 @@ private async Task ProcessQueue()
{
while (m_WorkAssets.TryDequeue(out var work))
{
await DecompressAssetBundleAsync(work.Path, work.Hash);
await DecompressAssetBundleAsync(work);
}
}
finally
Expand All @@ -155,35 +188,47 @@ private async Task ProcessQueue()
}
}

private async Task DecompressAssetBundleAsync(string path, byte[] hash)
private async Task DecompressAssetBundleAsync(WorkAsset workAsset)
{
var metadata = new Metadata()
{
OriginalAssetBundleHash = HashingHelper.HashToString(hash),
OriginalAssetBundleHash = HashingHelper.HashToString(workAsset.Hash),
LastAccessTime = DateTime.Now,
};
var originalFileName = Path.GetFileNameWithoutExtension(path);
var originalFileName = Path.GetFileNameWithoutExtension(workAsset.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
await FileHelper.RetryUntilFileIsClosedAsync(path, 5);
await FileHelper.RetryUntilFileIsClosedAsync(workAsset.Path, 5);
await AsyncHelper.SwitchToMainThread();

var op = AssetBundle.RecompressAssetBundleAsync(path, outputPath,
BuildCompression.UncompressedRuntime, 0, UnityEngine.ThreadPriority.Normal);
var op = AssetBundle.RecompressAssetBundleAsync(workAsset.Path, outputPath,
BuildCompression.UncompressedRuntime, 0, ThreadPriority.Normal);

await op.WaitCompletionAsync();

// we are in main thread, load results locally to make unity happy
var result = op.result;
var humanReadableResult = op.humanReadableResult;
var success = op.success;

await AsyncHelper.SwitchToThreadPool();

if (op.result is not AssetBundleLoadResult.Success || !op.success)
// delete temp bundle if needed
if (workAsset.DeleteBundleAfterOperation)
{
Patcher.Logger.LogWarning($"Failed to decompress a assetbundle at \"{path}\"\nResult: {op.result}, {op.humanReadableResult}");
FileHelper.TryDeleteFile(workAsset.Path, out _);
}

if (result is not AssetBundleLoadResult.Success || !success)
{
Patcher.Logger.LogWarning($"Failed to decompress a assetbundle at \"{workAsset.Path}\"\nResult: {result}, {humanReadableResult}");
return;
}

// check if unity returned the same assetbundle (means that assetbundle is already decompressed)
if (hash.AsSpan().SequenceEqual(HashingHelper.HashFile(outputPath)))
if (workAsset.Hash.AsSpan().SequenceEqual(HashingHelper.HashFile(outputPath)))
{
Patcher.Logger.LogDebug($"Assetbundle \"{originalFileName}\" is already uncompressed, adding to ignore list");

Expand All @@ -194,19 +239,23 @@ private async Task DecompressAssetBundleAsync(string path, byte[] hash)
return;
}

Patcher.Logger.LogDebug($"Assetbundle \"{originalFileName}\" is now uncompressed!");

metadata.UncompressedAssetBundleName = outputName;
Patcher.MetadataManager.SaveMetadata(metadata);
}

private struct WorkAsset
{
public WorkAsset(string path, byte[] hash)
public WorkAsset(string path, byte[] hash, bool deleteBundleAfterOperation)
{
Path = path;
Hash = hash;
DeleteBundleAfterOperation = deleteBundleAfterOperation;
}

public string Path { get; set; }
public byte[] Hash { get; set; }
public bool DeleteBundleAfterOperation { get; }
}
}
61 changes: 25 additions & 36 deletions BepInExFasterLoadAssetBundles/Patcher.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,17 +90,18 @@ private static void DeleteOldCache()
private static void LoadAssetBundleFromFileFast(ref string path)
{
// mod trying to load assetbundle at null path, buh
if (path == null)
if (string.IsNullOrEmpty(path))
{
return;
}

var tempPath = string.Copy(path);
try
{
if (AssetBundleManager.TryRecompressAssetBundle(ref tempPath))
using var bundleFileStream = new FileStream(path, FileMode.Open, FileAccess.ReadWrite);

if (HandleStreamBundle(bundleFileStream, out var newPath))
{
path = tempPath;
path = newPath;
}
}
catch (Exception ex)
Expand All @@ -109,55 +110,43 @@ private static void LoadAssetBundleFromFileFast(ref string path)
}
}

private static void LoadAssetBundleFromStreamFast(ref Stream stream)
private static bool LoadAssetBundleFromStreamFast(Stream stream, ref AssetBundle? __result)
{
if (stream is not FileStream fileStream)
if (HandleStreamBundle(stream, out var path))
{
return;
__result = AssetBundle.LoadFromFile_Internal(path, 0, 0);
return false;
}

var previousPosition = fileStream.Position;

try
{
if (AssetBundleManager.TryRecompressAssetBundle(fileStream, out var path))
{
stream = File.OpenRead(path);
return;
}
}
catch (Exception ex)
{
Logger.LogError($"Failed to decompress assetbundle\n{ex}");
}

fileStream.Position = previousPosition;

return true;
}

private static bool LoadAssetBundleFromStreamAsyncFast(Stream stream, out AssetBundleCreateRequest? __result)
private static bool LoadAssetBundleFromStreamAsyncFast(Stream stream, ref AssetBundleCreateRequest? __result)
{
__result = null;
if (stream is not FileStream fileStream)
if (HandleStreamBundle(stream, out var path))
{
return true;
__result = AssetBundle.LoadFromFileAsync_Internal(path, 0, 0);
return false;
}

return true;
}

var previousPosition = fileStream.Position;
private static bool HandleStreamBundle(Stream stream, out string path)
{
var previousPosition = stream.Position;

try
{
if (AssetBundleManager.TryRecompressAssetBundle(fileStream, out var path))
{
__result = AssetBundle.LoadFromFileAsync_Internal(path, 0, 0);
return false;
}
return AssetBundleManager.TryRecompressAssetBundle(stream, out path);
}
catch (Exception ex)
{
Logger.LogError($"Failed to decompress assetbundle\n{ex}");
}

fileStream.Position = previousPosition;
return true;
stream.Position = previousPosition;
path = null!;
return false;
}
}
6 changes: 6 additions & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.5.0] 2024-04-24
### Changed
- All bundle loading by stream are now recompressed.
## Fixed
- Array leaking from the pool.

## [0.4.0] 2024-04-04
### Added
- Check of drive space before trying to decompress.
Expand Down

0 comments on commit 5840e19

Please sign in to comment.