Make XAML hot reload JIT on UI load

This means we don't have to JIT a bunch of UIs that you might not open, reducing memory usage and startup overhead.

One (1) UI is always JITed in another thread before prototype UIs are loaded, so as to warm up the JIT machinery. Said type is DropDownDebugConsole which always gets used anyways so there's no harm in it.

In total, these changes save more than a second of startup time for me.

Author: PJB3005

RELEASE-NOTES.md

@@ -65,6 +65,7 @@ END TEMPLATE-->
 * ComponentNameSerializer will now ignore any components that have been ignored via `IComponentFactory.RegisterIgnore`.
 * Add pure to some SharedTransformSystem methods.
 * `Control.Stylesheet` does not do any work if assigning the value it already has.
+* XAML hot reload now JITs UIs when first opened rather than doing every single one at client startup. This reduces dev startup overhead significantly and probably helps with memory usage too.
 
 ### Internal
 

Robust.Client/GameController/GameController.cs

@@ -216,12 +216,12 @@ internal bool StartupContinue(DisplayMode displayMode)
 
             _loadscr.LoadingStep(_reload.Initialize, _reload);
             _loadscr.LoadingStep(_reflectionManager.Initialize, _reflectionManager);
+            _loadscr.LoadingStep(_xamlProxyManager.Initialize, _xamlProxyManager);
+            _loadscr.LoadingStep(_xamlHotReloadManager.Initialize, _xamlHotReloadManager);
             _loadscr.BeginLoadingSection(_prototypeManager);
             _prototypeManager.Initialize();
             _prototypeManager.LoadDefaultPrototypes();
             _loadscr.EndLoadingSection();
-            _loadscr.LoadingStep(_xamlProxyManager.Initialize, _xamlProxyManager);
-            _loadscr.LoadingStep(_xamlHotReloadManager.Initialize, _xamlHotReloadManager);
             _loadscr.LoadingStep(_userInterfaceManager.Initialize, "UI init");
             _loadscr.LoadingStep(_eyeManager.Initialize, _eyeManager);
             _loadscr.LoadingStep(_entityManager.Initialize, _entityManager);

Robust.Client/UserInterface/UserInterfaceManager.cs

@@ -29,6 +29,11 @@ namespace Robust.Client.UserInterface
 {
     internal sealed partial class UserInterfaceManager : IUserInterfaceManagerInternal
     {
+        /// <summary>
+        /// A type that will always be instantiated anyways.
+        /// </summary>
+        public static readonly Type XamlHotReloadWarmupType = typeof(DropDownDebugConsole);
+
         [Dependency] private readonly IDependencyCollection _rootDependencies = default!;
         [Dependency] private readonly IInputManager _inputManager = default!;
         [Dependency] private readonly IFontManager _fontManager = default!;

Robust.Client/UserInterface/XAML/Proxy/XamlImplementationStorage.cs

@@ -2,7 +2,10 @@
 using System;
 using System.Collections.Generic;
 using System.Reflection;
+using System.Threading;
 using Robust.Shared.Log;
+using Robust.Shared.Timing;
+using Robust.Shared.Utility;
 using Robust.Xaml;
 
 namespace Robust.Client.UserInterface.XAML.Proxy;
@@ -36,6 +39,8 @@ internal sealed class XamlImplementationStorage
     /// </summary>
     private readonly Dictionary<string, Type> _fileType = new();
 
+    private readonly Dictionary<Type, string> _fileTypeReverse = new();
+
     /// <summary>
     /// For each type, store the JIT-compiled implementation of Populate.
     /// </summary>
@@ -50,6 +55,8 @@ internal sealed class XamlImplementationStorage
     private readonly ISawmill _sawmill;
     private readonly XamlJitDelegate _jitDelegate;
 
+    private readonly Lock _compileLock = new();
+
     /// <summary>
     /// Create the storage.
     /// </summary>
@@ -102,6 +109,8 @@ public XamlImplementationStorage(ISawmill sawmill, XamlJitDelegate jitDelegate)
     /// <param name="assembly">an assembly</param>
     public void Add(Assembly assembly)
     {
+        using var _ = _compileLock.EnterScope();
+
         foreach (var (type, metadata) in TypesWithXamlMetadata(assembly))
         {
             // this can fail, but if it does, that means something is _really_ wrong
@@ -132,6 +141,8 @@ public void Add(Assembly assembly)
                     $"{fileName}. ({type.FullName} and {_fileType[fileName].FullName}). this is a bug in XamlAotCompiler"
                 );
             }
+
+            _fileTypeReverse.Add(type, fileName);
         }
     }
 
@@ -145,6 +156,8 @@ public void Add(Assembly assembly)
     /// </remarks>
     public void ForceReloadAll()
     {
+        using var _ = _compileLock.EnterScope();
+
         foreach (var (fileName, fileContent) in _fileContent)
         {
             SetImplementation(fileName, fileContent, true);
@@ -161,9 +174,19 @@ public void ForceReloadAll()
     /// <returns>true if not a no-op</returns>
     public bool CanSetImplementation(string fileName)
     {
+        using var _ = _compileLock.EnterScope();
         return _fileType.ContainsKey(fileName);
     }
 
+    public MethodInfo? CompileType(Type type)
+    {
+        if (_fileTypeReverse.TryGetValue(type, out var fileName))
+            return SetImplementation(fileName, _fileContent[fileName], quiet: true);
+
+        _sawmill.Warning($"Type {type} has no XAML file!");
+        return null;
+    }
+
     /// <summary>
     /// Replace the implementation of <paramref name="fileName"/> by JIT-ing
     /// <paramref name="fileContent"/>.
@@ -174,12 +197,14 @@ public bool CanSetImplementation(string fileName)
     /// <param name="fileName">the name of the file whose implementation should be replaced</param>
     /// <param name="fileContent">the new implementation</param>
     /// <param name="quiet">if true, then don't bother to log</param>
-    public void SetImplementation(string fileName, string fileContent, bool quiet)
+    public MethodInfo? SetImplementation(string fileName, string fileContent, bool quiet)
     {
+        using var _ = _compileLock.EnterScope();
+
         if (!_fileType.TryGetValue(fileName, out var type))
         {
             _sawmill.Warning($"SetImplementation called with {fileName}, but no types care about its contents");
-            return;
+            return null;
         }
 
         var uri =
@@ -190,12 +215,14 @@ public void SetImplementation(string fileName, string fileContent, bool quiet)
         {
             _sawmill.Debug($"replacing {fileName} for {type}");
         }
+
         var impl = _jitDelegate(type, uri, fileName, fileContent);
         if (impl != null)
         {
             _populateImplementations[type] = impl;
         }
         _fileContent[fileName] = fileContent;
+        return impl;
     }
 
     /// <summary>
@@ -210,8 +237,12 @@ public bool Populate(Type t, object o)
     {
         if (!_populateImplementations.TryGetValue(t, out var implementation))
         {
-            // pop out if we never JITed anything
-            return false;
+            // JIT if needed.
+            implementation = CompileType(t);
+
+            // pop out if we never JITed anything/couldn't JIT
+            if (implementation == null)
+                return false;
         }
 
         implementation.Invoke(null, [null, o]);

Robust.Client/UserInterface/XAML/Proxy/XamlProxyManager.cs

@@ -2,6 +2,9 @@
 using System;
 using System.Collections.Generic;
 using System.Reflection;
+using System.Threading;
+using Robust.Shared;
+using Robust.Shared.Configuration;
 using Robust.Shared.IoC;
 using Robust.Shared.Log;
 using Robust.Shared.Reflection;
@@ -17,6 +20,7 @@ public sealed class XamlProxyManager: IXamlProxyManager
     ISawmill _sawmill = null!;
     [Dependency] IReflectionManager _reflectionManager = null!;
     [Dependency] ILogManager _logManager = null!;
+    [Dependency] private readonly IConfigurationManager _cfg = null!;
 
     XamlImplementationStorage _xamlImplementationStorage = null!;
 
@@ -31,8 +35,21 @@ public void Initialize()
         _sawmill = _logManager.GetSawmill("xamlhotreload");
         _xamlImplementationStorage = new XamlImplementationStorage(_sawmill, Compile);
 
-        AddAssemblies();
+        var preload = _cfg.GetCVar(CVars.UIXamlJitPreload);
+
+        AddAssemblies(reload: preload);
         _reflectionManager.OnAssemblyAdded += (_, _) => { AddAssemblies(); };
+
+        if (!preload)
+        {
+            // Compile any type at all on another thread, so we don't hold up main thread init with loading
+            // the entire XAML compiler machinery.
+            // In my testing, it took like 0.5s on debug to run the first XAML compile. Yeah.
+            ThreadPool.QueueUserWorkItem(_ =>
+            {
+                _xamlImplementationStorage.CompileType(UserInterfaceManager.XamlHotReloadWarmupType);
+            });
+        }
     }
 
     /// <summary>
@@ -61,7 +78,7 @@ public void SetImplementation(string fileName, string fileContent)
     /// Add all the types from all known assemblies, then force-JIT everything
     /// again.
     /// </summary>
-    private void AddAssemblies()
+    private void AddAssemblies(bool reload = true)
     {
         foreach (var a in _reflectionManager.Assemblies)
         {
@@ -74,8 +91,8 @@ private void AddAssemblies()
             }
         }
 
-        // Always use the JITed versions on debug builds
-        _xamlImplementationStorage.ForceReloadAll();
+        if (reload)
+            _xamlImplementationStorage.ForceReloadAll();
     }
 
     /// <summary>

Robust.Shared/CVars.cs

@@ -1925,6 +1925,13 @@ internal static readonly CVarDef<string>
         public static readonly CVarDef<string> XamlHotReloadMarkerName =
             CVarDef.Create("ui.xaml_hot_reload_marker_name", "SpaceStation14.sln", CVar.CLIENTONLY);
 
+        /// <summary>
+        /// If true, all XAML UIs will be JITed for hot reload on client startup.
+        /// If false, they will be JITed on demand.
+        /// </summary>
+        public static readonly CVarDef<bool> UIXamlJitPreload =
+            CVarDef.Create("ui.xaml_jit_preload", false, CVar.CLIENTONLY);
+
         /*
          * FONT
          */