From e59ffbf54913830c1a186348d7ee453ad67dfd43 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 2858e6ed3ba..d30c767df69 100644
--- a/src/backend/distributed/utils/foreign_key_relationship.c
+++ b/src/backend/distributed/utils/foreign_key_relationship.c
@@ -28,6 +28,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"
 #include "common/hashfn.h"
@@ -96,6 +97,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
@@ -321,17 +324,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 = AllocSetContextCreateInternal(
-		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));
@@ -631,22 +653,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 ef2c5be33a8..fbbee831e7a 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 bool OidVisited(HTAB *oidVisitedMap, Oid oid);
 extern void VisitOid(HTAB *oidVisitedMap, Oid oid);