From 0cba392b66aaa97f5f60b9515be44ca546515726 Mon Sep 17 00:00:00 2001
From: Timothy Baldridge <tbaldridge@gmail.com>
Date: Mon, 9 Oct 2023 14:25:10 -0600
Subject: [PATCH] Several tricks to get memory usage under control during list
 validation.

---
 Wabbajack.CLI/Verbs/ValidateLists.cs                | 13 +++++++++++--
 Wabbajack.Installer/StandardInstaller.cs            |  4 ++--
 .../AChunkedBufferingStream.cs                      |  6 ++++++
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/Wabbajack.CLI/Verbs/ValidateLists.cs b/Wabbajack.CLI/Verbs/ValidateLists.cs
index 27a97559..45bf40b2 100644
--- a/Wabbajack.CLI/Verbs/ValidateLists.cs
+++ b/Wabbajack.CLI/Verbs/ValidateLists.cs
@@ -53,6 +53,8 @@ public class ValidateLists
     private readonly Networking.WabbajackClientApi.Client _wjClient;
     private readonly HttpClient _httpClient;
     private readonly IResource<HttpClient> _httpLimiter;
+    private readonly AsyncLock _imageProcessLock;
+
 
     public ValidateLists(ILogger<ValidateLists> logger, Networking.WabbajackClientApi.Client wjClient,
         Client gitHubClient, TemporaryFileManager temporaryFileManager,
@@ -73,6 +75,7 @@ public class ValidateLists
         _random = new Random();
         _httpClient = httpClient;
         _httpLimiter = httpLimiter;
+        _imageProcessLock = new AsyncLock();
     }
 
     public static VerbDefinition Definition = new("validate-lists",
@@ -139,6 +142,9 @@ public class ValidateLists
                 _logger.LogInformation("Loading Modlist");
                 modListData =
                     await StandardInstaller.Load(_dtos, _dispatcher, modList, token);
+                // Clear out the directives to save memory
+                modListData.Directives = Array.Empty<Directive>();
+                GC.Collect();
             }
             catch (Exception ex)
             {
@@ -146,8 +152,10 @@ public class ValidateLists
                 validatedList.Status = ListStatus.ForcedDown;
                 return validatedList;
             }
+            
 
-            _logger.LogInformation("Verifying {Count} archives", modListData.Archives.Length);
+
+            _logger.LogInformation("Verifying {Count} archives from {Name}", modListData.Archives.Length, modList.NamespacedName);
 
             var archives = await modListData.Archives.PMapAll(async archive =>
             {
@@ -248,7 +256,6 @@ public class ValidateLists
 
         await ExportReports(reports, validatedLists, token);
         
-        
         var usedMirroredFiles = validatedLists.SelectMany(a => a.Archives)
             .Where(m => m.Status == ArchiveStatus.Mirrored)
             .Select(m => m.Original.Hash)
@@ -261,6 +268,7 @@ public class ValidateLists
     private async Task<(RelativePath SmallImage, RelativePath LargeImage)> ProcessModlistImage(AbsolutePath reports, ModlistMetadata validatedList,
         CancellationToken token)
     {
+        using var _ = await _imageProcessLock.WaitAsync();
         _logger.LogInformation("Processing Modlist Image for {MachineUrl}", validatedList.NamespacedName);
         var baseFolder = reports.Combine(validatedList.NamespacedName);
         baseFolder.CreateDirectory();
@@ -401,6 +409,7 @@ public class ValidateLists
         "skyrimvr.ini",
     }.Select(f => f.ToRelativePath()).ToHashSet();
 
+
     private async Task<Dictionary<RelativePath, byte[]>> GetFiles(ModList modlist, ModlistMetadata metadata, CancellationToken token)
     {
         var archive = new Archive
diff --git a/Wabbajack.Installer/StandardInstaller.cs b/Wabbajack.Installer/StandardInstaller.cs
index 35c2487c..b839c09d 100644
--- a/Wabbajack.Installer/StandardInstaller.cs
+++ b/Wabbajack.Installer/StandardInstaller.cs
@@ -587,10 +587,10 @@ public class StandardInstaller : AInstaller<StandardInstaller>
             Hash = metadata.DownloadMetadata.Hash
         };
 
-        var stream = await dispatcher.ChunkedSeekableStream(archive, token);
+        await using var stream = await dispatcher.ChunkedSeekableStream(archive, token);
         await using var reader = new ZipReader(stream);
         var entry = (await reader.GetFiles()).First(e => e.FileName == "modlist");
-        var ms = new MemoryStream();
+        using var ms = new MemoryStream();
         await reader.Extract(entry, ms, token);
         ms.Position = 0;
         return JsonSerializer.Deserialize<ModList>(ms, dtos.Options)!;
diff --git a/Wabbajack.Networking.Http/AChunkedBufferingStream.cs b/Wabbajack.Networking.Http/AChunkedBufferingStream.cs
index 5c5e6e02..72bb95c9 100644
--- a/Wabbajack.Networking.Http/AChunkedBufferingStream.cs
+++ b/Wabbajack.Networking.Http/AChunkedBufferingStream.cs
@@ -23,6 +23,12 @@ public abstract class AChunkedBufferingStream : Stream
         _maxChunks = maxChunks;
         _chunks = new Dictionary<ulong, byte[]>();
     }
+    
+    protected override void Dispose(bool disposing)
+    {
+        _chunks.Clear();
+        base.Dispose(disposing);
+    }
 
     public abstract Task<byte[]> LoadChunk(long offset, int size);