Fix bug: rejectmap entries should not be removed by other databases. …

…(#279) This commit fixed two bugs: - Previously, refresh_rejectmap() cleared all entries in rejectmap, including other databases' entries, which causes hardlimit can not to work correctly. - soft-limit rejectmap entries should not be added into disk_quota_reject_map on segments, otherwise, these entries may remain in segments and trigger the soft-limit incorrectly. Co-authored-by: Chen Mulong <[email protected]>
arenadata · Dec 12, 2022 · 8cbdeb6 · 8cbdeb6
1 parent 73114d8
commit 8cbdeb6
Show file tree

Hide file tree

Showing 4 changed files with 168 additions and 13 deletions.
diff --git a/quotamodel.c b/quotamodel.c
@@ -811,8 +811,13 @@ refresh_disk_quota_usage(bool is_init)
 		flush_to_table_size();
 		/* copy local reject map back to shared reject map */
 		bool reject_map_changed = flush_local_reject_map();
-		/* Dispatch rejectmap entries to segments to perform hard-limit. */
-		if (diskquota_hardlimit && (reject_map_changed || hasActiveTable))
+		/*
+		 * Dispatch rejectmap entries to segments to perform hard-limit.
+		 * If the bgworker is in init mode, the rejectmap should be refreshed anyway.
+		 * Otherwise, only when the rejectmap is changed or the active_table_list is
+		 * not empty the rejectmap should be dispatched to segments.
+		 */
+		if (is_init || (diskquota_hardlimit && (reject_map_changed || hasActiveTable)))
 			dispatch_rejectmap(local_active_table_stat_map);
 		hash_destroy(local_active_table_stat_map);
 	}
@@ -1795,7 +1800,8 @@ refresh_rejectmap(PG_FUNCTION_ARGS)
 	int16                 elem_width;
 	bool                  elem_type_by_val;
 	char                  elem_alignment_code;
-	int                   count;
+	int                   reject_array_count;
+	int                   active_array_count;
 	HeapTupleHeader       lt;
 	bool                  segexceeded;
 	GlobalRejectMapEntry *rejectmapentry;
@@ -1832,8 +1838,8 @@ refresh_rejectmap(PG_FUNCTION_ARGS)
 	local_rejectmap = hash_create("local_rejectmap", 1024, &hashctl, HASH_ELEM | HASH_CONTEXT | HASH_FUNCTION);
 	get_typlenbyvalalign(rejectmap_elem_type, &elem_width, &elem_type_by_val, &elem_alignment_code);
 	deconstruct_array(rejectmap_array_type, rejectmap_elem_type, elem_width, elem_type_by_val, elem_alignment_code,
-	                  &datums, &nulls, &count);
-	for (int i = 0; i < count; ++i)
+	                  &datums, &nulls, &reject_array_count);
+	for (int i = 0; i < reject_array_count; ++i)
 	{
 		RejectMapEntry keyitem;
 		bool           isnull;
@@ -1864,8 +1870,8 @@ refresh_rejectmap(PG_FUNCTION_ARGS)
 	 */
 	get_typlenbyvalalign(active_oid_elem_type, &elem_width, &elem_type_by_val, &elem_alignment_code);
 	deconstruct_array(active_oid_array_type, active_oid_elem_type, elem_width, elem_type_by_val, elem_alignment_code,
-	                  &datums, &nulls, &count);
-	for (int i = 0; i < count; ++i)
+	                  &datums, &nulls, &active_array_count);
+	for (int i = 0; i < active_array_count; ++i)
 	{
 		Oid       active_oid = InvalidOid;
 		HeapTuple tuple;
@@ -2040,22 +2046,28 @@ refresh_rejectmap(PG_FUNCTION_ARGS)
 	/* Clear rejectmap entries. */
 	hash_seq_init(&hash_seq, disk_quota_reject_map);
 	while ((rejectmapentry = hash_seq_search(&hash_seq)) != NULL)
+	{
+		if (rejectmapentry->keyitem.relfilenode.dbNode != MyDatabaseId &&
+		    rejectmapentry->keyitem.databaseoid != MyDatabaseId)
+			continue;
 		hash_search(disk_quota_reject_map, &rejectmapentry->keyitem, HASH_REMOVE, NULL);
+	}
 
 	/* Flush the content of local_rejectmap to the global rejectmap. */
 	hash_seq_init(&hash_seq, local_rejectmap);
 	while ((rejectmapentry = hash_seq_search(&hash_seq)) != NULL)
 	{
 		bool                  found;
 		GlobalRejectMapEntry *new_entry;
-		new_entry = hash_search(disk_quota_reject_map, &rejectmapentry->keyitem, HASH_ENTER_NULL, &found);
+
 		/*
-		 * We don't perform soft-limit on segment servers, so we don't flush the
-		 * rejectmap entry with a valid targetoid to the global rejectmap on segment
-		 * servers.
+		 * Skip soft limit reject entry. We don't perform soft-limit on segment servers, so we don't flush the
+		 * rejectmap entry with a valid targetoid to the global rejectmap on segment servers.
 		 */
-		if (!found && new_entry && !OidIsValid(rejectmapentry->keyitem.targetoid))
-			memcpy(new_entry, rejectmapentry, sizeof(GlobalRejectMapEntry));
+		if (OidIsValid(rejectmapentry->keyitem.targetoid)) continue;
+
+		new_entry = hash_search(disk_quota_reject_map, &rejectmapentry->keyitem, HASH_ENTER_NULL, &found);
+		if (!found && new_entry) memcpy(new_entry, rejectmapentry, sizeof(GlobalRejectMapEntry));
 	}
 	LWLockRelease(diskquota_locks.reject_map_lock);
 

diff --git a/tests/regress/diskquota_schedule b/tests/regress/diskquota_schedule
@@ -31,6 +31,7 @@ test: test_fetch_table_stat
 test: test_appendonly
 test: test_rejectmap
 test: test_clean_rejectmap_after_drop
+test: test_rejectmap_mul_db
 test: test_ctas_pause
 test: test_ctas_role
 test: test_ctas_schema

diff --git a/tests/regress/expected/test_rejectmap_mul_db.out b/tests/regress/expected/test_rejectmap_mul_db.out
@@ -0,0 +1,89 @@
+-- One db's rejectmap update should not impact on other db's rejectmap
+CREATE DATABASE tjmu1;
+CREATE DATABASE tjmu2;
+-- start_ignore
+\! gpconfig -c "diskquota.hard_limit" -v "on" > /dev/null
+\! gpstop -u > /dev/null
+-- end_ignore
+\c tjmu1
+CREATE EXTENSION diskquota;
+SELECT diskquota.set_schema_quota('public', '1MB');
+ set_schema_quota 
+------------------
+
+(1 row)
+
+CREATE TABLE b (t TEXT) DISTRIBUTED BY (t);
+SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t
+(1 row)
+
+-- Trigger hard limit to dispatch rejectmap for tjmu1
+INSERT INTO b SELECT generate_series(1, 100000000); -- fail
+ERROR:  schema's disk space quota exceeded with name: 2200  (seg1 127.0.0.1:6003 pid=3985762)
+-- NOTE: Pause to avoid tjmu1's worker clear the active table. Since the naptime is 0 on CI, this might be flaky.
+SELECT diskquota.pause();
+ pause 
+-------
+
+(1 row)
+
+-- The rejectmap should contain entries with dbnode = 0 and dbnode = tjmu1_oid. count = 1
+SELECT COUNT(DISTINCT r.dbnode) FROM (SELECT (diskquota.show_rejectmap()).* FROM gp_dist_random('gp_id')) as r where r.dbnode != 0;
+ count 
+-------
+     1
+(1 row)
+
+\c tjmu2
+CREATE EXTENSION diskquota;
+SELECT diskquota.set_schema_quota('public', '1MB');
+ set_schema_quota 
+------------------
+
+(1 row)
+
+CREATE TABLE b (t TEXT) DISTRIBUTED BY (t);
+SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t
+(1 row)
+
+-- Trigger hard limit to dispatch rejectmap for tjmu2
+INSERT INTO b SELECT generate_series(1, 100000000); -- fail
+ERROR:  schema's disk space quota exceeded with name: 2200  (seg1 127.0.0.1:6003 pid=4001721)
+SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t
+(1 row)
+
+SELECT diskquota.pause();
+ pause 
+-------
+
+(1 row)
+
+--\c tjmu1
+-- The rejectmap should contain entris with dbnode = 0 and dbnode = tjmu1_oid and tjmu2_oid. count = 2
+-- The entries for tjmu1 should not be cleared
+SELECT COUNT(DISTINCT r.dbnode) FROM (SELECT (diskquota.show_rejectmap()).* FROM gp_dist_random('gp_id')) as r where r.dbnode != 0;
+ count 
+-------
+     2
+(1 row)
+
+-- start_ignore
+\! gpconfig -c "diskquota.hard_limit" -v "off" > /dev/null
+\! gpstop -u > /dev/null
+-- end_ignore
+\c tjmu1
+DROP EXTENSION diskquota;
+\c tjmu2
+DROP EXTENSION diskquota;
+\c contrib_regression
+DROP DATABASE tjmu1;
+DROP DATABASE tjmu2;
diff --git a/tests/regress/sql/test_rejectmap_mul_db.sql b/tests/regress/sql/test_rejectmap_mul_db.sql
@@ -0,0 +1,53 @@
+-- One db's rejectmap update should not impact on other db's rejectmap
+CREATE DATABASE tjmu1;
+CREATE DATABASE tjmu2;
+
+-- start_ignore
+\! gpconfig -c "diskquota.hard_limit" -v "on" > /dev/null
+-- increase the naptime to avoid active table gets cleared by tjmu1's worker
+\! gpconfig -c "diskquota.naptime" -v 1 > /dev/null
+\! gpstop -u > /dev/null
+-- end_ignore
+
+\c tjmu1
+CREATE EXTENSION diskquota;
+SELECT diskquota.set_schema_quota('public', '1MB');
+CREATE TABLE b (t TEXT) DISTRIBUTED BY (t);
+SELECT diskquota.wait_for_worker_new_epoch();
+-- Trigger hard limit to dispatch rejectmap for tjmu1
+INSERT INTO b SELECT generate_series(1, 100000000); -- fail
+-- NOTE: Pause to avoid tjmu1's worker clear the active table. Since the naptime is 0 on CI, this might be flaky.
+SELECT diskquota.pause();
+-- The rejectmap should contain entries with dbnode = 0 and dbnode = tjmu1_oid. count = 1
+SELECT COUNT(DISTINCT r.dbnode) FROM (SELECT (diskquota.show_rejectmap()).* FROM gp_dist_random('gp_id')) as r where r.dbnode != 0;
+
+\c tjmu2
+CREATE EXTENSION diskquota;
+SELECT diskquota.set_schema_quota('public', '1MB');
+CREATE TABLE b (t TEXT) DISTRIBUTED BY (t);
+SELECT diskquota.wait_for_worker_new_epoch();
+-- Trigger hard limit to dispatch rejectmap for tjmu2
+INSERT INTO b SELECT generate_series(1, 100000000); -- fail
+SELECT diskquota.wait_for_worker_new_epoch();
+SELECT diskquota.pause();
+
+--\c tjmu1
+-- The rejectmap should contain entris with dbnode = 0 and dbnode = tjmu1_oid and tjmu2_oid. count = 2
+-- The entries for tjmu1 should not be cleared
+SELECT COUNT(DISTINCT r.dbnode) FROM (SELECT (diskquota.show_rejectmap()).* FROM gp_dist_random('gp_id')) as r where r.dbnode != 0;
+
+-- start_ignore
+\! gpconfig -c "diskquota.hard_limit" -v "off" > /dev/null
+\! gpconfig -c "diskquota.naptime" -v 0 > /dev/null
+\! gpstop -u > /dev/null
+-- end_ignore
+
+\c tjmu1
+DROP EXTENSION diskquota;
+\c tjmu2
+DROP EXTENSION diskquota;
+
+\c contrib_regression
+DROP DATABASE tjmu1;
+DROP DATABASE tjmu2;
+