Add ptrack_get_and_clear function.

stalkerg · stalkerg · commit c75895c61b89 · 2016-05-10T18:39:57.000+03:00
diff --git a/src/backend/access/heap/ptrack.c b/src/backend/access/heap/ptrack.c
@@ -15,9 +15,9 @@
 #include "storage/smgr.h"
 #include "utils/inval.h"
 #include "utils/array.h"
+#include "utils/relfilenodemap.h"
 #include <unistd.h>
 
-
 /* Effective data size */
 #define MAPSIZE (BLCKSZ - MAXALIGN(SizeOfPageHeaderData))
 
@@ -252,7 +252,7 @@ ptrack_clear(void)
 			char *map = PageGetContents(page);
 			LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
 			START_CRIT_SECTION();
-			MemSet(map, 0, BLCKSZ-MAXALIGN(SizeOfPageHeaderData));
+			MemSet(map, 0, MAPSIZE);
 			MarkBufferDirty(buf);
 			END_CRIT_SECTION_WITHOUT_TRACK();
 			LockBuffer(buf, BUFFER_LOCK_UNLOCK);
@@ -270,6 +270,79 @@ ptrack_clear(void)
 	SetPtrackClearLSN(false);
 }
 
+/* Get ptrack file as bytea and clear him */
+bytea *
+ptrack_get_and_clear(Oid tablespace_oid, Oid table_oid)
+{
+	bytea *result = NULL;
+	BlockNumber nblock;
+	if (table_oid == InvalidOid)
+	{
+		elog(WARNING, "InvalidOid");
+		goto full_end;
+	}
+	Relation rel = RelationIdGetRelation(RelidByRelfilenode(tablespace_oid,
+															table_oid));
+
+	if (rel == InvalidRelation)
+	{
+		elog(WARNING, "InvalidRelation");
+		goto full_end;
+	}
+
+	RelationOpenSmgr(rel);
+	if (rel->rd_smgr == NULL)
+		goto end_rel;
+
+	LockSmgrForExtension(rel->rd_smgr, ExclusiveLock);
+
+	if (rel->rd_smgr->smgr_ptrack_nblocks == InvalidBlockNumber)
+	{
+		if (smgrexists(rel->rd_smgr, PAGESTRACK_FORKNUM))
+			rel->rd_smgr->smgr_ptrack_nblocks = smgrnblocks(rel->rd_smgr,
+															PAGESTRACK_FORKNUM);
+		else
+			rel->rd_smgr->smgr_ptrack_nblocks = 0;
+	}
+	if (rel->rd_smgr->smgr_ptrack_nblocks == 0)
+	{
+		UnlockSmgrForExtension(rel->rd_smgr, ExclusiveLock);
+		goto end_rel;
+	}
+	result = (bytea *) palloc(rel->rd_smgr->smgr_ptrack_nblocks*MAPSIZE + VARHDRSZ);
+	SET_VARSIZE(result, rel->rd_smgr->smgr_ptrack_nblocks*MAPSIZE + VARHDRSZ);
+
+	for(nblock = 0; nblock < rel->rd_smgr->smgr_ptrack_nblocks; nblock++)
+	{
+		Buffer	buf = ReadBufferExtended(rel, PAGESTRACK_FORKNUM,
+			   nblock, RBM_ZERO_ON_ERROR, NULL);
+		Page page = BufferGetPage(buf);
+		char *map = PageGetContents(page);
+		LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
+		START_CRIT_SECTION();
+		memcpy(VARDATA(result) + nblock*MAPSIZE, map, MAPSIZE);
+		MemSet(map, 0, MAPSIZE);
+		MarkBufferDirty(buf);
+		END_CRIT_SECTION_WITHOUT_TRACK();
+		LockBuffer(buf, BUFFER_LOCK_UNLOCK);
+		ReleaseBuffer(buf);
+	}
+
+	UnlockSmgrForExtension(rel->rd_smgr, ExclusiveLock);
+	end_rel:
+		RelationClose(rel);
+
+	SetPtrackClearLSN(false);
+	full_end:
+	if (result == NULL)
+	{
+		result = palloc0(VARHDRSZ);
+		SET_VARSIZE(result, VARHDRSZ);
+	}
+
+	return result;
+}
+
 void
 SetPtrackClearLSN(bool set_invalid)
 {
diff --git a/src/backend/access/transam/xlogfuncs.c b/src/backend/access/transam/xlogfuncs.c
@@ -115,6 +115,17 @@ pg_ptrack_clear(PG_FUNCTION_ARGS)
 	PG_RETURN_VOID();
 }
 
+Datum
+pg_ptrack_get_and_clear(PG_FUNCTION_ARGS)
+{
+	if (!superuser() && !has_rolreplication(GetUserId()))
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+		 (errmsg("must be superuser or replication role to clear ptrack files"))));
+
+	PG_RETURN_BYTEA_P(ptrack_get_and_clear(PG_GETARG_OID(0), PG_GETARG_OID(1)));
+}
+
 /*
  * pg_switch_xlog: switch to next xlog file
  */
diff --git a/src/include/access/ptrack.h b/src/include/access/ptrack.h
@@ -16,6 +16,7 @@ extern void ptrack_save(void);
 extern void ptrack_add_block(BlockNumber block_number, RelFileNode rel);
 
 extern void ptrack_clear(void);
+extern bytea *ptrack_get_and_clear(Oid tablespace_oid, Oid table_oid);
 extern void assign_ptrack_enable(bool newval, void *extra);
 
-#endif   /* VISIBILITYMAP_H */
+#endif   /* PTRACK_H */
diff --git a/src/include/catalog/pg_proc.h b/src/include/catalog/pg_proc.h
@@ -3147,6 +3147,8 @@ DATA(insert OID = 6016 ( pg_ptrack_clear		PGNSP PGUID 12 1 0 0 0 f f f f t f v 0
 DESCR("clear ptrack fork files");
 DATA(insert OID = 6017 ( pg_ptrack_test		PGNSP PGUID 12 1 0 0 0 f f f f t f v 1 0 2277 "26" _null_ _null_ _null_ _null_ _null_ pg_ptrack_test _null_ _null_ _null_ ));
 DESCR("test ptrack fork relation");
+DATA(insert OID = 6018 ( pg_ptrack_get_and_clear		PGNSP PGUID 12 1 0 0 0 f f f f t f v 2 0 17 "26 26" _null_ _null_ _null_ _null_ _null_ pg_ptrack_get_and_clear _null_ _null_ _null_ ));
+DESCR("get ptrack file as bytea and clear him");
 DATA(insert OID = 3098 ( pg_create_restore_point	PGNSP PGUID 12 1 0 0 0 f f f f t f v 1 0 3220 "25" _null_ _null_ _null_ _null_ _null_ pg_create_restore_point _null_ _null_ _null_ ));
 DESCR("create a named restore point");
 DATA(insert OID = 2849 ( pg_current_xlog_location	PGNSP PGUID 12 1 0 0 0 f f f f t f v 0 0 3220 "" _null_ _null_ _null_ _null_ _null_ pg_current_xlog_location _null_ _null_ _null_ ));
