From bb36a6292b2d46e9c230918e1c00f37b1c3ab325 Mon Sep 17 00:00:00 2001
From: Dennis Kovalenko <kds@arenadata.io>
Date: Fri, 23 Aug 2024 13:09:17 +0000
Subject: [PATCH] Track startup memory for resource group

---
 src/backend/utils/mmgr/vmem_tracker.c |  4 ++++
 src/backend/utils/resgroup/resgroup.c | 20 ++++++++++++++++++++
 src/include/utils/resgroup.h          |  3 +++
 3 files changed, 27 insertions(+)

diff --git a/src/backend/utils/mmgr/vmem_tracker.c b/src/backend/utils/mmgr/vmem_tracker.c
index d2b729c97b38..d8b5e594800a 100644
--- a/src/backend/utils/mmgr/vmem_tracker.c
+++ b/src/backend/utils/mmgr/vmem_tracker.c
@@ -670,6 +670,8 @@ VmemTracker_RegisterStartupMemory(int64 bytes)
 	pg_atomic_add_fetch_u32((pg_atomic_uint32 *) segmentVmemChunks,
 							startupChunks);
 
+	selfAddStartupChunks(startupChunks);
+
 	/*
 	 * Step 2, check if an OOM error should be raised by allocating 0 chunk.
 	 */
@@ -692,6 +694,8 @@ VmemTracker_UnregisterStartupMemory(void)
 		pg_atomic_sub_fetch_u32((pg_atomic_uint32 *) &MySessionState->sessionVmem,
 								startupChunks);
 
+	selfSubStartupChunks(startupChunks);
+
 	trackedBytes -= startupBytes;
 	trackedVmemChunks -= startupChunks;
 
diff --git a/src/backend/utils/resgroup/resgroup.c b/src/backend/utils/resgroup/resgroup.c
index 65e52db46ca9..0382d3c168e5 100644
--- a/src/backend/utils/resgroup/resgroup.c
+++ b/src/backend/utils/resgroup/resgroup.c
@@ -1594,6 +1594,26 @@ selfDetachResGroup(ResGroupData *group, ResGroupSlotData *slot)
 	selfUnsetGroup();
 }
 
+/* 
+ * Add startup memory before a resgroup is assigned. This memory
+ * will later be added to resgroup via selfAttachResGroup
+ */
+void
+selfAddStartupChunks(int32 chunks)
+{
+	self->memUsage += chunks;
+}
+
+/* 
+ * Sub startup memory at cleanup. This memory should already been
+ * subtracted from a resource group via selfDetachResGroup
+ */
+void
+selfSubStartupChunks(int32 chunks)
+{
+	self->memUsage -= chunks;
+}
+
 /*
  * Initialize the members of a slot
  */
diff --git a/src/include/utils/resgroup.h b/src/include/utils/resgroup.h
index 87b57d65c66b..3dbd828b1f23 100644
--- a/src/include/utils/resgroup.h
+++ b/src/include/utils/resgroup.h
@@ -235,6 +235,9 @@ extern Oid ResGroupGetGroupIdBySessionId(int sessionId);
 extern char *getCpuSetByRole(const char *cpuset);
 extern void checkCpuSetByRole(const char *cpuset);
 
+extern void selfAddStartupChunks(int32 chunks);
+extern void selfSubStartupChunks(int32 chunks);
+
 #define LOG_RESGROUP_DEBUG(...) \
 	do {if (Debug_resource_group) elog(__VA_ARGS__); } while(false);