add function pg_clean_all_cache,pg_clean_rel_cache in pg_prewarm exte…

…nsion to clean cache
ApsaraDB · Nov 20, 2023 · e56626c · e56626c
1 parent dd60f45
commit e56626c
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 4 deletions.
diff --git a/contrib/pg_prewarm/pg_prewarm--1.1.sql b/contrib/pg_prewarm/pg_prewarm--1.1.sql
@@ -12,3 +12,14 @@ CREATE FUNCTION pg_prewarm(regclass,
 RETURNS int8
 AS 'MODULE_PATHNAME', 'pg_prewarm'
 LANGUAGE C PARALLEL SAFE;
+
+
+CREATE OR REPLACE FUNCTION pg_clean_all_cache()
+RETURNS VOID
+AS 'MODULE_PATHNAME', 'pg_clean_all_cache'
+LANGUAGE C;
+
+CREATE OR REPLACE FUNCTION pg_clean_rel_cache(relation regclass, fork text default null)
+RETURNS VOID
+AS 'MODULE_PATHNAME', 'pg_clean_rel_cache'
+LANGUAGE C;
diff --git a/contrib/pg_prewarm/pg_prewarm.c b/contrib/pg_prewarm/pg_prewarm.c
@@ -25,9 +25,17 @@
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
 
+#include "utils/syscache.h"
+#include "access/htup_details.h"
+#include "miscadmin.h"
+#include "catalog/pg_class.h"
+
+
 PG_MODULE_MAGIC;
 
 PG_FUNCTION_INFO_V1(pg_prewarm);
+PG_FUNCTION_INFO_V1(pg_clean_all_cache);
+PG_FUNCTION_INFO_V1(pg_clean_rel_cache);
 
 typedef enum
 {
@@ -192,10 +200,7 @@ pg_prewarm(PG_FUNCTION_ARGS)
 			Buffer		buf;
 
 			CHECK_FOR_INTERRUPTS();
-			//buf = ReadBufferExtended(rel, forkNumber, block, RBM_NORMAL, NULL);
-			buf = polar_bulk_read_buffer_extended(rel, MAIN_FORKNUM, block,
-														RBM_NORMAL, NULL,
-														64);
+			buf = ReadBufferExtended(rel, forkNumber, block, RBM_NORMAL, NULL);
 			ReleaseBuffer(buf);
 			++blocks_done;
 		}
@@ -206,3 +211,46 @@ pg_prewarm(PG_FUNCTION_ARGS)
 
 	PG_RETURN_INT64(blocks_done);
 }
+
+
+Datum
+pg_clean_all_cache(PG_FUNCTION_ARGS)
+{
+	DropDatabaseBuffers(MyDatabaseId);
+	PG_RETURN_VOID();
+}
+
+Datum
+pg_clean_rel_cache(PG_FUNCTION_ARGS)
+{
+	Oid			relid = PG_GETARG_OID(0);
+	int			forkNum;
+	HeapTuple	tp;
+	RelFileNodeBackend	rnode;
+
+	tp = SearchSysCache1(RELOID, ObjectIdGetDatum(relid));
+	if (HeapTupleIsValid(tp))
+	{
+		Form_pg_class reltup = (Form_pg_class) GETSTRUCT(tp);
+
+		rnode.node.relNode = reltup->relfilenode;
+		rnode.node.spcNode = (reltup->reltablespace == InvalidOid) ?
+			MyDatabaseTableSpace :
+			reltup->reltablespace;
+		rnode.node.dbNode = MyDatabaseId;
+		rnode.backend = InvalidBackendId;
+
+		ReleaseSysCache(tp);
+	}
+	else
+	{
+		elog(ERROR, "Unable to get cache for relation %u", relid);
+		PG_RETURN_VOID();
+	}
+
+	forkNum = PG_ARGISNULL(1) ? 0 : forkname_to_number(text_to_cstring(PG_GETARG_TEXT_P(1)));
+	for (; forkNum <= MAX_FORKNUM; ++forkNum)
+		DropRelFileNodeBuffers(rnode, forkNum, 0);
+
+	PG_RETURN_VOID();
+}