Fix LargeTextureStore race conditions

osu.Framework/Graphics/Textures/LargeTextureStore.cs

···
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Threading;
+using JetBrains.Annotations;
 using osu.Framework.IO.Stores;
 using osuTK.Graphics.ES30;
-using osu.Framework.Graphics.OpenGL.Textures;
 
 namespace osu.Framework.Graphics.Textures
 {
···
         {
         }
 
-        /// <summary>
-        /// Retrieves a texture.
-        /// This texture should only be assigned once, as reference counting is being used internally.
-        /// If you wish to use the same texture multiple times, call this method an equal number of times.
-        /// </summary>
-        /// <param name="name">The name of the texture.</param>
-        /// <param name="wrapModeS">The horizontal wrap mode of the texture.</param>
-        /// <param name="wrapModeT">The vertical wrap mode of the texture.</param>
-        /// <returns>The texture.</returns>
-        public override Texture Get(string name, WrapMode wrapModeS, WrapMode wrapModeT)
+        protected override bool TryGetCached(string lookupKey, out Texture texture)
         {
-            var tex = base.Get(name, wrapModeS, wrapModeT);
+            lock (referenceCountLock)
+            {
+                if (base.TryGetCached(lookupKey, out var tex))
+                {
+                    texture = createTextureWithRefCount(lookupKey, tex);
+                    return true;
+                }
 
-            if (tex?.TextureGL == null)
+                texture = null;
+                return false;
+            }
+        }
+
+        protected override Texture CacheAndReturnTexture(string lookupKey, Texture texture)
+        {
+            lock (referenceCountLock)
+                return createTextureWithRefCount(lookupKey, base.CacheAndReturnTexture(lookupKey, texture));
+        }
+
+        private TextureWithRefCount createTextureWithRefCount([NotNull] string lookupKey, [CanBeNull] Texture baseTexture)
+        {
+            if (baseTexture == null)
                 return null;
 
             lock (referenceCountLock)
             {
-                if (!referenceCounts.TryGetValue(tex.LookupKey, out TextureWithRefCount.ReferenceCount count))
-                    referenceCounts[tex.LookupKey] = count = new TextureWithRefCount.ReferenceCount(referenceCountLock, () => onAllReferencesLost(tex));
+                if (!referenceCounts.TryGetValue(lookupKey, out TextureWithRefCount.ReferenceCount count))
+                    referenceCounts[lookupKey] = count = new TextureWithRefCount.ReferenceCount(referenceCountLock, () => onAllReferencesLost(baseTexture));
 
-                return new TextureWithRefCount(tex.TextureGL, count);
+                return new TextureWithRefCount(baseTexture.TextureGL, count);
             }
         }
 

osu.Framework/Graphics/Textures/TextureStore.cs

···
 using osu.Framework.IO.Stores;
 using System.Collections.Generic;
 using System.Threading.Tasks;
+using JetBrains.Annotations;
 using osu.Framework.Logging;
 using osuTK.Graphics.ES30;
 
···
             if (Atlas != null)
             {
                 if ((glTexture = Atlas.Add(upload.Width, upload.Height, wrapModeS, wrapModeT)) == null)
-                    Logger.Log($"Texture requested ({upload.Width}x{upload.Height}) which exceeds {nameof(TextureStore)}'s atlas size ({max_atlas_size}x{max_atlas_size}) - bypassing atlasing. Consider using {nameof(LargeTextureStore)}.", LoggingTarget.Performance);
+                {
+                    Logger.Log(
+                        $"Texture requested ({upload.Width}x{upload.Height}) which exceeds {nameof(TextureStore)}'s atlas size ({max_atlas_size}x{max_atlas_size}) - bypassing atlasing. Consider using {nameof(LargeTextureStore)}.",
+                        LoggingTarget.Performance);
+                }
             }
 
             glTexture ??= new TextureGLSingle(upload.Width, upload.Height, manualMipmaps, filteringMode, wrapModeS, wrapModeT);
···
         /// <param name="wrapModeS">The texture wrap mode in horizontal direction.</param>
         /// <param name="wrapModeT">The texture wrap mode in vertical direction.</param>
         /// <returns>The texture.</returns>
-        public Task<Texture> GetAsync(string name, WrapMode wrapModeT, WrapMode wrapModeS) => Task.Run(() => Get(name, wrapModeS, wrapModeT)); // TODO: best effort. need to re-think textureCache data structure to fix this.
+        public Task<Texture> GetAsync(string name, WrapMode wrapModeT, WrapMode wrapModeS) =>
+            Task.Run(() => Get(name, wrapModeS, wrapModeT)); // TODO: best effort. need to re-think textureCache data structure to fix this.
 
         /// <summary>
         /// Retrieves a texture from the store and adds it to the atlas.
···
             this.LogIfNonBackgroundThread(key);
 
             // Check if the texture exists in the cache.
-            lock (textureCache)
-            {
-                if (textureCache.TryGetValue(key, out var tex))
-                    return tex;
-            }
+            if (TryGetCached(key, out var cached))
+                return cached;
 
             // Take an exclusive lock on retrieval of the texture.
             lock (retrievalLock)
             {
                 // If another retrieval of the texture happened before us, we should check if the texture exists in the cache again.
-                lock (textureCache)
-                {
-                    if (textureCache.TryGetValue(key, out var tex))
-                        return tex;
-                }
+                if (TryGetCached(key, out cached))
+                    return cached;
 
                 try
                 {
···
                     if (tex != null)
                         tex.LookupKey = key;
 
-                    lock (textureCache)
-                        textureCache[key] = tex;
-
-                    return tex;
+                    return CacheAndReturnTexture(key, tex);
                 }
                 catch (TextureTooLargeForGLException)
                 {
···
             }
 
             return null;
+        }
+
+        /// <summary>
+        /// Attempts to retrieve an existing cached texture.
+        /// </summary>
+        /// <param name="lookupKey">The lookup key that uniquely identifies textures in the cache.</param>
+        /// <param name="texture">The returned texture. Null if the texture did not exist in the cache.</param>
+        /// <returns>Whether a cached texture was retrieved.</returns>
+        protected virtual bool TryGetCached([NotNull] string lookupKey, [CanBeNull] out Texture texture)
+        {
+            lock (textureCache)
+                return textureCache.TryGetValue(lookupKey, out texture);
+        }
+
+        /// <summary>
+        /// Caches and returns the given texture.
+        /// </summary>
+        /// <param name="lookupKey">The lookup key that uniquely identifies textures in the cache.</param>
+        /// <param name="texture">The texture to be cached and returned.</param>
+        /// <returns>The texture to be returned.</returns>
+        [CanBeNull]
+        protected virtual Texture CacheAndReturnTexture([NotNull] string lookupKey, [CanBeNull] Texture texture)
+        {
+            lock (textureCache)
+                return textureCache[lookupKey] = texture;
         }
 
         /// <summary>