Invalidate diskquota.table_size entries during startup

arenadata · Nov 16, 2023 · 1c26307 · 1c26307
1 parent 87b5dc3
commit 1c26307
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 5 deletions.
diff --git a/src/diskquota.h b/src/diskquota.h
@@ -265,7 +265,7 @@ extern bool  diskquota_hardlimit;
 extern int      SEGCOUNT;
 extern int      worker_spi_get_extension_version(int *major, int *minor);
 extern void     truncateStringInfo(StringInfo str, int nchars);
-extern List    *get_rel_oid_list(void);
+extern List    *get_rel_oid_list(bool is_init);
 extern int64    calculate_relation_size_all_forks(RelFileNodeBackend *rnode, char relstorage, Oid relam);
 extern Relation diskquota_relation_open(Oid relid);
 extern bool     get_rel_name_namespace(Oid relid, Oid *nsOid, char *relname);

diff --git a/src/diskquota_utility.c b/src/diskquota_utility.c
@@ -113,7 +113,7 @@ 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);
 
-List *get_rel_oid_list(void);
+List *get_rel_oid_list(bool is_init);
 
 /* ---- Help Functions to set quota limit. ---- */
 /*
@@ -1299,12 +1299,18 @@ worker_spi_get_extension_version(int *major, int *minor)
  */
 
 List *
-get_rel_oid_list(void)
+get_rel_oid_list(bool is_init)
 {
 	List *oidlist = NIL;
 	int   ret;
 
-	ret = SPI_execute_with_args("select oid from pg_class where oid >= $1 and (relkind='r' or relkind='m')", 1,
+#define SELECT_FROM_PG_CATALOG_PG_CLASS "select oid from pg_catalog.pg_class where oid >= $1 and (relkind='r' or relkind='m')"
+#define SELECT_FROM_DISKQUOTA_TABLE_SIZE "select tableid from diskquota.table_size where segid = -1"
+
+	ret = SPI_execute_with_args(is_init ?
+	                            SELECT_FROM_PG_CATALOG_PG_CLASS " union distinct " SELECT_FROM_DISKQUOTA_TABLE_SIZE :
+	                            SELECT_FROM_PG_CATALOG_PG_CLASS,
+	                            1,
 	                            (Oid[]){
 	                                    OIDOID,
 	                            },

diff --git a/src/quotamodel.c b/src/quotamodel.c
@@ -247,6 +247,8 @@ static bool get_table_size_entry_flag(TableSizeEntry *entry, TableSizeEntryFlag
 static void reset_table_size_entry_flag(TableSizeEntry *entry, TableSizeEntryFlag flag);
 static void set_table_size_entry_flag(TableSizeEntry *entry, TableSizeEntryFlag flag);
 
+static void delete_from_table_size_map(char *str);
+
 /* add a new entry quota or update the old entry quota */
 static void
 update_size_for_quota(int64 size, QuotaType type, Oid *keys, int16 segid)
@@ -923,6 +925,10 @@ calculate_table_disk_usage(bool is_init, HTAB *local_active_table_stat_map)
 	TableEntryKey              active_table_key;
 	List                      *oidlist;
 	ListCell                  *l;
+	int                        delete_entries_num = 0;
+	StringInfoData             delete_statement;
+
+	initStringInfo(&delete_statement);
 
 	/*
 	 * unset is_exist flag for tsentry in table_size_map this is used to
@@ -939,7 +945,7 @@ calculate_table_disk_usage(bool is_init, HTAB *local_active_table_stat_map)
 	 * calculate the file size for active table and update namespace_size_map
 	 * and role_size_map
 	 */
-	oidlist = get_rel_oid_list();
+	oidlist = get_rel_oid_list(is_init);
 
 	oidlist = merge_uncommitted_table_to_oidlist(oidlist);
 
@@ -973,6 +979,25 @@ calculate_table_disk_usage(bool is_init, HTAB *local_active_table_stat_map)
 			{
 				elog(WARNING, "cache lookup failed for relation %u", relOid);
 				LWLockRelease(diskquota_locks.relation_cache_lock);
+
+				if (!is_init)
+					continue;
+
+				for (int i = -1; i < SEGCOUNT; i++)
+				{
+					appendStringInfo(&delete_statement, "%s(%u,%d)", (delete_entries_num == 0) ? " " : ", ",
+					                 relOid, i);
+
+					delete_entries_num++;
+
+					if (delete_entries_num > SQL_MAX_VALUES_NUMBER)
+					{
+						delete_from_table_size_map(delete_statement.data);
+						resetStringInfo(&delete_statement);
+						delete_entries_num = 0;
+					}
+				}
+
 				continue;
 			}
 			relnamespace  = relation_entry->namespaceoid;
@@ -1112,6 +1137,9 @@ calculate_table_disk_usage(bool is_init, HTAB *local_active_table_stat_map)
 		}
 	}
 
+	if (delete_entries_num) delete_from_table_size_map(delete_statement.data);
+
+	pfree(delete_statement.data);
 	list_free(oidlist);
 
 	/*