CachedResources for GasTileDangerousTemperatureOverlay (#43032)

Content.Client/Atmos/Overlays/GasTileDangerousTemperatureOverlay.cs

@@ -1,4 +1,5 @@
 using Content.Client.Atmos.EntitySystems;
+using Content.Client.Graphics;
 using Content.Shared.Atmos;
 using Content.Shared.Atmos.Components;
 using Content.Shared.Atmos.EntitySystems;
@@ -25,7 +26,8 @@ public sealed class GasTileDangerousTemperatureOverlay : Overlay
     private readonly SharedTransformSystem _xformSys;
     private EntityQuery<GasTileOverlayComponent> _overlayQuery;
 
-    private IRenderTexture? _temperatureTarget;
+    private readonly OverlayResourceCache<CachedResources> _resources = new();
+    private List<Entity<MapGridComponent>> _grids = new();
 
     // Cache used to transform ThermalByte into Color for overlay
     private readonly Color[] _colorCache = new Color[256];
@@ -152,10 +154,11 @@ protected override bool BeforeDraw(in OverlayDrawArgs args)
 
         var target = args.Viewport.RenderTarget;
 
-        if (_temperatureTarget?.Texture.Size != target.Size)
+        var res = _resources.GetForViewport(args.Viewport, static _ => new CachedResources());
+        if (res.TemperatureTarget is null || res.TemperatureTarget.Texture.Size != target.Size)
         {
-            _temperatureTarget?.Dispose();
-            _temperatureTarget = _clyde.CreateRenderTarget(
+            res.TemperatureTarget?.Dispose();
+            res.TemperatureTarget = _clyde.CreateRenderTarget(
                 target.Size,
                 new RenderTargetFormatParameters(RenderTargetColorFormat.Rgba8Srgb),
                 name: nameof(GasTileDangerousTemperatureOverlay));
@@ -167,16 +170,13 @@ protected override bool BeforeDraw(in OverlayDrawArgs args)
         var mapId = args.MapId;
         var worldToViewportLocal = args.Viewport.GetWorldToLocalMatrix();
 
-        var anyGasDrawn = false;
-        List<Entity<MapGridComponent>> grids = new();
-
-        drawHandle.RenderInRenderTarget(_temperatureTarget,
+        drawHandle.RenderInRenderTarget(res.TemperatureTarget,
             () =>
             {
-                grids.Clear();
-                _mapManager.FindGridsIntersecting(mapId, worldAABB, ref grids);
+                _grids.Clear();
+                _mapManager.FindGridsIntersecting(mapId, worldAABB, ref _grids);
 
-                foreach (var grid in grids)
+                foreach (var grid in _grids)
                 {
                     if (!_overlayQuery.TryGetComponent(grid.Owner, out var comp))
                         continue;
@@ -211,8 +211,6 @@ protected override bool BeforeDraw(in OverlayDrawArgs args)
                             if (gasColor.A <= 0f)
                                 continue;
 
-                            anyGasDrawn = true;
-
                             drawHandle.DrawRect(
                                 Box2.CenteredAround(tilePosition + gridTileCenterVec, gridTileSizeVec),
                                 gasColor
@@ -225,29 +223,31 @@ protected override bool BeforeDraw(in OverlayDrawArgs args)
 
         drawHandle.SetTransform(Matrix3x2.Identity);
 
-        if (!anyGasDrawn)
-        {
-            _temperatureTarget?.Dispose();
-            _temperatureTarget = null;
-            return false;
-        }
-
         return true;
     }
 
     protected override void Draw(in OverlayDrawArgs args)
     {
-        if (_temperatureTarget is null)
-            return;
+        var res = _resources.GetForViewport(args.Viewport, static _ => new CachedResources());
 
-        args.WorldHandle.DrawTextureRect(_temperatureTarget.Texture, args.WorldBounds);
+        if (res.TemperatureTarget != null)
+            args.WorldHandle.DrawTextureRect(res.TemperatureTarget.Texture, args.WorldBounds);
         args.WorldHandle.SetTransform(Matrix3x2.Identity);
     }
 
     protected override void DisposeBehavior()
     {
-        _temperatureTarget?.Dispose();
-        _temperatureTarget = null;
+        _resources.Dispose();
         base.DisposeBehavior();
     }
+
+    private sealed class CachedResources : IDisposable
+    {
+        public IRenderTexture? TemperatureTarget;
+
+        public void Dispose()
+        {
+            TemperatureTarget?.Dispose();
+        }
+    }
 }

Content.Client/Graphics/OverlayResourceCache.cs

@@ -0,0 +1,90 @@
+using Robust.Client.Graphics;
+
+namespace Content.Client.Graphics;
+
+/// <summary>
+/// A cache for <see cref="Overlay"/>s to store per-viewport render resources, such as render targets.
+/// </summary>
+/// <typeparam name="T">The type of data stored in the cache.</typeparam>
+public sealed class OverlayResourceCache<T> : IDisposable where T : class, IDisposable
+{
+    private readonly Dictionary<long, CacheEntry> _cache = new();
+
+    /// <summary>
+    /// Get the data for a specific viewport, creating a new entry if necessary.
+    /// </summary>
+    /// <remarks>
+    /// The cached data may be cleared at any time if <see cref="IClydeViewport.ClearCachedResources"/> gets invoked.
+    /// </remarks>
+    /// <param name="viewport">The viewport for which to retrieve cached data.</param>
+    /// <param name="factory">A delegate used to create the cached data, if necessary.</param>
+    public T GetForViewport(IClydeViewport viewport, Func<IClydeViewport, T> factory)
+    {
+        return GetForViewport(viewport, out _, factory);
+    }
+
+    /// <summary>
+    /// Get the data for a specific viewport, creating a new entry if necessary.
+    /// </summary>
+    /// <remarks>
+    /// The cached data may be cleared at any time if <see cref="IClydeViewport.ClearCachedResources"/> gets invoked.
+    /// </remarks>
+    /// <param name="viewport">The viewport for which to retrieve cached data.</param>
+    /// <param name="wasCached">True if the data was pulled from cache, false if it was created anew.</param>
+    /// <param name="factory">A delegate used to create the cached data, if necessary.</param>
+    public T GetForViewport(IClydeViewport viewport, out bool wasCached, Func<IClydeViewport, T> factory)
+    {
+        if (_cache.TryGetValue(viewport.Id, out var entry))
+        {
+            wasCached = true;
+            return entry.Data;
+        }
+
+        wasCached = false;
+
+        entry = new CacheEntry
+        {
+            Data = factory(viewport),
+            Viewport = new WeakReference<IClydeViewport>(viewport),
+        };
+        _cache.Add(viewport.Id, entry);
+
+        viewport.ClearCachedResources += ViewportOnClearCachedResources;
+
+        return entry.Data;
+    }
+
+    private void ViewportOnClearCachedResources(ClearCachedViewportResourcesEvent ev)
+    {
+        if (!_cache.Remove(ev.ViewportId, out var entry))
+        {
+            // I think this could theoretically happen if you manually dispose the cache *after* a leaked viewport got
+            // GC'd, but before its ClearCachedResources got invoked.
+            return;
+        }
+
+        entry.Data.Dispose();
+
+        if (ev.Viewport != null)
+            ev.Viewport.ClearCachedResources -= ViewportOnClearCachedResources;
+    }
+
+    public void Dispose()
+    {
+        foreach (var entry in _cache)
+        {
+            if (entry.Value.Viewport.TryGetTarget(out var viewport))
+                viewport.ClearCachedResources -= ViewportOnClearCachedResources;
+
+            entry.Value.Data.Dispose();
+        }
+
+        _cache.Clear();
+    }
+
+    private struct CacheEntry
+    {
+        public T Data;
+        public WeakReference<IClydeViewport> Viewport;
+    }
+}

Content.Client/Light/AfterLightTargetOverlay.cs

@@ -30,6 +30,7 @@ protected override void Draw(in OverlayDrawArgs args)
             return;
 
         var lightOverlay = _overlay.GetOverlay<BeforeLightTargetOverlay>();
+        var lightRes = lightOverlay.GetCachedForViewport(args.Viewport);
         var bounds = args.WorldBounds;
 
         // at 1-1 render scale it's mostly fine but at 4x4 it's way too fkn big
@@ -38,7 +39,7 @@ protected override void Draw(in OverlayDrawArgs args)
 
         var localMatrix =
             viewport.LightRenderTarget.GetWorldToLocalMatrix(viewport.Eye, newScale);
-        var diff = (lightOverlay.EnlargedLightTarget.Size - viewport.LightRenderTarget.Size);
+        var diff = (lightRes.EnlargedLightTarget.Size - viewport.LightRenderTarget.Size);
         var halfDiff = diff / 2;
 
         // Pixels -> Metres -> Half distance.
@@ -53,7 +54,7 @@ protected override void Draw(in OverlayDrawArgs args)
                     viewport.LightRenderTarget.Size.Y + halfDiff.Y);
 
                 worldHandle.SetTransform(localMatrix);
-                worldHandle.DrawTextureRectRegion(lightOverlay.EnlargedLightTarget.Texture, bounds, subRegion: subRegion);
+                worldHandle.DrawTextureRectRegion(lightRes.EnlargedLightTarget.Texture, bounds, subRegion: subRegion);
             }, Color.Transparent);
     }
 }

Content.Client/Light/AmbientOcclusionOverlay.cs

@@ -1,4 +1,5 @@
 using System.Numerics;
+using Content.Client.Graphics;
 using Content.Shared.CCVar;
 using Content.Shared.Maps;
 using Robust.Client.Graphics;
@@ -27,11 +28,7 @@ public sealed class AmbientOcclusionOverlay : Overlay
 
     public override OverlaySpace Space => OverlaySpace.WorldSpaceBelowEntities;
 
-    private IRenderTexture? _aoTarget;
-    private IRenderTexture? _aoBlurBuffer;
-
-    // Couldn't figure out a way to avoid this so if you can then please do.
-    private IRenderTexture? _aoStencilTarget;
+    private readonly OverlayResourceCache<CachedResources> _resources = new ();
 
     public AmbientOcclusionOverlay()
     {
@@ -69,30 +66,32 @@ protected override void Draw(in OverlayDrawArgs args)
         var turfSystem = _entManager.System<TurfSystem>();
         var invMatrix = args.Viewport.GetWorldToLocalMatrix();
 
-        if (_aoTarget?.Texture.Size != target.Size)
+        var res = _resources.GetForViewport(args.Viewport, static _ => new CachedResources());
+
+        if (res.AOTarget?.Texture.Size != target.Size)
         {
-            _aoTarget?.Dispose();
-            _aoTarget = _clyde.CreateRenderTarget(target.Size, new RenderTargetFormatParameters(RenderTargetColorFormat.Rgba8Srgb), name: "ambient-occlusion-target");
+            res.AOTarget?.Dispose();
+            res.AOTarget = _clyde.CreateRenderTarget(target.Size, new RenderTargetFormatParameters(RenderTargetColorFormat.Rgba8Srgb), name: "ambient-occlusion-target");
         }
 
-        if (_aoBlurBuffer?.Texture.Size != target.Size)
+        if (res.AOBlurBuffer?.Texture.Size != target.Size)
         {
-            _aoBlurBuffer?.Dispose();
-            _aoBlurBuffer = _clyde.CreateRenderTarget(target.Size, new RenderTargetFormatParameters(RenderTargetColorFormat.Rgba8Srgb), name: "ambient-occlusion-blur-target");
+            res.AOBlurBuffer?.Dispose();
+            res.AOBlurBuffer = _clyde.CreateRenderTarget(target.Size, new RenderTargetFormatParameters(RenderTargetColorFormat.Rgba8Srgb), name: "ambient-occlusion-blur-target");
         }
 
-        if (_aoStencilTarget?.Texture.Size != target.Size)
+        if (res.AOStencilTarget?.Texture.Size != target.Size)
         {
-            _aoStencilTarget?.Dispose();
-            _aoStencilTarget = _clyde.CreateRenderTarget(target.Size, new RenderTargetFormatParameters(RenderTargetColorFormat.Rgba8Srgb), name: "ambient-occlusion-stencil-target");
+            res.AOStencilTarget?.Dispose();
+            res.AOStencilTarget = _clyde.CreateRenderTarget(target.Size, new RenderTargetFormatParameters(RenderTargetColorFormat.Rgba8Srgb), name: "ambient-occlusion-stencil-target");
         }
 
         // Draw the texture data to the texture.
-        args.WorldHandle.RenderInRenderTarget(_aoTarget,
+        args.WorldHandle.RenderInRenderTarget(res.AOTarget,
             () =>
             {
                 worldHandle.UseShader(_proto.Index(UnshadedShader).Instance());
-                var invMatrix = _aoTarget.GetWorldToLocalMatrix(viewport.Eye!, scale);
+                var invMatrix = res.AOTarget.GetWorldToLocalMatrix(viewport.Eye!, scale);
 
                 foreach (var entry in query.QueryAabb(mapId, worldBounds))
                 {
@@ -106,11 +105,11 @@ protected override void Draw(in OverlayDrawArgs args)
                 }
             }, Color.Transparent);
 
-        _clyde.BlurRenderTarget(viewport, _aoTarget, _aoBlurBuffer, viewport.Eye!, 14f);
+        _clyde.BlurRenderTarget(viewport, res.AOTarget, res.AOBlurBuffer, viewport.Eye!, 14f);
 
         // Need to do stencilling after blur as it will nuke it.
         // Draw stencil for the grid so we don't draw in space.
-        args.WorldHandle.RenderInRenderTarget(_aoStencilTarget,
+        args.WorldHandle.RenderInRenderTarget(res.AOStencilTarget,
             () =>
             {
                 // Don't want lighting affecting it.
@@ -136,13 +135,36 @@ protected override void Draw(in OverlayDrawArgs args)
 
         // Draw the stencil texture to depth buffer.
         worldHandle.UseShader(_proto.Index(StencilMaskShader).Instance());
-        worldHandle.DrawTextureRect(_aoStencilTarget!.Texture, worldBounds);
+        worldHandle.DrawTextureRect(res.AOStencilTarget!.Texture, worldBounds);
 
         // Draw the Blurred AO texture finally.
         worldHandle.UseShader(_proto.Index(StencilEqualDrawShader).Instance());
-        worldHandle.DrawTextureRect(_aoTarget!.Texture, worldBounds, color);
+        worldHandle.DrawTextureRect(res.AOTarget!.Texture, worldBounds, color);
 
         args.WorldHandle.SetTransform(Matrix3x2.Identity);
         args.WorldHandle.UseShader(null);
     }
+
+    protected override void DisposeBehavior()
+    {
+        _resources.Dispose();
+
+        base.DisposeBehavior();
+    }
+
+    private sealed class CachedResources : IDisposable
+    {
+        public IRenderTexture? AOTarget;
+        public IRenderTexture? AOBlurBuffer;
+
+        // Couldn't figure out a way to avoid this so if you can then please do.
+        public IRenderTexture? AOStencilTarget;
+
+        public void Dispose()
+        {
+            AOTarget?.Dispose();
+            AOBlurBuffer?.Dispose();
+            AOStencilTarget?.Dispose();
+        }
+    }
 }

Content.Client/Light/BeforeLightTargetOverlay.cs

@@ -1,4 +1,4 @@
-using System.Numerics;
+using Content.Client.Graphics;
 using Robust.Client.Graphics;
 using Robust.Shared.Enums;
 
@@ -13,7 +13,8 @@ public sealed class BeforeLightTargetOverlay : Overlay
 
     [Dependency] private readonly IClyde _clyde = default!;
 
-    public IRenderTexture EnlargedLightTarget = default!;
+    private readonly OverlayResourceCache<CachedResources> _resources = new();
+
     public Box2Rotated EnlargedBounds;
 
     /// <summary>
@@ -36,16 +37,42 @@ protected override void Draw(in OverlayDrawArgs args)
         var size = args.Viewport.LightRenderTarget.Size + (int) (_skirting * EyeManager.PixelsPerMeter);
         EnlargedBounds = args.WorldBounds.Enlarged(_skirting / 2f);
 
+        var res = _resources.GetForViewport(args.Viewport, static _ => new CachedResources());
+
         // This just exists to copy the lightrendertarget and write back to it.
-        if (EnlargedLightTarget?.Size != size)
+        if (res.EnlargedLightTarget?.Size != size)
         {
-            EnlargedLightTarget = _clyde
+            res.EnlargedLightTarget = _clyde
                 .CreateRenderTarget(size, new RenderTargetFormatParameters(RenderTargetColorFormat.Rgba8Srgb), name: "enlarged-light-copy");
         }
 
-        args.WorldHandle.RenderInRenderTarget(EnlargedLightTarget,
+        args.WorldHandle.RenderInRenderTarget(res.EnlargedLightTarget,
             () =>
             {
             }, _clyde.GetClearColor(args.MapUid));
     }
+
+    internal CachedResources GetCachedForViewport(IClydeViewport viewport)
+    {
+        return _resources.GetForViewport(viewport,
+            static _ => throw new InvalidOperationException(
+                "Expected BeforeLightTargetOverlay to have created its resources"));
+    }
+
+    protected override void DisposeBehavior()
+    {
+        _resources.Dispose();
+
+        base.DisposeBehavior();
+    }
+
+    internal sealed class CachedResources : IDisposable
+    {
+        public IRenderTexture EnlargedLightTarget = default!;
+
+        public void Dispose()
+        {
+            EnlargedLightTarget?.Dispose();
+        }
+    }
 }