ADBDEV-8953: Fix shared memory allocation and usage in diskquota

src/diskquota.c

@@ -91,7 +91,7 @@ static DiskQuotaWorkerEntry *volatile MyWorkerInfo = NULL;
 // how many database diskquota are monitoring on
 static int num_db = 0;
 
-static DiskquotaLauncherShmemStruct *DiskquotaLauncherShmem;
+DiskquotaLauncherShmemStruct *DiskquotaLauncherShmem = NULL;
 
 #define MIN_SLEEPTIME 100         /* milliseconds */
 #define BGWORKER_LOG_TIME 3600000 /* milliseconds */
@@ -1130,7 +1130,7 @@ process_extension_ddl_message()
 	ExtensionDDLMessage local_extension_ddl_message;
 
 	LWLockAcquire(diskquota_locks.extension_ddl_message_lock, LW_SHARED);
-	memcpy(&local_extension_ddl_message, extension_ddl_message, sizeof(ExtensionDDLMessage));
+	memcpy(&local_extension_ddl_message, extension_ddl_message, EXTENSION_DDL_MESSAGE_SIZE);
 	LWLockRelease(diskquota_locks.extension_ddl_message_lock);
 
 	/* create/drop extension message must be valid */
@@ -1146,7 +1146,7 @@ process_extension_ddl_message()
 
 	/* Send createdrop extension diskquota result back to QD */
 	LWLockAcquire(diskquota_locks.extension_ddl_message_lock, LW_EXCLUSIVE);
-	memset(extension_ddl_message, 0, sizeof(ExtensionDDLMessage));
+	memset(extension_ddl_message, 0, EXTENSION_DDL_MESSAGE_SIZE);
 	extension_ddl_message->launcher_pid = MyProcPid;
 	extension_ddl_message->result       = (int)code;
 	LWLockRelease(diskquota_locks.extension_ddl_message_lock);
@@ -1732,8 +1732,8 @@ void
 init_launcher_shmem()
 {
 	bool found;
-	DiskquotaLauncherShmem = (DiskquotaLauncherShmemStruct *)ShmemInitStruct("Diskquota launcher Data",
-	                                                                         diskquota_launcher_shmem_size(), &found);
+	DiskquotaLauncherShmem = (DiskquotaLauncherShmemStruct *)DiskquotaShmemInitStruct(
+	        "Diskquota launcher Data", diskquota_launcher_shmem_size(), &found);
 	memset(DiskquotaLauncherShmem, 0, diskquota_launcher_shmem_size());
 	if (!found)
 	{

src/diskquota.h

@@ -75,6 +75,13 @@ extern int diskquota_worker_timeout;
 #define DiskquotaGetRelstorage(classForm) (0)
 #endif /* GP_VERSION_NUM */
 
+#define EXTENSION_DDL_MESSAGE_SIZE sizeof(ExtensionDDLMessage)
+#define ACTIVE_TABLES_MAP_ENTRY_SIZE sizeof(DiskQuotaActiveTableFileEntry)
+#define RELATION_CACHE_ENTRY_SIZE sizeof(DiskQuotaRelationCacheEntry)
+#define RELID_CACHE_ENTRY_SIZE sizeof(DiskQuotaRelidCacheEntry)
+#define ALTERED_RELOID_CACHE_ENTRY_SIZE sizeof(Oid)
+#define MONITORED_DBID_CACHE_ENTRY_SIZE sizeof(struct MonitorDBEntryStruct)
+
 typedef enum
 {
 	NAMESPACE_QUOTA = 0,
@@ -316,6 +323,7 @@ extern HTAB        *diskquota_hash_create(const char *tabname, long nelem, HASHC
                                           DiskquotaHashFunction hashFunction);
 extern HTAB *DiskquotaShmemInitHash(const char *name, long init_size, long max_size, HASHCTL *infoP, int hash_flags,
                                     DiskquotaHashFunction hash_function);
+extern void *DiskquotaShmemInitStruct(const char *name, Size size, bool *foundPtr);
 extern void  refresh_monitored_dbid_cache(void);
 extern HASHACTION check_hash_fullness(HTAB *hashp, int max_size, const char *warning_message,
                                       TimestampTz *last_overflow_report);

src/diskquota_utility.c

@@ -114,6 +114,11 @@ static float4 get_per_segment_ratio(Oid spcoid);
 static bool   to_delete_quota(QuotaType type, int64 quota_limit_mb, float4 segratio);
 static void   check_role(Oid roleoid, char *rolname, int64 quota_limit_mb);
 
+#ifdef USE_ASSERT_CHECKING
+extern DiskquotaLauncherShmemStruct *DiskquotaLauncherShmem;
+extern void                          diskquota_shmem_size_sub(Size size);
+#endif
+
 /* ---- Help Functions to set quota limit. ---- */
 /*
  * Initialize table diskquota.table_size.
@@ -1642,6 +1647,11 @@ DiskquotaShmemInitHash(const char           *name,       /* table string name fo
                        int                   hash_flags, /* info about infoP */
                        DiskquotaHashFunction hashFunction)
 {
+#ifdef USE_ASSERT_CHECKING
+	if (!DiskquotaLauncherShmem || !DiskquotaLauncherShmem->isDynamicWorker)
+		diskquota_shmem_size_sub(hash_estimate_size(max_size, infoP->entrysize));
+#endif
+
 #if GP_VERSION_NUM < 70000
 	if (hashFunction == DISKQUOTA_TAG_HASH)
 		infoP->hash = tag_hash;
@@ -1655,6 +1665,16 @@ DiskquotaShmemInitHash(const char           *name,       /* table string name fo
 #endif /* GP_VERSION_NUM */
 }
 
+void *
+DiskquotaShmemInitStruct(const char *name, Size size, bool *foundPtr)
+{
+#ifdef USE_ASSERT_CHECKING
+	if (!DiskquotaLauncherShmem || !DiskquotaLauncherShmem->isDynamicWorker) diskquota_shmem_size_sub(size);
+#endif
+
+	return ShmemInitStruct(name, size, foundPtr);
+}
+
 /*
  * Returns HASH_FIND if hash table is full and HASH_ENTER otherwise.
  * It can be used only under lock.

src/gp_activetable.c

@@ -109,14 +109,14 @@ init_shm_worker_active_tables(void)
 	HASHCTL ctl;
 
 	memset(&ctl, 0, sizeof(ctl));
-	ctl.keysize       = sizeof(DiskQuotaActiveTableFileEntry);
-	ctl.entrysize     = sizeof(DiskQuotaActiveTableFileEntry);
+	ctl.keysize       = ACTIVE_TABLES_MAP_ENTRY_SIZE;
+	ctl.entrysize     = ACTIVE_TABLES_MAP_ENTRY_SIZE;
 	active_tables_map = DiskquotaShmemInitHash("active_tables", diskquota_max_active_tables,
 	                                           diskquota_max_active_tables, &ctl, HASH_ELEM, DISKQUOTA_TAG_HASH);
 
 	memset(&ctl, 0, sizeof(ctl));
-	ctl.keysize          = sizeof(Oid);
-	ctl.entrysize        = sizeof(Oid);
+	ctl.keysize          = ALTERED_RELOID_CACHE_ENTRY_SIZE;
+	ctl.entrysize        = ALTERED_RELOID_CACHE_ENTRY_SIZE;
 	altered_reloid_cache = DiskquotaShmemInitHash("altered_reloid_cache", diskquota_max_active_tables,
 	                                              diskquota_max_active_tables, &ctl, HASH_ELEM, DISKQUOTA_OID_HASH);
 }
@@ -322,7 +322,7 @@ report_active_table_helper(const RelFileNodeBackend *relFileNode)
 	}
 	found = false;
 
-	MemSet(&item, 0, sizeof(DiskQuotaActiveTableFileEntry));
+	MemSet(&item, 0, ACTIVE_TABLES_MAP_ENTRY_SIZE);
 	item.dbid          = relFileNode->node.dbNode;
 	item.relfilenode   = relFileNode->node.relNode;
 	item.tablespaceoid = relFileNode->node.spcNode;
@@ -730,15 +730,15 @@ get_active_tables_oid(void)
 	refresh_monitored_dbid_cache();
 
 	memset(&ctl, 0, sizeof(ctl));
-	ctl.keysize                 = sizeof(DiskQuotaActiveTableFileEntry);
-	ctl.entrysize               = sizeof(DiskQuotaActiveTableFileEntry);
+	ctl.keysize                 = ACTIVE_TABLES_MAP_ENTRY_SIZE;
+	ctl.entrysize               = ACTIVE_TABLES_MAP_ENTRY_SIZE;
 	ctl.hcxt                    = CurrentMemoryContext;
 	local_active_table_file_map = diskquota_hash_create("local active table map with relfilenode info", 1024, &ctl,
 	                                                    HASH_ELEM | HASH_CONTEXT, DISKQUOTA_TAG_HASH);
 
 	memset(&ctl, 0, sizeof(ctl));
-	ctl.keysize                = sizeof(Oid);
-	ctl.entrysize              = sizeof(Oid);
+	ctl.keysize                = ALTERED_RELOID_CACHE_ENTRY_SIZE;
+	ctl.entrysize              = ALTERED_RELOID_CACHE_ENTRY_SIZE;
 	ctl.hcxt                   = CurrentMemoryContext;
 	local_altered_reloid_cache = diskquota_hash_create("local_altered_reloid_cache", 1024, &ctl,
 	                                                   HASH_ELEM | HASH_CONTEXT, DISKQUOTA_OID_HASH);

src/monitored_db.c

@@ -322,14 +322,14 @@ dump_monitored_dbid_cache(long *nitems)
 	HASH_SEQ_STATUS seq;
 	MonitorDBEntry  curEntry;
 	int             count = *nitems = hash_get_num_entries(monitored_dbid_cache);
-	MonitorDBEntry  entries = curEntry = (MonitorDBEntry)palloc(sizeof(struct MonitorDBEntryStruct) * count);
+	MonitorDBEntry  entries = curEntry = (MonitorDBEntry)palloc(MONITORED_DBID_CACHE_ENTRY_SIZE * count);
 
 	hash_seq_init(&seq, monitored_dbid_cache);
 	MonitorDBEntry entry;
 	while ((entry = hash_seq_search(&seq)) != NULL)
 	{
 		Assert(count > 0);
-		memcpy(curEntry, entry, sizeof(struct MonitorDBEntryStruct));
+		memcpy(curEntry, entry, MONITORED_DBID_CACHE_ENTRY_SIZE);
 		curEntry++;
 		count--;
 	}