restructure content manager to better handle asset disposal (#352)

3 files changed, 108 insertions, 25 deletions

diff --git a/src/StardewModdingAPI/Framework/ContentManagerShim.cs b/src/StardewModdingAPI/Framework/ContentManagerShim.cs
new file mode 100644
index 00000000..d46f23a3
--- /dev/null
+++ b/src/StardewModdingAPI/Framework/ContentManagerShim.cs
@@ -0,0 +1,50 @@
+using StardewValley;
+
+namespace StardewModdingAPI.Framework
+{
+    /// <summary>A minimal content manager which defers to SMAPI's main content manager.</summary>
+    internal class ContentManagerShim : LocalizedContentManager
+    {
+        /*********
+        ** Properties
+        *********/
+        /// <summary>SMAPI's underlying content manager.</summary>
+        private readonly SContentManager ContentManager;
+
+
+        /*********
+        ** Accessors
+        *********/
+        /// <summary>The content manager's name for logs (if any).</summary>
+        public string Name { get; }
+
+
+        /*********
+        ** Public methods
+        *********/
+        /// <summary>Construct an instance.</summary>
+        /// <param name="contentManager">SMAPI's underlying content manager.</param>
+        /// <param name="name">The content manager's name for logs (if any).</param>
+        public ContentManagerShim(SContentManager contentManager, string name)
+                : base(contentManager.ServiceProvider, contentManager.RootDirectory, contentManager.CurrentCulture, contentManager.LanguageCodeOverride)
+        {
+            this.ContentManager = contentManager;
+            this.Name = name;
+        }
+
+        /// <summary>Load an asset that has been processed by the content pipeline.</summary>
+        /// <typeparam name="T">The type of asset to load.</typeparam>
+        /// <param name="assetName">The asset path relative to the loader root directory, not including the <c>.xnb</c> extension.</param>
+        public override T Load<T>(string assetName)
+        {
+            return this.ContentManager.LoadFor<T>(assetName, this);
+        }
+
+        /// <summary>Dispose held resources.</summary>
+        /// <param name="disposing">Whether the content manager is disposing (rather than finalising).</param>
+        protected override void Dispose(bool disposing)
+        {
+            this.ContentManager.DisposeFor(this);
+        }
+    }
+}
diff --git a/src/StardewModdingAPI/Framework/SContentManager.cs b/src/StardewModdingAPI/Framework/SContentManager.cs
index 25775291..4c4eee58 100644
--- a/src/StardewModdingAPI/Framework/SContentManager.cs
+++ b/src/StardewModdingAPI/Framework/SContentManager.cs
@@ -48,6 +48,9 @@ namespace StardewModdingAPI.Framework
         /// <summary>The assets currently being intercepted by <see cref="IAssetLoader"/> instances. This is used to prevent infinite loops when a loader loads a new asset.</summary>
         private readonly ContextHash<string> AssetsBeingLoaded = new ContextHash<string>();
 
+        /// <summary>A lookup of the content managers which loaded each asset.</summary>
+        private readonly IDictionary<string, HashSet<ContentManager>> AssetLoaders = new Dictionary<string, HashSet<ContentManager>>();
+
 
         /*********
         ** Accessors
@@ -98,7 +101,6 @@ namespace StardewModdingAPI.Framework
             // get asset data
             this.CoreAssets = new CoreAssets(this.NormaliseAssetName);
             this.KeyLocales = this.GetKeyLocales(reflection);
-
         }
 
         /// <summary>Normalise path separators in a file path. For asset keys, see <see cref="NormaliseAssetName"/> instead.</summary>
@@ -135,11 +137,23 @@ namespace StardewModdingAPI.Framework
         /// <param name="assetName">The asset path relative to the loader root directory, not including the <c>.xnb</c> extension.</param>
         public override T Load<T>(string assetName)
         {
+            return this.LoadFor<T>(assetName, this);
+        }
+
+        /// <summary>Load an asset that has been processed by the content pipeline.</summary>
+        /// <typeparam name="T">The type of asset to load.</typeparam>
+        /// <param name="assetName">The asset path relative to the loader root directory, not including the <c>.xnb</c> extension.</param>
+        /// <param name="instance">The content manager instance for which to load the asset.</param>
+        public T LoadFor<T>(string assetName, ContentManager instance)
+        {
             assetName = this.NormaliseAssetName(assetName);
 
             // skip if already loaded
             if (this.IsNormalisedKeyLoaded(assetName))
+            {
+                this.TrackAssetLoader(assetName, instance);
                 return base.Load<T>(assetName);
+            }
 
             // load asset
             T data;
@@ -162,6 +176,7 @@ namespace StardewModdingAPI.Framework
 
             // update cache & return data
             this.Cache[assetName] = data;
+            this.TrackAssetLoader(assetName, instance);
             return data;
         }
 
@@ -172,8 +187,8 @@ namespace StardewModdingAPI.Framework
         public void Inject<T>(string assetName, T value)
         {
             assetName = this.NormaliseAssetName(assetName);
-
             this.Cache[assetName] = value;
+            this.TrackAssetLoader(assetName, this);
         }
 
         /// <summary>Get the current content locale.</summary>
@@ -229,8 +244,9 @@ namespace StardewModdingAPI.Framework
 
         /// <summary>Purge matched assets from the cache.</summary>
         /// <param name="predicate">Matches the asset keys to invalidate.</param>
+        /// <param name="dispose">Whether to dispose invalidated assets. This should only be <c>true</c> when they're being invalidated as part of a dispose, to avoid crashing the game.</param>
         /// <returns>Returns whether any cache entries were invalidated.</returns>
-        public bool InvalidateCache(Func<string, Type, bool> predicate)
+        public bool InvalidateCache(Func<string, Type, bool> predicate, bool dispose = true)
         {
             // find matching asset keys
             HashSet<string> purgeCacheKeys = new HashSet<string>(StringComparer.InvariantCultureIgnoreCase);
@@ -246,9 +262,14 @@ namespace StardewModdingAPI.Framework
                 }
             }
 
-            // purge from cache
+            // purge assets
             foreach (string key in purgeCacheKeys)
+            {
+                if (dispose && this.Cache[key] is IDisposable disposable)
+                    disposable.Dispose();
                 this.Cache.Remove(key);
+                this.AssetLoaders.Remove(key);
+            }
 
             // reload core game assets
             int reloaded = 0;
@@ -268,6 +289,19 @@ namespace StardewModdingAPI.Framework
             return false;
         }
 
+        /// <summary>Dispose assets for the given content manager shim.</summary>
+        /// <param name="shim">The content manager whose assets to dispose.</param>
+        internal void DisposeFor(ContentManagerShim shim)
+        {
+            this.Monitor.Log($"Content manager '{shim.Name}' disposed, disposing assets that aren't needed by any other asset loader.", LogLevel.Trace);
+            HashSet<string> keys = new HashSet<string>(
+                from entry in this.AssetLoaders
+                where entry.Value.Count == 1 && entry.Value.First() == shim
+                select entry.Key
+            );
+            this.InvalidateCache((key, type) => keys.Contains(key));
+        }
+
 
         /*********
         ** Private methods
@@ -280,6 +314,16 @@ namespace StardewModdingAPI.Framework
                 || this.Cache.ContainsKey($"{normalisedAssetName}.{this.GetKeyLocale.Invoke<string>()}"); // translated asset
         }
 
+        /// <summary>Track that a content manager loaded an asset.</summary>
+        /// <param name="key">The asset key that was loaded.</param>
+        /// <param name="manager">The content manager that loaded the asset.</param>
+        private void TrackAssetLoader(string key, ContentManager manager)
+        {
+            if (!this.AssetLoaders.TryGetValue(key, out HashSet<ContentManager> hash))
+                hash = this.AssetLoaders[key] = new HashSet<ContentManager>();
+            hash.Add(manager);
+        }
+
         /// <summary>Get the locale codes (like <c>ja-JP</c>) used in asset keys.</summary>
         /// <param name="reflection">Simplifies access to private game code.</param>
         private IDictionary<string, LanguageCode> GetKeyLocales(Reflector reflection)
@@ -463,23 +507,12 @@ namespace StardewModdingAPI.Framework
             }
         }
 
-        /// <summary>Dispose all game resources.</summary>
+        /// <summary>Dispose held resources.</summary>
         /// <param name="disposing">Whether the content manager is disposing (rather than finalising).</param>
         protected override void Dispose(bool disposing)
         {
-            if (!disposing)
-                return;
-
-            // Clear cache & reload all assets. While that may seem perverse during disposal, it's
-            // necessary due to limitations in the way SMAPI currently intercepts content assets.
-            // 
-            // The game uses multiple content managers while SMAPI needs one and only one. The game
-            // only disposes some of its content managers when returning to title, which means SMAPI
-            // can't know which assets are meant to be disposed. Here we remove current assets from
-            // the cache, but don't dispose them to avoid crashing any code that still references
-            // them. The garbage collector will eventually clean up any unused assets.
-            this.Monitor.Log("Content manager disposed, resetting cache.", LogLevel.Trace);
-            this.InvalidateCache((key, type) => true);
+            this.Monitor.Log("Disposing SMAPI's main content manager. It will no longer be usable after this point.", LogLevel.Trace);
+            base.Dispose(disposing);
         }
     }
 }
diff --git a/src/StardewModdingAPI/Framework/SGame.cs b/src/StardewModdingAPI/Framework/SGame.cs
index 2bfbc06a..76c106d7 100644
--- a/src/StardewModdingAPI/Framework/SGame.cs
+++ b/src/StardewModdingAPI/Framework/SGame.cs
@@ -38,9 +38,6 @@ namespace StardewModdingAPI.Framework
         /// <summary>Encapsulates monitoring and logging.</summary>
         private readonly IMonitor Monitor;
 
-        /// <summary>SMAPI's content manager.</summary>
-        private readonly SContentManager SContentManager;
-
         /// <summary>The maximum number of consecutive attempts SMAPI should make to recover from a draw error.</summary>
         private readonly Countdown DrawCrashTimer = new Countdown(60); // 60 ticks = roughly one second
 
@@ -177,6 +174,9 @@ namespace StardewModdingAPI.Framework
         /*********
         ** Accessors
         *********/
+        /// <summary>SMAPI's content manager.</summary>
+        public SContentManager SContentManager { get; }
+
         /// <summary>Whether SMAPI should log more information about the game context.</summary>
         public bool VerboseLogging { get; set; }
 
@@ -201,9 +201,9 @@ namespace StardewModdingAPI.Framework
             // override content manager
             this.Monitor?.Log("Overriding content manager...", LogLevel.Trace);
             this.SContentManager = new SContentManager(this.Content.ServiceProvider, this.Content.RootDirectory, Thread.CurrentThread.CurrentUICulture, null, this.Monitor);
-            this.Content = this.SContentManager;
-            Game1.content = this.SContentManager;
-            reflection.GetPrivateField<LocalizedContentManager>(typeof(Game1), "_temporaryContent").SetValue(null); // regenerate value with new content manager
+            this.Content = new ContentManagerShim(this.SContentManager, "SGame.Content");
+            Game1.content = new ContentManagerShim(this.SContentManager, "Game1.content");
+            reflection.GetPrivateField<LocalizedContentManager>(typeof(Game1), "_temporaryContent").SetValue(new ContentManagerShim(this.SContentManager, "Game1._temporaryContent")); // regenerate value with new content manager
         }
 
         /****
@@ -226,7 +226,7 @@ namespace StardewModdingAPI.Framework
                 throw new InvalidOperationException("SMAPI uses a single content manager internally. You can't get a new content manager with a different service provider.");
             if (rootDirectory != this.Content.RootDirectory)
                 throw new InvalidOperationException($"SMAPI uses a single content manager internally. You can't get a new content manager with a different root directory (current is {this.Content.RootDirectory}, requested {rootDirectory}).");
-            return this.SContentManager;
+            return new ContentManagerShim(this.SContentManager, "(generated instance)");
         }
 
         /// <summary>The method called when the game is updating its state. This happens roughly 60 times per second.</summary>