Refactor active table logic.

Put all the active table code into one file activetable.c
Replace relfilenode convert to Oid code by built-in function
Add answer files.

Co-authored-by: Haozhou Wang <hawang@pivotal.io>
Co-authored-by: Hubert Zhang <hzhang@pivotal.io>

Author: 2 people

gpcontrib/diskquota/activetable.c

@@ -6,8 +6,6 @@
  *
  * Copyright (C) 2013, PostgreSQL Global Development Group
  *
- * IDENTIFICATION
- *		contrib/diskquota/activetable.c
  *
  * -------------------------------------------------------------------------
  */
@@ -24,20 +22,21 @@
 #include "storage/smgr.h"
 #include "utils/builtins.h"
 #include "utils/fmgroids.h"
+#include "utils/relfilenodemap.h"
 
 #include "activetable.h"
-#include "diskquota.h"
 
 HTAB *active_tables_map = NULL;
 static SmgrStat_hook_type prev_SmgrStat_hook = NULL;
-static ScanKeyData relfilenode_skey[2];
 
 static void report_active_table_SmgrStat(SMgrRelation reln);
-HTAB* get_active_tables(void);
+static HTAB* get_active_tables_stats(void);
+static HTAB* get_all_tables_stats(void);
+
 void init_active_table_hook(void);
 void init_shm_worker_active_tables(void);
 void init_lock_active_tables(void);
-void init_relfilenode_key(void);
+HTAB* pg_fetch_active_tables(bool);
 
 /*
  * Register smgr hook to detect active table.
@@ -86,39 +85,84 @@ void init_lock_active_tables(void)
 }
 
 /*
- * Init relfilenode key to index search table oid 
- * given relfilenode and tablespace.
+ * Fetch active table file size statistics.
+ * If force is true, then fetch all the tables.
  */
-void
-init_relfilenode_key(void)
+HTAB* pg_fetch_active_tables(bool force)
+{
+	if (force)
+	{
+		return get_all_tables_stats();
+	}
+	else
+	{
+		return get_active_tables_stats();
+	}
+}
+
+/*
+ * Get the table size statistics for all the tables
+ */
+static HTAB* 
+get_all_tables_stats()
 {
-	int			i;
+	HTAB *local_table_stats_map = NULL;
+	HASHCTL ctl;
+	HeapTuple tuple;
+	Relation classRel;
+	HeapScanDesc relScan;
+
+	memset(&ctl, 0, sizeof(ctl));
+	ctl.keysize = sizeof(Oid);
+	ctl.entrysize = sizeof(DiskQuotaActiveTableEntry);
+	ctl.hcxt = CurrentMemoryContext;
+	ctl.hash = oid_hash;
 
-	/* build skey */
-	MemSet(&relfilenode_skey, 0, sizeof(relfilenode_skey));
+	local_table_stats_map = hash_create("local table map with table size info",
+								1024,
+								&ctl,
+								HASH_ELEM | HASH_CONTEXT | HASH_FUNCTION);
+	
+	classRel = heap_open(RelationRelationId, AccessShareLock);
+    relScan = heap_beginscan_catalog(classRel, 0, NULL);
 
-	for (i = 0; i < 2; i++)
+	while ((tuple = heap_getnext(relScan, ForwardScanDirection)) != NULL)
 	{
-		fmgr_info_cxt(F_OIDEQ,
-					  &relfilenode_skey[i].sk_func,
-					  CacheMemoryContext);
-		relfilenode_skey[i].sk_strategy = BTEqualStrategyNumber;
-		relfilenode_skey[i].sk_subtype = InvalidOid;
-		relfilenode_skey[i].sk_collation = InvalidOid;
+		Oid relOid;
+		DiskQuotaActiveTableEntry *entry;
+
+		Form_pg_class classForm = (Form_pg_class) GETSTRUCT(tuple);
+		if (classForm->relkind != RELKIND_RELATION &&
+				classForm->relkind != RELKIND_MATVIEW)
+			continue;
+		relOid = HeapTupleGetOid(tuple);
+
+		/* ignore system table*/
+		if (relOid < FirstNormalObjectId)
+			continue;
+
+		entry = (DiskQuotaActiveTableEntry *) hash_search(local_table_stats_map, &relOid, HASH_ENTER, NULL);
+
+		entry->tableoid = relOid;
+		entry->tablesize = (Size) DatumGetInt64(DirectFunctionCall1(pg_total_relation_size,
+					ObjectIdGetDatum(relOid)));
+
 	}
 
-	relfilenode_skey[0].sk_attno = Anum_pg_class_reltablespace;
-	relfilenode_skey[1].sk_attno = Anum_pg_class_relfilenode;
-}
+	heap_endscan(relScan);
+	heap_close(classRel, AccessShareLock);
 
+    return local_table_stats_map;	
+}
 /*
  * Get local active table with table oid and table size info.
  * This function first copies active table map from shared memory 
  * to local active table map with refilenode info. Then traverses
  * the local map and find corresponding table oid and table file 
  * size. Finnaly stores them into local active table map and return.
  */
-HTAB* get_active_tables()
+static HTAB* 
+get_active_tables_stats()
 {
 	HASHCTL ctl;
 	HTAB *local_active_table_file_map = NULL;
@@ -127,9 +171,6 @@ HTAB* get_active_tables()
 	DiskQuotaActiveTableFileEntry *active_table_file_entry;
 	DiskQuotaActiveTableEntry *active_table_entry;
 
-	Relation relation;
-	HeapTuple tuple;
-	SysScanDesc relScan;
 	Oid relOid;
 
 	memset(&ctl, 0, sizeof(ctl));
@@ -178,56 +219,34 @@ HTAB* get_active_tables()
 								&ctl,
 								HASH_ELEM | HASH_CONTEXT | HASH_FUNCTION);
 
-	relation = heap_open(RelationRelationId, AccessShareLock);
 	/* traverse local active table map and calculate their file size. */
 	hash_seq_init(&iter, local_active_table_file_map);
 	/* scan whole local map, get the oid of each table and calculate the size of them */
 	while ((active_table_file_entry = (DiskQuotaActiveTableFileEntry *) hash_seq_search(&iter)) != NULL)
 	{
 		Size tablesize;
 		bool found;
-		ScanKeyData skey[2];
-		Oid reltablespace;
 
-		reltablespace = active_table_file_entry->tablespaceoid;
-
-		/* pg_class will show 0 when the value is actually MyDatabaseTableSpace */
-		if (reltablespace == MyDatabaseTableSpace)
-			reltablespace = 0;
-
-		/* set scan arguments */
-		memcpy(skey, relfilenode_skey, sizeof(skey));
-		skey[0].sk_argument = ObjectIdGetDatum(reltablespace);
-		skey[1].sk_argument = ObjectIdGetDatum(active_table_file_entry->relfilenode);
-		relScan = systable_beginscan(relation,
-									ClassTblspcRelfilenodeIndexId,
-									true,
-									NULL,
-									2,
-									skey);
-
-		tuple = systable_getnext(relScan);
-
-		if (!HeapTupleIsValid(tuple))
+		relOid = RelidByRelfilenode(active_table_file_entry->tablespaceoid, active_table_file_entry->relfilenode);
+
+		//TODO replace DirectFunctionCall1 by a new total relation size function, which could handle Invalid relOid
+		/* avoid to generate ERROR if relOid is not existed (i.e. table has been droped) */
+		PG_TRY();
 		{
-			systable_endscan(relScan);
-			continue;
+			tablesize = (Size) DatumGetInt64(DirectFunctionCall1(pg_total_relation_size,
+						ObjectIdGetDatum(relOid)));
 		}
-		relOid = HeapTupleGetOid(tuple);
-
-		/* Call function directly to get size of table by oid */
-		tablesize = (Size) DatumGetInt64(DirectFunctionCall1(pg_total_relation_size, ObjectIdGetDatum(relOid)));
-
-		active_table_entry = hash_search(local_active_table_stats_map, &relOid, HASH_ENTER, &found);
-		if (active_table_entry)
+		PG_CATCH();
 		{
-			active_table_entry->tableoid = relOid;
-			active_table_entry->tablesize = tablesize;
+			FlushErrorState();
+			tablesize = 0;
 		}
-		systable_endscan(relScan);
+		PG_END_TRY();
+		active_table_entry = hash_search(local_active_table_stats_map, &relOid, HASH_ENTER, &found);
+		active_table_entry->tableoid = relOid;
+		active_table_entry->tablesize = tablesize;
 	}
 	elog(DEBUG1, "active table number is:%ld", hash_get_num_entries(local_active_table_file_map));
-	heap_close(relation, AccessShareLock);
 	hash_destroy(local_active_table_file_map);
 	return local_active_table_stats_map;
 }

gpcontrib/diskquota/activetable.h

@@ -1,7 +1,6 @@
 #ifndef ACTIVE_TABLE_H
 #define ACTIVE_TABLE_H
 
-#include "storage/lwlock.h"
 #include "diskquota.h"
 
 /* Cache to detect the active table list */
@@ -19,9 +18,8 @@ typedef struct DiskQuotaActiveTableEntry
 } DiskQuotaActiveTableEntry;
 
 
-extern HTAB* get_active_tables(void);
+extern HTAB* pg_fetch_active_tables(bool);
 extern void init_active_table_hook(void);
-extern void init_relfilenode_key(void);
 extern void init_shm_worker_active_tables(void);
 extern void init_lock_active_tables(void);
 

gpcontrib/diskquota/data/generate_insert.py

@@ -9,8 +9,6 @@
  *
  * Copyright (C) 2013, PostgreSQL Global Development Group
  *
- * IDENTIFICATION
- *		contrib/diskquota/diskquota.c
  *
  * -------------------------------------------------------------------------
  */
@@ -266,9 +264,6 @@ disk_quota_launcher_main(Datum main_arg)
 	/* We're now ready to receive signals */
 	BackgroundWorkerUnblockSignals();
 
-	/* Connect to our database */
-	//BackgroundWorkerInitializeConnection("postgres", NULL, 0);
-
 	memset(&hash_ctl, 0, sizeof(hash_ctl));
 	hash_ctl.keysize = NAMEDATALEN;
 	hash_ctl.entrysize = sizeof(DiskQuotaWorkerEntry);
@@ -279,7 +274,7 @@ disk_quota_launcher_main(Datum main_arg)
 										  HASH_ELEM);
 
 	dblist = get_database_list();
-
+	elog(LOG,"diskquota launcher started");
 	foreach(cell, dblist)
 	{
 		char *db_name;
@@ -468,6 +463,7 @@ start_worker(char* dbname)
 	bool found;
 	DiskQuotaWorkerEntry* workerentry;
 
+	memset(&worker, 0, sizeof(BackgroundWorker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
 		BGWORKER_BACKEND_DATABASE_CONNECTION;
 	worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
@@ -477,6 +473,7 @@ start_worker(char* dbname)
 	snprintf(worker.bgw_name, BGW_MAXLEN, "%s", dbname);
 	/* set bgw_notify_pid so that we can use WaitForBackgroundWorkerStartup */
 	worker.bgw_notify_pid = MyProcPid;
+	worker.bgw_main_arg = (Datum) 0;
 
 	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
 		return -1;

gpcontrib/diskquota/data/generate_table.py

@@ -7,8 +7,6 @@
  *
  * Copyright (C) 2013, PostgreSQL Global Development Group
  *
- * IDENTIFICATION
- *		contrib/diskquota/enforcement.c
  *
  * -------------------------------------------------------------------------
  */

gpcontrib/diskquota/diskquota.c

@@ -0,0 +1,4 @@
+drop table badquota.t1;
+drop role testbody;
+drop schema badquota;
+drop extension diskquota;

gpcontrib/diskquota/enforcement.c

@@ -0,0 +1,36 @@
+create extension diskquota;
+select pg_sleep(1);
+ pg_sleep 
+----------
+ 
+(1 row)
+
+\! pg_ctl -D /tmp/pg_diskquota_test/data reload
+server signaled
+\! cp data/csmall.txt /tmp/csmall.txt
+select pg_sleep(5);
+ pg_sleep 
+----------
+ 
+(1 row)
+
+-- prepare a schema that has reached quota limit
+create schema badquota;
+select diskquota.set_schema_quota('badquota', '1 MB');
+ set_schema_quota 
+------------------
+ 
+(1 row)
+
+create role testbody;
+create table badquota.t1(i int);
+alter table badquota.t1 owner to testbody;
+insert into badquota.t1 select generate_series(0, 50000);
+select pg_sleep(5);
+ pg_sleep 
+----------
+ 
+(1 row)
+
+insert into badquota.t1 select generate_series(0, 10);
+ERROR:  schema's disk space quota exceeded with name:badquota