ADBDEV-8922: Restore global size limits for table_size_map and quota_info_map

src/diskquota.c

@@ -91,7 +91,18 @@ static DiskQuotaWorkerEntry *volatile MyWorkerInfo = NULL;
 // how many database diskquota are monitoring on
 static int num_db = 0;
 
-DiskquotaLauncherShmemStruct *DiskquotaLauncherShmem = NULL;
+static DiskquotaLauncherShmemStruct *DiskquotaLauncherShmem;
+
+/* how many TableSizeEntry are maintained in all the table_size_map in shared memory */
+pg_atomic_uint32 *diskquota_table_size_entry_num;
+
+/* how many QuotaInfoEntry are maintained in all the quota_info_map in shared memory */
+pg_atomic_uint32 *diskquota_quota_info_entry_num;
+
+#ifdef USE_ASSERT_CHECKING
+extern pg_atomic_flag *diskquota_table_size_flag;
+extern pg_atomic_flag *diskquota_quota_info_flag;
+#endif
 
 #define MIN_SLEEPTIME 100         /* milliseconds */
 #define BGWORKER_LOG_TIME 3600000 /* milliseconds */
@@ -1772,6 +1783,25 @@ init_launcher_shmem()
 			DiskquotaLauncherShmem->dbArray[i].workerId = INVALID_WORKER_ID;
 		}
 	}
+	/* init TableSizeEntry counter */
+	diskquota_table_size_entry_num =
+	        DiskquotaShmemInitStruct("diskquota TableSizeEntry counter", DISKQUOTA_TABLE_SIZE_ENTRY_NUM_SIZE, &found);
+	if (!found) pg_atomic_init_u32(diskquota_table_size_entry_num, 0);
+
+	/* init QuotaInfoEntry counter */
+	diskquota_quota_info_entry_num =
+	        DiskquotaShmemInitStruct("diskquota QuotaInfoEntry counter", DISKQUOTA_QUOTA_INFO_ENTRY_NUM_SIZE, &found);
+	if (!found) pg_atomic_init_u32(diskquota_quota_info_entry_num, 0);
+
+#ifdef USE_ASSERT_CHECKING
+	diskquota_table_size_flag =
+	        DiskquotaShmemInitStruct("diskquota TableSizeEntry flag", DISKQUOTA_TABLE_SIZE_FLAG_SIZE, &found);
+	if (!found) pg_atomic_init_flag(diskquota_table_size_flag);
+
+	diskquota_quota_info_flag =
+	        DiskquotaShmemInitStruct("diskquota QuotaInfoEntry flag", DISKQUOTA_QUOTA_INFO_FLAG_SIZE, &found);
+	if (!found) pg_atomic_init_flag(diskquota_quota_info_flag);
+#endif
 }
 
 /*

src/diskquota.h

@@ -42,9 +42,10 @@
 #define MAX_NUM_KEYS_QUOTA_MAP 8
 /* init number of QuotaInfoEntry in quota_info_map */
 #define INIT_QUOTA_MAP_ENTRIES 128
-#define AVG_QUOTA_MAP_ENTRIES (diskquota_max_quota_probes / diskquota_max_monitored_databases)
 /* max number of QuotaInfoEntry in quota_info_map */
-#define MAX_QUOTA_MAP_ENTRIES (AVG_QUOTA_MAP_ENTRIES < 1024 ? 1024 : AVG_QUOTA_MAP_ENTRIES)
+#define MAX_QUOTA_MAP_ENTRIES                                                                                         \
+	(diskquota_max_quota_probes < 1024 * diskquota_max_monitored_databases ? 1024 * diskquota_max_monitored_databases \
+	                                                                       : diskquota_max_quota_probes)
 
 typedef enum
 {
@@ -81,6 +82,14 @@ extern int diskquota_worker_timeout;
 #define RELID_CACHE_ENTRY_SIZE sizeof(DiskQuotaRelidCacheEntry)
 #define ALTERED_RELOID_CACHE_ENTRY_SIZE sizeof(Oid)
 #define MONITORED_DBID_CACHE_ENTRY_SIZE sizeof(struct MonitorDBEntryStruct)
+#define DISKQUOTA_TABLE_SIZE_ENTRY_NUM_SIZE sizeof(pg_atomic_uint32)
+#define DISKQUOTA_QUOTA_INFO_ENTRY_NUM_SIZE sizeof(pg_atomic_uint32)
+
+#ifdef USE_ASSERT_CHECKING
+#define DISKQUOTA_TABLE_SIZE_FLAG_SIZE sizeof(pg_atomic_flag)
+#define LOCAL_DISK_QUOTA_REJECT_FLAG_SIZE sizeof(pg_atomic_flag)
+#define DISKQUOTA_QUOTA_INFO_FLAG_SIZE sizeof(pg_atomic_flag)
+#endif
 
 typedef enum
 {
@@ -322,9 +331,11 @@ extern void         update_monitordb_status(Oid dbid, uint32 status);
 extern HTAB        *diskquota_hash_create(const char *tabname, long nelem, HASHCTL *info, int flags,
                                           DiskquotaHashFunction hashFunction);
 extern HTAB *DiskquotaShmemInitHash(const char *name, long init_size, long max_size, HASHCTL *infoP, int hash_flags,
-                                    DiskquotaHashFunction hash_function);
+                                    DiskquotaHashFunction hash_function, pg_atomic_flag *foundPtr);
 extern void *DiskquotaShmemInitStruct(const char *name, Size size, bool *foundPtr);
 extern void  refresh_monitored_dbid_cache(void);
+extern HASHACTION check_hash_fullness_num(HTAB *hashp, int num_entries, int max_size, const char *warning_message,
+                                          TimestampTz *last_overflow_report);
 extern HASHACTION check_hash_fullness(HTAB *hashp, int max_size, const char *warning_message,
                                       TimestampTz *last_overflow_report);
 bool              SPI_push_cond_and_connect(void);

src/diskquota_utility.c

@@ -115,8 +115,7 @@ static bool   to_delete_quota(QuotaType type, int64 quota_limit_mb, float4 segra
 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);
+extern void diskquota_shmem_size_sub(Size size);
 #endif
 
 /* ---- Help Functions to set quota limit. ---- */
@@ -1645,12 +1644,9 @@ DiskquotaShmemInitHash(const char           *name,       /* table string name fo
                        long                  max_size,   /* max size of the table */
                        HASHCTL              *infoP,      /* info about key and bucket size */
                        int                   hash_flags, /* info about infoP */
-                       DiskquotaHashFunction hashFunction)
+                       DiskquotaHashFunction hashFunction, pg_atomic_flag *foundPtr)
 {
-#ifdef USE_ASSERT_CHECKING
-	if (!DiskquotaLauncherShmem || !DiskquotaLauncherShmem->isDynamicWorker)
-		diskquota_shmem_size_sub(hash_estimate_size(max_size, infoP->entrysize));
-#endif
+	HTAB *hashp;
 
 #if GP_VERSION_NUM < 70000
 	if (hashFunction == DISKQUOTA_TAG_HASH)
@@ -1659,31 +1655,37 @@ DiskquotaShmemInitHash(const char           *name,       /* table string name fo
 		infoP->hash = oid_hash;
 	else
 		infoP->hash = string_hash;
-	return ShmemInitHash(name, init_size, max_size, infoP, hash_flags | HASH_FUNCTION);
+	hashp = ShmemInitHash(name, init_size, max_size, infoP, hash_flags | HASH_FUNCTION);
 #else
-	return ShmemInitHash(name, init_size, max_size, infoP, hash_flags | HASH_BLOBS);
+	hashp = ShmemInitHash(name, init_size, max_size, infoP, hash_flags | HASH_BLOBS);
 #endif /* GP_VERSION_NUM */
+
+#ifdef USE_ASSERT_CHECKING
+	AssertImply(IsBackgroundWorker, foundPtr != NULL); /* foundPtr only used in background worker */
+
+	if (foundPtr == NULL || pg_atomic_test_set_flag(foundPtr))
+		diskquota_shmem_size_sub(hash_estimate_size(max_size, infoP->entrysize));
+#endif
+
+	return hashp;
 }
 
 void *
 DiskquotaShmemInitStruct(const char *name, Size size, bool *foundPtr)
 {
+	void *structPtr = ShmemInitStruct(name, size, foundPtr);
+
 #ifdef USE_ASSERT_CHECKING
-	if (!DiskquotaLauncherShmem || !DiskquotaLauncherShmem->isDynamicWorker) diskquota_shmem_size_sub(size);
+	if (!*foundPtr) diskquota_shmem_size_sub(size);
 #endif
 
-	return ShmemInitStruct(name, size, foundPtr);
+	return structPtr;
 }
 
-/*
- * Returns HASH_FIND if hash table is full and HASH_ENTER otherwise.
- * It can be used only under lock.
- */
 HASHACTION
-check_hash_fullness(HTAB *hashp, int max_size, const char *warning_message, TimestampTz *last_overflow_report)
+check_hash_fullness_num(HTAB *hashp, int num_entries, int max_size, const char *warning_message,
+                        TimestampTz *last_overflow_report)
 {
-	long num_entries = hash_get_num_entries(hashp);
-
 	if (num_entries < max_size) return HASH_ENTER;
 
 	if (num_entries == max_size)
@@ -1701,6 +1703,16 @@ check_hash_fullness(HTAB *hashp, int max_size, const char *warning_message, Time
 	return HASH_FIND;
 }
 
+/*
+ * Returns HASH_FIND if hash table is full and HASH_ENTER otherwise.
+ * It can be used only under lock.
+ */
+HASHACTION
+check_hash_fullness(HTAB *hashp, int max_size, const char *warning_message, TimestampTz *last_overflow_report)
+{
+	return check_hash_fullness_num(hashp, hash_get_num_entries(hashp), max_size, warning_message, last_overflow_report);
+}
+
 bool
 SPI_push_cond_and_connect(void)
 {

src/gp_activetable.c

@@ -112,13 +112,14 @@ init_shm_worker_active_tables(void)
 	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);
+	                                           diskquota_max_active_tables, &ctl, HASH_ELEM, DISKQUOTA_TAG_HASH, NULL);
 
 	memset(&ctl, 0, sizeof(ctl));
-	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);
+	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, NULL);
 }
 
 /*