From 5ba7a841e4c333934cbd4a095555094baad8b6a0 Mon Sep 17 00:00:00 2001
From: Steve Streeting <steve@stevestreeting.com>
Date: Tue, 9 Nov 2021 16:05:12 +0000
Subject: [PATCH] Fixed some GC problems, we need to implement FGCObject

---
 .../Private/StevesTextureRenderTargetPool.cpp | 22 ++++++++++++++++++-
 .../Public/StevesTextureRenderTargetPool.h    | 12 +++++++---
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/Source/StevesUEHelpers/Private/StevesTextureRenderTargetPool.cpp b/Source/StevesUEHelpers/Private/StevesTextureRenderTargetPool.cpp
index 4b9d556..d79d4c9 100644
--- a/Source/StevesUEHelpers/Private/StevesTextureRenderTargetPool.cpp
+++ b/Source/StevesUEHelpers/Private/StevesTextureRenderTargetPool.cpp
@@ -29,6 +29,7 @@ void FStevesTextureRenderTargetPool::ReleaseReservation(UTextureRenderTarget2D*
 			UE_LOG(LogStevesUEHelpers, Verbose, TEXT("FStevesTextureRenderTargetPool: Released texture reservation on %s"), *Tex->GetName());
 			UnreservedTextures.Add(R.Key, Tex);
 			Reservations.RemoveAtSwap(i);
+			ReservedTextures.Remove(Tex);
 			return;
 		}
 	}
@@ -37,8 +38,20 @@ void FStevesTextureRenderTargetPool::ReleaseReservation(UTextureRenderTarget2D*
 
 }
 
+FStevesTextureRenderTargetPool::~FStevesTextureRenderTargetPool()
+{
+	DrainPool(true);
+}
+
+void FStevesTextureRenderTargetPool::AddReferencedObjects(FReferenceCollector& Collector)
+{
+	// We need to hold on to the texture references
+	Collector.AddReferencedObjects(ReservedTextures);
+	Collector.AddReferencedObjects(UnreservedTextures);
+}
+
 FStevesTextureRenderTargetReservationPtr FStevesTextureRenderTargetPool::ReserveTexture(FIntPoint Size,
-	ETextureRenderTargetFormat Format, const UObject* Owner)
+                                                                                        ETextureRenderTargetFormat Format, const UObject* Owner)
 {
 	const FTextureKey Key {Size, Format};
 	UTextureRenderTarget2D* Tex = nullptr;
@@ -63,6 +76,10 @@ FStevesTextureRenderTargetReservationPtr FStevesTextureRenderTargetPool::Reserve
 
 	// Record reservation
 	Reservations.Add(FReservationInfo(Key, Owner, Tex));
+
+	// Reservation doesn't keep the texture alive; if caller doesn't hold a strong pointer to it, it'll be destroyed
+	// So we need to hold it ourselves
+	ReservedTextures.Add(Tex);
 	
 	return MakeShared<FStevesTextureRenderTargetReservation>(Tex, this->AsShared(), Owner);
 }
@@ -78,6 +95,7 @@ void FStevesTextureRenderTargetPool::RevokeReservations(const UObject* ForOwner)
 			{
 				UE_LOG(LogStevesUEHelpers, Verbose, TEXT("FStevesTextureRenderTargetPool: Revoked texture reservation on %s"), *R.Texture->GetName());
 				UnreservedTextures.Add(R.Key, R.Texture.Get());
+				ReservedTextures.Remove(R.Texture.Get());
 			}
 			// Can't use RemoveAtSwap because it'll change order
 			Reservations.RemoveAt(i);
@@ -97,4 +115,6 @@ void FStevesTextureRenderTargetPool::DrainPool(bool bForceAndRevokeReservations)
 		UKismetRenderingLibrary::ReleaseRenderTarget2D(TexPair.Value);
 	}
 	UnreservedTextures.Empty();
+	ReservedTextures.Empty();
+
 }
diff --git a/Source/StevesUEHelpers/Public/StevesTextureRenderTargetPool.h b/Source/StevesUEHelpers/Public/StevesTextureRenderTargetPool.h
index 9a62d08..0deeb7d 100644
--- a/Source/StevesUEHelpers/Public/StevesTextureRenderTargetPool.h
+++ b/Source/StevesUEHelpers/Public/StevesTextureRenderTargetPool.h
@@ -40,8 +40,9 @@ public:
  * these textures at runtime.
  * A pool needs to be owned by a UObject, which will in turn own the textures and so will ultimately control the
  * ultimate lifecycle of textures if not released specifically.
+ * See FCompElementRenderTargetPool for inspiration
  */
-struct STEVESUEHELPERS_API FStevesTextureRenderTargetPool : public TSharedFromThis<FStevesTextureRenderTargetPool>
+struct STEVESUEHELPERS_API FStevesTextureRenderTargetPool : public FGCObject, public TSharedFromThis<FStevesTextureRenderTargetPool>
 {
 
 protected:
@@ -73,6 +74,7 @@ protected:
 
 	TWeakObjectPtr<UObject> PoolOwner;
 	TMultiMap<FTextureKey, UTextureRenderTarget2D*> UnreservedTextures;
+	TSet<UTextureRenderTarget2D*> ReservedTextures;
 
 	/// Weak reverse tracking of reservations, mostly for debugging
 	struct FReservationInfo
@@ -94,8 +96,7 @@ protected:
 	friend struct FStevesTextureRenderTargetReservation;
 	/// Release a reservation on a texture, allowing it back into the pool
 	/// Protected because only FStevesTextureRenderTargetReservation will need to do this.
-	void ReleaseReservation(UTextureRenderTarget2D* Tex);	
-
+	void ReleaseReservation(UTextureRenderTarget2D* Tex);
 public:
 
 	explicit FStevesTextureRenderTargetPool(const FName& InName, UObject* InOwner)
@@ -103,8 +104,13 @@ public:
 	{
 	}
 
+	virtual ~FStevesTextureRenderTargetPool() override;
+
 	const FName& GetName() const { return Name; }
 
+	// FGCObject
+	virtual void AddReferencedObjects(FReferenceCollector& Collector) override;
+	
 	/**
 	 * Reserve a texture for use as a render target. This will create a new texture target if needed. 
 	 * @param Size The dimensions of the texture