From cea5ea40968741a1db4069dcd0f4f753fc47e356 Mon Sep 17 00:00:00 2001
From: Omar Emara <mail@OmarEmara.dev>
Date: Mon, 16 Oct 2023 19:31:06 +0300
Subject: [PATCH] Fix #113768: GPU compositor is leaking memory

The experimental GPU compositor is leaking memory in any setup.

This is because the current implementation of the render texture pool
always created a new texture and only freed the textures upon deletion.
This was a temporary implementation until a proper implementation that
uses the DRW textures pool was used.

This patch implements a small texture pool as a temporary fix until the
aforementioned DRW texture pool implementation is done.
---
 source/blender/render/intern/compositor.cc | 91 ++++++++++++++++------
 1 file changed, 69 insertions(+), 22 deletions(-)

diff --git a/source/blender/render/intern/compositor.cc b/source/blender/render/intern/compositor.cc
index 51e7b0d36f7..256df63f996 100644
--- a/source/blender/render/intern/compositor.cc
+++ b/source/blender/render/intern/compositor.cc
@@ -32,32 +32,77 @@
 
 namespace blender::render {
 
-/* Render Texture Pool */
-
+/**
+ * Render Texture Pool
+ *
+ * TODO: should share pool with draw manager. It needs some globals initialization figured out
+ * there first.
+ */
 class TexturePool : public realtime_compositor::TexturePool {
- public:
-  Vector<GPUTexture *> textures_;
+ private:
+  /** Textures that are not yet used and are available to be acquired. After evaluation, any
+   * texture in this map should be freed because it was not acquired in the evaluation and is thus
+   * unused. Textures removed from this map should be moved to the textures_in_use_ map when
+   * acquired. */
+  Map<realtime_compositor::TexturePoolKey, Vector<GPUTexture *>> available_textures_;
+  /** Textures that were acquired in this compositor evaluation. After evaluation, those textures
+   * are moved to the available_textures_ map to be acquired in the next evaluation. */
+  Map<realtime_compositor::TexturePoolKey, Vector<GPUTexture *>> textures_in_use_;
 
+ public:
   virtual ~TexturePool()
   {
-    for (GPUTexture *texture : textures_) {
-      GPU_texture_free(texture);
+    for (Vector<GPUTexture *> &available_textures : available_textures_.values()) {
+      for (GPUTexture *texture : available_textures) {
+        GPU_texture_free(texture);
+      }
+    }
+
+    for (Vector<GPUTexture *> &textures_in_use : textures_in_use_.values()) {
+      for (GPUTexture *texture : textures_in_use) {
+        GPU_texture_free(texture);
+      }
     }
   }
 
   GPUTexture *allocate_texture(int2 size, eGPUTextureFormat format) override
   {
-    /* TODO: should share pool with draw manager. It needs some globals
-     * initialization figured out there first. */
-#if 0
-    DrawEngineType *owner = (DrawEngineType *)this;
-    return DRW_texture_pool_query_2d(size.x, size.y, format, owner);
-#else
-    GPUTexture *texture = GPU_texture_create_2d(
-        "compositor_texture_pool", size.x, size.y, 1, format, GPU_TEXTURE_USAGE_GENERAL, nullptr);
-    textures_.append(texture);
+    const realtime_compositor::TexturePoolKey key(size, format);
+    Vector<GPUTexture *> &available_textures = available_textures_.lookup_or_add_default(key);
+    GPUTexture *texture = nullptr;
+    if (available_textures.is_empty()) {
+      texture = GPU_texture_create_2d("compositor_texture_pool",
+                                      size.x,
+                                      size.y,
+                                      1,
+                                      format,
+                                      GPU_TEXTURE_USAGE_GENERAL,
+                                      nullptr);
+    }
+    else {
+      texture = available_textures.pop_last();
+    }
+
+    textures_in_use_.lookup_or_add_default(key).append(texture);
     return texture;
-#endif
+  }
+
+  /** Should be called after compositor evaluation to free unused textures and reset the texture
+   * pool. */
+  void free_unused_and_reset()
+  {
+    /* Free all textures in the available textures vectors. The fact that they still exist in those
+     * vectors after evaluation means they were not acquired during the evaluation, and are thus
+     * consequently no longer used. */
+    for (Vector<GPUTexture *> &available_textures : available_textures_.values()) {
+      for (GPUTexture *texture : available_textures) {
+        GPU_texture_free(texture);
+      }
+    }
+
+    /* Move all textures in-use to be available textures for the next evaluation. */
+    available_textures_ = textures_in_use_;
+    textures_in_use_.clear();
   }
 };
 
@@ -102,14 +147,12 @@ class Context : public realtime_compositor::Context {
   /* Viewer output texture. */
   GPUTexture *viewer_output_texture_ = nullptr;
 
-  /* Texture pool. */
-  TexturePool &render_texture_pool_;
+  /* Cached textures that the compositor took ownership of. */
+  Vector<GPUTexture *> textures_;
 
  public:
   Context(const ContextInputData &input_data, TexturePool &texture_pool)
-      : realtime_compositor::Context(texture_pool),
-        input_data_(input_data),
-        render_texture_pool_(texture_pool)
+      : realtime_compositor::Context(texture_pool), input_data_(input_data)
   {
   }
 
@@ -117,6 +160,9 @@ class Context : public realtime_compositor::Context {
   {
     GPU_TEXTURE_FREE_SAFE(output_texture_);
     GPU_TEXTURE_FREE_SAFE(viewer_output_texture_);
+    for (GPUTexture *texture : textures_) {
+      GPU_texture_free(texture);
+    }
   }
 
   void update_input_data(const ContextInputData &input_data)
@@ -235,7 +281,7 @@ class Context : public realtime_compositor::Context {
             if (input_texture) {
               /* Don't assume render keeps texture around, add our own reference. */
               GPU_texture_ref(input_texture);
-              render_texture_pool_.textures_.append(input_texture);
+              textures_.append(input_texture);
             }
           }
         }
@@ -421,6 +467,7 @@ class RealtimeCompositor {
     context_->output_to_render_result();
     context_->viewer_output_to_viewer_image();
     DRW_render_context_disable(&render_);
+    texture_pool_->free_unused_and_reset();
   }
 };