Invalidate diskquota.table_size entries during startup

Diskquota calculates sizes and stores information in the diskquota.table_size
table periodically with a pause in diskquota.naptime, 2 seconds by default.
If we restart the cluster during this pause, then diskquota will lose all
changes that have occurred since the last save to the diskquota.table_size
table. We could create temporary tables, wait when it will be flushed to
diskquota.table_size table, restart the cluster, and diskquota would remember
the information about the temporary tables. Or we could delete the tables,
restart the cluster, and again diskquota will remember information about the
deleted tables. This happens because at the start of the cluster, diskquota
remembers all the information written to the diskquota.table_size table,
but does not check that some tables may have already been deleted.

As a solution, we invalidate diskquota.table_size during diskquota
worker start in addition to pg_class validation.

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);

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,19 @@ 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,
 	                            },

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,23 @@ 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 +1135,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);
 
 	/*

tests/isolation2/expected/test_dropped.out

@@ -0,0 +1,43 @@
+1: CREATE SCHEMA dropped_schema;
+CREATE
+1: SET search_path TO dropped_schema;
+SET
+1: SELECT diskquota.set_schema_quota('dropped_schema', '1 MB');
+ set_schema_quota 
+------------------
+
+(1 row)
+1: SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t                         
+(1 row)
+1: CREATE TABLE dropped_table(id int) DISTRIBUTED BY (id);
+CREATE
+1: INSERT INTO dropped_table SELECT generate_series(1, 100000);
+INSERT 100000
+1: SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t                         
+(1 row)
+1: DROP TABLE dropped_table;
+DROP
+1q: ... <quitting>
+
+!\retcode gpstop -afr;
+-- start_ignore
+-- end_ignore
+(exited with code 0)
+
+1: SELECT pg_table_size('dropped_table'::regclass::oid) size;
+ERROR:  relation "dropped_table" does not exist
+LINE 1: SELECT pg_table_size('dropped_table'::regclass::oid) size;
+                             ^
+1: SELECT size FROM diskquota.table_size WHERE NOT EXISTS (SELECT 1 FROM pg_class WHERE tableid = oid) AND segid = -1;
+ size 
+------
+(0 rows)
+1: DROP SCHEMA dropped_schema CASCADE;
+DROP
+1q: ... <quitting>

tests/isolation2/expected/test_temporary.out

@@ -0,0 +1,41 @@
+1: CREATE SCHEMA temporary_schema;
+CREATE
+1: SET search_path TO temporary_schema;
+SET
+1: SELECT diskquota.set_schema_quota('temporary_schema', '1 MB');
+ set_schema_quota 
+------------------
+
+(1 row)
+1: SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t                         
+(1 row)
+1: CREATE TEMPORARY TABLE temporary_table(id int) DISTRIBUTED BY (id);
+CREATE
+1: INSERT INTO temporary_table SELECT generate_series(1, 100000);
+INSERT 100000
+1: SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t                         
+(1 row)
+1q: ... <quitting>
+
+!\retcode gpstop -afr;
+-- start_ignore
+-- end_ignore
+(exited with code 0)
+
+1: SELECT pg_table_size('temporary_table'::regclass::oid) size;
+ERROR:  relation "temporary_table" does not exist
+LINE 1: SELECT pg_table_size('temporary_table'::regclass::oid) size;
+                             ^
+1: SELECT size FROM diskquota.table_size WHERE NOT EXISTS (SELECT 1 FROM pg_class WHERE tableid = oid) AND segid = -1;
+ size 
+------
+(0 rows)
+1: DROP SCHEMA temporary_schema CASCADE;
+DROP
+1q: ... <quitting>

tests/isolation2/isolation2_schedule

@@ -5,6 +5,8 @@ test: test_relation_size
 test: test_rejectmap
 test: test_vacuum
 test: test_truncate
+test: test_temporary
+test: test_dropped
 test: test_postmaster_restart
 test: test_worker_timeout
 test: test_per_segment_config

tests/isolation2/sql/test_dropped.sql

@@ -0,0 +1,16 @@
+1: CREATE SCHEMA dropped_schema;
+1: SET search_path TO dropped_schema;
+1: SELECT diskquota.set_schema_quota('dropped_schema', '1 MB');
+1: SELECT diskquota.wait_for_worker_new_epoch();
+1: CREATE TABLE dropped_table(id int) DISTRIBUTED BY (id);
+1: INSERT INTO dropped_table SELECT generate_series(1, 100000);
+1: SELECT diskquota.wait_for_worker_new_epoch();
+1: DROP TABLE dropped_table;
+1q:
+
+!\retcode gpstop -afr;
+
+1: SELECT pg_table_size('dropped_table'::regclass::oid) size;
+1: SELECT size FROM diskquota.table_size WHERE NOT EXISTS (SELECT 1 FROM pg_class WHERE tableid = oid) AND segid = -1;
+1: DROP SCHEMA dropped_schema CASCADE;
+1q:

tests/isolation2/sql/test_temporary.sql

@@ -0,0 +1,15 @@
+1: CREATE SCHEMA temporary_schema;
+1: SET search_path TO temporary_schema;
+1: SELECT diskquota.set_schema_quota('temporary_schema', '1 MB');
+1: SELECT diskquota.wait_for_worker_new_epoch();
+1: CREATE TEMPORARY TABLE temporary_table(id int) DISTRIBUTED BY (id);
+1: INSERT INTO temporary_table SELECT generate_series(1, 100000);
+1: SELECT diskquota.wait_for_worker_new_epoch();
+1q:
+
+!\retcode gpstop -afr;
+
+1: SELECT pg_table_size('temporary_table'::regclass::oid) size;
+1: SELECT size FROM diskquota.table_size WHERE NOT EXISTS (SELECT 1 FROM pg_class WHERE tableid = oid) AND segid = -1;
+1: DROP SCHEMA temporary_schema CASCADE;
+1q: