Add weak WorkingBeatmap cache (#5163)

Add weak WorkingBeatmap cache

Co-authored-by: Dan Balasescu <1329837+smoogipoo@users.noreply.github.com>

osu.Game.Tests/WaveformTestBeatmap.cs

@@ -30,9 +30,9 @@ public WaveformTestBeatmap(AudioManager audioManager)
             trackStore = audioManager.GetTrackStore(reader);
         }
 
-        public override void Dispose()
+        protected override void Dispose(bool isDisposing)
         {
-            base.Dispose();
+            base.Dispose(isDisposing);
             stream?.Dispose();
             reader?.Dispose();
             trackStore?.Dispose();

osu.Game/Beatmaps/BeatmapManager.cs

@@ -13,6 +13,7 @@
 using osu.Framework.Audio.Track;
 using osu.Framework.Extensions;
 using osu.Framework.Graphics.Textures;
+using osu.Framework.Lists;
 using osu.Framework.Logging;
 using osu.Framework.Platform;
 using osu.Framework.Threading;
@@ -159,6 +160,8 @@ private void validateOnlineIds(BeatmapSetInfo beatmapSet)
         /// <param name="beatmap">The beatmap difficulty to restore.</param>
         public void Restore(BeatmapInfo beatmap) => beatmaps.Restore(beatmap);
 
+        private readonly WeakList<WorkingBeatmap> workingCache = new WeakList<WorkingBeatmap>();
+
         /// <summary>
         /// Retrieve a <see cref="WorkingBeatmap"/> instance for the provided <see cref="BeatmapInfo"/>
         /// </summary>
@@ -173,12 +176,18 @@ public WorkingBeatmap GetWorkingBeatmap(BeatmapInfo beatmapInfo, WorkingBeatmap
             if (beatmapInfo?.BeatmapSet == null || beatmapInfo == DefaultBeatmap?.BeatmapInfo)
                 return DefaultBeatmap;
 
+            var cached = workingCache.FirstOrDefault(w => w.BeatmapInfo?.ID == beatmapInfo.ID);
+
+            if (cached != null)
+                return cached;
+
             if (beatmapInfo.Metadata == null)
                 beatmapInfo.Metadata = beatmapInfo.BeatmapSet.Metadata;
 
             WorkingBeatmap working = new BeatmapManagerWorkingBeatmap(Files.Store, new LargeTextureStore(host?.CreateTextureLoaderStore(Files.Store)), beatmapInfo, audioManager);
 
             previous?.TransferTo(working);
+            workingCache.Add(working);
 
             return working;
         }

osu.Game/Beatmaps/WorkingBeatmap.cs

@@ -11,6 +11,7 @@
 using System.IO;
 using System.Linq;
 using System.Threading;
+using System.Threading.Tasks;
 using osu.Framework.Audio;
 using osu.Game.IO.Serialization;
 using osu.Game.Rulesets;
@@ -38,19 +39,6 @@ protected WorkingBeatmap(BeatmapInfo beatmapInfo, AudioManager audioManager)
             BeatmapSetInfo = beatmapInfo.BeatmapSet;
             Metadata = beatmapInfo.Metadata ?? BeatmapSetInfo?.Metadata ?? new BeatmapMetadata();
 
-            beatmap = new RecyclableLazy<IBeatmap>(() =>
-            {
-                var b = GetBeatmap() ?? new Beatmap();
-
-                // The original beatmap version needs to be preserved as the database doesn't contain it
-                BeatmapInfo.BeatmapVersion = b.BeatmapInfo.BeatmapVersion;
-
-                // Use the database-backed info for more up-to-date values (beatmap id, ranked status, etc)
-                b.BeatmapInfo = BeatmapInfo;
-
-                return b;
-            });
-
             track = new RecyclableLazy<Track>(() => GetTrack() ?? GetVirtualTrack());
             background = new RecyclableLazy<Texture>(GetBackground, BackgroundStillValid);
             waveform = new RecyclableLazy<Waveform>(GetWaveform);
@@ -58,6 +46,11 @@ protected WorkingBeatmap(BeatmapInfo beatmapInfo, AudioManager audioManager)
             skin = new RecyclableLazy<Skin>(GetSkin);
         }
 
+        ~WorkingBeatmap()
+        {
+            Dispose(false);
+        }
+
         protected virtual Track GetVirtualTrack()
         {
             const double excess_length = 1000;
@@ -153,10 +146,26 @@ public IBeatmap GetPlayableBeatmap(RulesetInfo ruleset, IReadOnlyList<Mod> mods)
 
         public override string ToString() => BeatmapInfo.ToString();
 
-        public bool BeatmapLoaded => beatmap.IsResultAvailable;
-        public IBeatmap Beatmap => beatmap.Value;
+        public bool BeatmapLoaded => beatmapLoadTask?.IsCompleted ?? false;
+
+        public Task<IBeatmap> LoadBeatmapAsync() => (beatmapLoadTask ?? (beatmapLoadTask = Task.Factory.StartNew(() =>
+        {
+            var b = GetBeatmap() ?? new Beatmap();
+
+            // The original beatmap version needs to be preserved as the database doesn't contain it
+            BeatmapInfo.BeatmapVersion = b.BeatmapInfo.BeatmapVersion;
+
+            // Use the database-backed info for more up-to-date values (beatmap id, ranked status, etc)
+            b.BeatmapInfo = BeatmapInfo;
+
+            return b;
+        }, beatmapCancellation.Token, TaskCreationOptions.LongRunning, TaskScheduler.Default)));
+
+        public IBeatmap Beatmap => LoadBeatmapAsync().Result;
+
+        private readonly CancellationTokenSource beatmapCancellation = new CancellationTokenSource();
         protected abstract IBeatmap GetBeatmap();
-        private readonly RecyclableLazy<IBeatmap> beatmap;
+        private Task<IBeatmap> beatmapLoadTask;
 
         public bool BackgroundLoaded => background.IsResultAvailable;
         public Texture Background => background.Value;
@@ -195,20 +204,33 @@ public virtual void TransferTo(WorkingBeatmap other)
                 other.track = track;
         }
 
-        public virtual void Dispose()
-        {
-            background.Recycle();
-            waveform.Recycle();
-            storyboard.Recycle();
-            skin.Recycle();
-        }
-
         /// <summary>
         /// Eagerly dispose of the audio track associated with this <see cref="WorkingBeatmap"/> (if any).
         /// Accessing track again will load a fresh instance.
         /// </summary>
         public virtual void RecycleTrack() => track.Recycle();
 
+        #region Disposal
+
+        public void Dispose()
+        {
+            Dispose(true);
+            GC.SuppressFinalize(this);
+        }
+
+        protected virtual void Dispose(bool isDisposing)
+        {
+            // recycling logic is not here for the time being, as components which use
+            // retrieved objects from WorkingBeatmap may not hold a reference to the WorkingBeatmap itself.
+            // this should be fine as each retrieved comopnent do have their own finalizers.
+
+            // cancelling the beatmap load is safe for now since the retrieval is a synchronous
+            // operation. if we add an async retrieval method this may need to be reconsidered.
+            beatmapCancellation.Cancel();
+        }
+
+        #endregion
+
         public class RecyclableLazy<T>
         {
             private Lazy<T> lazy;

osu.Game/OsuGame.cs

@@ -295,6 +295,8 @@ private void beatmapChanged(ValueChangedEvent<WorkingBeatmap> beatmap)
             var nextBeatmap = beatmap.NewValue;
             if (nextBeatmap?.Track != null)
                 nextBeatmap.Track.Completed += currentTrackCompleted;
+
+            nextBeatmap?.LoadBeatmapAsync();
         }
 
         private void currentTrackCompleted()

osu.Game/Tests/Visual/OsuTestScene.cs

@@ -137,9 +137,9 @@ public ClockBackedTestWorkingBeatmap(IBeatmap beatmap, IFrameBasedClock referenc
                     track = audio?.Tracks.GetVirtual(length);
             }
 
-            public override void Dispose()
+            protected override void Dispose(bool isDisposing)
             {
-                base.Dispose();
+                base.Dispose(isDisposing);
                 store?.Dispose();
             }