Sandboxing performance improvements.

Don't leave file handles dangling. Prefetch verifying assembly images to speed stuff up.
space-wizards · Jul 15, 2023 · 177ca6b · 177ca6b
1 parent 3c262af
commit 177ca6b
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 7 deletions.
diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
@@ -50,6 +50,7 @@ END TEMPLATE-->
 
 * BQL `with` now includes paused entities.
 * The game loop now times more accurately and avoids sleeping more than necessary.
+* Sandboxing (and thus, client startup) should be much faster when ran from the launcher.
 
 ### Internal
 

diff --git a/Robust.Shared/ContentPack/AssemblyTypeChecker.cs b/Robust.Shared/ContentPack/AssemblyTypeChecker.cs
@@ -59,7 +59,7 @@ public AssemblyTypeChecker(IResourceManager res, ISawmill sawmill)
             _config = Task.Run(() => LoadConfig(sawmill));
         }
 
-        private Resolver CreateResolver()
+        internal Resolver CreateResolver()
         {
             var dotnetDir = Path.GetDirectoryName(typeof(int).Assembly.Location)!;
             var ourPath = typeof(AssemblyTypeChecker).Assembly.Location;
@@ -100,6 +100,19 @@ private Resolver CreateResolver()
         /// <param name="assembly">Assembly to load.</param>
         /// <returns></returns>
         public bool CheckAssembly(Stream assembly)
+        {
+            using var resolver = CreateResolver();
+
+            return CheckAssembly(assembly, resolver);
+        }
+
+        /// <summary>
+        ///     Check the assembly for any illegal types. Any types not on the white list
+        ///     will cause the assembly to be rejected.
+        /// </summary>
+        /// <param name="assembly">Assembly to load.</param>
+        /// <returns></returns>
+        public bool CheckAssembly(Stream assembly, Resolver resolver)
         {
             if (WouldNoOp)
             {
@@ -110,8 +123,7 @@ public bool CheckAssembly(Stream assembly)
             _sawmill.Debug("Checking assembly...");
             var fullStopwatch = Stopwatch.StartNew();
 
-            var resolver = CreateResolver();
-            using var peReader = ModLoader.MakePEReader(assembly, leaveOpen: true);
+            using var peReader = ModLoader.MakePEReader(assembly, leaveOpen: true, PEStreamOptions.PrefetchEntireImage);
             var reader = peReader.GetMetadataReader();
 
             var asmName = reader.GetString(reader.GetAssemblyDefinition().Name);
@@ -856,7 +868,7 @@ public bool CheckAssembly(string diskPath)
             return handle.IsNil ? null : reader.GetString(handle);
         }
 
-        private sealed class Resolver : IResolver
+        internal sealed class Resolver : IResolver, IDisposable
         {
             private readonly ConcurrentDictionary<string, PEReader?> _dictionary = new();
             private readonly AssemblyTypeChecker _parent;
@@ -910,6 +922,14 @@ public Resolver(AssemblyTypeChecker parent, string[] diskLoadPaths, ResPath[] re
             {
                 return _dictionary.GetOrAdd(simpleName, ResolveCore);
             }
+
+            public void Dispose()
+            {
+                foreach (var reader in _dictionary.Values)
+                {
+                    reader?.Dispose();
+                }
+            }
         }
 
         internal sealed class TypeProvider : ISignatureTypeProvider<MType, int>

diff --git a/Robust.Shared/ContentPack/ModLoader.cs b/Robust.Shared/ContentPack/ModLoader.cs
@@ -118,13 +118,14 @@ public bool TryLoadModules(IEnumerable<ResPath> paths)
                 var checkerSw = Stopwatch.StartNew();
 
                 var typeChecker = MakeTypeChecker();
+                var resolver = typeChecker.CreateResolver();
 
                 Parallel.ForEach(files, pair =>
                 {
                     var (name, (path, _)) = pair;
 
                     using var stream = _res.ContentFileRead(path);
-                    if (!typeChecker.CheckAssembly(stream))
+                    if (!typeChecker.CheckAssembly(stream, resolver))
                     {
                         throw new TypeCheckFailedException($"Assembly {name} failed type checks.");
                     }
@@ -426,12 +427,15 @@ private AssemblyTypeChecker MakeTypeChecker()
             };
         }
 
-        internal static PEReader MakePEReader(Stream stream, bool leaveOpen=false)
+        internal static PEReader MakePEReader(Stream stream, bool leaveOpen=false, PEStreamOptions options=PEStreamOptions.Default)
         {
             if (!stream.CanSeek)
                 stream = leaveOpen ? stream.CopyToMemoryStream() : stream.ConsumeToMemoryStream();
 
-            return new PEReader(stream, leaveOpen ? PEStreamOptions.LeaveOpen : default);
+            if (leaveOpen)
+                options |= PEStreamOptions.LeaveOpen;
+
+            return new PEReader(stream, options);
         }
     }
 }