From d67cf907a2ed20c9f8a1d190d1eeec69d3d9a75d Mon Sep 17 00:00:00 2001
From: Onder Kalaci <onderkalaci@gmail.com>
Date: Wed, 27 Jul 2022 18:03:53 +0200
Subject: [PATCH] Detach relation access tracking from connection management

---
 .../connection/placement_connection.c         |  3 --
 src/backend/distributed/shared_library_init.c |  1 +
 .../transaction/relation_access_tracking.c    | 34 ++++++++++++++++---
 .../distributed/relation_access_tracking.h    |  2 +-
 4 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/backend/distributed/connection/placement_connection.c b/src/backend/distributed/connection/placement_connection.c
index d3929f4b9..2aa3994a0 100644
--- a/src/backend/distributed/connection/placement_connection.c
+++ b/src/backend/distributed/connection/placement_connection.c
@@ -1089,9 +1089,6 @@ InitPlacementConnectionManagement(void)
 
 	ConnectionShardHash = hash_create("citus connection cache (shardid)",
 									  64, &info, hashFlags);
-
-	/* (relationId) = [relationAccessMode] hash */
-	AllocateRelationAccessHash();
 }
 
 
diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c
index 8ea60fe74..600b62b69 100644
--- a/src/backend/distributed/shared_library_init.c
+++ b/src/backend/distributed/shared_library_init.c
@@ -390,6 +390,7 @@ _PG_init(void)
 	InitializeBackendManagement();
 	InitializeConnectionManagement();
 	InitPlacementConnectionManagement();
+	InitRelationAccessHash();
 	InitializeCitusQueryStats();
 	InitializeSharedConnectionStats();
 	InitializeLocallyReservedSharedConnections();
diff --git a/src/backend/distributed/transaction/relation_access_tracking.c b/src/backend/distributed/transaction/relation_access_tracking.c
index 9fdb226e1..f69de6f8a 100644
--- a/src/backend/distributed/transaction/relation_access_tracking.c
+++ b/src/backend/distributed/transaction/relation_access_tracking.c
@@ -47,6 +47,8 @@ bool EnforceForeignKeyRestrictions = true;
 									(1 << (PLACEMENT_ACCESS_DDL + \
 										   PARALLEL_MODE_FLAG_OFFSET)))
 
+MemoryContext RelationAcessContext = NULL;
+
 
 /*
  * Hash table mapping relations to the
@@ -84,8 +86,8 @@ typedef struct RelationAccessHashEntry
 
 static HTAB *RelationAccessHash;
 
-
 /* functions related to access recording */
+static void AllocateRelationAccessHash(void);
 static void RecordRelationAccessBase(Oid relationId, ShardPlacementAccessType accessType);
 static void RecordPlacementAccessToCache(Oid relationId,
 										 ShardPlacementAccessType accessType);
@@ -120,6 +122,18 @@ static bool HoldsConflictingLockWithReferencedRelations(Oid relationId,
 														conflictingAccessMode);
 
 
+/*
+ * InitRelationAccessHash performs initialization of the
+ * infrastructure in this file at backend start.
+ */
+void
+InitRelationAccessHash(void)
+{
+	/* allocate (relationId) = [relationAccessMode] hash */
+	AllocateRelationAccessHash();
+}
+
+
 /*
  * Empty RelationAccessHash, without destroying the hash table itself.
  */
@@ -133,19 +147,29 @@ ResetRelationAccessHash()
 /*
  * Allocate RelationAccessHash.
  */
-void
+static void
 AllocateRelationAccessHash(void)
 {
-	HASHCTL info;
+	/*
+	 * Create a single context for relation access related memory
+	 * management. Doing so, instead of allocating in TopMemoryContext, makes
+	 * it easier to associate used memory.
+	 */
+	RelationAcessContext = AllocSetContextCreateExtended(TopMemoryContext,
+														 "Relation Access Context",
+														 ALLOCSET_DEFAULT_MINSIZE,
+														 ALLOCSET_DEFAULT_INITSIZE,
+														 ALLOCSET_DEFAULT_MAXSIZE);
 
+	HASHCTL info;
 	memset(&info, 0, sizeof(info));
 	info.keysize = sizeof(RelationAccessHashKey);
 	info.entrysize = sizeof(RelationAccessHashEntry);
 	info.hash = tag_hash;
-	info.hcxt = ConnectionContext;
+	info.hcxt = RelationAcessContext;
 	uint32 hashFlags = (HASH_ELEM | HASH_BLOBS | HASH_CONTEXT);
 
-	RelationAccessHash = hash_create("citus connection cache (relationid)",
+	RelationAccessHash = hash_create("citus relation access cache (relationid)",
 									 8, &info, hashFlags);
 }
 
diff --git a/src/include/distributed/relation_access_tracking.h b/src/include/distributed/relation_access_tracking.h
index deacdec94..295d54351 100644
--- a/src/include/distributed/relation_access_tracking.h
+++ b/src/include/distributed/relation_access_tracking.h
@@ -34,7 +34,7 @@ typedef enum RelationAccessMode
 	RELATION_PARALLEL_ACCESSED
 } RelationAccessMode;
 
-extern void AllocateRelationAccessHash(void);
+extern void InitRelationAccessHash(void);
 extern void ResetRelationAccessHash(void);
 extern void RecordRelationAccessIfNonDistTable(Oid relationId,
 											   ShardPlacementAccessType accessType);