From 72dbd541e8efdbaa7340df871823da71f3c6e0fa Mon Sep 17 00:00:00 2001
From: reshke kirill <reshke@double.cloud>
Date: Mon, 25 Nov 2024 18:03:44 +0000
Subject: [PATCH v2] Use stream read interface for pgstattuple routines.

Patch implements new streaming read API for pgstattuple contrib
extension.

This is not perfomance speedup patch.
This patch enables this extension to benefit from stream API in the future
without introducing any performance improvements.
---
 contrib/pgstattuple/pgstattuple.c | 64 +++++++++++++++++++------------
 1 file changed, 39 insertions(+), 25 deletions(-)

diff --git a/contrib/pgstattuple/pgstattuple.c b/contrib/pgstattuple/pgstattuple.c
index 48cb8f59c4f..79b0157b9c0 100644
--- a/contrib/pgstattuple/pgstattuple.c
+++ b/contrib/pgstattuple/pgstattuple.c
@@ -61,22 +61,18 @@ typedef struct pgstattuple_type
 	uint64		free_space;		/* free/reusable space in bytes */
 } pgstattuple_type;
 
-typedef void (*pgstat_page) (pgstattuple_type *, Relation, BlockNumber,
-							 BufferAccessStrategy);
+typedef void (*pgstat_page) (pgstattuple_type *, Relation, Buffer);
 
 static Datum build_pgstattuple_type(pgstattuple_type *stat,
 									FunctionCallInfo fcinfo);
 static Datum pgstat_relation(Relation rel, FunctionCallInfo fcinfo);
 static Datum pgstat_heap(Relation rel, FunctionCallInfo fcinfo);
 static void pgstat_btree_page(pgstattuple_type *stat,
-							  Relation rel, BlockNumber blkno,
-							  BufferAccessStrategy bstrategy);
+							  Relation rel, Buffer buf);
 static void pgstat_hash_page(pgstattuple_type *stat,
-							 Relation rel, BlockNumber blkno,
-							 BufferAccessStrategy bstrategy);
+							 Relation rel, Buffer buf);
 static void pgstat_gist_page(pgstattuple_type *stat,
-							 Relation rel, BlockNumber blkno,
-							 BufferAccessStrategy bstrategy);
+							 Relation rel, Buffer buf);
 static Datum pgstat_index(Relation rel, BlockNumber start,
 						  pgstat_page pagefn, FunctionCallInfo fcinfo);
 static void pgstat_index_page(pgstattuple_type *stat, Page page,
@@ -405,13 +401,10 @@ pgstat_heap(Relation rel, FunctionCallInfo fcinfo)
  * pgstat_btree_page -- check tuples in a btree page
  */
 static void
-pgstat_btree_page(pgstattuple_type *stat, Relation rel, BlockNumber blkno,
-				  BufferAccessStrategy bstrategy)
+pgstat_btree_page(pgstattuple_type *stat, Relation rel, Buffer buf)
 {
-	Buffer		buf;
 	Page		page;
 
-	buf = ReadBufferExtended(rel, MAIN_FORKNUM, blkno, RBM_NORMAL, bstrategy);
 	LockBuffer(buf, BT_READ);
 	page = BufferGetPage(buf);
 
@@ -449,13 +442,15 @@ pgstat_btree_page(pgstattuple_type *stat, Relation rel, BlockNumber blkno,
  * pgstat_hash_page -- check tuples in a hash page
  */
 static void
-pgstat_hash_page(pgstattuple_type *stat, Relation rel, BlockNumber blkno,
-				 BufferAccessStrategy bstrategy)
+pgstat_hash_page(pgstattuple_type *stat, Relation rel, Buffer buf)
 {
-	Buffer		buf;
 	Page		page;
 
-	buf = _hash_getbuf_with_strategy(rel, blkno, HASH_READ, 0, bstrategy);
+	LockBuffer(buf, HASH_READ);
+
+	/* ref count and lock type are correct */
+
+	_hash_checkpage(rel, buf, 0);
 	page = BufferGetPage(buf);
 
 	if (PageGetSpecialSize(page) == MAXALIGN(sizeof(HashPageOpaqueData)))
@@ -491,13 +486,10 @@ pgstat_hash_page(pgstattuple_type *stat, Relation rel, BlockNumber blkno,
  * pgstat_gist_page -- check tuples in a gist page
  */
 static void
-pgstat_gist_page(pgstattuple_type *stat, Relation rel, BlockNumber blkno,
-				 BufferAccessStrategy bstrategy)
+pgstat_gist_page(pgstattuple_type *stat, Relation rel, Buffer buf)
 {
-	Buffer		buf;
 	Page		page;
 
-	buf = ReadBufferExtended(rel, MAIN_FORKNUM, blkno, RBM_NORMAL, bstrategy);
 	LockBuffer(buf, GIST_SHARE);
 	gistcheckpage(rel, buf);
 	page = BufferGetPage(buf);
@@ -523,14 +515,24 @@ pgstat_index(Relation rel, BlockNumber start, pgstat_page pagefn,
 			 FunctionCallInfo fcinfo)
 {
 	BlockNumber nblocks;
-	BlockNumber blkno;
 	BufferAccessStrategy bstrategy;
 	pgstattuple_type stat = {0};
+	Buffer		buf;
+	BlockRangeReadStreamPrivate p;
+	ReadStream *stream = NULL;
 
 	/* prepare access strategy for this index */
 	bstrategy = GetAccessStrategy(BAS_BULKREAD);
 
-	blkno = start;
+	p.current_blocknum = start;
+	stream = read_stream_begin_relation(READ_STREAM_FULL,
+										bstrategy,
+										rel,
+										MAIN_FORKNUM,
+										block_range_read_stream_cb,
+										&p,
+										0);
+
 	for (;;)
 	{
 		/* Get the current relation length */
@@ -539,21 +541,33 @@ pgstat_index(Relation rel, BlockNumber start, pgstat_page pagefn,
 		UnlockRelationForExtension(rel, ExclusiveLock);
 
 		/* Quit if we've scanned the whole relation */
-		if (blkno >= nblocks)
+		if (p.current_blocknum >= nblocks)
 		{
 			stat.table_len = (uint64) nblocks * BLCKSZ;
 
 			break;
 		}
 
-		for (; blkno < nblocks; blkno++)
+		p.last_exclusive = nblocks;
+
+		while (BufferIsValid(buf = read_stream_next_buffer(stream, NULL)))
 		{
 			CHECK_FOR_INTERRUPTS();
 
-			pagefn(&stat, rel, blkno, bstrategy);
+			pagefn(&stat, rel, buf);
 		}
+
+		Assert(read_stream_next_buffer(stream, NULL) == InvalidBuffer);
+
+		/*
+		 * After reaching the end we have to reset stream to use it again.
+		 * Extra restart in case of just one iteration does not cost us much.
+		 */
+		read_stream_reset(stream);
 	}
 
+	read_stream_end(stream);
+
 	relation_close(rel, AccessShareLock);
 
 	return build_pgstattuple_type(&stat, fcinfo);
-- 
2.34.1