disk_quota_hooks_v2.patch

From 3d275c78b304b308d288bd227f6dcab45dc5f595 Mon Sep 17 00:00:00 2001
From: Hubert Zhang <hubertzhang@apache.org>
Date: Tue, 6 Nov 2018 06:51:22 +0000
Subject: [PATCH] Add hooks for diskquota extension.

Add ReadBufferExtend_hook() and smgr*_hook()
hook points to extend logic of storage
management.

Co-authored-by: Haozhou Wang <hawang@pivotal.io>
Co-authored-by: Hubert Zhang <hzhang@pivotal.io>
Co-authored-by: Hao Wu <gfphoenix78@gmail.com>
---
 src/backend/storage/buffer/bufmgr.c | 14 ++++++++++++++
 src/backend/storage/smgr/smgr.c     | 33 +++++++++++++++++++++++++++++++++
 src/include/storage/bufmgr.h        | 10 ++++++++++
 src/include/storage/smgr.h          | 18 ++++++++++++++++++
 4 files changed, 75 insertions(+)

diff --git a/src/backend/storage/buffer/bufmgr.c b/src/backend/storage/buffer/bufmgr.c
index 01eabe5706..5499495506 100644
--- a/src/backend/storage/buffer/bufmgr.c
+++ b/src/backend/storage/buffer/bufmgr.c
@@ -104,6 +104,13 @@ typedef struct CkptTsStatus
 	int			index;
 } CkptTsStatus;
 
+/*
+ * Hook for plugins to check permissions when doing a buffer extend.
+ * One example is to check whether there is additional disk quota for
+ * the table to be inserted.
+ */
+ReadBufferExtended_hook_type ReadBufferExtended_hook = NULL;
+
 /* GUC variables */
 bool		zero_damaged_pages = false;
 int			bgwriter_lru_maxpages = 100;
@@ -661,6 +668,13 @@ ReadBufferExtended(Relation reln, ForkNumber forkNum, BlockNumber blockNum,
 	 * miss.
 	 */
 	pgstat_count_buffer_read(reln);
+
+	/* hook function for doing a buffer extend */
+	if (blockNum == P_NEW && ReadBufferExtended_hook)
+	{
+		(*ReadBufferExtended_hook)(reln, forkNum, blockNum, mode, strategy);
+	}
+
 	buf = ReadBuffer_common(reln->rd_smgr, reln->rd_rel->relpersistence,
 							forkNum, blockNum, mode, strategy, &hit);
 	if (hit)
diff --git a/src/backend/storage/smgr/smgr.c b/src/backend/storage/smgr/smgr.c
index 189342ef86..fa36a18e15 100644
--- a/src/backend/storage/smgr/smgr.c
+++ b/src/backend/storage/smgr/smgr.c
@@ -90,6 +90,16 @@ static const f_smgr smgrsw[] = {
 
 static const int NSmgr = lengthof(smgrsw);
 
+/*
+ * Hook for plugins to extend smgr functions.
+ * for example, collect statistics from smgr functions
+ * via recording the active relfilenode information.
+ */
+smgrcreate_hook_type smgrcreate_hook = NULL;
+smgrextend_hook_type smgrextend_hook = NULL;
+smgrtruncate_hook_type smgrtruncate_hook = NULL;
+smgrdounlinkall_hook_type smgrdounlinkall_hook = NULL;
+
 
 /*
  * Each backend has a hashtable that stores all extant SMgrRelation objects.
@@ -397,6 +407,11 @@ smgrcreate(SMgrRelation reln, ForkNumber forknum, bool isRedo)
 	if (isRedo && reln->md_num_open_segs[forknum] > 0)
 		return;
 
+	if (smgrcreate_hook)
+	{
+		(*smgrcreate_hook)(reln, forknum, isRedo);
+	}
+
 	/*
 	 * We may be using the target table space for the first time in this
 	 * database, so create a per-database subdirectory if needed.
@@ -411,6 +426,7 @@ smgrcreate(SMgrRelation reln, ForkNumber forknum, bool isRedo)
 							isRedo);
 
 	smgrsw[reln->smgr_which].smgr_create(reln, forknum, isRedo);
+
 }
 
 /*
@@ -492,6 +508,11 @@ smgrdounlinkall(SMgrRelation *rels, int nrels, bool isRedo)
 	if (nrels == 0)
 		return;
 
+	if (smgrdounlinkall_hook)
+	{
+		(*smgrdounlinkall_hook)(rels, nrels, isRedo);
+	}
+
 	/*
 	 * create an array which contains all relations to be dropped, and close
 	 * each relation's forks at the smgr level while at it
@@ -615,8 +636,14 @@ void
 smgrextend(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 		   char *buffer, bool skipFsync)
 {
+	if (smgrextend_hook)
+	{
+		(*smgrextend_hook)(reln, forknum, blocknum, buffer, skipFsync);
+	}
+
 	smgrsw[reln->smgr_which].smgr_extend(reln, forknum, blocknum,
 										 buffer, skipFsync);
+
 }
 
 /*
@@ -698,6 +725,11 @@ smgrnblocks(SMgrRelation reln, ForkNumber forknum)
 void
 smgrtruncate(SMgrRelation reln, ForkNumber forknum, BlockNumber nblocks)
 {
+	if (smgrtruncate_hook)
+	{
+		(*smgrtruncate_hook)(reln, forknum, nblocks);
+	}
+
 	/*
 	 * Get rid of any buffers for the about-to-be-deleted blocks. bufmgr will
 	 * just drop them without bothering to write the contents.
@@ -720,6 +752,7 @@ smgrtruncate(SMgrRelation reln, ForkNumber forknum, BlockNumber nblocks)
 	 * Do the truncation.
 	 */
 	smgrsw[reln->smgr_which].smgr_truncate(reln, forknum, nblocks);
+
 }
 
 /*
diff --git a/src/include/storage/bufmgr.h b/src/include/storage/bufmgr.h
index 3cce3906a0..f1dcc77bf7 100644
--- a/src/include/storage/bufmgr.h
+++ b/src/include/storage/bufmgr.h
@@ -159,6 +159,16 @@ extern PGDLLIMPORT int32 *LocalRefCount;
  */
 #define BufferGetPage(buffer) ((Page)BufferGetBlock(buffer))
 
+/*
+ * Hook for plugins to add external logic when doing a buffer extend.
+ * One example is to check whether there is additional disk quota for
+ * the table to be inserted.
+ */
+typedef bool (*ReadBufferExtended_hook_type) (Relation reln,
+					ForkNumber forkNum, BlockNumber blockNum,
+					ReadBufferMode mode, BufferAccessStrategy strategy);
+extern PGDLLIMPORT ReadBufferExtended_hook_type ReadBufferExtended_hook;
+
 /*
  * prototypes for functions in bufmgr.c
  */
diff --git a/src/include/storage/smgr.h b/src/include/storage/smgr.h
index c843bbc969..d070b3d573 100644
--- a/src/include/storage/smgr.h
+++ b/src/include/storage/smgr.h
@@ -144,5 +144,23 @@ extern void RememberFsyncRequest(RelFileNode rnode, ForkNumber forknum,
 extern void ForgetRelationFsyncRequests(RelFileNode rnode, ForkNumber forknum);
 extern void ForgetDatabaseFsyncRequests(Oid dbid);
 extern void DropRelationFiles(RelFileNode *delrels, int ndelrels, bool isRedo);
+/*
+ * Hook for plugins to extend smgr functions.
+ * for example, collect statistics from smgr functions
+ * via recording the active relfilenode information.
+ */
+typedef void (*smgrcreate_hook_type)(SMgrRelation reln, ForkNumber forknum,
+									 bool isRedo);
+extern PGDLLIMPORT smgrcreate_hook_type smgrcreate_hook;
+typedef void (*smgrextend_hook_type)(SMgrRelation reln, ForkNumber forknum,
+									 BlockNumber blocknum,
+									 char *buffer, bool skipFsync);
+extern PGDLLIMPORT smgrextend_hook_type smgrextend_hook;
+typedef void (*smgrtruncate_hook_type)(SMgrRelation reln, ForkNumber forknum,
+									   BlockNumber nblocks);
+extern PGDLLIMPORT smgrtruncate_hook_type smgrtruncate_hook;
+typedef void (*smgrdounlinkall_hook_type)(SMgrRelation *rels, int nrels,
+										  bool isRedo);
+extern PGDLLIMPORT smgrdounlinkall_hook_type smgrdounlinkall_hook;
 
 #endif							/* SMGR_H */
-- 
2.16.2