From 6e0ae3f9db2af0498fc6e3da8434972034e5e0c3 Mon Sep 17 00:00:00 2001
From: zingballyhoo <zingballyhoo@live.co.uk>
Date: Mon, 18 Nov 2024 16:12:51 +0000
Subject: [PATCH] add ParallelCDNIndexLoading option

---
 TACTLib/Client/ClientCreateArgs.cs  |  2 ++
 TACTLib/Protocol/CDNIndexHandler.cs | 47 ++++++++++++++++++++++-------
 2 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/TACTLib/Client/ClientCreateArgs.cs b/TACTLib/Client/ClientCreateArgs.cs
index d4313d0..a55f45c 100644
--- a/TACTLib/Client/ClientCreateArgs.cs
+++ b/TACTLib/Client/ClientCreateArgs.cs
@@ -116,6 +116,8 @@ public bool Online {
         
         public bool LoadCDNIndices { get; set; } = true;
         public ClientHandler? TryShareCDNIndexWithHandler { get; set; } = null;
+        public bool ParallelCDNIndexLoading { get; set; } = false;
+        public int MaxCDNIndexLoadingParallelism { get; set; } = 4;
         
         public bool LoadRoot { get; set; } = true;
         public bool LoadVFS { get; set; } = true;
diff --git a/TACTLib/Protocol/CDNIndexHandler.cs b/TACTLib/Protocol/CDNIndexHandler.cs
index 1eb328f..0cfe334 100644
--- a/TACTLib/Protocol/CDNIndexHandler.cs
+++ b/TACTLib/Protocol/CDNIndexHandler.cs
@@ -1,9 +1,12 @@
 using System;
+using System.Collections.Concurrent;
+using System.Collections.Frozen;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.Runtime.InteropServices;
+using System.Threading.Tasks;
 using CommunityToolkit.HighPerformance;
 using Microsoft.Win32.SafeHandles;
 using TACTLib.Client;
@@ -25,7 +28,7 @@ public struct IndexEntry
     public class CDNIndexHandler
     {
         private readonly ClientHandler Client;
-        private readonly Dictionary<FullKey, IndexEntry> CDNIndexData = new Dictionary<FullKey, IndexEntry>(CASCKeyComparer.Instance);
+        private readonly IDictionary<FullKey, IndexEntry> CDNIndexData;
 
         private FixedFooter ArchiveGroupFooter;
         private SafeFileHandle ArchiveGroupFileHandle = new SafeFileHandle();
@@ -46,6 +49,13 @@ public static CDNIndexHandler Initialize(ClientHandler clientHandler)
         private CDNIndexHandler(ClientHandler client)
         {
             Client = client;
+            if (client.CreateArgs.ParallelCDNIndexLoading) 
+            {
+                CDNIndexData = new ConcurrentDictionary<FullKey, IndexEntry>(CASCKeyComparer.Instance);
+            } else
+            {
+                CDNIndexData = new Dictionary<FullKey, IndexEntry>(CASCKeyComparer.Instance);
+            }
             
             // load loose files index so we dont have to hit the cdn just to get 404'd
             OpenOrDownloadIndexFile(client.ConfigHandler.CDNConfig.Values["file-index"][0], ARCHIVE_ID_LOOSE);
@@ -68,11 +78,30 @@ private CDNIndexHandler(ClientHandler client)
                 // only loose files will be available
                 return;
             }
-            for (var index = 0; index < client.ConfigHandler.CDNConfig.Archives.Count; index++)
+            if (client.CreateArgs.ParallelCDNIndexLoading)
+            {
+                Parallel.ForEach(client.ConfigHandler.CDNConfig.Archives, new ParallelOptions {
+                    MaxDegreeOfParallelism = client.CreateArgs.MaxCDNIndexLoadingParallelism
+                }, (archive, _, index) => {
+                    OpenOrDownloadIndexFile(archive, (int)index);
+                });
+            } else
             {
-                string archive = client.ConfigHandler.CDNConfig.Archives[index];
-                OpenOrDownloadIndexFile(archive, index);
+                for (var index = 0; index < client.ConfigHandler.CDNConfig.Archives.Count; index++)
+                {
+                    var archive = client.ConfigHandler.CDNConfig.Archives[index];
+                    OpenOrDownloadIndexFile(archive, index);
+                }
             }
+            
+            // todo: still not very happy about this system
+            // we create a giant dictionary (with no initial size) which can contain a lot of empty space (90+MB)
+            // converting to a frozen dictionary afterwards helps a bit but the arrays are still wasteful
+            // also means the peak memory is higher during conversion
+            // using IDictionary is also worse for lookup perf, but this keeps the code simpler for now
+            CDNIndexData = CDNIndexData.ToFrozenDictionary(CASCKeyComparer.Instance);
+            // we could load each index into an array of entries and then merge sort into one giant array...
+            // (also means higher building memory cost but maybe that's inevitable)
         }
 
         private bool LoadGroupIndexFile(string hash) {
@@ -229,8 +258,7 @@ private void DownloadIndexFile(string archive, int i)
         {
             try
             {
-                var cdn = Client.CDNClient!;
-                var indexData = cdn.FetchIndexFile(archive);
+                var indexData = Client.CDNClient!.FetchIndexFile(archive);
                 if (indexData == null) throw new Exception($"failed to fetch archive index data for {archive} (index {i})");
                 using var indexDataStream = new MemoryStream(indexData);
                 ParseIndex(indexDataStream, i);
@@ -344,11 +372,8 @@ public bool IsLooseFile(FullKey key) {
 
         public byte[]? OpenIndexEntry(IndexEntry entry)
         {
-            var archive = Client.ConfigHandler.CDNConfig.Archives[entry.Index];
-
-            var cdn = Client.CDNClient!;
-            var stream = cdn.FetchIndexEntry(archive, entry);
-            return stream;
+            var archiveKey = Client.ConfigHandler.CDNConfig.Archives[entry.Index];
+            return Client.CDNClient!.FetchIndexEntry(archiveKey, entry);
         }
         
         [StructLayout(LayoutKind.Sequential, Pack = 1)]