From 76655957fd8d69ab6fcd67ced214e16a40cb5854 Mon Sep 17 00:00:00 2001
From: Nils Dijk <nils@citusdata.com>
Date: Mon, 9 Oct 2023 13:05:51 +0200
Subject: [PATCH] Fix leaking of memory and memory contexts in Foreign
 Constraint Graphs (#7236)

DESCRIPTION: Fix leaking of memory and memory contexts in Foreign
Constraint Graphs

Previously, every time we (re)created the Foreign Constraint
Relationship Graph, we created a new Memory Context while loosing a
reference to the previous context. This old context could still have
left over memory in there causing a memory leak.

With this patch we statically have one memory context that we lazily
initialize the first time we create our foreign constraint relationship
graph. On every subsequent creation, beside destroying our previous
hashmap we also reset our memory context to remove any left over
references.
---
 .../utils/foreign_key_relationship.c          | 57 ++++++++++---------
 .../distributed/foreign_key_relationship.h    |  1 -
 2 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/src/backend/distributed/utils/foreign_key_relationship.c b/src/backend/distributed/utils/foreign_key_relationship.c
index 84ff21f8c57..10215254346 100644
--- a/src/backend/distributed/utils/foreign_key_relationship.c
+++ b/src/backend/distributed/utils/foreign_key_relationship.c
@@ -27,6 +27,7 @@
 #include "distributed/version_compat.h"
 #include "nodes/pg_list.h"
 #include "storage/lockdefs.h"
+#include "utils/catcache.h"
 #include "utils/fmgroids.h"
 #include "utils/hsearch.h"
 #if PG_VERSION_NUM >= PG_VERSION_13
@@ -97,6 +98,8 @@ static List * GetConnectedListHelper(ForeignConstraintRelationshipNode *node,
 									 bool isReferencing);
 static List * GetForeignConstraintRelationshipHelper(Oid relationId, bool isReferencing);
 
+MemoryContext ForeignConstraintRelationshipMemoryContext = NULL;
+
 
 /*
  * GetForeignKeyConnectedRelationIdList returns a list of relation id's for
@@ -324,17 +327,36 @@ CreateForeignConstraintRelationshipGraph()
 		return;
 	}
 
-	ClearForeignConstraintRelationshipGraphContext();
+	/*
+	 * Lazily create our memory context once and reset on every reuse.
+	 * Since we have cleared and invalidated the fConstraintRelationshipGraph, right
+	 * before we can simply reset the context if it was already existing.
+	 */
+	if (ForeignConstraintRelationshipMemoryContext == NULL)
+	{
+		/* make sure we've initialized CacheMemoryContext */
+		if (CacheMemoryContext == NULL)
+		{
+			CreateCacheMemoryContext();
+		}
+
+		ForeignConstraintRelationshipMemoryContext = AllocSetContextCreate(
+			CacheMemoryContext,
+			"Foreign Constraint Relationship Graph Context",
+			ALLOCSET_DEFAULT_MINSIZE,
+			ALLOCSET_DEFAULT_INITSIZE,
+			ALLOCSET_DEFAULT_MAXSIZE);
+	}
+	else
+	{
+		fConstraintRelationshipGraph = NULL;
+		MemoryContextReset(ForeignConstraintRelationshipMemoryContext);
+	}
 
-	MemoryContext fConstraintRelationshipMemoryContext = AllocSetContextCreateExtended(
-		CacheMemoryContext,
-		"Forign Constraint Relationship Graph Context",
-		ALLOCSET_DEFAULT_MINSIZE,
-		ALLOCSET_DEFAULT_INITSIZE,
-		ALLOCSET_DEFAULT_MAXSIZE);
+	Assert(fConstraintRelationshipGraph == NULL);
 
 	MemoryContext oldContext = MemoryContextSwitchTo(
-		fConstraintRelationshipMemoryContext);
+		ForeignConstraintRelationshipMemoryContext);
 
 	fConstraintRelationshipGraph = (ForeignConstraintRelationshipGraph *) palloc(
 		sizeof(ForeignConstraintRelationshipGraph));
@@ -668,22 +690,3 @@ CreateOrFindNode(HTAB *adjacencyLists, Oid relid)
 
 	return node;
 }
-
-
-/*
- * ClearForeignConstraintRelationshipGraphContext clear all the allocated memory obtained
- * for foreign constraint relationship graph. Since all the variables of relationship
- * graph was obtained within the same context, destroying hash map is enough as
- * it deletes the context.
- */
-void
-ClearForeignConstraintRelationshipGraphContext()
-{
-	if (fConstraintRelationshipGraph == NULL)
-	{
-		return;
-	}
-
-	hash_destroy(fConstraintRelationshipGraph->nodeMap);
-	fConstraintRelationshipGraph = NULL;
-}
diff --git a/src/include/distributed/foreign_key_relationship.h b/src/include/distributed/foreign_key_relationship.h
index 6e476162ec2..5e85b7edba9 100644
--- a/src/include/distributed/foreign_key_relationship.h
+++ b/src/include/distributed/foreign_key_relationship.h
@@ -20,7 +20,6 @@ extern bool ShouldUndistributeCitusLocalTable(Oid relationId);
 extern List * ReferencedRelationIdList(Oid relationId);
 extern List * ReferencingRelationIdList(Oid relationId);
 extern void SetForeignConstraintRelationshipGraphInvalid(void);
-extern void ClearForeignConstraintRelationshipGraphContext(void);
 extern HTAB * CreateOidVisitedHashSet(void);
 extern bool OidVisited(HTAB *oidVisitedMap, Oid oid);
 extern void VisitOid(HTAB *oidVisitedMap, Oid oid);